This is an automated email from the ASF dual-hosted git repository.
yaxingcai pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 902c2e2db7 [TVMScript] Support SizeVar Roundtripping (#14227)
902c2e2db7 is described below
commit 902c2e2db70b75c36f9bd8c253707b1e8761cc18
Author: Junru Shao <[email protected]>
AuthorDate: Wed Mar 8 01:49:09 2023 -0800
[TVMScript] Support SizeVar Roundtripping (#14227)
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 +++++++
.../python/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 8d8b0b42ba..a0343b0395 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<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) \
diff --git a/include/tvm/tir/var.h b/include/tvm/tir/var.h
index 0dadd3dc71..52827f706a 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 7826278431..d65f9adea8 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_gen(name: str):
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 @@ float64x64 = func_gen(("Float64x64"))
# 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 8d6c51be3a..aee0b4bb62 100644
--- a/src/script/ir_builder/tir/ir.cc
+++ b/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;
diff --git a/src/script/printer/tir/expr.cc b/src/script/printer/tir/expr.cc
index dda2c73b8e..9c4f62eb1c 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<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;
diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc
index d5caeab539..db09ac17e6 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<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);
});
diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py
b/tests/python/unittest/test_tvmscript_printer_tir.py
index 87ec98e9a2..171d49b619 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""",
)
