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 b5fda2d93a [TIR] Ramp and Broadcast lanes fixed to int32 dtype (#16795) b5fda2d93a is described below commit b5fda2d93ab91b08753b23eff92916f097bd2620 Author: Anirudh Sundar Subramaniam <quic_sanir...@quicinc.com> AuthorDate: Mon Apr 1 17:37:11 2024 +0530 [TIR] Ramp and Broadcast lanes fixed to int32 dtype (#16795) * [TIR] Ramp and Broadcast lanes fixed to int32 dtype When Ramp and Broadcast nodes are created with fixed length lanes, they're fixed to int32 dtype since DLDataType always supports only uint16 lanes. * Add test cases for int64 type lanes * Update test case with int64 iterators --- src/tir/ir/expr.cc | 8 ++++++-- tests/python/arith/test_arith_rewrite_simplify.py | 15 +++++++++++++++ tests/python/tir-base/test_tir_nodes.py | 10 ++++++++++ 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc index c2baad2096..90dad72039 100644 --- a/src/tir/ir/expr.cc +++ b/src/tir/ir/expr.cc @@ -449,16 +449,18 @@ Ramp::Ramp(PrimExpr base, PrimExpr stride, PrimExpr lanes, Span span) { int lanes = static_cast<int>(lanes_as_int->value); ICHECK_GT(lanes, 1); node->dtype = base.dtype().with_lanes(lanes); + // Stick to int32 lanes for fixed length vectors + node->lanes = lanes; } else { /* scalable vector */ std::optional<int> vscale_factor = arith::ExtractVscaleFactor(lanes); ICHECK(vscale_factor) << "Invalid expression for scalable lanes " << lanes; node->dtype = base.dtype().with_scalable_vscale_factor(vscale_factor.value()); lanes = Mul(Call(DataType::Int(32), tir::builtin::vscale(), {}), vscale_factor.value()); + node->lanes = lanes; } node->base = base; node->stride = stride; - node->lanes = lanes; node->span = std::move(span); data_ = std::move(node); } @@ -481,15 +483,17 @@ Broadcast::Broadcast(PrimExpr value, PrimExpr lanes, Span span) { int lanes = static_cast<int>(lanes_int->value); ICHECK_GT(lanes, 1); node->dtype = value.dtype().with_lanes(lanes); + // Stick to int32 lanes for fixed length vectors + node->lanes = lanes; } else { /* scalable vector */ std::optional<int> vscale_factor = arith::ExtractVscaleFactor(lanes); ICHECK(vscale_factor) << "Invalid expression for scalable lanes " << lanes; node->dtype = value.dtype().with_scalable_vscale_factor(vscale_factor.value()); lanes = Mul(Call(DataType::Int(32), tir::builtin::vscale(), {}), vscale_factor.value()); + node->lanes = lanes; } node->value = std::move(value); - node->lanes = lanes; node->span = std::move(span); data_ = node; } diff --git a/tests/python/arith/test_arith_rewrite_simplify.py b/tests/python/arith/test_arith_rewrite_simplify.py index 8645e5b26a..9cc44aa6a2 100644 --- a/tests/python/arith/test_arith_rewrite_simplify.py +++ b/tests/python/arith/test_arith_rewrite_simplify.py @@ -75,6 +75,7 @@ class BaseCompare: class TestVector(BaseCompare): x, y, z = te.var("x"), te.var("y"), te.var("z") + x64 = te.var("x", dtype="int64") vx = te.var("vx", dtype="int32x2") vc = te.var("vc", dtype="uint1") test_case = tvm.testing.parameter( @@ -88,6 +89,20 @@ class TestVector(BaseCompare): ), TestCase(y.astype("int32x2") + x.astype("int32x2"), (y + x).astype("int32x2")), TestCase(tvm.tir.Broadcast(0, 4) + y, tvm.tir.Broadcast(y, 4)), + # int64 lanes + TestCase( + tvm.tir.Broadcast(x, 4) + tvm.tir.Ramp(0, 1, tvm.tir.IntImm(dtype="int64", value=4)), + tvm.tir.Ramp(x, 1, 4), + ), + TestCase( + tvm.tir.Broadcast(x, tvm.tir.IntImm(dtype="int64", value=4)) + tvm.tir.Ramp(0, 1, 4), + tvm.tir.Ramp(x, 1, 4), + ), + # int64 iterators with int32 lanes + TestCase( + tvm.tir.Broadcast(x64, 4) + tvm.tir.Ramp(tvm.tir.IntImm(dtype="int64", value=0), 1, 4), + tvm.tir.Ramp(x64, 1, 4), + ), TestCase( tvm.tir.Broadcast(0, tir.vscale() * 8) + y, tvm.tir.Broadcast(y, tir.vscale() * 8) ), diff --git a/tests/python/tir-base/test_tir_nodes.py b/tests/python/tir-base/test_tir_nodes.py index 60f8278ec2..31a1317e68 100644 --- a/tests/python/tir-base/test_tir_nodes.py +++ b/tests/python/tir-base/test_tir_nodes.py @@ -409,6 +409,16 @@ def _create_broadcast(lanes): return tvm.tir.Broadcast(0, lanes) +@pytest.mark.parametrize("lanes", [(tvm.tir.IntImm(dtype="int64", value=11))]) +@pytest.mark.parametrize("node_func", [_create_ramp, _create_broadcast]) +def test_lane_types(lanes, node_func): + def _check_dtype(node): + assert node.lanes.dtype == "int32" + assert node.lanes == 11 + + _check_dtype(node_func(lanes)) + + @pytest.mark.parametrize("lanes", [(11 * tvm.tir.vscale()), (tvm.tir.vscale() * 11)]) @pytest.mark.parametrize("node_func", [_create_ramp, _create_broadcast]) def test_scalable_vec(lanes, node_func):