This is an automated email from the ASF dual-hosted git repository.
ruihangl pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 2232207afd [Unity] Fix ForceNarrowIndexToI32 so it ignores i16 (#14733)
2232207afd is described below
commit 2232207afd9d8f29858e80d9a2cc0afaa25a9b63
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Apr 27 14:14:01 2023 -0400
[Unity] Fix ForceNarrowIndexToI32 so it ignores i16 (#14733)
This PR fixes the forece narrow to i32 pass so it ignores
integer with smaller bitwidth
---
include/tvm/tir/data_type_rewriter.h | 3 +++
src/tir/ir/data_type_rewriter.cc | 9 +++++++--
src/tir/transforms/force_narrow_index_to_i32.cc | 4 ++--
...test_tir_transform_force_narrow_index_to_i32.py | 22 ++++++++++++++++++++++
4 files changed, 34 insertions(+), 4 deletions(-)
diff --git a/include/tvm/tir/data_type_rewriter.h
b/include/tvm/tir/data_type_rewriter.h
index 76b53b18ca..8bdcc097a2 100644
--- a/include/tvm/tir/data_type_rewriter.h
+++ b/include/tvm/tir/data_type_rewriter.h
@@ -153,6 +153,9 @@ class IndexDataTypeNormalizer : public
IndexDataTypeRewriter {
PrimExpr VisitExpr_(const VarNode* op) override;
PrimExpr VisitExpr_(const CastNode* op) override;
+ /*! \brief Specifies which data type we can rewrite */
+ virtual bool CanRewriteDType(DataType dtype) const;
+
DataType target_data_type_ = DataType::Int(64);
};
diff --git a/src/tir/ir/data_type_rewriter.cc b/src/tir/ir/data_type_rewriter.cc
index 97ad1b7cc3..619804b880 100644
--- a/src/tir/ir/data_type_rewriter.cc
+++ b/src/tir/ir/data_type_rewriter.cc
@@ -605,6 +605,10 @@ PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) {
return func;
}
+bool IndexDataTypeNormalizer::CanRewriteDType(DataType dtype) const {
+ return dtype.is_int() && dtype.bits() >= 32;
+}
+
PrimExpr IndexDataTypeNormalizer::VisitExpr_(const IntImmNode* op) {
if (is_enabled_) {
ICHECK_LE(op->value,
Downcast<Integer>(max_value(target_data_type_))->value);
@@ -614,7 +618,8 @@ PrimExpr IndexDataTypeNormalizer::VisitExpr_(const
IntImmNode* op) {
}
PrimExpr IndexDataTypeNormalizer::VisitExpr_(const VarNode* op) {
- if (is_enabled_ && op->dtype != target_data_type_ && !var_remap_.count(op)) {
+ if (is_enabled_ && CanRewriteDType(op->dtype) && op->dtype !=
target_data_type_ &&
+ !var_remap_.count(op)) {
var_remap_[op] = GetRef<Var>(op).copy_with_dtype(target_data_type_);
}
return DataTypeLegalizer::VisitExpr_(op);
@@ -624,7 +629,7 @@ PrimExpr IndexDataTypeNormalizer::VisitExpr_(const
CastNode* op) {
// Unwrap the cast only when the dtype of this cast is integer dtype.
// When the dtype of this cast is not integer dtype, it means that this cast
// has some other purpose, and we should not unwrap the cast.
- if (is_enabled_ && op->dtype.is_int()) {
+ if (is_enabled_ && CanRewriteDType(op->dtype)) {
PrimExpr value = IndexDataTypeNormalizer::VisitExpr(op->value);
return value->dtype == target_data_type_ ? value : Cast(target_data_type_,
value);
}
diff --git a/src/tir/transforms/force_narrow_index_to_i32.cc
b/src/tir/transforms/force_narrow_index_to_i32.cc
index c559360bf5..86f839c4f5 100644
--- a/src/tir/transforms/force_narrow_index_to_i32.cc
+++ b/src/tir/transforms/force_narrow_index_to_i32.cc
@@ -35,7 +35,7 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer {
static PrimFunc RewriteDataType(PrimFunc func) {
// Check if the integer parameter buffers have dtype other than int32.
for (auto it : func->buffer_map) {
- if (it.second->dtype.is_int() && it.second->dtype.bits() != 32) {
+ if (it.second->dtype.is_int() && it.second->dtype.bits() > 32) {
LOG(FATAL) << "The buffer " << it.second << " in the function buffer
map has dtype "
<< it.second->dtype << ". The function is " << func;
}
@@ -62,7 +62,7 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer {
Block block_ = Downcast<Block>(IndexDataTypeNormalizer::VisitStmt_(block));
// Check if the allocated integer buffers have dtype other than int32.
for (const Buffer& buf : block_->alloc_buffers) {
- if (buf->dtype.is_int() && buf->dtype.bits() != 32) {
+ if (buf->dtype.is_int() && buf->dtype.bits() > 32) {
LOG(FATAL) << "The buffer " << buf << " allocated in the function has
dtype " << buf->dtype
<< ". The function is " << func_;
}
diff --git
a/tests/python/unittest/test_tir_transform_force_narrow_index_to_i32.py
b/tests/python/unittest/test_tir_transform_force_narrow_index_to_i32.py
index 6c12229b5c..8a2a286671 100644
--- a/tests/python/unittest/test_tir_transform_force_narrow_index_to_i32.py
+++ b/tests/python/unittest/test_tir_transform_force_narrow_index_to_i32.py
@@ -182,6 +182,28 @@ def test_block():
tvm.ir.assert_structural_equal(func, expected)
+def test_i16_buffer():
+ @T.prim_func
+ def before(A: T.Buffer((128,), "int16"), B: T.Buffer((128,), "int16")):
+ for i in T.serial(0, T.int64(16)):
+ for j in T.serial(0, T.int64(16)):
+ with T.block():
+ vi = T.axis.spatial(T.int64(128), i * 8 + j)
+ B[vi] = A[vi] + T.int16(1)
+
+ @T.prim_func
+ def expected(A: T.Buffer((128,), "int16"), B: T.Buffer((128,), "int16")):
+ for i in T.serial(0, 16):
+ for j in T.serial(0, 16):
+ with T.block():
+ vi = T.axis.spatial(128, i * 8 + j)
+ B[vi] = A[vi] + T.int16(1)
+
+ mod = tvm.IRModule.from_expr(before)
+ after = tvm.tir.transform.ForceNarrowIndexToInt32()(mod)["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
+
def test_fail_on_buffer_map():
@T.prim_func
def func(A: T.Buffer((128,), "int64"), B: T.Buffer((128,), "int64")):
