| import tvm |
| |
| def test_inline(): |
| m = tvm.var('m') |
| A = tvm.placeholder((m,), name='A') |
| T = tvm.compute((m,), lambda i,: A[i] + 10, name='T') |
| stmt = tvm.make.Evaluate(T[10] + 11 * T[100]) |
| stmt = tvm.ir_pass.Inline( |
| stmt, T.op, [x.var for x in T.op.axis], T.op.body[0]) |
| print(stmt) |
| assert(tvm.ir_pass.VerifySSA(stmt)) |
| |
| try: |
| # pass in int array(wrong argument type) |
| # must raise an error |
| stmt = tvm.ir_pass.Inline( |
| T.op, [1,2,3], T.op.body, stmt) |
| assert False |
| except tvm.TVMError: |
| pass |
| |
| def test_inline2(): |
| m = tvm.var('m') |
| A = tvm.placeholder((m,), name='A') |
| T = tvm.compute((m,), lambda i,: A[i] + 10, name='T') |
| stmt = tvm.make.Evaluate(tvm.exp(T[10]) + 11 * T[100]) |
| stmt = tvm.ir_pass.Inline( |
| stmt, T.op, [x.var for x in T.op.axis], T.op.body[0]) |
| def check(op): |
| if isinstance(op, tvm.expr.Call): |
| assert op.func != T.op |
| tvm.ir_pass.PostOrderVisit(stmt, check) |
| |
| |
| if __name__ == "__main__": |
| test_inline2() |
| test_inline() |
