[TIR] Fix data dependent indexing when lowering TE to TIR (apache#8217)

A conversion pass was missing the recursive VisitExpr statement.

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<te::Tensor>(op->producer);
+    auto visited_op = Downcast<ProducerLoad>(StmtExprMutator::VisitExpr_(op));
+    te::Tensor tensor = Downcast<te::Tensor>(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:

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()