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 dc35fb049 [FIX] Rebuild Ramp types during dtype legalization (#43)
dc35fb049 is described below
commit dc35fb049bc963b98b6ce851aa161073e5ead797
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 24 14:32:05 2026 -0400
[FIX] Rebuild Ramp types during dtype legalization (#43)
Rebuild Ramp expressions when unsupported-dtype legalization promotes
their operands, so the derived vector result type follows the promoted
base and stride. Preserve lane expressions, source locations, and
unchanged nodes at the existing legalization boundary.
This prevents a BF16 Ramp from retaining a BF16 vector result after its
operands become float32, which otherwise reaches LLVM with an
unsupported vector element type.
---
src/tirx/transform/unsupported_dtype_legalize.cc | 15 +++++++++++++++
tests/python/tirx/transform/test_bf16_legalize.py | 22 ++++++++++++++++++++++
2 files changed, 37 insertions(+)
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc
b/src/tirx/transform/unsupported_dtype_legalize.cc
index c675a5357..af578f7e8 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -209,6 +209,21 @@ class ComputeLegalizer : public StmtExprMutator {
}
}
+ UnchangedOr<PrimExpr> Mutate_(const prim::RampObj* op, InplaceMode
inplace_mode) final {
+ PrimExpr base =
+ PromoteToTarget(this->Mutate(op->base,
inplace_mode).ValueOrUnchanged(op->base));
+ PrimExpr stride =
+ PromoteToTarget(this->Mutate(op->stride,
inplace_mode).ValueOrUnchanged(op->stride));
+ auto lanes_result = this->Mutate(op->lanes, inplace_mode);
+ bool lanes_unchanged = lanes_result.UnchangedOrSameAs(op->lanes);
+ PrimExpr lanes = std::move(lanes_result).ValueOrUnchanged(op->lanes);
+ if (base.same_as(op->base) && stride.same_as(op->stride) &&
lanes_unchanged &&
+ !MatchType(op->ty)) {
+ return ffi::Unchanged();
+ }
+ return prim::Ramp(base, stride, lanes, op->loc);
+ }
+
UnchangedOr<PrimExpr> Mutate_(const prim::BroadcastObj* op, InplaceMode
inplace_mode) final {
PrimExpr value =
PromoteToTarget(this->Mutate(op->value,
inplace_mode).ValueOrUnchanged(op->value));
diff --git a/tests/python/tirx/transform/test_bf16_legalize.py
b/tests/python/tirx/transform/test_bf16_legalize.py
index aa42d9c2c..6e8af995d 100644
--- a/tests/python/tirx/transform/test_bf16_legalize.py
+++ b/tests/python/tirx/transform/test_bf16_legalize.py
@@ -47,6 +47,28 @@ def f32tobf16(v):
return T.reinterpret(f32tou16(v), ty="bfloat16")
+def test_bf16_ramp_will_legalize():
+ @I.ir_module
+ class Before:
+ @T.function
+ def main():
+ T.evaluate(T.Ramp(T.bfloat16(0), T.bfloat16(1), 2))
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main():
+ T.evaluate(T.Ramp(T.float32(0), T.float32(1), 2))
+
+ target = Target("llvm")
+ before = BindTarget(target)(Before)
+ expected = BindTarget(target)(Expected)
+ after_compute = tvm_tirx.tirx.transform.BF16ComputeLegalize()(before)
+ tvm_tirx.ir.assert_structural_equal(after_compute, expected)
+ after_storage =
tvm_tirx.tirx.transform.BF16StorageLegalize()(after_compute)
+ tvm_tirx.ir.assert_structural_equal(after_storage, expected)
+
+
def test_bf16_simple_store_will_legalize():
def get_before():
@tvm_tirx.script.ir_module
