From 086f676f0c27eed40e1e7d43d468a0c06a01713c Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Tue, 8 Jun 2021 22:53:04 -0600 Subject: [PATCH] [TIR] Fix data dependent indexing when lowering TE to TIR (#8217) A conversion pass was missing the recursive VisitExpr statement. --- src/te/operation/create_primfunc.cc | 5 +++-- tests/python/unittest/test_te_create_primfunc.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index ca2a1305434b..190892b2283f 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -35,11 +35,12 @@ class ProducerToBufferTransformer : public StmtExprMutator { : tensor2buffers_(tensor2buffers) {} PrimExpr VisitExpr_(const ProducerLoadNode* op) final { - te::Tensor tensor = Downcast(op->producer); + auto visited_op = Downcast(StmtExprMutator::VisitExpr_(op)); + te::Tensor tensor = Downcast(visited_op->producer); auto it = tensor2buffers_.find(tensor); ICHECK(it != tensor2buffers_.end()) << "IndexError: Cannot find the tensor " << tensor; const Buffer& buffer = it->second; - return BufferLoad(buffer, op->indices); + return BufferLoad(buffer, visited_op->indices); } private: diff --git a/tests/python/unittest/test_te_create_primfunc.py b/tests/python/unittest/test_te_create_primfunc.py index 6fbf95d24101..52dc4ccd9fef 100644 --- a/tests/python/unittest/test_te_create_primfunc.py +++ b/tests/python/unittest/test_te_create_primfunc.py @@ -300,6 +300,21 @@ def test_constant(): tvm.testing.assert_allclose(a_np + 2, c.numpy()) +def test_data_dependent_access(): + A = te.placeholder((10,), name="A") + B = te.placeholder((10,), name="B", dtype="int32") + C = te.compute((10,), lambda i: A[B[i]]) + + func = te.create_prim_func([C, A, B]) + func = tvm.build(func) + + a_np = np.random.uniform(size=(10,)).astype(A.dtype) + b_np = np.arange(10, dtype=B.dtype) + c = tvm.nd.array(np.zeros(10, dtype=C.dtype)) + func(c, tvm.nd.array(a_np), tvm.nd.array(b_np)) + tvm.testing.assert_allclose(a_np[b_np], c.numpy()) + + if __name__ == "__main__": test_unique_name() test_matmul()