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 a7b48136a3 [REFACTOR][S-TIR] Localize index datatype normalization
(#20688)
a7b48136a3 is described below
commit a7b48136a32e745fb7dd4af074eb10a7cab6010a
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 19:18:41 2026 -0400
[REFACTOR][S-TIR] Localize index datatype normalization (#20688)
Give S-TIR its own private datatype legalization and index-context
helpers so TE function creation and tensorization no longer inherit the
private TIRx datatype rewriter. Preserve the existing S-TIR
normalization interface and behavior.
---
src/s_tir/ir/data_type_rewriter.cc | 532 +++++++++++++++++++++++++++++++++++++
src/s_tir/ir/data_type_rewriter.h | 116 +++++++-
2 files changed, 644 insertions(+), 4 deletions(-)
diff --git a/src/s_tir/ir/data_type_rewriter.cc
b/src/s_tir/ir/data_type_rewriter.cc
index 25bb577f0e..9e770981c7 100644
--- a/src/s_tir/ir/data_type_rewriter.cc
+++ b/src/s_tir/ir/data_type_rewriter.cc
@@ -19,15 +19,521 @@
#include "data_type_rewriter.h"
+#include <tvm/ffi/cast.h>
+#include <tvm/ir/op.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/tirx/op/region.h>
+#include <algorithm>
#include <functional>
+#include <utility>
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
using namespace tvm::prim;
+UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode
inplace_mode) {
+ auto result = StmtExprMutator::Mutate_(op, inplace_mode);
+ if (!result.IsUnchanged()) {
+ op = ffi::AnyView(result).as<ForNode>();
+ TVM_FFI_ICHECK(op != nullptr) << "Expected type to be ForNode, but get "
+ << ffi::AnyView(result).GetTypeKey();
+ if (!op->unique()) inplace_mode = InplaceMode::kDisallow;
+ }
+ PrimExpr e = Mutate(op->loop_var,
inplace_mode).ValueOrUnchanged(op->loop_var);
+ Var var = e.as_or_throw<Var>();
+ PrimType var_ty = var->ty.as_or_throw<PrimType>();
+ if (inplace_mode == InplaceMode::kAllow) {
+ auto* n = const_cast<ForNode*>(op);
+ n->min = prim::cast(var_ty, op->min);
+ n->extent = prim::cast(var_ty, op->extent);
+ if (op->step.has_value()) {
+ n->step = prim::cast(var_ty, *op->step);
+ }
+ return result;
+ }
+ auto n = ffi::make_object<ForNode>(*op);
+ n->min = prim::cast(var_ty, op->min);
+ n->extent = prim::cast(var_ty, op->extent);
+ if (op->step.has_value()) {
+ n->step = prim::cast(var_ty, *op->step);
+ }
+ return For(n);
+}
+
+UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const RegionStmtNode* op,
InplaceMode inplace_mode) {
+ Stmt result = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ const auto* region = result.as<RegionStmtNode>();
+ if (!region || !region->op.same_as(tirx::launch_thread_op())) return result;
+ PrimVar var = region->body_params[0].as_or_throw<PrimVar>();
+ PrimExpr extent = region->args[1].as_or_throw<PrimExpr>();
+ if (var.ty() == extent.ty()) return result;
+ RegionStmt updated = ffi::GetRef<RegionStmt>(region);
+ updated.CopyOnWrite()->args.Set(1, prim::cast(var.ty(), extent));
+ return updated;
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::LetNode* op,
+ InplaceMode inplace_mode) {
+ auto value_result = this->Mutate(op->value, inplace_mode);
+ bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
+ PrimExpr value = std::move(value_result).ValueOrUnchanged(op->value);
+ Var var = Mutate(op->var).ValueOrUnchanged(op->var).as_or_throw<Var>();
+ if (value.ty() != var->ty.as_or_throw<PrimType>()) {
+ if (var.same_as(op->var)) {
+ var = op->var.CopyWithDType(value.ty());
+ VarRemapSet(op->var, var);
+ } else {
+ value = prim::cast(var->ty.as_or_throw<PrimType>(), value);
+ value_unchanged = false;
+ }
+ }
+ auto new_body_result = this->Mutate(op->body, inplace_mode);
+ bool new_body_unchanged = new_body_result.UnchangedOrSameAs(op->body);
+ PrimExpr new_body = std::move(new_body_result).ValueOrUnchanged(op->body);
+
+ if (value_unchanged && new_body_unchanged && var.same_as(op->var)) {
+ return ffi::Unchanged();
+ } else {
+ return prim::Let(var, value, new_body, op->loc);
+ }
+}
+
+UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const BindNode* op, InplaceMode
inplace_mode) {
+ if (!op->value.as<PrimExpr>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ auto value_result = this->Mutate(op->value, inplace_mode);
+ bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
+ Expr value = std::move(value_result).ValueOrUnchanged(op->value);
+ Var var = op->var;
+
+ if (auto prim_value = value.as<PrimExpr>()) {
+ if (prim_value.value().ty() != op->var->ty.as_or_throw<PrimType>()) {
+ var = op->var.CopyWithDType(prim_value.value().ty());
+ VarRemapSet(op->var, var);
+ }
+ }
+
+ if (value_unchanged && var.same_as(op->var)) {
+ return ffi::Unchanged();
+ } else {
+ return Bind(var, value, op->loc);
+ }
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::SelectNode* op,
+ InplaceMode inplace_mode) {
+ auto condition_result = this->Mutate(op->condition, inplace_mode);
+ bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition);
+ PrimExpr condition =
std::move(condition_result).ValueOrUnchanged(op->condition);
+ auto true_value_result = this->Mutate(op->true_value, inplace_mode);
+ bool true_value_unchanged =
true_value_result.UnchangedOrSameAs(op->true_value);
+ PrimExpr true_value =
std::move(true_value_result).ValueOrUnchanged(op->true_value);
+ auto false_value_result = this->Mutate(op->false_value, inplace_mode);
+ bool false_value_unchanged =
false_value_result.UnchangedOrSameAs(op->false_value);
+ PrimExpr false_value =
std::move(false_value_result).ValueOrUnchanged(op->false_value);
+ if (condition_unchanged && true_value_unchanged && false_value_unchanged &&
+ true_value.ty() == false_value.ty()) {
+ return ffi::Unchanged();
+ } else {
+ PrimType true_dtype = true_value.ty();
+ PrimType false_dtype = false_value.ty();
+ int bits = std::max(true_dtype.bits(), false_dtype.bits());
+ PrimType dtype = true_dtype.WithBits(bits);
+ if (true_dtype != dtype) true_value = prim::cast(dtype, true_value);
+ if (false_dtype != dtype) false_value = prim::cast(dtype, false_value);
+ return prim::Select(condition, true_value, false_value);
+ }
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::RampNode* op,
+ InplaceMode inplace_mode) {
+ auto base_result = Mutate(op->base, inplace_mode);
+ bool base_unchanged = base_result.UnchangedOrSameAs(op->base);
+ PrimExpr base = std::move(base_result).ValueOrUnchanged(op->base);
+ auto stride_result = Mutate(op->stride, inplace_mode);
+ bool stride_unchanged = stride_result.UnchangedOrSameAs(op->stride);
+ PrimExpr stride = std::move(stride_result).ValueOrUnchanged(op->stride);
+ if (base_unchanged && stride_unchanged && base.ty() == stride.ty()) {
+ return ffi::Unchanged();
+ } 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());
+ PrimType dtype = base_dtype.WithBits(bits);
+ if (base_dtype->dtype != dtype->dtype) base = prim::cast(dtype, base);
+ if (stride_dtype->dtype != dtype->dtype) stride = prim::cast(dtype,
stride);
+ return prim::Ramp(base, stride, op->lanes);
+ }
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::BroadcastNode* op,
+ InplaceMode inplace_mode) {
+ auto value = Mutate(op->value, inplace_mode);
+ auto lanes = Mutate(op->lanes, inplace_mode);
+ if (value.UnchangedOrSameAs(op->value) &&
lanes.UnchangedOrSameAs(op->lanes)) {
+ return ffi::Unchanged();
+ }
+ // Construction re-infers the dtype; the shared structural hook retains the
old dtype.
+ return prim::Broadcast(std::move(value).ValueOrUnchanged(op->value),
+ std::move(lanes).ValueOrUnchanged(op->lanes));
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::ShuffleNode* op,
+ InplaceMode inplace_mode) {
+ auto vectors = Mutate(op->vectors,
inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
+ auto indices = Mutate(op->indices,
inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
+ bool unchanged = vectors.UnchangedOrSameAs(op->vectors) &&
indices.UnchangedOrSameAs(op->indices);
+ if (unchanged && inplace_mode == InplaceMode::kDisallow) return
ffi::Unchanged();
+ // Arrays can change in place; reconstruct to infer the dtype from the final
vectors.
+ PrimExpr updated =
prim::Shuffle(std::move(vectors).ValueOrUnchanged(op->vectors),
+
std::move(indices).ValueOrUnchanged(op->indices));
+ if (unchanged && updated.ty() == op->ty.as_or_throw<PrimType>()) return
ffi::Unchanged();
+ return updated;
+}
+
+#define TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC)
\
+ UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const OP* op, InplaceMode
inplace_mode) { \
+ auto a_result = this->Mutate(op->a, inplace_mode);
\
+ bool a_unchanged = a_result.UnchangedOrSameAs(op->a);
\
+ PrimExpr a = std::move(a_result).ValueOrUnchanged(op->a);
\
+ auto b_result = this->Mutate(op->b, inplace_mode);
\
+ bool b_unchanged = b_result.UnchangedOrSameAs(op->b);
\
+ PrimExpr b = std::move(b_result).ValueOrUnchanged(op->b);
\
+ if (a_unchanged && b_unchanged && a.ty() == b.ty()) {
\
+ return ffi::Unchanged();
\
+ } else {
\
+ return FUNC(a, b);
\
+ }
\
+ }
+
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::AddNode, operator+);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::SubNode, operator-);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MulNode, operator*);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::DivNode, div);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::ModNode, truncmod);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::FloorDivNode, floordiv);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::FloorModNode, floormod);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MinNode, min);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MaxNode, max);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::EQNode, operator==);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::NENode, operator!=);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LENode, operator<=);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LTNode, operator<); //
NOLINT(*)
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GTNode, operator>); //
NOLINT(*)
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GENode, operator>=);
+
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseAndNode, bitwise_and);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseOrNode, bitwise_or);
+TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseXorNode, bitwise_xor);
+
+#undef TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::LShiftNode* op,
+ InplaceMode inplace_mode) {
+ PrimType before_dtype = op->a.ty();
+ // Preserve the original operands while computing the narrowed shift.
+ PrimExpr lhs = Mutate(op->a, InplaceMode::kDisallow).ValueOrUnchanged(op->a);
+ PrimExpr rhs = Mutate(op->b, InplaceMode::kDisallow).ValueOrUnchanged(op->b);
+ PrimType after_dtype = lhs.ty();
+ if (ShouldClampShiftAmounts() && before_dtype.code() ==
DLDataTypeCode::kDLInt &&
+ after_dtype.code() == DLDataTypeCode::kDLInt && before_dtype.bits() >
after_dtype.bits()) {
+ // Values fit in the narrowed dtype. Clamp lane-wise to keep dynamic and
+ // vector shift amounts below its width, preserving representable results.
+ rhs = min(rhs, MakeConst(rhs.ty(), after_dtype.bits() - 1, op->loc),
op->loc);
+ }
+ if (lhs.same_as(op->a) && rhs.same_as(op->b) && lhs.ty() == rhs.ty()) {
+ return ffi::Unchanged();
+ }
+ return left_shift(lhs, rhs, op->loc);
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::RShiftNode* op,
+ InplaceMode inplace_mode) {
+ PrimType before_dtype = op->a.ty();
+ // Preserve the original operands while computing the narrowed shift.
+ PrimExpr lhs = Mutate(op->a, InplaceMode::kDisallow).ValueOrUnchanged(op->a);
+ PrimExpr rhs = Mutate(op->b, InplaceMode::kDisallow).ValueOrUnchanged(op->b);
+ PrimType after_dtype = lhs.ty();
+ if (ShouldClampShiftAmounts() && before_dtype.code() ==
DLDataTypeCode::kDLInt &&
+ after_dtype.code() == DLDataTypeCode::kDLInt && before_dtype.bits() >
after_dtype.bits()) {
+ // Values fit in the narrowed dtype. Clamp lane-wise to keep dynamic and
+ // vector shift amounts below its width, preserving representable results.
+ rhs = min(rhs, MakeConst(rhs.ty(), after_dtype.bits() - 1, op->loc),
op->loc);
+ }
+ if (lhs.same_as(op->a) && rhs.same_as(op->b) && lhs.ty() == rhs.ty()) {
+ return ffi::Unchanged();
+ }
+ return right_shift(lhs, rhs, op->loc);
+}
+
+UnchangedOr<PrimExpr> DataTypeLegalizer::Mutate_(const prim::BitwiseNotNode*
op,
+ InplaceMode inplace_mode) {
+ auto a = Mutate(op->a, inplace_mode);
+ if (a.UnchangedOrSameAs(op->a)) return ffi::Unchanged();
+ return prim::BitwiseNot(std::move(a).ValueOrUnchanged(op->a), op->loc);
+}
+
+UnchangedOr<Expr> DataTypeLegalizer::Mutate_(const CallNode* op, InplaceMode
inplace_mode) {
+ Call before = ffi::GetRef<Call>(op);
+ // Keep the original argument dtype available for clz correction below.
+ Expr e =
+ StmtExprMutator::Mutate_(op,
InplaceMode::kDisallow).ValueOrUnchanged(ffi::GetRef<Expr>(op));
+ op = e.as<CallNode>();
+ TVM_FFI_ICHECK(op != nullptr) << "Expected type to be CallNode"
+ << ", but get " << e->GetTypeKey();
+ if (!op->ty.as<PrimTypeNode>()) {
+ return e;
+ }
+ PrimExpr prim_e = e.as_or_throw<PrimExpr>();
+ static const Op pow_op = Op::Get("prim.pow");
+ static const Op& clz_op = prim::clz_op();
+ if (op->op.same_as(pow_op)) {
+ return pow(op->args[0].as_or_throw<PrimExpr>(),
op->args[1].as_or_throw<PrimExpr>());
+ } else if (op->op.same_as(prim::if_then_else_op())) {
+ return Call(op->ty.as_or_throw<PrimType>(), op->op,
+ {op->args[0].as_or_throw<PrimExpr>(),
op->args[1].as_or_throw<PrimExpr>(),
+ op->args[2].as_or_throw<PrimExpr>()},
+ op->attrs, {}, op->loc)
+ .as_or_throw<PrimExpr>();
+ } else if (op->op.same_as(clz_op)) {
+ PrimType before_dtype = before->args[0].as_or_throw<PrimExpr>().ty();
+ PrimType after_dtype = op->args[0].as_or_throw<PrimExpr>().ty();
+ TVM_FFI_ICHECK((before_dtype.code() == DLDataTypeCode::kDLInt ||
+ before_dtype.code() == DLDataTypeCode::kDLUInt) &&
+ (before_dtype.bits() == 32 || before_dtype.bits() == 64))
+ << "clz only supports 32 or 64 bit integer types, but get type before
legalizing: "
+ << before_dtype;
+ TVM_FFI_ICHECK((after_dtype.code() == DLDataTypeCode::kDLInt ||
+ after_dtype.code() == DLDataTypeCode::kDLUInt) &&
+ (after_dtype.bits() == 32 || after_dtype.bits() == 64))
+ << "clz only supports 32 or 64 bit integer types, but get type after
legalizing: "
+ << after_dtype;
+ return prim_e - after_dtype.bits() + before_dtype.bits();
+ }
+ return prim_e;
+}
+
+UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const RegionStmtNode* op,
+ InplaceMode inplace_mode) {
+ bool was_enabled = is_enabled_;
+ if (op->op.same_as(tirx::launch_thread_op())) is_enabled_ = true;
+ auto result = DataTypeLegalizer::Mutate_(op, inplace_mode);
+ is_enabled_ = was_enabled;
+ return result;
+}
+
+UnchangedOr<ffi::Any> IndexDataTypeRewriter::Mutate(ffi::AnyView value,
InplaceMode inplace_mode) {
+ bool is_enabled = is_enabled_;
+ if (value.as<TensorTypeNode>()) is_enabled_ = true;
+ auto result = DataTypeLegalizer::Mutate(value, inplace_mode);
+ is_enabled_ = is_enabled;
+ return result;
+}
+
+UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const TensorStoreNode* op,
+ InplaceMode inplace_mode) {
+ TensorStore store = ffi::GetRef<TensorStore>(op);
+
+ TensorVar new_tensor = Mutate(op->dest.as_or_throw<TensorVar>(),
inplace_mode)
+ .as_or_throw<UnchangedOr<TensorVar>>()
+
.ValueOrUnchanged(op->dest.as_or_throw<TensorVar>());
+ auto value_result = this->Mutate(op->value, inplace_mode);
+ bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
+ auto value = std::move(value_result).ValueOrUnchanged(op->value);
+ PrimType value_dtype = value.ty();
+ if (new_tensor->dtype != value_dtype && value_dtype.IsScalar()) {
+ value = prim::cast(new_tensor->dtype, value);
+ value_unchanged = false;
+ }
+ auto indices = VisitIndices(op->indices, inplace_mode);
+
+ if (!new_tensor.same_as(op->dest.as_or_throw<TensorVar>()) ||
!value_unchanged ||
+ !indices.same_as(op->indices)) {
+ auto writer = store.CopyOnWrite();
+ writer->dest = new_tensor;
+ writer->value = value;
+ writer->indices = indices;
+ }
+
+ return store;
+}
+
+UnchangedOr<PrimExpr> IndexDataTypeRewriter::Mutate_(const TensorLoadNode* op,
+ InplaceMode inplace_mode)
{
+ TensorLoad load = ffi::GetRef<TensorLoad>(op);
+
+ TensorVar new_tensor =
+ Mutate(op->source,
inplace_mode).ValueOrUnchanged(op->source).as_or_throw<TensorVar>();
+ auto indices = VisitIndices(op->indices, inplace_mode);
+
+ if (!new_tensor.same_as(op->source.as_or_throw<tvm::tirx::TensorVar>()) ||
+ !indices.same_as(op->indices)) {
+ return MakeTensorLoad(new_tensor, indices, op->loc);
+ }
+
+ return load;
+}
+
+ffi::Array<PrimExpr> IndexDataTypeRewriter::VisitIndices(const
ffi::Array<PrimExpr>& indices,
+ InplaceMode
inplace_mode) {
+ bool is_enabled = is_enabled_;
+ is_enabled_ = true;
+ auto result = Mutate(indices, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>()
+ .ValueOrUnchanged(indices);
+ is_enabled_ = is_enabled;
+ return result;
+}
+
+UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const IfNode* op, InplaceMode
inplace_mode) {
+ bool is_condition = is_condition_;
+ is_condition_ = true;
+ auto cond_result = Mutate(op->condition, inplace_mode);
+ bool cond_unchanged = cond_result.UnchangedOrSameAs(op->condition);
+ PrimExpr cond = std::move(cond_result).ValueOrUnchanged(op->condition);
+ is_condition_ = is_condition;
+ auto then_case_result = Mutate(op->then_case, inplace_mode);
+ bool then_case_unchanged = then_case_result.UnchangedOrSameAs(op->then_case);
+ SeqStmt then_case =
std::move(then_case_result).ValueOrUnchanged(op->then_case);
+ ffi::Optional<SeqStmt> else_case =
+ op->else_case.has_value()
+ ? ffi::Optional<SeqStmt>{Mutate(op->else_case.value(), inplace_mode)
+
.ValueOrUnchanged(op->else_case.value())}
+ : std::nullopt;
+ if (!cond_unchanged || !then_case_unchanged ||
!else_case.same_as(op->else_case)) {
+ If new_stmt = ffi::GetRef<If>(op);
+ auto* n = new_stmt.CopyOnWrite();
+ n->condition = std::move(cond);
+ n->then_case = std::move(then_case);
+ n->else_case = std::move(else_case);
+ return new_stmt;
+ }
+ return ffi::Unchanged();
+}
+
+UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const ForNode* op,
InplaceMode inplace_mode) {
+ bool is_enabled = is_enabled_;
+ is_enabled_ = true;
+ PrimVar new_loop_var =
+ Mutate(op->loop_var,
inplace_mode).ValueOrUnchanged(op->loop_var).as_or_throw<PrimVar>();
+ auto min_result = Mutate(op->min, inplace_mode);
+ bool min_unchanged = min_result.UnchangedOrSameAs(op->min);
+ PrimExpr min = std::move(min_result).ValueOrUnchanged(op->min);
+ auto extent_result = Mutate(op->extent, inplace_mode);
+ bool extent_unchanged = extent_result.UnchangedOrSameAs(op->extent);
+ PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent);
+ is_enabled_ = is_enabled;
+ auto new_body_result = Mutate(op->body, inplace_mode);
+ bool new_body_unchanged = new_body_result.UnchangedOrSameAs(op->body);
+ Stmt new_body = std::move(new_body_result).ValueOrUnchanged(op->body);
+
+ if (!new_loop_var.same_as(op->loop_var) || !min_unchanged ||
!extent_unchanged ||
+ !new_body_unchanged) {
+ For new_for = ffi::GetRef<For>(op);
+ auto* n = new_for.CopyOnWrite();
+ n->loop_var = new_loop_var;
+ n->min = prim::cast(new_loop_var.ty(), min);
+ n->extent = prim::cast(new_loop_var.ty(), extent);
+ n->body = new_body;
+ return new_for;
+
+ } else {
+ return ffi::Unchanged();
+ }
+}
+
+UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const BindNode* op,
InplaceMode inplace_mode) {
+ if (!op->value.as<PrimExpr>()) {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ auto mapped = VarRemapGet(op->var);
+ if (mapped == nullptr || mapped.type_index() ==
ffi::TypeIndex::kTVMFFIUnchanged) {
+ return DataTypeLegalizer::Mutate_(op, inplace_mode);
+ }
+ Var var = mapped.as_or_throw<Var>();
+ bool is_enabled = is_enabled_;
+ is_enabled_ = true;
+ PrimExpr value =
+ Mutate(op->value,
inplace_mode).ValueOrUnchanged(op->value).as_or_throw<PrimExpr>();
+ is_enabled_ = is_enabled;
+ // The collected index requirement need not apply to every variable in the
RHS.
+ return Bind(var, prim::cast(var->ty.as_or_throw<PrimType>(), value),
op->loc);
+}
+
+#define TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC)
\
+ UnchangedOr<PrimExpr> IndexDataTypeRewriter::Mutate_(const OP* op,
InplaceMode inplace_mode) { \
+ bool is_enabled = is_enabled_;
\
+ is_enabled_ = is_condition_ &&
op->a.ty().MatchesCode(DLDataTypeCode::kDLInt) && \
+ op->b.ty().MatchesCode(DLDataTypeCode::kDLInt);
\
+ auto result = Parent::Mutate_(op, inplace_mode);
\
+ is_enabled_ = is_enabled;
\
+ return result;
\
+ }
+
+TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::EQNode, operator==);
+TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::NENode, operator!=);
+TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LENode, operator<=);
+TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LTNode, operator<); //
NOLINT(*)
+TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GTNode, operator>); //
NOLINT(*)
+TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GENode, operator>=);
+
+UnchangedOr<Expr> IndexDataTypeRewriter::Mutate_(const CallNode* op,
InplaceMode inplace_mode) {
+ // handle if_then_else condition
+ if (op->op.same_as(prim::if_then_else_op())) {
+ bool is_condition = is_condition_;
+ is_condition_ = true;
+ PrimExpr cond =
Mutate(op->args[0]).ValueOrUnchanged(op->args[0]).as_or_throw<PrimExpr>();
+ is_condition_ = is_condition;
+ PrimExpr true_value =
Mutate(op->args[1]).ValueOrUnchanged(op->args[1]).as_or_throw<PrimExpr>();
+ PrimExpr false_value =
+
Mutate(op->args[2]).ValueOrUnchanged(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 = prim::cast(dtype, true_value);
+ if (false_dtype != dtype) false_value = prim::cast(dtype, false_value);
+ return Call(dtype, op->op, {cond, true_value, false_value}, op->attrs, {},
op->loc)
+ .as_or_throw<PrimExpr>();
+ }
+ return Parent::Mutate_(op, inplace_mode);
+}
+
+UnchangedOr<PrimExpr> IndexDataTypeRewriter::Mutate_(const prim::SelectNode*
op,
+ InplaceMode inplace_mode)
{
+ bool is_condition = true;
+ std::swap(is_condition_, is_condition);
+ auto condition_result = this->Mutate(op->condition, inplace_mode);
+ bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition);
+ PrimExpr condition =
std::move(condition_result).ValueOrUnchanged(op->condition);
+ std::swap(is_condition_, is_condition);
+ auto true_value_result = this->Mutate(op->true_value, inplace_mode);
+ bool true_value_unchanged =
true_value_result.UnchangedOrSameAs(op->true_value);
+ PrimExpr true_value =
std::move(true_value_result).ValueOrUnchanged(op->true_value);
+ auto false_value_result = this->Mutate(op->false_value, inplace_mode);
+ bool false_value_unchanged =
false_value_result.UnchangedOrSameAs(op->false_value);
+ PrimExpr false_value =
std::move(false_value_result).ValueOrUnchanged(op->false_value);
+
+ if (condition_unchanged && true_value_unchanged && false_value_unchanged &&
+ true_value.ty() == false_value.ty()) {
+ return ffi::Unchanged();
+ } else {
+ PrimType true_dtype = true_value.ty();
+ PrimType false_dtype = false_value.ty();
+ int bits = std::max(true_dtype.bits(), false_dtype.bits());
+ PrimType dtype = true_dtype.WithBits(bits);
+ if (true_dtype->dtype != dtype->dtype) true_value = prim::cast(dtype,
true_value);
+ if (false_dtype->dtype != dtype->dtype) false_value = prim::cast(dtype,
false_value);
+ return prim::Select(condition, true_value, false_value);
+ }
+}
+
+#undef TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH
+
Function IndexDataTypeNormalizer::Rewrite(Function func) {
// Keep this short setup local so its collector uses S-TIR block semantics
// without adding a dialect-specific collector hook to the TIRX normalizer.
@@ -238,5 +744,31 @@ TensorRegion
IndexDataTypeNormalizer::VisitTensorRegion(const TensorRegion& tens
}
}
+bool IndexDataTypeNormalizer::CanRewriteDType(PrimType dtype) const {
+ return dtype.code() == DLDataTypeCode::kDLInt && dtype.bits() >= 32;
+}
+
+UnchangedOr<PrimExpr> IndexDataTypeNormalizer::Mutate_(const IntImmNode* op,
+ InplaceMode
inplace_mode) {
+ if (is_enabled_ && CanRewriteDType(op->ty.as_or_throw<PrimType>())) {
+ TVM_FFI_ICHECK_LE(op->value,
max_value(target_data_type_).as_or_throw<IntImm>()->value);
+ return prim::cast(target_data_type_, ffi::GetRef<IntImm>(op));
+ }
+ return ffi::Unchanged();
+}
+
+UnchangedOr<PrimExpr> IndexDataTypeNormalizer::Mutate_(const prim::CastNode*
op,
+ InplaceMode
inplace_mode) {
+ // Unwrap the cast only when the dtype of this cast is integer dtype.
+ // When the dtype of this cast is not integer dtype, it means that this cast
+ // has some other purpose, and we should not unwrap the cast.
+ PrimType dtype = op->ty.as_or_throw<PrimType>();
+ if (is_enabled_ && CanRewriteDType(dtype)) {
+ PrimExpr value = this->Mutate(op->value,
inplace_mode).ValueOrUnchanged(op->value);
+ return value.ty() == target_data_type_ ? value :
prim::Cast(target_data_type_, value);
+ }
+ return IndexDataTypeRewriter::Mutate_(op, inplace_mode);
+}
+
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/ir/data_type_rewriter.h
b/src/s_tir/ir/data_type_rewriter.h
index 755e298924..2fc3cd7d6f 100644
--- a/src/s_tir/ir/data_type_rewriter.h
+++ b/src/s_tir/ir/data_type_rewriter.h
@@ -21,27 +21,135 @@
#define TVM_S_TIR_IR_DATA_TYPE_REWRITER_H_
#include <tvm/s_tir/stmt.h>
+#include <tvm/tirx/stmt_functor.h>
-#include "../../tirx/ir/data_type_rewriter.h"
+#include <utility>
namespace tvm {
namespace s_tir {
+/*!
+ * \brief Legalize the data types of expressions to make sure they are
consistent with other
+ * parts of the program.
+ *
+ * It enforces the following rules:
+ * - The data type of the index variable in a loop must be consistent with the
data type of the loop
+ * bounds.
+ * - The data type of the binary and ternary expressions must be consistent
with the data types of
+ * each of their operands.
+ *
+ * Usually we enforce the consistency of data types when constructing the IR
nodes. However, such
+ * inconsistency may happen as a result of IR mutation in some passes. This
class can be used as
+ * base class of such passes to ensure the consistency of data types.
+ */
+class DataTypeLegalizer : public tirx::StmtExprMutator {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(DataTypeLegalizer,
tirx::StmtExprMutator)
+ using tirx::StmtExprMutator::Mutate;
+ using tirx::StmtExprMutator::Mutate_;
+
+ protected:
+ explicit DataTypeLegalizer(const VTable* vtable) :
tirx::StmtExprMutator(vtable) {}
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::SelectNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::RampNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::BroadcastNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::ShuffleNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LShiftNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::RShiftNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::BitwiseAndNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::BitwiseOrNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::BitwiseXorNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::BitwiseNotNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::AddNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::SubNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::MulNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::DivNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::ModNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::FloorDivNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::FloorModNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::MinNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::MaxNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::EQNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::NENode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LTNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LENode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::GTNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::GENode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LetNode* op, InplaceMode
inplace_mode) override;
+
+ /*! \brief Whether to clamp shift amounts after narrowing signed integers. */
+ virtual bool ShouldClampShiftAmounts() const { return false; }
+};
+
+/*!
+ * \brief Data type rewriter for tensor indices.
+ *
+ * Detect the components of tensor indices that should be considered for data
type rewriting.
+ * This class doesn't perform actual rewriting of data types. During recursive
visiting, the
+ * internal flags `is_enabled_` and `is_conditional_` are used to indicate
whether the current
+ * expression is a tensor index or a conditional expression, which can be used
in the sub-classes to
+ * implement different rewriting rules.
+ */
+class IndexDataTypeRewriter : public DataTypeLegalizer {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IndexDataTypeRewriter,
DataTypeLegalizer)
+ using DataTypeLegalizer::Mutate;
+ using DataTypeLegalizer::Mutate_;
+
+ protected:
+ explicit IndexDataTypeRewriter(const VTable* vtable) :
DataTypeLegalizer(vtable) {}
+ using Parent = DataTypeLegalizer;
+ UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode)
override;
+ UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) override;
+ ffi::Array<PrimExpr> VisitIndices(const ffi::Array<PrimExpr>& indices,
InplaceMode inplace_mode);
+ UnchangedOr<Stmt> Mutate_(const IfNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode)
override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::EQNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::NENode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LTNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::LENode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::GTNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::GENode* op, InplaceMode
inplace_mode) override;
+
+ UnchangedOr<PrimExpr> Mutate_(const prim::SelectNode* op, InplaceMode
inplace_mode) override;
+
+ // indicator of index expr to rewrite
+ bool is_enabled_{false};
+ // indicator of condition
+ bool is_condition_{false};
+};
+
// Explicit normalization for schedulable TE functions and tensor intrinsic
bodies.
// Ordinary TIRX dtype passes run after block lowering.
-class IndexDataTypeNormalizer : public tirx::IndexDataTypeNormalizer {
+class IndexDataTypeNormalizer : public IndexDataTypeRewriter {
public:
- using Parent = tirx::IndexDataTypeNormalizer;
+ using Parent = IndexDataTypeRewriter;
using Parent::Mutate;
using Parent::Mutate_;
explicit IndexDataTypeNormalizer(PrimType target_data_type)
- : Parent(std::move(target_data_type), GlobalVTable()) {}
+ : Parent(GlobalVTable()), target_data_type_(std::move(target_data_type))
{}
tirx::Function Rewrite(tirx::Function func);
virtual UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode
inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode
inplace_mode);
protected:
+ UnchangedOr<PrimExpr> Mutate_(const IntImmNode* op, InplaceMode
inplace_mode) override;
+ UnchangedOr<PrimExpr> Mutate_(const prim::CastNode* op, InplaceMode
inplace_mode) override;
+
+ /*! \brief Specifies which data type we can rewrite */
+ virtual bool CanRewriteDType(PrimType dtype) const;
+
+ PrimType target_data_type_ = PrimType::Int(64);
+
static void InitVTable(VTable* table) {
Parent::InitVTable(table);
SetDispatch<IndexDataTypeNormalizer, SBlockNode>(table);
