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 91fb997563 [REFACTOR][IR] Use CamelCase Var copy helpers (#20008)
91fb997563 is described below
commit 91fb997563e61ba21550357e11862468b6ec3336
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Jul 16 07:43:00 2026 +0800
[REFACTOR][IR] Use CamelCase Var copy helpers (#20008)
---
include/tvm/ir/expr.h | 4 +--
include/tvm/tirx/var.h | 4 +--
src/ir/expr.cc | 6 ++--
.../schedule/primitive/loop_transformation.cc | 6 ++--
src/s_tir/transform/lift_thread_binding.cc | 2 +-
src/tirx/ir/data_type_rewriter.cc | 8 ++---
src/tirx/ir/expr.cc | 4 +--
src/tirx/transform/unsupported_dtype_legalize.cc | 8 ++---
tests/cpp/expr_test.cc | 34 ++++++++++++++++++++++
9 files changed, 55 insertions(+), 21 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 30fe958b95..0af9a8a103 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -289,13 +289,13 @@ class Var : public Expr {
Span span = Span());
/*! \brief Return a fresh ordinary Var with the same type and a new name. */
- TVM_DLL Var copy_with_name(const ffi::String& name) const;
+ TVM_DLL Var CopyWithName(const ffi::String& name) const;
/*! \brief Return a fresh ordinary Var with a suffix appended to its name. */
TVM_DLL Var CopyWithSuffix(const ffi::String& suffix) const;
/*! \brief Return a fresh ordinary Var with a new primitive type. */
- TVM_DLL Var copy_with_dtype(PrimType dtype) const;
+ TVM_DLL Var CopyWithDType(PrimType dtype) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode);
};
diff --git a/include/tvm/tirx/var.h b/include/tvm/tirx/var.h
index e295358b83..467b2f7bee 100644
--- a/include/tvm/tirx/var.h
+++ b/include/tvm/tirx/var.h
@@ -61,8 +61,8 @@ class PrimVar : public PrimExpr {
PrimVar CopyWithSuffix(const ffi::String& suffix) const {
return
this->as_or_throw<Var>().CopyWithSuffix(suffix).as_or_throw<PrimVar>();
}
- PrimVar copy_with_dtype(PrimType dtype) const {
- return
this->as_or_throw<Var>().copy_with_dtype(dtype).as_or_throw<PrimVar>();
+ PrimVar CopyWithDType(PrimType dtype) const {
+ return
this->as_or_throw<Var>().CopyWithDType(dtype).as_or_throw<PrimVar>();
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimVar, PrimExpr, VarNode);
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 21b345587c..2330e643cd 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -261,7 +261,7 @@ Var::Var(ffi::String name_hint, ffi::Optional<Type>
ty_annotation, Span span) {
data_ = std::move(n);
}
-Var Var::copy_with_name(const ffi::String& name) const {
+Var Var::CopyWithName(const ffi::String& name) const {
TVM_FFI_CHECK_EQ(type_index(), VarNode::RuntimeTypeIndex(), TypeError)
<< "Cannot copy a Var runtime subtype as an ordinary Var";
ffi::ObjectPtr<VarNode> copy = ffi::make_object<VarNode>(*get());
@@ -270,10 +270,10 @@ Var Var::copy_with_name(const ffi::String& name) const {
}
Var Var::CopyWithSuffix(const ffi::String& suffix) const {
- return copy_with_name(get()->name_hint + suffix);
+ return CopyWithName(get()->name_hint + suffix);
}
-Var Var::copy_with_dtype(PrimType dtype) const {
+Var Var::CopyWithDType(PrimType dtype) const {
TVM_FFI_CHECK_EQ(type_index(), VarNode::RuntimeTypeIndex(), TypeError)
<< "Cannot copy a Var runtime subtype as an ordinary Var";
ffi::ObjectPtr<VarNode> copy = ffi::make_object<VarNode>(*get());
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc
b/src/s_tir/schedule/primitive/loop_transformation.cc
index 7a007058e9..b3573843b4 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -424,7 +424,7 @@ ffi::Array<StmtSRef> Split(ScheduleState self, const
StmtSRef& loop_sref,
new_loop_vars.reserve(n);
for (int i = 0; i < n; i++) {
const PrimExpr& factor = factors[i];
- Var var = loop->loop_var.CopyWithSuffix("_" +
std::to_string(i)).copy_with_dtype(dtype);
+ Var var = loop->loop_var.CopyWithSuffix("_" +
std::to_string(i)).CopyWithDType(dtype);
substitute_value = substitute_value * factor + var.as_or_throw<PrimExpr>();
analyzer->Bind(var, Range::FromMinExtent(IntImm(dtype, 0),
tvm::cast(dtype, factor)));
new_loop_vars.emplace_back(std::move(var));
@@ -675,7 +675,7 @@ ffi::Array<StmtSRef> LoopPartition(ScheduleState self,
const StmtSRef& loop_sref
// Iterate over each pair of factors and create partition
for (int i = 0; i < n; i++) {
extent_value = analyzer->Simplify(factors[i]);
- Var new_loop_var =
loop->loop_var.CopyWithSuffix(std::to_string(i)).copy_with_dtype(dtype);
+ Var new_loop_var =
loop->loop_var.CopyWithSuffix(std::to_string(i)).CopyWithDType(dtype);
Stmt loop_body = tirx::Substitute(loop->body, {{loop->loop_var,
new_loop_var}});
// Create new block with new reference to each variable/stmt/expr in the
existing block
@@ -929,7 +929,7 @@ StmtSRef Fuse(ScheduleState self, const
ffi::Array<StmtSRef>& loop_srefs,
}
suffix += "_fused";
- Var fused_var =
loops[0]->loop_var.CopyWithSuffix(suffix).copy_with_dtype(PrimType::Int(bits));
+ Var fused_var =
loops[0]->loop_var.CopyWithSuffix(suffix).CopyWithDType(PrimType::Int(bits));
ffi::Array<PrimExpr> substitute_value;
substitute_value.resize(loops.size());
PrimExpr lower = 1;
diff --git a/src/s_tir/transform/lift_thread_binding.cc
b/src/s_tir/transform/lift_thread_binding.cc
index 6387b8595d..5447ecf5ad 100644
--- a/src/s_tir/transform/lift_thread_binding.cc
+++ b/src/s_tir/transform/lift_thread_binding.cc
@@ -62,7 +62,7 @@ FindLoopLCA(const Stmt& root) {
iter_var = IterVar(Range::FromMinExtent(loop->min, loop->extent), //
loop->loop_var
.as_or_throw<Var>() //
- .copy_with_name(thread_tag) //
+ .CopyWithName(thread_tag) //
.as_or_throw<PrimVar>(), //
loop->thread_binding.value()->iter_type, //
thread_tag);
diff --git a/src/tirx/ir/data_type_rewriter.cc
b/src/tirx/ir/data_type_rewriter.cc
index 9fb408039d..d20346d905 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -132,7 +132,7 @@ Expr DataTypeLegalizer::VisitExpr_(const LetNode* op) {
Var var = op->var;
if (value.ty() != op->var->ty.as_or_throw<PrimType>()) {
- var = op->var.copy_with_dtype(value.ty());
+ var = op->var.CopyWithDType(value.ty());
var_remap_[op->var.get()] = var;
}
@@ -151,7 +151,7 @@ Stmt DataTypeLegalizer::VisitStmt_(const BindNode* op) {
if (auto prim_value = value.as<PrimExpr>()) {
if (prim_value.value().ty() != op->var->ty.as_or_throw<PrimType>()) {
- var = op->var.copy_with_dtype(prim_value.value().ty());
+ var = op->var.CopyWithDType(prim_value.value().ty());
var_remap_[op->var.get()] = var;
}
}
@@ -539,7 +539,7 @@ Stmt IndexDataTypeRewriter::VisitStmt_(const ForNode* op) {
if (op->thread_binding.has_value()) {
auto old_thread_binding = op->thread_binding.value();
auto* ptr = old_thread_binding.CopyOnWrite();
- ptr->var = old_thread_binding->var.copy_with_dtype(new_loop_var.ty());
+ ptr->var = old_thread_binding->var.CopyWithDType(new_loop_var.ty());
n->thread_binding =
ffi::Optional<IterVar>(std::move(old_thread_binding));
}
n->body = new_body;
@@ -674,7 +674,7 @@ Expr IndexDataTypeNormalizer::VisitExpr_(const VarNode* op)
{
PrimType dtype = dtype_opt.value();
if (is_enabled_ && CanRewriteDType(dtype) && dtype->dtype !=
target_data_type_->dtype &&
!var_remap_.count(op)) {
- var_remap_[op] = ffi::GetRef<Var>(op).copy_with_dtype(target_data_type_);
+ var_remap_[op] = ffi::GetRef<Var>(op).CopyWithDType(target_data_type_);
}
return DataTypeLegalizer::VisitExpr_(op);
}
diff --git a/src/tirx/ir/expr.cc b/src/tirx/ir/expr.cc
index 19516d5ff0..c7a044a5aa 100644
--- a/src/tirx/ir/expr.cc
+++ b/src/tirx/ir/expr.cc
@@ -623,8 +623,8 @@ CommReducer::CommReducer(ffi::Array<PrimVar> lhs,
ffi::Array<PrimVar> rhs,
var_map.reserve(n_group * 2);
for (int i = 0; i < static_cast<int>(n_group); ++i) {
PrimType dtype = identity_element[i].ty();
- PrimVar l = lhs[i].copy_with_dtype(dtype);
- PrimVar r = rhs[i].copy_with_dtype(dtype);
+ PrimVar l = lhs[i].CopyWithDType(dtype);
+ PrimVar r = rhs[i].CopyWithDType(dtype);
var_map[lhs[i].get()] = l;
var_map[rhs[i].get()] = r;
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc
b/src/tirx/transform/unsupported_dtype_legalize.cc
index a1e580cc50..2d2e8ed74a 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -309,7 +309,7 @@ class ComputeLegalizer : public StmtExprMutator {
PrimExpr value = PromoteToTarget(op->value);
Var var = op->var;
if (value.ty() != op->value.ty()) {
- var = op->var.copy_with_dtype(op->value.ty());
+ var = op->var.CopyWithDType(op->value.ty());
var_remap_[op->var] = var;
}
@@ -343,7 +343,7 @@ class ComputeLegalizer : public StmtExprMutator {
PrimExpr value = PromoteToTarget(prim_value.value());
Var var = op->var;
if (value.ty() != prim_value.value().ty()) {
- var = op->var.copy_with_dtype(prim_value.value().ty());
+ var = op->var.CopyWithDType(prim_value.value().ty());
var_remap_[op->var] = var;
}
@@ -407,11 +407,11 @@ class ComputeLegalizer : public StmtExprMutator {
for (size_t i = 0; i < legalized_identity_elements.size(); i++) {
Var lhs_var = reducer->lhs[i];
if (lhs_var->ty.as_or_throw<PrimType>() !=
legalized_identity_elements[i].ty()) {
- var_remap_[lhs_var] =
lhs_var.copy_with_dtype(legalized_identity_elements[i].ty());
+ var_remap_[lhs_var] =
lhs_var.CopyWithDType(legalized_identity_elements[i].ty());
}
Var rhs_var = reducer->rhs[i];
if (rhs_var->ty.as_or_throw<PrimType>() !=
legalized_identity_elements[i].ty()) {
- var_remap_[rhs_var] =
rhs_var.copy_with_dtype(legalized_identity_elements[i].ty());
+ var_remap_[rhs_var] =
rhs_var.CopyWithDType(legalized_identity_elements[i].ty());
}
}
diff --git a/tests/cpp/expr_test.cc b/tests/cpp/expr_test.cc
index a4fbb82878..6ad0f81191 100644
--- a/tests/cpp/expr_test.cc
+++ b/tests/cpp/expr_test.cc
@@ -20,9 +20,12 @@
#include <gtest/gtest.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ir/source_map.h>
#include <tvm/runtime/logging.h>
#include <tvm/te/operation.h>
+#include <type_traits>
+
TEST(Expr, Basic) {
using namespace tvm;
using namespace tvm::tirx;
@@ -46,6 +49,37 @@ TEST(Expr, VarTypeAnnotation) {
TVM_FFI_ICHECK(checker(x->ty, y->ty));
}
+TEST(Expr, VarCopyHelpers) {
+ using namespace tvm;
+ using namespace tvm::tirx;
+
+ Span span(SourceName::Get("test.cc"), 1, 1, 1, 10);
+ Type pointer_type = PointerType(PrimType::Float(32), "global");
+ Var var("x", pointer_type, span);
+
+ Var renamed = var.CopyWithName("y");
+ EXPECT_FALSE(renamed.same_as(var));
+ EXPECT_EQ(renamed->name_hint, "y");
+ EXPECT_TRUE(renamed->ty.same_as(pointer_type));
+ EXPECT_TRUE(renamed->span.same_as(span));
+
+ PrimType dtype = PrimType::Int(64);
+ Var retyped = var.CopyWithDType(dtype);
+ EXPECT_FALSE(retyped.same_as(var));
+ EXPECT_EQ(retyped->name_hint, "x");
+ EXPECT_TRUE(retyped->ty.same_as(dtype));
+ EXPECT_TRUE(retyped->span.same_as(span));
+
+ PrimVar prim_var("i", PrimType::Int(32), span);
+
static_assert(std::is_same_v<decltype(prim_var.CopyWithDType(PrimType::Float(32))),
PrimVar>);
+ PrimType prim_dtype = PrimType::Float(32);
+ PrimVar retyped_prim_var = prim_var.CopyWithDType(prim_dtype);
+ EXPECT_FALSE(retyped_prim_var.same_as(prim_var));
+ EXPECT_EQ(retyped_prim_var->name_hint, "i");
+ EXPECT_TRUE(retyped_prim_var.ty().same_as(prim_dtype));
+ EXPECT_TRUE(retyped_prim_var->span.same_as(span));
+}
+
TEST(Expr, PrimTypeBoolLanes) {
using namespace tvm;
PrimType boolx4 = PrimType::Bool(4);
