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-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 77e9059d8 [FIX] Restrict Ramp elements to 32-bit and 64-bit integers
(#64)
77e9059d8 is described below
commit 77e9059d8585635abe13c0e72baa769074cc9aaf
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 25 11:19:40 2026 -0400
[FIX] Restrict Ramp elements to 32-bit and 64-bit integers (#64)
---
include/tvm/ir/prim/vector_expr.h | 5 ++++-
python/tvm_tirx/ir/prim/vector_expr.py | 6 +++++-
python/tvm_tirx/support/popen_worker.py | 6 ++++--
src/ir/prim/op.cc | 3 ++-
src/ir/prim/vector_expr.cc | 12 ++++++++++--
src/tirx/ir/data_type_rewriter.cc | 7 ++++---
src/tirx/transform/vectorize_loop.cc | 14 ++++++++++----
tests/python/tirx/transform/test_bf16_legalize.py | 4 ++--
8 files changed, 41 insertions(+), 16 deletions(-)
diff --git a/include/tvm/ir/prim/vector_expr.h
b/include/tvm/ir/prim/vector_expr.h
index 7a7881096..45fe28945 100644
--- a/include/tvm/ir/prim/vector_expr.h
+++ b/include/tvm/ir/prim/vector_expr.h
@@ -33,7 +33,10 @@ namespace prim {
/*!
* \brief Construct a vector with lanes elements
* where its i-th element equals base + i * stride.
- * This is useful to construct a index for a continuous vector load.
+ * This is useful to construct an index for a continuous vector load.
+ * Base and stride must be scalar int32, int64, uint32, or uint64. The stride
+ * is cast to the base type when their types differ. Other element types can
+ * be expressed by explicitly casting the resulting vector.
*
* Examples:
* - ramp(0, 1, 3) = [0, 1, 2]
diff --git a/python/tvm_tirx/ir/prim/vector_expr.py
b/python/tvm_tirx/ir/prim/vector_expr.py
index 827a0887f..83e9b5df7 100644
--- a/python/tvm_tirx/ir/prim/vector_expr.py
+++ b/python/tvm_tirx/ir/prim/vector_expr.py
@@ -25,7 +25,11 @@ from . import _ffi_api as _prim_ffi_api
@tvm_ffi.register_object("prim.Ramp")
class Ramp(ExprWithOp):
- """Construct a vector whose lane ``i`` equals ``base + i * stride``.
+ """Construct an index vector whose lane ``i`` equals ``base + i * stride``.
+
+ Base and stride must be scalar ``int32``, ``int64``, ``uint32``, or
+ ``uint64``. If their types differ, stride is cast to the base type.
+ Use an explicit vector cast to produce other element types.
Parameters
----------
diff --git a/python/tvm_tirx/support/popen_worker.py
b/python/tvm_tirx/support/popen_worker.py
index 4ae7f9d89..21b7081e2 100644
--- a/python/tvm_tirx/support/popen_worker.py
+++ b/python/tvm_tirx/support/popen_worker.py
@@ -109,8 +109,10 @@ def main():
msg = traceback.format_exc()
ret_value = (StatusKind.EXCEPTION, type(exception)(msg))
- if timeout is not None:
- watcher.cancel()
+ finally:
+ # SystemExit must cancel the timer before interpreter shutdown too.
+ if timeout is not None:
+ watcher.cancel()
lock.acquire()
if status.status == StatusKind.RUNNING:
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index bf9203c92..d53568f6b 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -401,7 +401,8 @@ PrimExpr cast(PrimExpr value, PrimType t, Location loc) {
const auto* base = ramp->base.as<IntImmObj>();
const auto* stride = ramp->stride.as<IntImmObj>();
if (dtype.IsFixedLengthVector() && base && stride &&
- dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
+ dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)
&&
+ (dtype.bits() == 32 || dtype.bits() == 64)) {
ffi::BigInt last = base->value + stride->value * (dtype.lanes() - 1);
if (detail::IsRepresentable(base->value, elem_ty) &&
detail::IsRepresentable(stride->value, elem_ty) &&
diff --git a/src/ir/prim/vector_expr.cc b/src/ir/prim/vector_expr.cc
index 180c5dacb..55e15de77 100644
--- a/src/ir/prim/vector_expr.cc
+++ b/src/ir/prim/vector_expr.cc
@@ -213,8 +213,16 @@ Ramp::Ramp(PrimExpr base, PrimExpr stride, PrimExpr lanes,
Location loc)
TVM_FFI_ICHECK(stride.defined());
PrimType base_ty = base.ty();
PrimType stride_ty = stride.ty();
- TVM_FFI_ICHECK(base_ty.IsScalar());
- TVM_FFI_ICHECK(stride_ty.IsScalar());
+ TVM_FFI_CHECK(base_ty.IsScalar() &&
+ base_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) &&
+ (base_ty.bits() == 32 || base_ty.bits() == 64),
+ TypeError)
+ << "Ramp base must be scalar int32, int64, uint32, or uint64, but got "
<< base_ty;
+ TVM_FFI_CHECK(stride_ty.IsScalar() &&
+ stride_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) &&
+ (stride_ty.bits() == 32 || stride_ty.bits() == 64),
+ TypeError)
+ << "Ramp stride must be scalar int32, int64, uint32, or uint64, but got
" << stride_ty;
if (stride_ty != base_ty) {
stride = prim::cast(stride, base_ty);
}
diff --git a/src/tirx/ir/data_type_rewriter.cc
b/src/tirx/ir/data_type_rewriter.cc
index 84de59cbd..78f325056 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -163,9 +163,10 @@ UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const
prim::RampObj* op,
} else {
PrimType base_dtype = base.ty();
PrimType stride_dtype = stride.ty();
- TVM_FFI_ICHECK(base_dtype.MatchesCode(DLDataTypeCode::kDLInt) &&
- stride_dtype.MatchesCode(DLDataTypeCode::kDLInt));
- int bits = std::max(base_dtype.bits(), stride_dtype.bits());
+ TVM_FFI_ICHECK(base_dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) &&
+ stride_dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt));
+ // Index narrowing may produce narrow children, but Ramp requires at least
32 bits.
+ int bits = std::max(32, std::max(base_dtype.bits(), stride_dtype.bits()));
PrimType dtype = base_dtype.WithBits(bits);
if (base_dtype->dtype != dtype->dtype) base = prim::cast(base, dtype);
if (stride_dtype->dtype != dtype->dtype) stride = prim::cast(stride,
dtype);
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index c65e61c1d..786e3515e 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -475,8 +475,7 @@ class Vectorizer : public StmtExprMutator {
Vectorizer(Var var, PrimExpr var_lanes, ffi::Optional<Target> target)
: var_(var), var_lanes_(var_lanes), target_(target) {
- PrimType var_ty = var->ty.as_or_throw<PrimType>();
- ramp_ = prim::Ramp(IntImm(var_ty, 0), IntImm(var_ty, 1), var_lanes);
+ ramp_ = MakeIndexRamp(1, var_lanes);
VarRemapSet(var_, ramp_);
}
@@ -1004,8 +1003,7 @@ class Vectorizer : public StmtExprMutator {
} else {
PrimExpr prev_ramp = ramp_;
PrimExpr prev_var_lanes = var_lanes_;
- PrimType var_ty = var_->ty.as_or_throw<PrimType>();
- ramp_ = prim::Ramp(IntImm(var_ty, 0), IntImm(var_ty, 2), new_vec_length);
+ ramp_ = MakeIndexRamp(2, new_vec_length);
var_lanes_ = tvm::IntImm(var_lanes_.ty(), new_vec_length);
lane_vectors = 0;
vectors = MutateArray(op->vectors, &lane_vectors, inplace_mode);
@@ -1187,6 +1185,14 @@ class Vectorizer : public StmtExprMutator {
}
private:
+ PrimExpr MakeIndexRamp(int stride, PrimExpr lanes) const {
+ PrimType var_ty = var_->ty.as_or_throw<PrimType>();
+ PrimType index_ty = var_ty.WithBits(std::max(32, var_ty.bits()));
+ PrimExpr ramp = prim::Ramp(IntImm(index_ty, 0), IntImm(index_ty, stride),
lanes);
+ // Preserve narrow loop-variable types with an explicit vector cast.
+ return prim::cast(ramp,
ramp.ty().WithCode(var_ty.code()).WithBits(var_ty.bits()));
+ }
+
// analyzer
sym::Analyzer analyzer_;
// deep equal
diff --git a/tests/python/tirx/transform/test_bf16_legalize.py
b/tests/python/tirx/transform/test_bf16_legalize.py
index 6e8af995d..449ca8ef0 100644
--- a/tests/python/tirx/transform/test_bf16_legalize.py
+++ b/tests/python/tirx/transform/test_bf16_legalize.py
@@ -52,13 +52,13 @@ def test_bf16_ramp_will_legalize():
class Before:
@T.function
def main():
- T.evaluate(T.Ramp(T.bfloat16(0), T.bfloat16(1), 2))
+ T.evaluate(T.Cast("bfloat16x2", T.Ramp(0, 1, 2)))
@I.ir_module
class Expected:
@T.function
def main():
- T.evaluate(T.Ramp(T.float32(0), T.float32(1), 2))
+ T.evaluate(T.Cast("float32x2", T.Ramp(0, 1, 2)))
target = Target("llvm")
before = BindTarget(target)(Before)
