[usmp] Also remap VarNode to USMP-allocated buffer

src/tir/usmp/analysis/extract_buffer_info.cc

@@ -429,15 +429,17 @@ void BufferInfoExtractor::VisitExpr_(const VarNode* op) {
 
 Array<Var> static GetMatchedBuffers(const PrimFunc& func) {
   Array<Var> buffer_vars;
-  for (unsigned int i = 0; i < func->params.size() - 1; i++) {
-    Var param = func->params[i];
-    buffer_vars.push_back(func->buffer_map[param]->data);
-  }
-  Var last_param = func->params.back();
-  // Checks whether last var is present in the buffer map
-  // because it could be the resource handle
-  if (func->buffer_map.find(last_param) != func->buffer_map.end()) {
-    buffer_vars.push_back(func->buffer_map[last_param]->data);
+  if (func->params.size() > 0) {
+    for (unsigned int i = 0; i < func->params.size() - 1; i++) {
+      Var param = func->params[i];
+      buffer_vars.push_back(func->buffer_map[param]->data);
+    }
+    Var last_param = func->params.back();
+    // Checks whether last var is present in the buffer map
+    // because it could be the resource handle
+    if (func->buffer_map.find(last_param) != func->buffer_map.end()) {
+      buffer_vars.push_back(func->buffer_map[last_param]->data);
+    }
   }
   return buffer_vars;
 }

src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc

@@ -96,6 +96,7 @@ class PoolAllocationToOffsetConverter : public StmtExprMutator {
  private:
   PrimExpr VisitExpr_(const CallNode* op) override;
   Stmt VisitStmt_(const AllocateNode* op) override;
+  PrimExpr VisitExpr_(const VarNode* op) override;
   PrimExpr VisitExpr_(const BufferLoadNode* op) override;
   Stmt VisitStmt_(const BufferStoreNode* op) override;
 
@@ -395,6 +396,15 @@ PrimExpr PoolAllocationToOffsetConverter::VisitExpr_(const BufferLoadNode* op) {
   return std::move(load);
 }
 
+PrimExpr PoolAllocationToOffsetConverter::VisitExpr_(const VarNode* op) {
+  auto it = allocate_var_to_let_var_.find(GetRef<Var>(op));
+  if (it != allocate_var_to_let_var_.end()) {
+    return (*it).second;
+  }
+
+  return StmtExprMutator::VisitExpr_(op);
+}
+
 Buffer PoolAllocationToOffsetConverter::GetRemappedBuffer(Buffer original) {
   {
     auto it = original_buf_to_let_buf_.find(original);

tests/python/unittest/test_tir_usmp_transform_convert_pool_allocations_to_offsets.py

@@ -600,5 +600,98 @@ def test_resnet_subgraph():
         tvm.ir.assert_structural_equal(actual_func, ref_func)
 
 
+@tvm.script.ir_module
+class TensorIntrinStructure:
+    @T.prim_func
+    def tensor_intrin_primfunc() -> None:
+        dense_data = T.allocate([10], "int32", "global")
+        T.evaluate(
+            T.call_extern(
+                "intrin_function",
+                T.tvm_access_ptr(
+                    T.type_annotation(dtype="int32"), dense_data, 0, 1, 2, dtype="handle"
+                ),
+                dtype="int32",
+            )
+        )
+
+        dense = T.buffer_decl([10], "int32", data=dense_data)
+        dense[0] = T.q_multiply_shift(dense[0], 1608879842, 31, -7, dtype="int32")
+
+    @T.prim_func
+    def __tvm_main__(input: T.handle, output: T.handle) -> None:
+        T.evaluate(T.call_extern("tensor_intrin_primfunc", dtype="int32"))
+
+
+@tvm.script.ir_module
+class TensorIntrinStructurePlanned:
+    @T.prim_func
+    def tensor_intrin_primfunc(global_workspace_1_var: T.Ptr[T.uint8]) -> None:
+        global_workspace_1_buffer_var = T.match_buffer(
+            global_workspace_1_var, [40], dtype="uint8", strides=[1], elem_offset=0, align=16
+        )
+        T.preflattened_buffer(
+            global_workspace_1_buffer_var, [40], dtype="uint8", strides=[1], elem_offset=0, align=16
+        )
+        dense_let = T.buffer_decl([10], "int32")
+        with T.let(dense_let.data, T.address_of(global_workspace_1_buffer_var[0], dtype="handle")):
+            T.evaluate(
+                T.call_extern(
+                    "intrin_function",
+                    T.tvm_access_ptr(
+                        T.type_annotation(dtype="int32"), dense_let.data, 0, 1, 2, dtype="handle"
+                    ),
+                    dtype="int32",
+                )
+            )
+            dense_let[0] = T.q_multiply_shift(dense_let[0], 1608879842, 31, -7, dtype="int32")
+
+    @T.prim_func
+    def __tvm_main__(
+        input: T.handle, global_workspace_1_var: T.Ptr[T.uint8], output: T.handle
+    ) -> None:
+        global_workspace_1_buffer_var = T.match_buffer(
+            global_workspace_1_var, [40], dtype="uint8", strides=[1], elem_offset=0, align=16
+        )
+        T.evaluate(
+            T.call_extern(
+                "tensor_intrin_primfunc", global_workspace_1_buffer_var.data, dtype="int32"
+            )
+        )
+
+
+def test_tensor_intrin():
+    target = Target("c")
+    global_workspace_pool = WorkspacePoolInfo(
+        "global_workspace",
+        [target],
+    )
+
+    tir_mod = TensorIntrinStructure
+    tir_mod = _assign_targets_to_primfuncs_irmodule(tir_mod, target)
+    tir_mod = assign_poolinfos_to_allocates_in_irmodule(tir_mod, [global_workspace_pool])
+    main_func = tir_mod["__tvm_main__"]
+    buffer_analysis = tvm.tir.usmp.analysis.extract_buffer_info(main_func, tir_mod)
+    buffer_info_map = buffer_analysis.buffer_info_stmts
+
+    fcreate_array_bi = tvm.get_global_func("tir.usmp.CreateArrayBufferInfo")
+    buffer_info_arr = fcreate_array_bi(buffer_info_map)
+    fusmp_algo_greedy_by_size = tvm.get_global_func("tir.usmp.algo.greedy_by_size")
+    buffer_pool_allocations = fusmp_algo_greedy_by_size(
+        buffer_info_arr, buffer_analysis.memory_pressure
+    )
+    fassign_stmt_pool_allocations = tvm.get_global_func("tir.usmp.AssignStmtPoolAllocations")
+    pool_allocations = fassign_stmt_pool_allocations(buffer_info_map, buffer_pool_allocations)
+    tir_mod_with_offsets = tvm.tir.usmp.transform.convert_pool_allocations_to_offsets(
+        pool_allocations, emit_tvmscript_printable=True
+    )(tir_mod)
+
+    expected = TensorIntrinStructurePlanned
+
+    for gv, ref_func in expected.functions.items():
+        actual_func = tir_mod_with_offsets[gv.name_hint]
+        tvm.ir.assert_structural_equal(actual_func, ref_func)
+
+
 if __name__ == "__main__":
     pytest.main([__file__] + sys.argv[1:])