From c99a06e0b3120b075d72cdd9119e34becf2668f5 Mon Sep 17 00:00:00 2001 From: Junru Shao Date: Tue, 7 Mar 2023 13:43:35 -0800 Subject: [PATCH] [TVMScript] Support SizeVar Roundtripping Previously the printer does not respect the difference between tir.Var and tir.SizeVar, which means SizeVar info will be lost during roundtripping, which is not ideal when dealing with dynamic shape workloads, where SizeVar provides additional bound information to the arithmetic analyzer. This PR extends Var printing to support SizeVar: ```python >>> a = tir.SizeVar("a", "int32") >>> print(a.script(verbose_expr=True)) a = T.int32(is_size_var=True) a ``` --- include/tvm/script/ir_builder/tir/ir.h | 25 +++++++++++---- include/tvm/tir/var.h | 7 +++++ python/tvm/script/ir_builder/tir/ir.py | 31 +++++++++++++------ src/script/ir_builder/tir/ir.cc | 15 ++------- src/script/printer/tir/expr.cc | 15 +++++++-- src/tir/ir/expr.cc | 9 ++++++ .../unittest/test_tvmscript_printer_tir.py | 2 +- 7 files changed, 72 insertions(+), 32 deletions(-) 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""", )