This is an automated email from the ASF dual-hosted git repository.
tqchen 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 7359313b40 [TIR] Fix Shuffle rewrite (#17030)
7359313b40 is described below
commit 7359313b40dd1927cd27e2c60539575ae08a4dc5
Author: Siyuan Feng <[email protected]>
AuthorDate: Mon May 27 21:25:06 2024 +0800
[TIR] Fix Shuffle rewrite (#17030)
This PR fixes the shuffle rewrite pass to handle the case where the
vector lanes are larger than the data type of the input vector.
---
src/target/source/codegen_c.cc | 4 +-
src/tir/transforms/storage_rewrite.cc | 2 +-
...est_tir_transform_pointer_value_type_rewrite.py | 46 ++++++++++++++++++++--
3 files changed, 47 insertions(+), 5 deletions(-)
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 009fc1672a..344d0392d4 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -932,7 +932,9 @@ void CodeGenC::VisitExpr_(const ShuffleNode* op,
std::ostream& os) { // NOLINT(
}
if (op->indices.size() == 1) {
// This is an extract element
- os << concat_vec[Downcast<IntImm>(op->indices[0])->value];
+ int64_t idx = Downcast<IntImm>(op->indices[0])->value;
+ ICHECK_LT(idx, concat_vec.size());
+ os << concat_vec[idx];
} else {
// Print the shuffle as vector constructor
// vec(e0, e1, e2, .. en)
diff --git a/src/tir/transforms/storage_rewrite.cc
b/src/tir/transforms/storage_rewrite.cc
index 2ebb767149..1c3f916a44 100644
--- a/src/tir/transforms/storage_rewrite.cc
+++ b/src/tir/transforms/storage_rewrite.cc
@@ -1493,7 +1493,7 @@ class VectorTypeRewriter : public StmtExprMutator {
arith::ModularSet me = analyzer_.modular_set(last_dim_index);
ICHECK(me->coeff == 0 || info.factor() % me->coeff == 0);
PrimExpr new_index = last_dim_index / make_const(last_dim_index.dtype(),
info.factor());
- shuffle_index = me->base;
+ shuffle_index = me->base % info.factor();
indices.Set(indices.size() - 1, new_index);
}
diff --git
a/tests/python/tir-transform/test_tir_transform_pointer_value_type_rewrite.py
b/tests/python/tir-transform/test_tir_transform_pointer_value_type_rewrite.py
index 7baa96c1a1..186f6bd02a 100644
---
a/tests/python/tir-transform/test_tir_transform_pointer_value_type_rewrite.py
+++
b/tests/python/tir-transform/test_tir_transform_pointer_value_type_rewrite.py
@@ -14,10 +14,10 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+# pylint: disable=invalid-name, missing-docstring
+
import tvm
import tvm.testing
-from tvm import te
-from tvm.driver.build_module import schedule_to_module
from tvm.script import tir as T
@@ -25,7 +25,7 @@ class BaseCompare(tvm.testing.CompareBeforeAfter):
transform = tvm.tir.transform.PointerValueTypeRewrite()
-class TestRewriteToShuffle(BaseCompare):
+class TestRewriteToShuffle0(BaseCompare):
@T.prim_func
def before(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32")):
A_local_data = T.allocate([16], "float32", scope="local")
@@ -50,6 +50,42 @@ class TestRewriteToShuffle(BaseCompare):
)
+class TestRewriteToShuffle1(BaseCompare):
+ @T.prim_func
+ def before(A: T.Buffer((8,), "float32"), B: T.Buffer((1,), "float32")):
+ A_local_data = T.allocate([8], "float32", scope="local")
+ A_local = T.Buffer((8,), "float32", data=A_local_data, scope="local")
+ A_local[0:4] = A[0:4]
+ A_local[4:8] = A[4:8]
+ B[0] = (
+ A_local[0]
+ + A_local[1]
+ + A_local[2]
+ + A_local[3]
+ + A_local[4]
+ + A_local[5]
+ + A_local[6]
+ + A_local[7]
+ )
+
+ @T.prim_func
+ def expected(A: T.Buffer((2,), "float32x4"), B: T.Buffer((1,), "float32")):
+ A_local_data = T.allocate([2], "float32x4", "local")
+ A_local = T.Buffer((2,), "float32x4", data=A_local_data, scope="local")
+ A_local[0] = A[0]
+ A_local[1] = A[1]
+ B[0] = (
+ T.Shuffle([A_local[0]], [0])
+ + T.Shuffle([A_local[0]], [1])
+ + T.Shuffle([A_local[0]], [2])
+ + T.Shuffle([A_local[0]], [3])
+ + T.Shuffle([A_local[1]], [0])
+ + T.Shuffle([A_local[1]], [1])
+ + T.Shuffle([A_local[1]], [2])
+ + T.Shuffle([A_local[1]], [3])
+ )
+
+
class TestAddressOf(BaseCompare):
@T.prim_func
def before(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
@@ -71,3 +107,7 @@ class TestScalarReadWithoutWrite(BaseCompare):
T.evaluate(A[i * 4])
expected = before
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
