[IR] Add storage scope to PointerType (apache#8017)

* Add storage scope to PointerType.

* Apply suggestions from code review

Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>

include/tvm/ir/type.h

@@ -153,8 +153,15 @@ class PointerTypeNode : public TypeNode {
    * \brief The type of the element which the pointer points to.
    */
   Type element_type;
+  /*!
+   * \brief The storage scope of the pointer
+   */
+  String storage_scope;
 
-  void VisitAttrs(AttrVisitor* v) { v->Visit("element_type", &element_type); }
+  void VisitAttrs(AttrVisitor* v) {
+    v->Visit("element_type", &element_type);
+    v->Visit("storage_scope", &storage_scope);
+  }
 
   bool SEqualReduce(const PointerTypeNode* other, SEqualReducer equal) const {
     return equal(element_type, other->element_type);
@@ -175,8 +182,9 @@ class PointerType : public Type {
   /*!
    * \brief Constructor
    * \param element_type The type of the element which the pointer points to.
+   * \param storage_scope The storage scope into which the pointer addresses
    */
-  TVM_DLL explicit PointerType(Type element_type);
+  TVM_DLL explicit PointerType(Type element_type, String storage_scope = "");
 
   TVM_DEFINE_OBJECT_REF_METHODS(PointerType, Type, PointerTypeNode);
 };

python/tvm/ir/type.py

@@ -71,10 +71,13 @@ class PointerType(Type):
     ----------
     element_type : tvm.ir.Type
         The type of pointer's element.
+
+    storage_scope : str
+        The storage scope into which the pointer addresses.
     """
 
-    def __init__(self, element_type):
-        self.__init_handle_by_constructor__(_ffi_api.PointerType, element_type)
+    def __init__(self, element_type, storage_scope=""):
+        self.__init_handle_by_constructor__(_ffi_api.PointerType, element_type, storage_scope)
 
 
 @tvm._ffi.register_object("TypeVar")

src/ir/type.cc

@@ -43,21 +43,26 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->stream << node->dtype;
     });
 
-PointerType::PointerType(Type element_type) {
+PointerType::PointerType(Type element_type, String storage_scope) {
   ObjectPtr<PointerTypeNode> n = make_object<PointerTypeNode>();
   n->element_type = std::move(element_type);
+  n->storage_scope = std::move(storage_scope);
   data_ = std::move(n);
 }
 
 TVM_REGISTER_NODE_TYPE(PointerTypeNode);
 
-TVM_REGISTER_GLOBAL("ir.PointerType").set_body_typed([](Type element_type) {
-  return PointerType(element_type);
-});
+TVM_REGISTER_GLOBAL("ir.PointerType")
+    .set_body_typed([](Type element_type, String storage_scope = "") {
+      return PointerType(element_type, storage_scope);
+    });
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
     .set_dispatch<PointerTypeNode>([](const ObjectRef& ref, ReprPrinter* p) {
       auto* node = static_cast<const PointerTypeNode*>(ref.get());
+      if (!node->storage_scope.empty()) {
+        p->stream << node->storage_scope << " ";
+      }
       p->Print(node->element_type);
       p->stream << '*';
     });

src/ir/type_functor.cc

@@ -194,7 +194,7 @@ Type TypeMutator::VisitType_(const PointerTypeNode* op) {
   if (element_type.same_as(op->element_type)) {
     return GetRef<Type>(op);
   } else {
-    return PointerType(element_type);
+    return PointerType(element_type, op->storage_scope);
   }
 }
 

src/printer/tir_text_printer.cc

@@ -620,7 +620,11 @@ Doc TIRTextPrinter::VisitType_(const PrimTypeNode* node) {
 
 Doc TIRTextPrinter::VisitType_(const PointerTypeNode* node) {
   Doc doc;
-  doc << "Pointer(" << Print(node->element_type) << ")";
+  doc << "Pointer(";
+  if (!node->storage_scope.empty()) {
+    doc << node->storage_scope << " ";
+  }
+  doc << Print(node->element_type) << ")";
   return doc;
 }
 

src/printer/tvmscript_printer.cc

@@ -764,7 +764,11 @@ Doc TVMScriptPrinter::VisitType_(const PrimTypeNode* node) {
 
 Doc TVMScriptPrinter::VisitType_(const PointerTypeNode* node) {
   Doc doc;
-  doc << "ty.Ptr[" << Print(node->element_type) << "]";
+  doc << "ty.Ptr[";
+  if (!node->storage_scope.empty()) {
+    doc << node->storage_scope << " ";
+  }
+  doc << Print(node->element_type) << "]";
   return doc;
 }
 

tests/python/unittest/test_tir_constructor.py

@@ -161,6 +161,15 @@ def test_stmt_constructor():
     assert x.buffer_var == buffer_var
     assert x.body == nop
 
+    storage_scope = "global.texture"
+    buffer_var = tvm.tir.Var("buf", tvm.ir.PointerType(tvm.ir.PrimType("float32"), storage_scope))
+    x = tvm.tir.Allocate(buffer_var, "float32", [10], tvm.tir.const(1, "uint1"), nop)
+    assert isinstance(x, tvm.tir.Allocate)
+    assert x.dtype == "float32"
+    assert x.buffer_var == buffer_var
+    assert x.buffer_var.type_annotation.storage_scope == storage_scope
+    assert x.body == nop
+
     x = tvm.tir.AttrStmt(buffer_var, "xyz", 1, nop)
     assert isinstance(x, tvm.tir.AttrStmt)
     assert x.node == buffer_var

tests/python/unittest/test_tir_nodes.py

@@ -338,6 +338,17 @@ def test_vars():
     assert isinstance(ptype.element_type, tvm.ir.PrimType)
 
 
+def test_scoped_storage_vars():
+    dtype = "float"
+    storage_scope = "global.texture"
+    ptype = tvm.ir.PointerType(tvm.ir.PrimType(dtype), storage_scope)
+    x = tvm.tir.Var("xyz", ptype)
+    assert x.dtype == "handle"
+    assert x.type_annotation == ptype
+    assert x.type_annotation.storage_scope == storage_scope
+    assert isinstance(ptype.element_type, tvm.ir.PrimType)
+
+
 def test_buffer_load_store():
     b = tvm.tir.decl_buffer((10,), "float32")
     x = tvm.tir.BufferLoad(b, [0])
@@ -460,6 +471,7 @@ def test_block_blockrealize():
     test_intimm_cond()
     test_buffer_load_store()
     test_vars()
+    test_scoped_storage_var()
     test_prim_func()
     test_cast()
     test_attr()