This is an automated email from the ASF dual-hosted git repository.
tlopex 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 44597a1d69 [Fix][TIRx] Preserve index semantics when narrowing to
int32 (#20129)
44597a1d69 is described below
commit 44597a1d69e636569e3758e95017372a6fca464f
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Sat Aug 29 20:30:07 2026 -0700
[Fix][TIRx] Preserve index semantics when narrowing to int32 (#20129)
This PR fixes two cases where `ForceNarrowIndexToInt32` could produce
invalid or semantically different index expressions. When rewriting
`if_then_else` and `Select`, one branch may remain `int64` while the
other is narrowed to `int32`. The pass now derives the result type from
the rewritten branches and casts the narrower branch to the common type.
The pass also adjusts signed left and right shifts when their input is
narrowed. Under the pass invariant that index values fit in `int32`,
scalar, dynamic, and vector shift amounts greater than the new sign-bit
position are clamped lane-wise to that position. This preserves
right-shift sign extraction and keeps left shifts valid without changing
representable results. The clamp is scoped to `ForceNarrowIndexToInt32`,
so datatype legalization used by substitution preserves the original
shift amount.
---
src/tirx/ir/data_type_rewriter.cc | 39 +++++-
src/tirx/ir/data_type_rewriter.h | 3 +
src/tirx/transform/force_narrow_index_to_i32.cc | 2 +
tests/cpp/ir_functor_test.cc | 22 ++++
...test_tir_transform_force_narrow_index_to_i32.py | 140 +++++++++++++++++++++
5 files changed, 200 insertions(+), 6 deletions(-)
diff --git a/src/tirx/ir/data_type_rewriter.cc
b/src/tirx/ir/data_type_rewriter.cc
index b2f66d06f6..4ead88b9c3 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -251,9 +251,32 @@ Expr DataTypeLegalizer::VisitExpr_(const CallNode* op) {
}
PrimExpr prim_e = e.as_or_throw<PrimExpr>();
if (op->op.same_as(builtin::shift_right())) {
- return op->args[0].as_or_throw<PrimExpr>() >>
op->args[1].as_or_throw<PrimExpr>();
+ PrimExpr lhs = op->args[0].as_or_throw<PrimExpr>();
+ PrimExpr rhs = op->args[1].as_or_throw<PrimExpr>();
+ PrimType before_dtype = before->args[0].as_or_throw<PrimExpr>().ty();
+ PrimType after_dtype = lhs.ty();
+ if (ShouldClampShiftAmounts() && before_dtype.code() ==
DLDataTypeCode::kDLInt &&
+ after_dtype.code() == DLDataTypeCode::kDLInt && before_dtype.bits() >
after_dtype.bits()) {
+ // Values are assumed to fit in the narrowed dtype. An arithmetic right
+ // shift at or beyond its sign bit therefore has the same value as a
shift
+ // by the new sign-bit position. Clamp lane-wise so dynamic and vector
+ // shift amounts remain valid for the narrowed dtype.
+ rhs = min(rhs, MakeConst(rhs.ty(), after_dtype.bits() - 1, op->span),
op->span);
+ }
+ return lhs >> rhs;
} else if (op->op.same_as(builtin::shift_left())) {
- return op->args[0].as_or_throw<PrimExpr>() <<
op->args[1].as_or_throw<PrimExpr>();
+ PrimExpr lhs = op->args[0].as_or_throw<PrimExpr>();
+ PrimExpr rhs = op->args[1].as_or_throw<PrimExpr>();
+ PrimType before_dtype = before->args[0].as_or_throw<PrimExpr>().ty();
+ PrimType after_dtype = lhs.ty();
+ if (ShouldClampShiftAmounts() && before_dtype.code() ==
DLDataTypeCode::kDLInt &&
+ after_dtype.code() == DLDataTypeCode::kDLInt && before_dtype.bits() >
after_dtype.bits()) {
+ // Keep dynamic and vector shift amounts valid for the narrowed dtype.
Under the pass's
+ // representability precondition, a left shift at or beyond the narrowed
width can only
+ // produce a representable result when lhs is zero, so clamping does not
alter valid cases.
+ rhs = min(rhs, MakeConst(rhs.ty(), after_dtype.bits() - 1, op->span),
op->span);
+ }
+ return lhs << rhs;
} else if (op->op.same_as(builtin::bitwise_and())) {
return op->args[0].as_or_throw<PrimExpr>() &
op->args[1].as_or_throw<PrimExpr>();
} else if (op->op.same_as(builtin::bitwise_or())) {
@@ -593,10 +616,14 @@ Expr IndexDataTypeRewriter::VisitExpr_(const CallNode*
op) {
is_condition_ = true;
PrimExpr cond = VisitPrimExpr(op->args[0].as_or_throw<PrimExpr>());
is_condition_ = is_condition;
- return Call(op->ty.as_or_throw<PrimType>(), op->op,
- {cond, VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>()),
- VisitPrimExpr(op->args[2].as_or_throw<PrimExpr>())},
- op->attrs, {}, op->span)
+ PrimExpr true_value = VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>());
+ PrimExpr false_value = VisitPrimExpr(op->args[2].as_or_throw<PrimExpr>());
+ PrimType true_dtype = true_value.ty();
+ PrimType false_dtype = false_value.ty();
+ PrimType dtype = true_dtype.WithBits(std::max(true_dtype.bits(),
false_dtype.bits()));
+ if (true_dtype != dtype) true_value = cast(dtype, true_value);
+ if (false_dtype != dtype) false_value = cast(dtype, false_value);
+ return Call(dtype, op->op, {cond, true_value, false_value}, op->attrs, {},
op->span)
.as_or_throw<PrimExpr>();
}
return Parent::VisitExpr_(op);
diff --git a/src/tirx/ir/data_type_rewriter.h b/src/tirx/ir/data_type_rewriter.h
index 883bde94b0..03d3945b39 100644
--- a/src/tirx/ir/data_type_rewriter.h
+++ b/src/tirx/ir/data_type_rewriter.h
@@ -76,6 +76,9 @@ class DataTypeLegalizer : public StmtExprMutator {
Expr VisitExpr_(const CastNode* op) override;
Expr VisitExpr_(const LetNode* op) override;
+ /*! \brief Whether to clamp shift amounts after narrowing signed integers. */
+ virtual bool ShouldClampShiftAmounts() const { return false; }
+
using StmtExprMutator::VisitExpr_;
using StmtExprMutator::VisitStmt_;
diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc
b/src/tirx/transform/force_narrow_index_to_i32.cc
index 2e444c273b..86c35fea99 100644
--- a/src/tirx/transform/force_narrow_index_to_i32.cc
+++ b/src/tirx/transform/force_narrow_index_to_i32.cc
@@ -54,6 +54,8 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer {
explicit Int32DTypeNarrower(PrimFunc func)
: IndexDataTypeNormalizer(PrimType::Int(32)), func_(std::move(func)) {}
+ bool ShouldClampShiftAmounts() const final { return true; }
+
Expr VisitExpr_(const IntImmNode* op) final {
// ignore the enabled condition and always rewrite i64
if (op->ty.as_or_throw<PrimType>() == PrimType::Int(64)) {
diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc
index 6ca291ce43..ea5db7893e 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -18,6 +18,7 @@
*/
#include <gtest/gtest.h>
+#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ir/module.h>
#include <tvm/ir/node_functor.h>
#include <tvm/runtime/logging.h>
@@ -382,3 +383,24 @@ TEST(IRF, Substitute) {
TVM_FFI_ICHECK(new_expr.same_as(expr));
}
}
+
+TEST(IRF, SubstituteWithDataTypeLegalizationPreservesShiftAmounts) {
+ using namespace tvm;
+ using namespace tvm::tirx;
+
+ PrimVar x("x", PrimType::Int(64));
+ PrimVar y("y", PrimType::Int(32));
+ auto f_subst = [&](const tirx::Var& var) -> ffi::Optional<PrimExpr> {
+ if (var.same_as(x)) return PrimExpr(y);
+ return std::nullopt;
+ };
+
+ PrimExpr shift_amount = IntImm::Int64(40);
+ PrimExpr widened_y = cast(PrimType::Int(64), y);
+ PrimExpr actual_left = SubstituteWithDataTypeLegalization(x << shift_amount,
f_subst);
+ PrimExpr actual_right = SubstituteWithDataTypeLegalization(x >>
shift_amount, f_subst);
+
+ ffi::StructuralEqual structural_equal;
+ EXPECT_TRUE(structural_equal(actual_left, widened_y << shift_amount));
+ EXPECT_TRUE(structural_equal(actual_right, widened_y >> shift_amount));
+}
diff --git
a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
index af5b719417..1016966729 100644
---
a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
+++
b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py
@@ -260,6 +260,48 @@ def test_pod_params_and_select():
tvm.ir.assert_structural_equal(Expected, after)
+def test_if_then_else_index():
+ @tvm.script.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((T.int64(4),), "float32"), B: T.Buffer((1,),
"float32"), n: T.int64):
+ B[0] = A[T.if_then_else(n < T.int64(0), n + T.int64(1), n)]
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((4,), "float32"), B: T.Buffer((1,), "float32"),
n: T.int32):
+ B[0] = A[T.if_then_else(n < 0, n + 1, n)]
+
+ after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+ tvm.ir.assert_structural_equal(Expected, after)
+
+
+def test_conditional_index_mixed_width_branches():
+ @tvm.script.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((T.int64(4),), "float32"), B: T.Buffer((4,),
"float32"), n: T.int64):
+ opaque_index = T.call_extern("opaque_index", n, dtype="int64")
+ B[0] = A[T.if_then_else(n < T.int64(0), opaque_index, n)]
+ B[1] = A[T.if_then_else(n < T.int64(0), n, opaque_index)]
+ B[2] = A[T.Select(n < T.int64(0), opaque_index, n)]
+ B[3] = A[T.Select(n < T.int64(0), n, opaque_index)]
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32"),
n: T.int32):
+ opaque_index = T.call_extern("opaque_index", n, dtype="int64")
+ B[0] = A[T.if_then_else(n < 0, opaque_index, T.Cast("int64", n))]
+ B[1] = A[T.if_then_else(n < 0, T.Cast("int64", n), opaque_index)]
+ B[2] = A[T.Select(n < 0, opaque_index, T.Cast("int64", n))]
+ B[3] = A[T.Select(n < 0, T.Cast("int64", n), opaque_index)]
+
+ after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+ tvm.ir.assert_structural_equal(Expected, after)
+
+
def test_clz():
@tvm.script.ir_module
class Before:
@@ -279,6 +321,104 @@ def test_clz():
tvm.ir.assert_structural_equal(Expected, after)
+def test_right_shift_preserves_sign_extension_after_narrowing():
+ @tvm.script.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((T.int64(6),), "float32"), B: T.Buffer((1,),
"float32"), n: T.int64):
+ B[0] = A[T.shift_right(T.truncmod(n - T.int64(8), T.int64(6)),
T.int64(63))]
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((6,), "float32"), B: T.Buffer((1,), "float32"),
n: T.int32):
+ B[0] = A[T.shift_right(T.truncmod(n - 8, 6), 31)]
+
+ # ForceNarrowIndexToInt32 assumes that index values fit in int32. Under
+ # that precondition, shifting the original int64 value by 63 and shifting
+ # the narrowed value by its sign-bit position have the same result.
+ after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+ tvm.ir.assert_structural_equal(Expected, after)
+
+
+def test_right_shift_dynamic_and_vector_amounts():
+ @tvm.script.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((T.int64(6),), "float32"),
+ B: T.Buffer((T.int64(5),), "float32"),
+ n: T.int64,
+ shift: T.int64,
+ ):
+ B[0] = A[T.shift_right(T.truncmod(n - T.int64(8), T.int64(6)),
shift)]
+ B[T.Ramp(T.int64(1), T.int64(1), 4)] = A[
+ T.shift_right(
+ T.Broadcast(T.truncmod(n - T.int64(8), T.int64(6)), 4),
+ T.Ramp(T.int64(30), T.int64(1), 4),
+ )
+ ]
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((6,), "float32"),
+ B: T.Buffer((5,), "float32"),
+ n: T.int32,
+ shift: T.int32,
+ ):
+ B[0] = A[T.shift_right(T.truncmod(n - 8, 6), T.min(shift, 31))]
+ B[T.Ramp(1, 1, 4)] = A[
+ T.shift_right(
+ T.Broadcast(T.truncmod(n - 8, 6), 4),
+ T.min(T.Ramp(30, 1, 4), T.Broadcast(31, 4)),
+ )
+ ]
+
+ after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+ tvm.ir.assert_structural_equal(Expected, after)
+
+
+def test_left_shift_dynamic_and_vector_amounts_remain_valid():
+ @tvm.script.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((T.int64(1),), "float32"),
+ B: T.Buffer((T.int64(5),), "float32"),
+ n: T.int64,
+ shift: T.int64,
+ ):
+ B[0] = A[T.shift_left(T.truncmod(n, T.int64(1)), shift)]
+ B[T.Ramp(T.int64(1), T.int64(1), 4)] = A[
+ T.shift_left(
+ T.Broadcast(T.truncmod(n, T.int64(1)), 4),
+ T.Ramp(T.int64(30), T.int64(1), 4),
+ )
+ ]
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((1,), "float32"),
+ B: T.Buffer((5,), "float32"),
+ n: T.int32,
+ shift: T.int32,
+ ):
+ B[0] = A[T.shift_left(T.truncmod(n, 1), T.min(shift, 31))]
+ B[T.Ramp(1, 1, 4)] = A[
+ T.shift_left(
+ T.Broadcast(T.truncmod(n, 1), 4),
+ T.min(T.Ramp(30, 1, 4), T.Broadcast(31, 4)),
+ )
+ ]
+
+ after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before)
+ tvm.ir.assert_structural_equal(Expected, after)
+
+
def test_let_binding():
@tvm.script.ir_module
class Before:
