diff --git a/include/tvm/script/ir_builder/tir/ir.h b/include/tvm/script/ir_builder/tir/ir.h index 8d8b0b42ba5c..a0343b03955b 100644 --- a/include/tvm/script/ir_builder/tir/ir.h +++ b/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 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 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) \ diff --git a/include/tvm/tir/var.h b/include/tvm/tir/var.h index 0dadd3dc712e..52827f706a56 100644 --- a/include/tvm/tir/var.h +++ b/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. diff --git a/python/tvm/script/ir_builder/tir/ir.py b/python/tvm/script/ir_builder/tir/ir.py index 782627843179..d65f9adea86f 100644 --- a/python/tvm/script/ir_builder/tir/ir.py +++ b/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,7 +1422,7 @@ 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 @@ -1428,15 +1430,18 @@ def boolean(expr: Optional[PrimExpr] = None) -> PrimExpr: 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") diff --git a/src/script/ir_builder/tir/ir.cc b/src/script/ir_builder/tir/ir.cc index 8d6c51be3a11..aee0b4bb6253 100644 --- a/src/script/ir_builder/tir/ir.cc +++ b/src/script/ir_builder/tir/ir.cc @@ -560,18 +560,9 @@ DeclBufferFrame DeclBuffer(Array 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; diff --git a/src/script/printer/tir/expr.cc b/src/script/printer/tir/expr.cc index dda2c73b8e97..9c4f62eb1c1b 100644 --- a/src/script/printer/tir/expr.cc +++ b/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 kwargs_keys; + Array kwargs_values; + + if (var->IsInstance()) { + kwargs_keys.push_back("is_size_var"); + kwargs_values.push_back(LiteralDoc::Boolean(true, NullOpt)); + } + if (const auto* ptr_type = type.as()) { const auto* prim_type = ptr_type->element_type.as(); 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; diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc index d5caeab53922..db09ac17e6eb 100644 --- a/src/tir/ir/expr.cc +++ b/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(); + 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); }); diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py index 87ec98e9a266..171d49b6191b 100644 --- a/tests/python/unittest/test_tvmscript_printer_tir.py +++ b/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""", )