[TVMScript] Support SizeVar Roundtripping

include/tvm/script/ir_builder/tir/ir.h

@@ -430,14 +430,27 @@ void Evaluate(PrimExpr value);
  * \brief Create a TIR var that represents a pointer
  * \param dtype The data type of the pointer.
  * \param storage_scope The storage scope of the pointer.
+ * \param is_size_var Whether the pointer is a size var.
  * \return The pointer.
  */
-Var Handle(runtime::DataType dtype = runtime::DataType::Void(), String storage_scope = "global");
-
-#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType)                             \
-  inline PrimExpr FuncName(Optional<PrimExpr> expr = NullOpt) {                        \
-    DataType dtype = DType;                                                            \
-    return expr.defined() ? tvm::cast(dtype, expr.value()) : tvm::tir::Var("", dtype); \
+inline Var Handle(runtime::DataType dtype = runtime::DataType::Void(),  //
+                  String storage_scope = "global",                      //
+                  bool is_size_var = false) {
+  Type type_annotation{nullptr};
+  if (dtype.is_void() && storage_scope == "global") {
+    type_annotation = PrimType(runtime::DataType::Handle());
+  } else {
+    type_annotation = PointerType(PrimType(dtype), storage_scope);
+  }
+  return is_size_var ? tvm::tir::SizeVar("", type_annotation) : tvm::tir::Var("", type_annotation);
+}
+
+#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType)                                \
+  inline PrimExpr FuncName(Optional<PrimExpr> expr = NullOpt, bool is_size_var = false) { \
+    DataType dtype = DType;                                                               \
+    return expr.defined()                                                                 \
+               ? tvm::cast(dtype, expr.value())                                           \
+               : (is_size_var ? tvm::tir::SizeVar("", dtype) : tvm::tir::Var("", dtype)); \
   }
 
 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType) \

include/tvm/tir/var.h

