[TIR][BugFix]Ensure the Var's scope is correct (#15406)

* [TIR][BugFix]Ensure Var's scope is correct

* add new testcases

* fix lint

---------

Co-authored-by: Bin Li <[email protected]>

Author: sisleyli

src/tir/transforms/unsupported_dtype_legalize.cc

@@ -79,7 +79,11 @@ class ComputeLegalizePlanner : public StmtExprVisitor {
     // remap all intermediate constant buffer to promote data types (fp16/fp32)
     if (MatchDType(op->dtype) && op->ConstantAllocationSize() != 0) {
       DataType dtype = promote_dtype_.with_lanes(op->dtype.lanes());
-      Var buffer_var = Var(op->buffer_var->name_hint, PointerType(PrimType(dtype)));
+      String storage_scope = "global";
+      if (auto* ptr_type = op->buffer_var->type_annotation.as<PointerTypeNode>()) {
+        storage_scope = ptr_type->storage_scope;
+      }
+      Var buffer_var = Var(op->buffer_var->name_hint, PointerType(PrimType(dtype), storage_scope));
       (*var_remap_)[op->buffer_var] = buffer_var;
     }
     return StmtExprVisitor::VisitStmt_(op);
@@ -496,7 +500,11 @@ class StorageLegalizer : public StmtExprMutator {
   Stmt VisitStmt_(const AllocateNode* op) final {
     if (MatchDType(op->dtype)) {
       DataType dtype = GetStorageUIntDType(op->dtype);
-      Var buffer_var = Var(op->buffer_var->name_hint, PointerType(PrimType(dtype)));
+      String storage_scope = "global";
+      if (auto* ptr_type = op->buffer_var->type_annotation.as<PointerTypeNode>()) {
+        storage_scope = ptr_type->storage_scope;
+      }
+      Var buffer_var = Var(op->buffer_var->name_hint, PointerType(PrimType(dtype), storage_scope));
       var_remap_[op->buffer_var] = buffer_var;
       return VisitStmt(Allocate(buffer_var, dtype, op->extents, op->condition, op->body));
     } else {
@@ -637,7 +645,8 @@ class StorageLegalizer : public StmtExprMutator {
         if (auto* elem_type = ptr_type->element_type.as<PrimTypeNode>()) {
           if (MatchDType(elem_type->dtype)) {
             Var new_var =
-                Var(var->name_hint, PointerType(PrimType(GetStorageUIntDType(elem_type->dtype))));
+                Var(var->name_hint, PointerType(PrimType(GetStorageUIntDType(elem_type->dtype)),
+                                                ptr_type->storage_scope));
             var_remap_[var] = new_var;
             return new_var;
           }

tests/python/unittest/test_tir_transform_bf16_legalize.py

@@ -114,5 +114,74 @@ def test_bf16_storage_legalize():
     tvm.ir.assert_structural_equal(after, expected)
 
 
+def test_bf16_storage_scope():
+    def get_before():
+        @tvm.script.ir_module
+        class Before:
+            @T.prim_func
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "bfloat16")
+                for i in T.grid(100):
+                    C[i] = A[i] + B[i]
+                    D[i] = T.exp(C[i])
+
+        return Before
+
+    def after_compute_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func
+            def main(
+                Aptr: T.handle("bfloat16", storage_scope="shared"),
+                Bptr: T.handle("bfloat16", storage_scope="local"),
+                Dptr: T.handle("bfloat16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "bfloat16", data=Aptr)
+                B = T.decl_buffer((100,), "bfloat16", data=Bptr)
+                D = T.decl_buffer((100,), "bfloat16", data=Dptr)
+                C = T.decl_buffer((100,), "float32")
+                for i in T.grid(100):
+                    C[i] = bf16tof32(A[i]) + bf16tof32(B[i])
+                    D[i] = f32tobf16(T.exp(C[i]))
+
+        return After
+
+    def after_storage_legalize():
+        @tvm.script.ir_module
+        class After:
+            @T.prim_func
+            def main(
+                Aptr: T.handle("uint16", storage_scope="shared"),
+                Bptr: T.handle("uint16", storage_scope="local"),
+                Dptr: T.handle("uint16"),
+            ):
+                T.func_attr({"global_symbol": "main"})
+                A = T.decl_buffer((100,), "uint16", data=Aptr)
+                B = T.decl_buffer((100,), "uint16", data=Bptr)
+                D = T.decl_buffer((100,), "uint16", data=Dptr)
+                C = T.decl_buffer((100,), "float32")
+                for i in T.grid(100):
+                    C[i] = u16tof32(A[i]) + u16tof32(B[i])
+                    D[i] = f32tou16(T.exp(C[i]))
+
+        return After
+
+    before = get_before()
+    after_compute = tvm.tir.transform.BF16ComputeLegalize()(before)
+    after_storage = tvm.tir.transform.BF16StorageLegalize()(after_compute)
+    tvm.ir.assert_structural_equal(after_compute, after_compute_legalize())
+    tvm.ir.assert_structural_equal(after_storage, after_storage_legalize())
+
+
 if __name__ == "__main__":
     test_bf16_storage_legalize()
+    test_bf16_storage_scope()