This is an automated email from the ASF dual-hosted git repository.
masahi 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 990c768e9c [Unity][Analysis] Reshape TIR detection with
iter-map-simplify (#15099)
990c768e9c is described below
commit 990c768e9c7ed8968da0bb442a84f94046dee9f5
Author: Ruihang Lai <[email protected]>
AuthorDate: Wed Jun 14 14:57:36 2023 -0700
[Unity][Analysis] Reshape TIR detection with iter-map-simplify (#15099)
This PR supports the reshape PrimFunc detection analysis with
iter-map-simplify, in order to better support TIR functions with dynamic
variables, which is not well handled before.
Meanwhie, this PR also keeps the existing analysis logic, so that we are
still able to cover the cases like reshape with raggedness.
Unit tests for reshape are added to ensure correctness.
---
src/relax/analysis/tir_op_pattern_kind.cc | 69 +++++++++++-
tests/python/relax/test_analysis.py | 122 +++++++++++++++++++++
.../test_transform_rewrite_dataflow_reshape.py | 12 +-
3 files changed, 194 insertions(+), 9 deletions(-)
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc
b/src/relax/analysis/tir_op_pattern_kind.cc
index 950e6a10e0..55b5d76213 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -17,6 +17,7 @@
* under the License.
*/
+#include <tvm/arith/iter_affine_map.h>
#include <tvm/relax/analysis.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/expr_functor.h>
@@ -420,7 +421,10 @@ bool HasReshapePattern(const PrimFunc& func) {
return;
}
- // Step 3. Calculate the flattened access index according to the
load/store pattern.
+ // Apply check 1: use iter_map_simplify
+ // This check requires at least one of the src/dst side is a trivial
buffer
+ // access (e.g., buf[ax0, ax1, ax2]).
+
auto f_calc_flattened_idx = [](const Buffer& buffer, const
Array<PrimExpr>& indices) {
ICHECK_EQ(indices.size(), buffer->shape.size());
int ndim = indices.size();
@@ -430,12 +434,71 @@ bool HasReshapePattern(const PrimFunc& func) {
}
return idx;
};
+
+ auto f_is_trivial_indices = [block, this](const Buffer& buffer,
+ const Array<PrimExpr>&
indices) {
+ if (indices.size() != block->iter_vars.size()) {
+ return false;
+ }
+ for (int i = 0; i < static_cast<int>(block->iter_vars.size()); ++i) {
+ if (!(indices[i].same_as(block->iter_vars[i]->var) &&
+ this->ana_.CanProveEqual(block->iter_vars[i]->dom->min,
+ IntImm(DataType::Int(64),
/*value=*/0)) &&
+ this->ana_.CanProveEqual(buffer->shape[i],
block->iter_vars[i]->dom->extent))) {
+ return false;
+ }
+ }
+ return true;
+ };
+
+ Array<PrimExpr> nontrivial_indices{nullptr};
+ Buffer nontrivial_buffer{nullptr};
+ if (f_is_trivial_indices(dst_buffer_, buffer_store->indices)) {
+ nontrivial_indices = buffer_load->indices;
+ nontrivial_buffer = src_buffer_;
+ } else if (f_is_trivial_indices(src_buffer_, buffer_load->indices)) {
+ nontrivial_indices = buffer_store->indices;
+ nontrivial_buffer = dst_buffer_;
+ }
+
+ if (nontrivial_indices.defined()) {
+ DataType dtype =
+ !block->iter_vars.empty() ? block->iter_vars[0]->var->dtype :
DataType::Int(64);
+ tir::Var fused_var("fused", dtype);
+ Map<tir::Var, PrimExpr> inverse_indices_map;
+ PrimExpr stride = IntImm(dtype, /*value=*/1);
+ for (int i = static_cast<int>(block->iter_vars.size()) - 1; i >= 0;
--i) {
+ inverse_indices_map.Set(
+ block->iter_vars[i]->var,
+ floormod(floordiv(fused_var, stride),
block->iter_vars[i]->dom->extent));
+ stride *= block->iter_vars[i]->dom->extent;
+ }
+ PrimExpr flattened_idx = f_calc_flattened_idx(nontrivial_buffer,
nontrivial_indices);
+ flattened_idx = Substitute(std::move(flattened_idx),
inverse_indices_map);
+
+ Array<PrimExpr> simplify_res = arith::IterMapSimplify(
+ /*indices=*/{flattened_idx},
+ /*input_iters=*/{{fused_var, Range(IntImm(dtype, /*value=*/0),
stride)}},
+ /*input_pred=*/Bool(true),
+ /*check_level=*/arith::IterMapLevel::Surjective,
+ /*analyzer=*/&this->ana_,
+ /*simplify_trivial_iterators=*/true);
+ ICHECK_EQ(simplify_res.size(), 1);
+
+ if (simplify_res[0].same_as(fused_var)) {
+ this->is_reshape_ = true;
+ return;
+ }
+ }
+
+ // Apply check 2 as followup when check 1 is not satisfied.
+ // Calculate the flattened access index according to the load/store
pattern.
PrimExpr src_idx = f_calc_flattened_idx(src_buffer_,
buffer_load->indices);
PrimExpr dst_idx = f_calc_flattened_idx(dst_buffer_,
buffer_store->indices);
-
- // Step 4. Check if we can prove the equality of flattened indices.
+ // Check if we can prove the equality of flattened indices.
if (ana_.CanProveEqual(src_idx, dst_idx)) {
this->is_reshape_ = true;
+ return;
}
}
diff --git a/tests/python/relax/test_analysis.py
b/tests/python/relax/test_analysis.py
index 31efb98646..fd47d25a47 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -357,6 +357,128 @@ def test_reshape_pattern_expand_dims():
assert has_reshape_pattern(expand_dims)
+def test_reshape_pattern_dyn_1():
+ @T.prim_func
+ def reshape(var_A: T.handle, var_T_reshape: T.handle):
+ n = T.int64()
+ A = T.match_buffer(var_A, (n, T.int64(32), T.int64(128)), "float16")
+ T_reshape = T.match_buffer(
+ var_T_reshape, (T.int64(1), n, T.int64(32), T.int64(128)),
"float16"
+ )
+ for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), n, T.int64(32),
T.int64(128)):
+ with T.block("T_reshape"):
+ v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1,
ax2, ax3])
+ T.reads(
+ A[
+ ((v_ax3 // T.int64(128) + v_ax2) // T.int64(32) +
v_ax0 * n + v_ax1) % n,
+ (v_ax3 // T.int64(128) + v_ax2) % T.int64(32),
+ v_ax3 % T.int64(128),
+ ]
+ )
+ T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
+ T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = A[
+ ((v_ax3 // T.int64(128) + v_ax2) // T.int64(32) + v_ax0 *
n + v_ax1) % n,
+ (v_ax3 // T.int64(128) + v_ax2) % T.int64(32),
+ v_ax3 % T.int64(128),
+ ]
+
+ assert has_reshape_pattern(reshape)
+
+
+def test_reshape_pattern_dyn_2():
+ @T.prim_func
+ def reshape(var_A: T.handle, var_T_reshape: T.handle):
+ n = T.int64()
+ A = T.match_buffer(var_A, (T.int64(1), n), "int32")
+ T_reshape = T.match_buffer(var_T_reshape, (n,), "int32")
+ for ax0 in range(n):
+ with T.block("T_reshape"):
+ v_ax0 = T.axis.spatial(n, ax0)
+ T.reads(A[T.int64(0), v_ax0 % n])
+ T.writes(T_reshape[v_ax0])
+ T_reshape[v_ax0] = A[T.int64(0), v_ax0 % n]
+
+ assert has_reshape_pattern(reshape)
+
+
+def test_reshape_pattern_dyn_3():
+ @T.prim_func
+ def reshape(var_A: T.handle, var_T_reshape: T.handle):
+ T.func_attr({"op_pattern": 8, "tir.noalias": T.bool(True)})
+ n = T.int64()
+ A = T.match_buffer(var_A, (n, T.int64(4096)), "float16")
+ T_reshape = T.match_buffer(var_T_reshape, (T.int64(1), n,
T.int64(4096)), "float16")
+ for ax0, ax1, ax2 in T.grid(T.int64(1), n, T.int64(4096)):
+ with T.block("T_reshape"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(A[(v_ax2 // T.int64(4096) + v_ax0 * n + v_ax1) % n,
v_ax2 % T.int64(4096)])
+ T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
+ T_reshape[v_ax0, v_ax1, v_ax2] = A[
+ (v_ax2 // T.int64(4096) + v_ax0 * n + v_ax1) % n, v_ax2 %
T.int64(4096)
+ ]
+
+ assert has_reshape_pattern(reshape)
+
+
+def test_reshape_pattern_dyn_4():
+ @T.prim_func
+ def reshape(var_A: T.handle, var_T_reshape: T.handle):
+ T.func_attr({"op_pattern": 8, "tir.noalias": T.bool(True)})
+ n = T.int64()
+ A = T.match_buffer(var_A, (T.int64(1), n, T.int64(4096)), "float16")
+ T_reshape = T.match_buffer(
+ var_T_reshape, (T.int64(1), n, T.int64(32), T.int64(128)),
"float16"
+ )
+ for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), n, T.int64(32),
T.int64(128)):
+ with T.block("T_reshape"):
+ v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1,
ax2, ax3])
+ T.reads(
+ A[
+ T.int64(0),
+ ((v_ax2 * T.int64(128) + v_ax3) // T.int64(4096) +
v_ax0 * n + v_ax1) % n,
+ (v_ax2 * T.int64(128) + v_ax3) % T.int64(4096),
+ ]
+ )
+ T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
+ T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = A[
+ T.int64(0),
+ ((v_ax2 * T.int64(128) + v_ax3) // T.int64(4096) + v_ax0 *
n + v_ax1) % n,
+ (v_ax2 * T.int64(128) + v_ax3) % T.int64(4096),
+ ]
+
+ assert has_reshape_pattern(reshape)
+
+
+def test_reshape_pattern_dyn_5():
+ @T.prim_func
+ def reshape(var_A: T.handle, var_T_reshape: T.handle):
+ T.func_attr({"op_pattern": 8, "tir.noalias": T.bool(True)})
+ n = T.int64()
+ A = T.match_buffer(var_A, (T.int64(1), n, T.int64(32), T.int64(128)),
"float16")
+ T_reshape = T.match_buffer(var_T_reshape, (T.int64(1), n,
T.int64(4096)), "float16")
+ # with T.block("root"):
+ for ax0, ax1, ax2 in T.grid(T.int64(1), n, T.int64(4096)):
+ with T.block("T_reshape"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(
+ A[
+ T.int64(0),
+ (v_ax2 // T.int64(4096) + v_ax0 * n + v_ax1) % n,
+ v_ax2 % T.int64(4096) // T.int64(128),
+ v_ax2 % T.int64(128),
+ ]
+ )
+ T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
+ T_reshape[v_ax0, v_ax1, v_ax2] = A[
+ T.int64(0),
+ (v_ax2 // T.int64(4096) + v_ax0 * n + v_ax1) % n,
+ v_ax2 % T.int64(4096) // T.int64(128),
+ v_ax2 % T.int64(128),
+ ]
+
+ assert has_reshape_pattern(reshape)
+
+
def test_reshape_pattern_with_raggedness():
@T.prim_func
def reshape_raggedness(
diff --git a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
index db737f82c0..ff802bfb56 100644
--- a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
+++ b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
@@ -35,13 +35,13 @@ def test_reshape_expand_dims():
T.reads(
rxplaceholder[
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
- (v_ax1 * 12 + v_ax2 * 3 + v_ax2) % T.int64(3),
+ (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
]
)
T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
- (v_ax1 * 12 + v_ax2 * 3 + v_ax2) % T.int64(3),
+ (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
]
@T.prim_func
@@ -87,13 +87,13 @@ def test_reshape_expand_dims():
T.reads(
rxplaceholder[
(v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2)
// T.int64(3),
- (v_ax1 * T.int64(12) + v_ax2 * T.int64(3) + v_ax2)
% T.int64(3),
+ (v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2)
% T.int64(3),
]
)
T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
(v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) //
T.int64(3),
- (v_ax1 * T.int64(12) + v_ax2 * T.int64(3) + v_ax2) %
T.int64(3),
+ (v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) %
T.int64(3),
]
@T.prim_func
@@ -236,13 +236,13 @@ def test_reshape_non_dataflow():
T.reads(
rxplaceholder[
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
- (v_ax1 * 12 + v_ax2 * 3 + v_ax2) % T.int64(3),
+ (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
]
)
T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
- (v_ax1 * 12 + v_ax2 * 3 + v_ax2) % T.int64(3),
+ (v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
]
@R.function