@@ -152,6 +152,13 @@ class SizeVar : public Var {
    */
   TVM_DLL explicit SizeVar(String name_hint = "s", DataType t = DataType::Int(32),
                            Span span = Span());
+  /*!
+   * \brief Constructor which provides a more detailed type annotation.
+   * \param name_hint variable name.
+   * \param type_annotation The type annotation.
+   * \param span The location of this object in the source code.
+   */
+  TVM_DLL explicit SizeVar(String name_hint, Type type_annotation, Span span = Span());
   /*!
    * \brief Get pointer to the internal value.
    * \return the corresponding Variable.

python/tvm/script/ir_builder/tir/ir.py

@@ -1333,11 +1333,13 @@ def func(
             Literal["inf", "-inf", "nan"],
             int,
             float,
-        ] = None
+        ] = None,
+        *,
+        is_size_var: bool = False,
     ) -> PrimExpr:
         if isinstance(expr, str):
             expr = float(expr)
-        return getattr(_ffi_api, name)(expr)
+        return getattr(_ffi_api, name)(expr, is_size_var)
 
     return func
 
@@ -1420,23 +1422,26 @@ def func(
 # pylint: enable=invalid-name
 
 
-def boolean(expr: Optional[PrimExpr] = None) -> PrimExpr:
+def boolean(expr: Optional[PrimExpr] = None, is_size_var: bool = False) -> PrimExpr:
     """Construct a new tir.Var with type boolean or cast expression to type boolean.
 
     Parameters
     ----------
     expr: PrimExpr
         The expression to be cast.
 
+    is_size_var: bool
+        Whether or not to return a SizeVar instead of Var.
+
     Returns
     -------
     res : PrimExpr
         The new tir.Var with type boolean or casted expression with type boolean.
     """
-    return _ffi_api.Boolean(expr)  # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Boolean(expr, is_size_var)  # type: ignore[attr-defined] # pylint: disable=no-member
 
 
-def handle(dtype: str = "void", storage_scope: str = "global") -> Var:
+def handle(dtype: str = "void", storage_scope: str = "global", *, is_size_var: bool = False) -> Var:
     """Create a TIR var that represents a pointer.
 
     Parameters
@@ -1447,15 +1452,18 @@ def handle(dtype: str = "void", storage_scope: str = "global") -> Var:
     storage_scope: str
         The storage scope of the pointer.
 
+    is_size_var: bool
+        Whether or not to return a SizeVar instead of Var.
+
     Returns
     -------
     res : PrimExpr
         The new tir.Var with type handle or casted expression with type handle.
     """
-    return _ffi_api.Handle(dtype, storage_scope)  # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Handle(dtype, storage_scope, is_size_var)  # type: ignore[attr-defined] # pylint: disable=no-member
 
 
-def void(expr: Optional[PrimExpr] = None) -> PrimExpr:
+def void(expr: Optional[PrimExpr] = None, *, is_size_var: bool = False) -> PrimExpr:
     """Construct a new tir.Var with type void or cast expression to type void.
 
     Parameters
@@ -1468,7 +1476,7 @@ def void(expr: Optional[PrimExpr] = None) -> PrimExpr:
     res : PrimExpr
         The new tir.Var with type void or casted expression with type void.
     """
-    return _ffi_api.Void(expr)  # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Void(expr, is_size_var)  # type: ignore[attr-defined] # pylint: disable=no-member
 
 
 @deprecated("T.var", "T.{dtype}")
@@ -1491,7 +1499,7 @@ def var(dtype: str, name: str = "") -> Var:
     return Var(name, dtype)  # pylint: disable=no-member
 
 
-def ptr(dtype: str, storage_scope: str = "global") -> Var:
+def ptr(dtype: str, storage_scope: str = "global", is_size_var: bool = False) -> Var:
     """The pointer declaration function.
 
     Parameters
@@ -1502,12 +1510,15 @@ def ptr(dtype: str, storage_scope: str = "global") -> Var:
     storage_scope : str
         The storage scope of the pointer.
 
+    is_size_var: bool
+        Whether or not to return a SizeVar instead of Var.
+
     Returns
     -------
     res : Var
         The pointer.
     """
-    return _ffi_api.Ptr(dtype, storage_scope)  # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Ptr(dtype, storage_scope, is_size_var)  # type: ignore[attr-defined] # pylint: disable=no-member
 
 
 @deprecated("T.buffer_var", "T.handle")

src/script/ir_builder/tir/ir.cc

@@ -560,18 +560,9 @@ DeclBufferFrame DeclBuffer(Array<PrimExpr> shape, DataType dtype, String buffer_
 
 void Evaluate(PrimExpr value) { AddToParent(tvm::tir::Evaluate(value)); }
 
-PrimExpr Ptr(runtime::DataType dtype, String storage_scope) {
-  return tvm::tir::Var("", tvm::PointerType(PrimType(dtype), storage_scope));
-}
-
-Var Handle(runtime::DataType dtype, String storage_scope) {
-  Type type_annotation{nullptr};
-  if (dtype.is_void() && storage_scope == "global") {
-    type_annotation = PrimType(runtime::DataType::Handle());
-  } else {
-    type_annotation = PointerType(PrimType(dtype), storage_scope);
-  }
-  return tvm::tir::Var("", type_annotation);
+PrimExpr Ptr(runtime::DataType dtype, String storage_scope = "global", bool is_size_var = false) {
+  PointerType type_annotation(PrimType(dtype), storage_scope);
+  return is_size_var ? tvm::tir::SizeVar("", type_annotation) : tvm::tir::Var("", type_annotation);
 }
 
 using tvm::script::ir_builder::details::Namer;

src/script/printer/tir/expr.cc

@@ -28,6 +28,14 @@ ExprDoc PrintVarCreation(const tir::Var& var, const ObjectPath& var_p, const IRD
   Type type = var->type_annotation;
   ObjectPath type_p = var_p->Attr("type_annotation");
   ExprDoc rhs{nullptr};
+  Array<String> kwargs_keys;
+  Array<ExprDoc> kwargs_values;
+
+  if (var->IsInstance<tir::SizeVarNode>()) {
+    kwargs_keys.push_back("is_size_var");
+    kwargs_values.push_back(LiteralDoc::Boolean(true, NullOpt));
+  }
+
   if (const auto* ptr_type = type.as<PointerTypeNode>()) {
     const auto* prim_type = ptr_type->element_type.as<PrimTypeNode>();
     ICHECK(prim_type);
@@ -36,16 +44,17 @@ ExprDoc PrintVarCreation(const tir::Var& var, const ObjectPath& var_p, const IRD
     rhs = TIR(d, "handle");
     rhs->source_paths.push_back(var_p->Attr("dtype"));
     if (ptr_type->storage_scope == "") {
-      rhs = rhs->Call({element_type});
+      rhs = rhs->Call({element_type}, kwargs_keys, kwargs_values);
     } else {
       rhs = rhs->Call({element_type,
                        LiteralDoc::Str(ptr_type->storage_scope,  //
-                                       type_p->Attr("storage_scope"))});
+                                       type_p->Attr("storage_scope"))},
+                      kwargs_keys, kwargs_values);
     }
   } else {
     rhs = TIR(d, DType2Str(var->dtype));
     rhs->source_paths.push_back(var_p->Attr("dtype"));
-    rhs = rhs->Call({});
+    rhs = rhs->Call({}, kwargs_keys, kwargs_values);
   }
   rhs->source_paths.push_back(type_p);
   return rhs;

src/tir/ir/expr.cc

@@ -126,6 +126,15 @@ SizeVar::SizeVar(String name_hint, DataType dtype, Span span) {
   data_ = std::move(n);
 }
 
+SizeVar::SizeVar(String name_hint, Type type_annotation, Span span) {
+  auto n = make_object<SizeVarNode>();
+  n->name_hint = std::move(name_hint);
+  n->dtype = GetRuntimeDataType(type_annotation);
+  n->type_annotation = std::move(type_annotation);
+  n->span = std::move(span);
+  data_ = std::move(n);
+}
+
 TVM_REGISTER_GLOBAL("tir.SizeVar").set_body_typed([](String s, DataType t, Span span) {
   return SizeVar(s, t, span);
 });

tests/python/unittest/test_tvmscript_printer_tir.py

@@ -468,7 +468,7 @@ def test_size_var():
     _assert_print(
         a,
         """
-a = T.float32()
+a = T.float32(is_size_var=True)
 a""",
     )