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 ded6ad8dd2 [Cleanup] Remove macOS Clang warnings (#19954)
ded6ad8dd2 is described below
commit ded6ad8dd212869c881efb5590f8a33fc972728e
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Jul 6 09:22:43 2026 +0800
[Cleanup] Remove macOS Clang warnings (#19954)
## Summary
- remove redundant parentheses from four VM dtype comparisons
- clean up pessimizing moves, missing overrides, unused state, and
hidden virtual overloads reported by macOS Clang
- preserve `PrimType` across source-codegen type-bearing paths and
remove transitional raw dtype adapters
- retain raw `DLDataType` only for explicit runtime launch metadata and
runtime-helper boundaries
- leave the deferred LLVM and CMake compatibility paths unchanged
`PrimType::operator==` already compares the represented dtype by code,
bits, and lanes, so no backend-specific non-identity equality regression
is needed.
## Validation
- compiled all original warning-producing translation units with Clang
22 and the relevant warning families promoted to errors
- built `tvm_runtime` and `tvm_compiler` with LLVM 15 after the codegen
migration
- passed 14 focused C-host, bool, OpenCL, Metal, device, and Vulkan
codegen tests, with 70 hardware-dependent skips
- generated OpenCL C, Metal, WGSL, and CUDA source from the same typed
kernel
- compiled the changed Hexagon and Vulkan translation units
independently
- passed pre-commit, Git whitespace/log checks, and a static audit
leaving raw codegen `DLDataType` declarations only at the WebGPU runtime
ABI boundary
---
src/backend/cuda/codegen/codegen_cuda.cc | 17 +++---
src/backend/cuda/codegen/codegen_cuda.h | 3 -
.../hexagon/codegen/llvm/codegen_hexagon.cc | 3 +-
src/backend/metal/codegen/codegen_metal.cc | 29 ++++------
src/backend/opencl/codegen/codegen_opencl.cc | 61 ++++++++++----------
src/backend/opencl/codegen/codegen_opencl.h | 11 ++--
src/backend/trn/codegen/codegen_trn.h | 1 +
src/backend/vulkan/codegen/codegen_spirv.cc | 66 ++++++++++------------
src/backend/webgpu/codegen/codegen_webgpu.cc | 63 ++++++++++-----------
src/relax/ir/expr.cc | 4 +-
src/relax/ir/py_expr_functor.cc | 42 +++++++-------
src/relax/op/tensor/create.cc | 6 +-
src/relax/op/tensor/manipulate.cc | 2 +-
src/relax/transform/attach_global_symbol.cc | 1 +
src/relax/transform/compute_prim_value.cc | 2 +
src/relax/transform/dead_code_elimination.cc | 1 +
src/relax/transform/fuse_tir.cc | 5 +-
src/runtime/vm/builtin.cc | 8 +--
src/target/source/codegen_c.cc | 6 +-
src/target/source/codegen_c.h | 3 -
src/target/source/codegen_c_host.cc | 25 ++++----
src/target/source/codegen_source_base.h | 1 -
src/tirx/transform/storage_rewrite.cc | 2 +
src/tirx/transform/tile_primitive_dispatch.cc | 2 +
24 files changed, 171 insertions(+), 193 deletions(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 8cc95bacae..5e65eaa034 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -418,7 +418,7 @@ void CodeGenCUDA::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*)
fail = true;
}
return;
- } else if (t->dtype == DLDataType{kDLBool, 8, 1}) {
+ } else if (t == PrimType::Bool()) {
os << "bool";
return;
} else if (t.MatchesCode(DLDataTypeCode::kDLBool) && lanes > 1) {
@@ -1641,11 +1641,9 @@ void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) {
if (scope.find("wmma.") == 0) {
if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") {
bool supported_wmma_input_dtype = dtype == PrimType::Float(16) || dtype
== PrimType::Int(8) ||
- dtype == PrimType::UInt(8) ||
- dtype == PrimType(DLDataType{kDLInt,
4, 1}) ||
- dtype == PrimType(DLDataType{kDLUInt,
4, 1}) ||
- dtype == PrimType(DLDataType{kDLInt,
1, 1}) ||
- dtype ==
PrimType(DLDataType{kDLBfloat, 16, 1});
+ dtype == PrimType::UInt(8) || dtype ==
PrimType::Int(4) ||
+ dtype == PrimType::UInt(4) || dtype ==
PrimType::Int(1) ||
+ dtype == PrimType::BFloat(16);
TVM_FFI_ICHECK(supported_wmma_input_dtype)
<< "Matrix_a and matrix_b only support half or char or unsigned char
"
<< "or uint4 or int4 or int1 type for now";
@@ -1689,9 +1687,8 @@ void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) {
if (scope.find("wmma.") == 0) {
constant_size = GetWmmaFragmentSize(scope, buffer, constant_size);
}
- bool is_packed_integer_dtype = dtype == PrimType(DLDataType{kDLInt, 4, 1})
||
- dtype == PrimType(DLDataType{kDLUInt, 4,
1}) ||
- dtype == PrimType(DLDataType{kDLInt, 1, 1});
+ bool is_packed_integer_dtype =
+ dtype == PrimType::Int(4) || dtype == PrimType::UInt(4) || dtype ==
PrimType::Int(1);
if (is_packed_integer_dtype && scope == "shared") {
constant_size = constant_size / (32 / dtype.bits());
}
@@ -1902,7 +1899,7 @@ void CodeGenCUDA::VisitExpr_(const SelectNode* op,
std::ostream& os) {
// The condition is stored as an ushort vector.
int lanes = op_ty.lanes();
- PrimType memory_ty(DLDataType{kDLUInt, 16, static_cast<uint16_t>(lanes)});
+ PrimType memory_ty = PrimType::UInt(16, lanes);
for (int i = 0; i < lanes; ++i) {
std::ostringstream item;
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index ce5f24e8e4..3435f0dc9c 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -118,9 +118,6 @@ class CodeGenCUDA final : public CodeGenC {
const std::string barrier_name_ = "barrier";
// The size of the barrier array in shared memory
std::unordered_map<int, int> barrier_count_;
- // The alignment of the barrier array in shared memory
- // Set to 16 to maintain minimum alignment requirements for async bulk copy
- const int barrier_alignment_bytes_ = 16;
// Functions to be added to the util functions during codegen
std::unordered_map<std::string, std::string> util_funcs_;
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index e7389f84c5..665067ea3a 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -202,8 +202,7 @@ llvm::Value* CodeGenHexagon::VisitExpr_(const
BufferLoadNode* op) {
if (!op->buffer.same_as(op->buffer->data)) {
// Check if we can generate a vector lookup.
if (!op->indices[0].as<RampNode>()) {
- if (auto* vlut = VectorLookupLoad(op->buffer,
PrimType(op->ty.as_or_throw<PrimType>()->dtype),
- op->indices)) {
+ if (auto* vlut = VectorLookupLoad(op->buffer,
op->ty.as_or_throw<PrimType>(), op->indices)) {
return vlut;
}
}
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 607b04ce21..af914f6821 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -166,14 +166,10 @@ void CodeGenMetal::AddFunction(const GlobalVar& gvar,
const PrimFunc& func) {
if (work_dim != 0) {
// use ushort by default for now
stream << " ";
- PrintType(DLDataType{kDLUInt, static_cast<uint8_t>(thread_index_bits_),
- static_cast<uint16_t>(work_dim)},
- stream);
+ PrintType(PrimType::UInt(thread_index_bits_, work_dim), stream);
stream << " blockIdx [[threadgroup_position_in_grid]],\n";
stream << " ";
- PrintType(DLDataType{kDLUInt, static_cast<uint8_t>(thread_index_bits_),
- static_cast<uint16_t>(work_dim)},
- stream);
+ PrintType(PrimType::UInt(thread_index_bits_, work_dim), stream);
stream << " threadIdx [[thread_position_in_threadgroup]]\n";
}
thread_work_dim_ = work_dim;
@@ -195,18 +191,16 @@ void CodeGenMetal::BindThreadIndex(const IterVar& iv) {
if (thread_work_dim_ <= 1) {
vname = vname.substr(0, iv->thread_tag.length() - 2);
}
- var_idmap_[iv->var.get()] = CastFromTo(
- vname, DLDataType{kDLUInt, static_cast<uint8_t>(thread_index_bits_), 1},
iv->var.ty()->dtype);
+ var_idmap_[iv->var.get()] = CastFromTo(vname,
PrimType::UInt(thread_index_bits_), iv->var.ty());
}
void CodeGenMetal::PrintType(const PrimType& t, std::ostream& os) { //
NOLINT(*)
- const DLDataType& raw_t = t->dtype;
int lanes = t.lanes();
if (t.IsVoid()) {
os << "void";
return;
}
- if (raw_t == DLDataType{kDLBool, 8, 1}) {
+ if (t == PrimType::Bool()) {
os << "bool";
return;
}
@@ -272,7 +266,7 @@ void CodeGenMetal::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*
os << "bfloat";
return;
}
- TVM_FFI_THROW(InternalError) << "Cannot convert type " <<
ffi::DLDataTypeToString(raw_t)
+ TVM_FFI_THROW(InternalError) << "Cannot convert type " <<
ffi::DLDataTypeToString(t->dtype)
<< " to Metal type";
}
@@ -329,14 +323,13 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
auto scope = GetPtrStorageScope(op->buffer->data);
alloc_storage_scope_[op->buffer->data.get()] = scope;
- DLDataType dtype = op->buffer->dtype->dtype;
+ const PrimType& dtype = op->buffer->dtype;
if (scope == "metal.simdgroup") {
- bool supported_simdgroup_dtype = dtype == DLDataType{kDLFloat, 16, 1} ||
- dtype == DLDataType{kDLFloat, 32, 1} ||
- dtype == DLDataType{kDLBfloat, 16, 1};
+ bool supported_simdgroup_dtype = dtype == PrimType::Float(16) || dtype ==
PrimType::Float(32) ||
+ dtype == PrimType::BFloat(16);
TVM_FFI_ICHECK(supported_simdgroup_dtype)
<< "Only float16, float32, and bfloat16 are supported, but got "
- << ffi::DLDataTypeToString(dtype);
+ << ffi::DLDataTypeToString(dtype->dtype);
TVM_FFI_ICHECK(constant_size % 64 == 0)
<< "Only 8x8 matrix is supported, but got " << constant_size << "
bytes\n";
@@ -365,7 +358,7 @@ void CodeGenMetal::VisitExpr_(const SelectNode* op,
std::ostream& os) { // NOLI
void CodeGenMetal::VisitExpr_(const BroadcastNode* op, std::ostream& os) { //
NOLINT(*)
std::string v = PrintExpr(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
os << "(";
for (int i = 0; i < lanes; ++i) {
if (i != 0) os << ", ";
@@ -432,7 +425,7 @@ void CodeGenMetal::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLINT
}
// generate as_type<TYPE>(ARG)
os << "(as_type<";
- this->PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ this->PrintType(op->ty.as_or_throw<PrimType>(), os);
os << ">(";
this->PrintExpr(op->args[0], os);
os << "))";
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc
b/src/backend/opencl/codegen/codegen_opencl.cc
index 93e925a6b4..d1d19a2fe4 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -43,6 +43,7 @@ class InferTextureAccess : public StmtExprVisitor {
static constexpr const uint8_t kWriteAccess = 2;
InferTextureAccess() {}
+ using StmtExprVisitor::VisitExpr_;
std::unordered_map<const VarNode*, std::string> Infer(const Stmt& n) {
StmtExprVisitor::VisitStmt(n);
std::unordered_map<const VarNode*, std::string> storage_scope_qualifiers;
@@ -57,7 +58,7 @@ class InferTextureAccess : public StmtExprVisitor {
}
return storage_scope_qualifiers;
}
- void VisitExpr_(const CallNode* op) {
+ void VisitExpr_(const CallNode* op) final {
if (op->op.same_as(builtin::texture2d_load())) {
var_access_map_[op->args[0].as<VarNode>()] |= kReadAccess;
} else if (op->op.same_as(builtin::texture2d_store())) {
@@ -190,17 +191,16 @@ void CodeGenOpenCL::BindThreadIndex(const IterVar& iv) {
} else {
os << "get_group_id(" << ts.dim_index << ")";
}
- var_idmap_[iv->var.get()] = CastFromTo(os.str(), DLDataType{kDLUInt, 64, 1},
iv->var.ty()->dtype);
+ var_idmap_[iv->var.get()] = CastFromTo(os.str(), PrimType::UInt(64),
iv->var.ty());
}
void CodeGenOpenCL::PrintType(const PrimType& t, std::ostream& os) { //
NOLINT(*)
- const DLDataType& raw_t = t->dtype;
int lanes = t.lanes();
if (t.IsVoid()) {
os << "void";
return;
}
- if (raw_t == DLDataType{kDLBool, 8, 1}) {
+ if (t == PrimType::Bool()) {
os << "bool";
return;
}
@@ -263,7 +263,7 @@ void CodeGenOpenCL::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(
return;
}
}
- TVM_FFI_THROW(InternalError) << "Cannot convert type " <<
ffi::DLDataTypeToString(raw_t)
+ TVM_FFI_THROW(InternalError) << "Cannot convert type " <<
ffi::DLDataTypeToString(t->dtype)
<< " to OpenCL type";
}
@@ -376,14 +376,15 @@ void CodeGenOpenCL::PrintRestrict(const Var& v,
std::ostream& os) {
}
}
-std::string CodeGenOpenCL::CastFromTo(std::string value, DLDataType from,
DLDataType target) {
+std::string CodeGenOpenCL::CastFromTo(std::string value, const PrimType& from,
+ const PrimType& target) {
if (from == target) return value;
return CastTo(value, target);
}
-std::string CodeGenOpenCL::CastTo(std::string value, DLDataType target) {
+std::string CodeGenOpenCL::CastTo(std::string value, const PrimType& target) {
std::ostringstream os;
- if (target == DLDataType{kDLBool, 8, 1}) {
+ if (target == PrimType::Bool()) {
os << "(";
os << "(";
this->PrintType(target, os);
@@ -422,9 +423,7 @@ void CodeGenOpenCL::VisitExpr_(const CallNode* op,
std::ostream& os) {
if (it != alloc_storage_scope_.end()) {
PrintStorageScope(it->second, os);
}
- this->PrintType(DLDataType{load->ty.as_or_throw<PrimType>()->dtype.code,
- load->ty.as_or_throw<PrimType>()->dtype.bits,
1},
- os);
+ this->PrintType(load->ty.as_or_throw<PrimType>().WithLanes(1), os);
os << " *)" << this->GetVarID(load->buffer->data.get()) << " + ";
this->PrintExpr(load->indices[0], os);
os << ')';
@@ -436,14 +435,13 @@ void CodeGenOpenCL::VisitExpr_(const CallNode* op,
std::ostream& os) {
const int channel_size = op->args[4].as_or_throw<IntImm>()->value;
TVM_FFI_ICHECK(channel_size == 64 || channel_size == 128)
<< "Unsupported Channel Size: " << channel_size;
- DLDataType channel_type = runtime::GetChannelType(channel_size);
+ PrimType channel_type(runtime::GetChannelType(channel_size));
- DLDataType buffer_type = ptr_type->element_type.as<PrimTypeNode>()->dtype;
+ PrimType buffer_type = ptr_type->element_type.as_or_throw<PrimType>();
std::stringstream ss;
this->PrintExpr(op->args[5].as_or_throw<PrimExpr>(), ss);
std::string value;
- value =
- this->SSAGetID(ss.str(), PrimType(buffer_type).WithLanes(channel_size
/ buffer_type.bits));
+ value = this->SSAGetID(ss.str(), buffer_type.WithLanes(channel_size /
buffer_type.bits()));
if (channel_size == 64) {
os << "write_imageh(";
} else if (channel_size == 128) {
@@ -475,7 +473,7 @@ void CodeGenOpenCL::VisitExpr_(const CallNode* op,
std::ostream& os) {
TVM_FFI_ICHECK(channel_size == 64 || channel_size == 128)
<< "Unsupported Channel Size: " << channel_size;
ss << "as_";
- this->PrintType(op_ty.WithLanes(data_lanes)->dtype, ss);
+ this->PrintType(op_ty.WithLanes(data_lanes), ss);
ss << "(";
if (channel_size == 64) {
ss << "READ_IMAGEH(";
@@ -505,10 +503,10 @@ void CodeGenOpenCL::VisitExpr_(const CallNode* op,
std::ostream& os) {
os << rhs;
} else if (*tirx::as_const_int(ramp->stride) == 1) {
os << "(*(";
-
this->PrintType(op_ty.WithLanes(*tirx::as_const_int(ramp->lanes))->dtype, os);
+ this->PrintType(op_ty.WithLanes(*tirx::as_const_int(ramp->lanes)), os);
os << "*)";
os << "((";
- this->PrintType(op_ty.WithLanes(1)->dtype, os);
+ this->PrintType(op_ty.WithLanes(1), os);
os << "*)&" << rhs << " + ";
this->PrintExpr(ramp->base, os);
os << "))";
@@ -517,7 +515,7 @@ void CodeGenOpenCL::VisitExpr_(const CallNode* op,
std::ostream& os) {
}
} else {
os << "((";
- this->PrintType(op_ty.WithLanes(1)->dtype, os);
+ this->PrintType(op_ty.WithLanes(1), os);
os << "*)&" << rhs << ")[";
this->PrintExpr(op->args.back().as_or_throw<PrimExpr>(), os);
os << "]";
@@ -548,7 +546,7 @@ void CodeGenOpenCL::VisitExpr_(const BroadcastNode* op,
std::ostream& os) { //
std::string v = PrintExpr(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
os << "((";
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
os << ")(";
for (int i = 0; i < lanes; ++i) {
if (i != 0) os << ", ";
@@ -559,7 +557,7 @@ void CodeGenOpenCL::VisitExpr_(const BroadcastNode* op,
std::ostream& os) { //
void CodeGenOpenCL::VisitExpr_(const RampNode* op, std::ostream& os) { //
NOLINT(*)
os << "((";
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
os << ")(";
int lanes = op->ty.as_or_throw<PrimType>().lanes();
for (int i = 0; i < lanes; i++) {
@@ -642,11 +640,11 @@ void CodeGenOpenCL::VisitExpr_(const AndNode* op,
std::ostream& os) {
std::ostringstream oss;
os << "(";
this->PrintExpr(op->a, oss);
- os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>());
oss.str("");
os << " && ";
this->PrintExpr(op->b, oss);
- os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>());
os << ")";
}
@@ -654,32 +652,31 @@ void CodeGenOpenCL::VisitExpr_(const OrNode* op,
std::ostream& os) {
std::ostringstream oss;
os << "(";
this->PrintExpr(op->a, oss);
- os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>());
oss.str("");
os << " || ";
this->PrintExpr(op->b, oss);
- os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastTo(oss.str(), op->ty.as_or_throw<PrimType>());
os << ")";
}
void CodeGenOpenCL::VisitExpr_(const SelectNode* op, std::ostream& os) {
+ PrimType op_ty = op->ty.as_or_throw<PrimType>();
std::ostringstream oss;
os << "select(";
PrintExpr(op->false_value, oss);
- os << CastFromTo(oss.str(), op->false_value.ty()->dtype,
op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastFromTo(oss.str(), op->false_value.ty(), op_ty);
oss.str("");
os << ", ";
PrintExpr(op->true_value, oss);
- os << CastFromTo(oss.str(), op->true_value.ty()->dtype,
op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastFromTo(oss.str(), op->true_value.ty(), op_ty);
oss.str("");
os << ", ";
PrintExpr(op->condition, oss);
- if (op->ty.as_or_throw<PrimType>().code() == DLDataTypeCode::kDLFloat) {
- os << CastTo(oss.str(),
- DLDataType{kDLInt,
static_cast<uint8_t>(op->ty.as_or_throw<PrimType>().bits()),
-
static_cast<uint16_t>(op->ty.as_or_throw<PrimType>().lanes())});
+ if (op_ty.code() == DLDataTypeCode::kDLFloat) {
+ os << CastTo(oss.str(), PrimType::Int(op_ty.bits(), op_ty.lanes()));
} else {
- os << CastFromTo(oss.str(), op->condition.ty()->dtype,
op->ty.as_or_throw<PrimType>()->dtype);
+ os << CastFromTo(oss.str(), op->condition.ty(), op_ty);
}
os << ")";
}
diff --git a/src/backend/opencl/codegen/codegen_opencl.h
b/src/backend/opencl/codegen/codegen_opencl.h
index 362409fa79..a537a2ede6 100644
--- a/src/backend/opencl/codegen/codegen_opencl.h
+++ b/src/backend/opencl/codegen/codegen_opencl.h
@@ -56,11 +56,12 @@ class CodeGenOpenCL final : public CodeGenC {
std::ostream& os) final; // NOLINT(*)
// the address of load/store
void PrintVecAddr(const BufferNode* buffer, const PrimType& t, PrimExpr base,
- std::ostream& os);
// NOLINT(*)
- void PrintRestrict(const Var& v, std::ostream& os) final;
// NOLINT(*)
- std::string CastFromTo(std::string value, DLDataType from, DLDataType
target); // NOLINT(*)
- std::string CastTo(std::string value, DLDataType target);
// NOLINT(*)
- void SetTextureScope(const std::unordered_map<const VarNode*,
std::string>&); // NOLINT(*)
+ std::ostream& os); // NOLINT(*)
+ void PrintRestrict(const Var& v, std::ostream& os) final; // NOLINT(*)
+ std::string CastFromTo(std::string value, const PrimType& from,
+ const PrimType& target) final;
// NOLINT(*)
+ std::string CastTo(std::string value, const PrimType& target);
// NOLINT(*)
+ void SetTextureScope(const std::unordered_map<const VarNode*,
std::string>&); // NOLINT(*)
// overload visitor
void VisitStmt_(const AllocBufferNode* op) final; //
NOLINT(*)
diff --git a/src/backend/trn/codegen/codegen_trn.h
b/src/backend/trn/codegen/codegen_trn.h
index 853a915faf..bcfd9451b2 100644
--- a/src/backend/trn/codegen/codegen_trn.h
+++ b/src/backend/trn/codegen/codegen_trn.h
@@ -49,6 +49,7 @@ struct NKIInstructionCtx {
class CodeGenTrainium final : public CodeGenC {
public:
explicit CodeGenTrainium(Target target);
+ using CodeGenC::PrintType;
using CodeGenC::VisitExpr_;
using CodeGenC::VisitStmt_;
// override print thread tag.
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 75ca467872..1165dbd78d 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -80,7 +80,7 @@ runtime::SPIRVShader CodeGenSPIRV::BuildFunction(const
PrimFunc& f, const std::s
<< "All handles passed to the Vulkan codegen must have a
type_annotation as a "
"PointerType, "
<< "and must point to a PrimType";
- PrimType value_storage_type(prim->dtype);
+ PrimType value_storage_type = ffi::GetRef<PrimType>(prim);
if (value_storage_type == PrimType::Bool()) {
// We need a physically addressable buffer type to support boolean
tensors.
// The loaded byte is cast to bool inside the LoadNode visitor below.
@@ -168,7 +168,7 @@ spirv::Value CodeGenSPIRV::GetThreadIndex(const IterVar&
iv, const PrimExpr& ext
} else {
v = builder_->GetWorkgroupID(ts.dim_index);
}
- return builder_->Cast(builder_->GetSType(PrimType(iv->var.ty()->dtype)), v);
+ return builder_->Cast(builder_->GetSType(iv->var.ty()), v);
}
spirv::Value CodeGenSPIRV::CreateStorageSync(const CallNode* op) {
@@ -212,13 +212,11 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const VarNode* op) {
}
spirv::Value CodeGenSPIRV::VisitExpr_(const IntImmNode* op) {
- return
builder_->IntImm(builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
- op->value);
+ return builder_->IntImm(builder_->GetSType(op->ty.as_or_throw<PrimType>()),
op->value);
}
spirv::Value CodeGenSPIRV::VisitExpr_(const FloatImmNode* op) {
- return
builder_->FloatImm(builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
- op->value);
+ return
builder_->FloatImm(builder_->GetSType(op->ty.as_or_throw<PrimType>()),
op->value);
}
spirv::Value CodeGenSPIRV::VisitExpr_(const StringImmNode* op) {
@@ -227,8 +225,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const StringImmNode*
op) {
}
spirv::Value CodeGenSPIRV::VisitExpr_(const CastNode* op) {
- return
builder_->Cast(builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
- MakeValue(op->value));
+ return builder_->Cast(builder_->GetSType(op->ty.as_or_throw<PrimType>()),
MakeValue(op->value));
}
spirv::Value CodeGenSPIRV::VisitExpr_(const AddNode* op) {
@@ -330,8 +327,8 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) {
for (size_t i = 1; i < op->args.size(); ++i) {
values.push_back(MakeValue(op->args[i]));
}
- return builder_->CallGLSL450(
- builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
inst_id, values);
+ return
builder_->CallGLSL450(builder_->GetSType(op->ty.as_or_throw<PrimType>()),
inst_id,
+ values);
} else if (op->op.same_as(builtin::bitwise_and())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
spirv::Value a = MakeValue(op->args[0]);
@@ -360,23 +357,20 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op)
{
TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
spirv::Value a = MakeValue(op->args[0]);
spirv::Value b = MakeValue(op->args[1]);
- if (PrimType(op->args[0].as_or_throw<PrimExpr>().ty()->dtype)
- .MatchesCode(DLDataTypeCode::kDLInt)) {
+ if
(op->args[0].as_or_throw<PrimExpr>().ty().MatchesCode(DLDataTypeCode::kDLInt)) {
return builder_->MakeValue(spv::OpShiftRightArithmetic, a.stype, a, b);
} else {
return builder_->MakeValue(spv::OpShiftRightLogical, a.stype, a, b);
}
} else if (op->op.same_as(builtin::reinterpret())) {
- return builder_->MakeValue(spv::OpBitcast,
-
builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
+ return builder_->MakeValue(spv::OpBitcast,
builder_->GetSType(op->ty.as_or_throw<PrimType>()),
MakeValue(op->args[0]));
} else if (op->op.same_as(builtin::large_uint_imm())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
uint64_t low = static_cast<uint64_t>(AsIntImmNode(op->args[0])->value);
uint64_t high = static_cast<uint64_t>(AsIntImmNode(op->args[1])->value);
uint64_t val = (high << 32U) | low;
- return
builder_->UIntImm(builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
- val);
+ return
builder_->UIntImm(builder_->GetSType(op->ty.as_or_throw<PrimType>()), val);
} else if (op->op.same_as(builtin::tvm_storage_sync())) {
return this->CreateStorageSync(op);
} else if (op->op.same_as(builtin::if_then_else())) {
@@ -404,8 +398,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) {
phi.SetIncoming(1, else_value, else_value_label);
return phi;
} else if (op->op.same_as(builtin::popcount())) {
- return builder_->MakeValue(spv::OpBitCount,
-
builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype)),
+ return builder_->MakeValue(spv::OpBitCount,
builder_->GetSType(op->ty.as_or_throw<PrimType>()),
MakeValue(op->args[0]));
} else if (op->op.same_as(builtin::call_pure_extern())) {
TVM_FFI_ICHECK_GE(op->args.size(), 1U);
@@ -415,7 +408,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) {
for (size_t i = 1; i < op->args.size(); ++i) {
values.push_back(MakeValue(op->args[i]));
}
- PrimType op_dtype(op->ty.as_or_throw<PrimType>()->dtype);
+ PrimType op_dtype = op->ty.as_or_throw<PrimType>();
return builder_->CallKHRIntegerDotProduct(builder_->GetSType(op_dtype),
values, op_dtype);
} else {
TVM_FFI_THROW(InternalError)
@@ -588,7 +581,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const BufferLoadNode*
op) {
Var buffer_var = op->buffer->data;
PrimExpr prim_index = op->indices[0];
- PrimType desired_read_type(op->ty.as_or_throw<PrimType>()->dtype);
+ PrimType desired_read_type = op->ty.as_or_throw<PrimType>();
if (desired_read_type == PrimType::Bool()) {
desired_read_type =
boolean_storage_type_.WithLanes(desired_read_type.lanes());
}
@@ -596,7 +589,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const BufferLoadNode*
op) {
auto it = storage_info_.find(buffer_var.get());
TVM_FFI_ICHECK(it != storage_info_.end());
StorageInfo& info = it->second;
- info.CheckContentType(desired_read_type,
PrimType(prim_index.ty()->dtype).lanes());
+ info.CheckContentType(desired_read_type, prim_index.ty().lanes());
spirv::SType content_type = builder_->GetSType(info.element_type);
spirv::Value buffer = MakeValue(buffer_var);
@@ -616,7 +609,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const BufferLoadNode*
op) {
spirv::Value loaded = builder_->MakeValue(spv::OpLoad, content_type, ptr,
mask);
// OpTypeBool have no physical address/storage. Here, cast from
// the storage type to an OpTypeBool.
- if (PrimType(op->ty.as_or_throw<PrimType>()->dtype) == PrimType::Bool()) {
+ if (op->ty.as_or_throw<PrimType>() == PrimType::Bool()) {
auto spirv_bool = builder_->GetSType(PrimType::Bool());
loaded = builder_->Cast(spirv_bool, loaded);
}
@@ -636,10 +629,8 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const
BufferLoadNode* op) {
} else {
TVM_FFI_THROW(InternalError) << "Cannot perform buffer access of buffer
variable '"
<< buffer_var->name_hint << "' with element
type "
- << info.element_type << " using index of type
"
- << PrimType(prim_index.ty()->dtype)
- << " to produce output of type "
- <<
PrimType(op->ty.as_or_throw<PrimType>()->dtype);
+ << info.element_type << " using index of type
" << prim_index.ty()
+ << " to produce output of type " <<
op->ty.as_or_throw<PrimType>();
return spirv::Value();
}
}
@@ -651,9 +642,10 @@ void CodeGenSPIRV::Scalarize(const PrimExpr& e,
std::function<void(int i, spirv:
f(i, MakeValue(offset));
}
} else {
- spirv::SType etype =
builder_->GetSType(PrimType(e.ty()->dtype).WithLanes(1));
+ PrimType e_type = e.ty();
+ spirv::SType etype = builder_->GetSType(e_type.WithLanes(1));
spirv::Value value = MakeValue(e);
- for (int i = 0; i < PrimType(e.ty()->dtype).lanes(); ++i) {
+ for (int i = 0; i < e_type.lanes(); ++i) {
f(i, builder_->MakeValue(spv::OpCompositeExtract, etype, value, i));
}
}
@@ -665,7 +657,7 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const ShuffleNode*
op) {
<< "of one vector with one index";
spirv::Value vector = MakeValue(op->vectors[0]);
int index = AsIntImmNode(op->indices[0])->value;
- spirv::SType etype =
builder_->GetSType(PrimType(op->ty.as_or_throw<PrimType>()->dtype));
+ spirv::SType etype = builder_->GetSType(op->ty.as_or_throw<PrimType>());
spirv::Value element = builder_->MakeValue(spv::OpCompositeExtract, etype,
vector, index);
return element;
}
@@ -679,7 +671,9 @@ void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
auto it = storage_info_.find(buffer_var.get());
TVM_FFI_ICHECK(it != storage_info_.end());
StorageInfo& info = it->second;
- info.CheckContentType(PrimType(op->value.ty()->dtype),
PrimType(prim_index.ty()->dtype).lanes());
+ PrimType value_type = op->value.ty();
+ PrimType index_type = prim_index.ty();
+ info.CheckContentType(value_type, index_type.lanes());
spirv::SType content_type = builder_->GetSType(info.element_type);
spirv::Value buffer = MakeValue(buffer_var);
@@ -691,16 +685,16 @@ void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
mask |= spv::MemoryAccessVolatileMask;
}
- if (PrimType(op->value.ty()->dtype) == info.element_type) {
+ if (value_type == info.element_type) {
// Requested store of a single value. This may be a scalar store
// or a vectorized store, based on the array element type.
- TVM_FFI_ICHECK_EQ(info.element_type, PrimType(op->value.ty()->dtype))
+ TVM_FFI_ICHECK_EQ(info.element_type, value_type)
<< "Vulkan only allow one type access to the same buffer";
spirv::Value index = MakeValue(prim_index);
spirv::Value ptr = builder_->StructArrayAccess(ptr_type, buffer, index);
builder_->MakeInst(spv::OpStore, ptr, value, mask);
- } else if (PrimType(op->value.ty()->dtype).WithLanes(1) ==
info.element_type) {
+ } else if (value_type.WithLanes(1) == info.element_type) {
// Requested store of several arbitrarily located values. Extract
// each value from the composite, then assign to the buffer.
auto f = [&](int i, spirv::Value index) {
@@ -711,10 +705,10 @@ void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
this->Scalarize(prim_index, f);
} else {
- TVM_FFI_THROW(InternalError) << "Cannot store value of type " <<
PrimType(op->value.ty()->dtype)
+ TVM_FFI_THROW(InternalError) << "Cannot store value of type " << value_type
<< " into buffer variable '" <<
buffer_var->name_hint
<< "' with element type " << info.element_type
- << " using index of type " <<
PrimType(prim_index.ty()->dtype);
+ << " using index of type " << index_type;
}
}
@@ -727,7 +721,7 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// loop step
spirv::Value step;
if (op->HasTrivialStep()) {
- step =
PrimType(op->loop_var.ty()->dtype).MatchesCode(DLDataTypeCode::kDLInt)
+ step = op->loop_var.ty().MatchesCode(DLDataTypeCode::kDLInt)
? builder_->IntImm(init_value.stype, 1)
: builder_->UIntImm(init_value.stype, 1);
} else {
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index c93d5e3780..cadd169b41 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -65,6 +65,8 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor {
}
private:
+ using StmtExprVisitor::VisitExpr_;
+
void VisitExpr_(const VarNode* op) final {
StmtExprVisitor::VisitExpr_(op);
Var buffer_var = ffi::GetRef<Var>(op);
@@ -292,7 +294,7 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const
PrimFunc& f, bool skip_re
void CodeGenWebGPU::BindThreadIndex(const IterVar& iv) {
TVM_FFI_ICHECK(!var_idmap_.count(iv->var.get()));
std::ostringstream os;
- PrintType(iv->var.ty()->dtype, os);
+ PrintType(iv->var.ty(), os);
if (iv->thread_tag == "blockIdx.x") {
// WebGPU have restriction to limit the maximum size of blockId.x to be
65535
// We allow runtime to spread the load out to blockIdx.z so it can be a
large number.
@@ -309,13 +311,12 @@ void CodeGenWebGPU::BindThreadIndex(const IterVar& iv) {
}
void CodeGenWebGPU::PrintType(const PrimType& t, std::ostream& os) { //
NOLINT(*)
- const DLDataType& raw_t = t->dtype;
int lanes = t.lanes();
if (t.IsVoid()) {
os << "void";
return;
}
- if (raw_t == DLDataType{kDLBool, 8, 1}) {
+ if (t == PrimType::Bool()) {
os << "bool";
return;
}
@@ -346,7 +347,7 @@ void CodeGenWebGPU::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(
os << "i" << t.bits();
} else {
TVM_FFI_THROW(InternalError) << "CodeGenWebGPU: Cannot convert type "
- << ffi::DLDataTypeToString(raw_t) << " to
WebGPU type";
+ << ffi::DLDataTypeToString(t->dtype) << " to
WebGPU type";
}
if (lanes != 1) {
os << ">";
@@ -389,7 +390,7 @@ void CodeGenWebGPU::PrintVecElemStore(const std::string&
vec, const PrimType& t,
void CodeGenWebGPU::VisitExpr_(const BroadcastNode* op, std::ostream& os) {
// NOLINT(*)
std::string v = PrintExpr(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
os << "(";
for (int i = 0; i < lanes; ++i) {
if (i != 0) os << ", ";
@@ -406,7 +407,7 @@ void CodeGenWebGPU::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLIN
if (op->op.same_as(builtin::reinterpret())) {
// generate bitcast<TYPE>(ARG)
os << "bitcast<";
- this->PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ this->PrintType(op->ty.as_or_throw<PrimType>(), os);
os << ">(";
this->PrintExpr(op->args[0], os);
os << ")";
@@ -430,7 +431,7 @@ void CodeGenWebGPU::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLIN
std::string cond = PrintExpr(op->args[0]);
this->PrintIndent();
this->stream << "var " << result << " : ";
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, this->stream);
+ PrintType(op->ty.as_or_throw<PrimType>(), this->stream);
this->stream << ";\n";
this->PrintIndent();
this->stream << "if (" << cond << ") {\n";
@@ -463,7 +464,7 @@ void CodeGenWebGPU::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLIN
}
void CodeGenWebGPU::VisitExpr_(const CastNode* op, std::ostream& os) { //
NOLINT(*)
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
os << "(" << PrintExpr(op->value) << ")";
}
@@ -505,7 +506,7 @@ void CodeGenWebGPU::VisitExpr_(const IntImmNode* op,
std::ostream& os) { // NOL
this->MarkConst(temp.str());
os << temp.str();
} else {
- this->PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ this->PrintType(op->ty.as_or_throw<PrimType>(), os);
os << "(" << op->value << ")";
}
}
@@ -535,12 +536,10 @@ void CodeGenWebGPU::VisitExpr_(const BufferLoadNode* op,
std::ostream& os) { //
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Load from non-flat memory not
supported.";
TVM_FFI_ICHECK(!op->predicate.defined()) << "Predicated buffer load is not
supported.";
- DLDataType value_dtype = op->ty.as_or_throw<PrimType>()->dtype;
- PrimType value_ty(value_dtype);
+ PrimType value_ty = op->ty.as_or_throw<PrimType>();
PrimExpr index = op->indices[0];
Var buffer_var = op->buffer->data;
- DLDataType element_dtype = op->buffer->dtype->dtype;
- PrimType element_ty(element_dtype);
+ const PrimType& element_ty = op->buffer->dtype;
int lanes = value_ty.lanes();
std::string buffer_vid = GetVarID(buffer_var.get());
@@ -548,29 +547,28 @@ void CodeGenWebGPU::VisitExpr_(const BufferLoadNode* op,
std::ostream& os) { //
if (value_ty.lanes() == element_ty.lanes()) {
// Direct buffer loading
// Special handle bool loading
- if (value_dtype == DLDataType{kDLBool, 8, 1}) {
- this->PrintType(value_dtype, os);
+ if (value_ty == PrimType::Bool()) {
+ this->PrintType(value_ty, os);
os << "(";
} else {
- TVM_FFI_ICHECK(value_dtype == element_dtype);
+ TVM_FFI_ICHECK(value_ty == element_ty);
}
TVM_FFI_ICHECK_EQ(index.ty().lanes(), 1);
os << buffer_vid << "[" << this->PrintExpr(index) << "]";
// Special handle bool loading
- if (value_dtype == DLDataType{kDLBool, 8, 1}) {
+ if (value_ty == PrimType::Bool()) {
os << ")";
}
} else {
// Vector load from scalar buffer
TVM_FFI_ICHECK_EQ(element_ty.lanes(), 1) << "Can only vector load scalar
array";
- DLDataType value_element_dtype{value_dtype.code, value_dtype.bits, 1};
- TVM_FFI_ICHECK(value_element_dtype == element_dtype)
+ TVM_FFI_ICHECK(value_ty.WithLanes(1) == element_ty)
<< "WebGPU vector loading requires base type to match";
arith::PVar<PrimExpr> base;
if (arith::ramp(base, 1, value_ty.lanes()).Match(index)) {
// vec3<f32>(buf[base + 0], buf[base + 1], buf[base + 2]);
std::string base_vid = SSAGetID(PrintExpr(base.Eval()),
base.Eval().ty());
- PrintType(element_ty.WithLanes(value_ty.lanes())->dtype, os);
+ PrintType(element_ty.WithLanes(value_ty.lanes()), os);
os << "(";
for (int i = 0; i < lanes; ++i) {
if (i != 0) os << ", ";
@@ -580,7 +578,7 @@ void CodeGenWebGPU::VisitExpr_(const BufferLoadNode* op,
std::ostream& os) { //
} else {
// vec3<f32>(buf[index[0]], buf[index[1]], buf[index[2]]);
std::string index_vid = SSAGetID(PrintExpr(index), index.ty());
- PrintType(element_ty.WithLanes(value_ty.lanes())->dtype, os);
+ PrintType(element_ty.WithLanes(value_ty.lanes()), os);
os << "(";
for (int i = 0; i < lanes; ++i) {
if (i != 0) os << ", ";
@@ -610,10 +608,8 @@ void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not
supported.";
TVM_FFI_ICHECK(!op->predicate.defined()) << "Predicated buffer store is not
supported.";
- DLDataType value_dtype = op->value.ty()->dtype;
- PrimType value_ty(value_dtype);
- DLDataType element_dtype = op->buffer->dtype->dtype;
- PrimType element_ty(element_dtype);
+ PrimType value_ty = op->value.ty();
+ const PrimType& element_ty = op->buffer->dtype;
PrimExpr index = op->indices[0];
Var buffer_var = op->buffer->data;
@@ -628,23 +624,22 @@ void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op)
{
this->PrintIndent();
stream << buffer_vid << "[" << index_vid << "] = ";
// special explicit conversion of bool
- if (value_dtype == DLDataType{kDLBool, 8, 1}) {
- PrintType(element_dtype, stream);
+ if (value_ty == PrimType::Bool()) {
+ PrintType(element_ty, stream);
stream << "(";
} else {
- TVM_FFI_ICHECK(value_dtype == element_dtype);
+ TVM_FFI_ICHECK(value_ty == element_ty);
}
stream << value_vid;
// Special handle bool store
- if (value_dtype == DLDataType{kDLBool, 8, 1}) {
+ if (value_ty == PrimType::Bool()) {
stream << ")";
}
stream << ";\n";
} else {
// Vector store into scalar buffer
TVM_FFI_ICHECK_EQ(element_ty.lanes(), 1) << "Can only vector load scalar
array";
- DLDataType value_element_dtype{value_dtype.code, value_dtype.bits, 1};
- TVM_FFI_ICHECK(value_element_dtype == element_dtype)
+ TVM_FFI_ICHECK(value_ty.WithLanes(1) == element_ty)
<< "WebGPU vector stire requires base type to match";
std::string value_vid = PrintExpr(op->value);
arith::PVar<PrimExpr> base;
@@ -684,12 +679,12 @@ void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op)
{
if (storage_scope.rank == runtime::StorageRank::kShared) {
this->decl_stream << "var<workgroup> " << vid << " : array<";
- PrintType(op->buffer->dtype->dtype, this->decl_stream);
+ PrintType(op->buffer->dtype, this->decl_stream);
this->decl_stream << ", " << constant_size << ">;\n";
} else if (storage_scope.rank == runtime::StorageRank::kLocal) {
this->PrintIndent();
this->stream << "var " << vid << " : array<";
- PrintType(op->buffer->dtype->dtype, this->stream);
+ PrintType(op->buffer->dtype, this->stream);
this->stream << ", " << constant_size << ">;\n";
} else {
TVM_FFI_THROW(InternalError) << "WebGPU: Do not support storage scope: "
@@ -705,7 +700,7 @@ void CodeGenWebGPU::VisitStmt_(const ForNode* op) {
std::string vid = AllocVarID(op->loop_var.get());
PrintIndent();
stream << "for (var " << vid << " : ";
- PrintType(op->loop_var.ty()->dtype, stream);
+ PrintType(op->loop_var.ty(), stream);
stream << " = " << begin_str << "; " << vid << " < " << end_str << "; " <<
vid;
if (step_str.empty()) {
stream << "++";
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 798f149c3c..0f35d34ba9 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -150,7 +150,7 @@ Var::Var(Id vid, ffi::Optional<Type> ty_annotation, Span
span) {
ffi::ObjectPtr<VarNode> n = ffi::make_object<VarNode>();
n->vid = std::move(vid);
if (ty_annotation.defined()) {
- n->ty = std::move(ty_annotation.value());
+ n->ty = ty_annotation.value();
}
n->span = std::move(span);
data_ = std::move(n);
@@ -189,7 +189,7 @@ DataflowVar::DataflowVar(Id vid, ffi::Optional<Type>
ty_annotation, Span span) {
ffi::ObjectPtr<DataflowVarNode> n = ffi::make_object<DataflowVarNode>();
n->vid = std::move(vid);
if (ty_annotation.defined()) {
- n->ty = std::move(ty_annotation.value());
+ n->ty = ty_annotation.value();
}
n->span = std::move(span);
data_ = std::move(n);
diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc
index 30f90eddc2..1257948663 100644
--- a/src/relax/ir/py_expr_functor.cc
+++ b/src/relax/ir/py_expr_functor.cc
@@ -97,7 +97,7 @@ class PyExprVisitorNode : public ffi::Object, public
ExprVisitor {
/*! \brief The packed function to the `VisitSpan(const Span& span)`
function. */
ffi::Function f_visit_span{nullptr};
- void VisitExpr(const Expr& expr) {
+ void VisitExpr(const Expr& expr) override {
if (f_visit_expr != nullptr) {
f_visit_expr(expr);
} else {
@@ -115,36 +115,36 @@ class PyExprVisitorNode : public ffi::Object, public
ExprVisitor {
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<Expr>(op), f_visit_expr_fallback_,
ExprVisitor::VisitExprFallback_(op));
- void VisitBinding(const Binding& binding)
+ void VisitBinding(const Binding& binding) override
PY_EXPR_VISITOR_DEFAULT(binding, f_visit_binding,
ExprVisitor::VisitBinding(binding));
- void VisitBinding_(const VarBindingNode* binding)
+ void VisitBinding_(const VarBindingNode* binding) override
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<VarBinding>(binding),
f_visit_var_binding_,
ExprVisitor::VisitBinding_(binding));
- void VisitBinding_(const MatchCastNode* binding)
+ void VisitBinding_(const MatchCastNode* binding) override
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<MatchCast>(binding),
f_visit_match_cast_,
ExprVisitor::VisitBinding_(binding));
- void VisitBindingBlock(const BindingBlock& block)
+ void VisitBindingBlock(const BindingBlock& block) override
PY_EXPR_VISITOR_DEFAULT(block, f_visit_binding_block,
ExprVisitor::VisitBindingBlock(block));
- void VisitBindingBlock_(const BindingBlockNode* block)
+ void VisitBindingBlock_(const BindingBlockNode* block) override
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<BindingBlock>(block),
f_visit_binding_block_,
ExprVisitor::VisitBindingBlock_(block));
- void VisitBindingBlock_(const DataflowBlockNode* block)
+ void VisitBindingBlock_(const DataflowBlockNode* block) override
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<DataflowBlock>(block),
f_visit_dataflow_block_,
ExprVisitor::VisitBindingBlock_(block));
- void VisitVarDef(const Var& var)
+ void VisitVarDef(const Var& var) override
PY_EXPR_VISITOR_DEFAULT(var, f_visit_var_def,
ExprVisitor::VisitVarDef(var));
- void VisitVarDef_(const VarNode* var)
+ void VisitVarDef_(const VarNode* var) override
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<Var>(var), f_visit_var_def_,
ExprVisitor::VisitVarDef_(var));
- void VisitVarDef_(const DataflowVarNode* var)
+ void VisitVarDef_(const DataflowVarNode* var) override
PY_EXPR_VISITOR_DEFAULT(ffi::GetRef<DataflowVar>(var),
f_visit_dataflow_var_def_,
ExprVisitor::VisitVarDef_(var));
- void VisitSpan(const Span& span)
+ void VisitSpan(const Span& span) override
PY_EXPR_VISITOR_DEFAULT(span, f_visit_span,
ExprVisitor::VisitSpan(span));
static void RegisterReflection() {
@@ -342,7 +342,7 @@ class PyExprMutatorNode : public ffi::Object, public
ExprMutator {
/*! \brief The packed function to the `VisitSpan(const Span& span)`
function. */
ffi::Function f_visit_span{nullptr};
- Expr VisitExpr(const Expr& expr) {
+ Expr VisitExpr(const Expr& expr) override {
if (f_visit_expr != nullptr) {
return builder_->Normalize(f_visit_expr(expr).cast<Expr>());
} else {
@@ -358,44 +358,44 @@ class PyExprMutatorNode : public ffi::Object, public
ExprMutator {
PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef<Expr>(op), f_visit_expr_fallback_,
ExprMutator::VisitExprFallback_(op), Expr);
- void VisitBinding(const Binding& binding) {
+ void VisitBinding(const Binding& binding) override {
if (f_visit_binding != nullptr)
f_visit_binding(binding);
else
ExprMutator::VisitBinding(binding);
}
- void VisitBinding_(const VarBindingNode* binding) {
+ void VisitBinding_(const VarBindingNode* binding) override {
if (f_visit_var_binding_ != nullptr)
f_visit_var_binding_(ffi::GetRef<VarBinding>(binding));
else
ExprMutator::VisitBinding_(binding);
}
- void VisitBinding_(const MatchCastNode* binding) {
+ void VisitBinding_(const MatchCastNode* binding) override {
if (f_visit_match_cast_ != nullptr)
f_visit_match_cast_(ffi::GetRef<MatchCast>(binding));
else
ExprMutator::VisitBinding_(binding);
}
- BindingBlock VisitBindingBlock(const BindingBlock& block)
+ BindingBlock VisitBindingBlock(const BindingBlock& block) override
PY_EXPR_MUTATOR_DEFAULT(block, f_visit_binding_block,
ExprMutator::VisitBindingBlock(block),
BindingBlock);
- BindingBlock VisitBindingBlock_(const BindingBlockNode* block)
+ BindingBlock VisitBindingBlock_(const BindingBlockNode* block) override
PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef<BindingBlock>(block),
f_visit_binding_block_,
ExprMutator::VisitBindingBlock_(block),
BindingBlock);
- BindingBlock VisitBindingBlock_(const DataflowBlockNode* block)
+ BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) override
PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef<DataflowBlock>(block),
f_visit_dataflow_block_,
ExprMutator::VisitBindingBlock_(block),
BindingBlock);
- Var VisitVarDef(const Var& var)
+ Var VisitVarDef(const Var& var) override
PY_EXPR_MUTATOR_DEFAULT(var, f_visit_var_def,
ExprMutator::VisitVarDef(var), Var);
- Var VisitVarDef_(const VarNode* var)
+ Var VisitVarDef_(const VarNode* var) override
PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef<Var>(var), f_visit_var_def_,
ExprMutator::VisitVarDef_(var), Var);
- Var VisitVarDef_(const DataflowVarNode* var)
+ Var VisitVarDef_(const DataflowVarNode* var) override
PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef<DataflowVar>(var),
f_visit_dataflow_var_def_,
ExprMutator::VisitVarDef_(var), Var);
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index d257852413..09f842e9d8 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -276,7 +276,7 @@ Type InferTypeEye(const Call& call, const BlockBuilder&
ctx) {
<< call->args.size() << " arguments";
}
- auto get_prim_value = [&ctx](const Expr& expr, std::string key) {
+ auto get_prim_value = [](const Expr& expr, std::string key) {
auto prim_value = expr.as<PrimExpr>();
if (!prim_value) {
TVM_FFI_VISIT_THROW(TypeError, expr)
@@ -360,7 +360,7 @@ Type InferTypeArange(const Call& call, const BlockBuilder&
ctx) {
<< call->args.size() << " arguments";
}
// TODO(Siyuan): Support indirect prim_values
- auto get_prim_value = [&ctx](const Expr& expr, std::string key) {
+ auto get_prim_value = [](const Expr& expr, std::string key) {
auto prim_value = expr.as<PrimExpr>();
if (!prim_value) {
TVM_FFI_VISIT_THROW(TypeError, expr)
@@ -421,7 +421,7 @@ Type InferTypeHammingWindow(const Call& call, const
BlockBuilder& ctx) {
TVM_FFI_VISIT_THROW(TypeError, call)
<< "Hamming Window expects the datatype to be float but got " << dtype;
}
- auto get_prim_value = [&ctx](const Expr& expr, std::string key) {
+ auto get_prim_value = [](const Expr& expr, std::string key) {
auto prim_value = expr.as<PrimExpr>();
if (!prim_value) {
TVM_FFI_VISIT_THROW(TypeError, expr) << "Hamming_window expects the `"
<< key
diff --git a/src/relax/op/tensor/manipulate.cc
b/src/relax/op/tensor/manipulate.cc
index 74727ac9a3..d9c1b35931 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -3016,7 +3016,7 @@ Type InferTypeSliceScatter(const Call& call, const
BlockBuilder& ctx) {
}
}
- auto get_prim_expr_from_arg = [&ctx, &call](const Expr& arg_expr,
std::string key) -> PrimExpr {
+ auto get_prim_expr_from_arg = [&call](const Expr& arg_expr, std::string key)
-> PrimExpr {
auto prim_value = arg_expr.as<PrimExpr>();
if (!prim_value) {
TVM_FFI_VISIT_THROW(TypeError, call)
diff --git a/src/relax/transform/attach_global_symbol.cc
b/src/relax/transform/attach_global_symbol.cc
index e8f5d368d0..cdefafb7f7 100644
--- a/src/relax/transform/attach_global_symbol.cc
+++ b/src/relax/transform/attach_global_symbol.cc
@@ -57,6 +57,7 @@ struct TirxGvarMutator : tirx::StmtExprMutator {
explicit TirxGvarMutator(ffi::Map<GlobalVar, GlobalVar> replacements)
: replacements(replacements) {}
+ using tirx::StmtExprMutator::VisitExpr_;
Expr VisitExpr_(const CallNode* node) override {
auto call =
tirx::StmtExprMutator::VisitExpr_(node).as_or_throw<tvm::Call>();
if (auto old_gvar = call->op.as<GlobalVar>()) {
diff --git a/src/relax/transform/compute_prim_value.cc
b/src/relax/transform/compute_prim_value.cc
index f82721cb9c..49beccc88f 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -44,6 +44,8 @@ class PrimExprComputeInjector : public ExprMutator {
IRModule Finalize() const { return builder_->Finalize(); }
private:
+ using ExprMutator::VisitExpr_;
+
Expr VisitExpr_(const CallNode* op) final {
Call call = ffi::GetRef<Call>(op);
if (auto prim_expr = call.as<PrimExpr>()) {
diff --git a/src/relax/transform/dead_code_elimination.cc
b/src/relax/transform/dead_code_elimination.cc
index e80ac2fbac..10b6a1d7c0 100644
--- a/src/relax/transform/dead_code_elimination.cc
+++ b/src/relax/transform/dead_code_elimination.cc
@@ -61,6 +61,7 @@ struct RelaxCalleeCollector : relax::ExprVisitor {
struct TIRxCalleeCollector : tirx::StmtExprVisitor {
std::vector<GlobalVar>* callees;
explicit TIRxCalleeCollector(std::vector<GlobalVar>* out) : callees(out) {}
+ using tirx::StmtExprVisitor::VisitExpr_;
void VisitExpr_(const CallNode* node) final {
tirx::StmtExprVisitor::VisitExpr_(node);
if (auto opt_gvar = node->op.as<GlobalVar>()) {
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index d28a0e0ea7..e3b00defe6 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -57,7 +57,8 @@ class SymbolicMatcher : ExprFunctor<void(const Expr& n, const
PrimExpr& other)>
}
private:
- void VisitExpr(const PrimExpr& node, const PrimExpr& other) {
+ void VisitExpr(const Expr& expr, const PrimExpr& other) final {
+ PrimExpr node = expr.as_or_throw<PrimExpr>();
if (node.same_as(other)) {
return;
} else if (node.ty().code() != other.ty().code()) {
@@ -65,7 +66,7 @@ class SymbolicMatcher : ExprFunctor<void(const Expr& n, const
PrimExpr& other)>
<< "Parameter expression " << node << " with dtype " <<
node.ty()->dtype
<< " cannot match to argument " << other << " with dtype " <<
other.ty()->dtype;
} else {
- ExprFunctor::VisitExpr(node, other);
+ ExprFunctor::VisitExpr(expr, other);
}
}
diff --git a/src/runtime/vm/builtin.cc b/src/runtime/vm/builtin.cc
index c23fbd868f..162cc59ba9 100644
--- a/src/runtime/vm/builtin.cc
+++ b/src/runtime/vm/builtin.cc
@@ -310,13 +310,13 @@ void CheckPrimValueInfo(ffi::AnyView arg, DLDataType
dtype, ffi::Optional<ffi::S
TVM_FFI_THROW(TypeError) << err_ctx.value_or("") << ", expected dtype " <<
dtype
<< ", but received ObjectRef of type "
<< opt_obj.value()->GetTypeKey();
- } else if (((dtype).code == kDLBool)) {
+ } else if (dtype.code == kDLBool) {
arg.cast<bool>();
- } else if (((dtype).code == kDLInt)) {
+ } else if (dtype.code == kDLInt) {
arg.cast<int64_t>();
- } else if (((dtype).code == kDLUInt)) {
+ } else if (dtype.code == kDLUInt) {
arg.cast<uint64_t>();
- } else if (((dtype).code == kDLFloat)) {
+ } else if (dtype.code == kDLFloat) {
arg.cast<double>();
} else if (dtype.code == kDLOpaqueHandle && !(dtype.bits == 0 && dtype.lanes
== 0)) {
arg.cast<void*>();
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index abec9a6b1a..8bf390a508 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -485,9 +485,9 @@ inline void PrintConst(const IntImmNode* op, std::ostream&
os, CodeGenC* p) { /
}
}
-inline void PrintUIntConst(DLDataType dtype, uint64_t val, std::ostream& os,
+inline void PrintUIntConst(const PrimType& dtype, uint64_t val, std::ostream&
os,
CodeGenC* p) { // NOLINT(*)
- if (dtype == DLDataType{kDLUInt, 32, 1}) {
+ if (dtype == PrimType::UInt(32)) {
std::ostringstream temp;
temp << val << "U";
p->MarkConst(temp.str());
@@ -704,7 +704,7 @@ void CodeGenC::VisitExpr_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
uint64_t low =
static_cast<uint64_t>(op->args[0].as_or_throw<IntImm>()->value);
uint64_t high =
static_cast<uint64_t>(op->args[1].as_or_throw<IntImm>()->value);
uint64_t val = (high << 32U) | low;
- PrintUIntConst(op->ty.as_or_throw<PrimType>()->dtype, val, os, this);
+ PrintUIntConst(op->ty.as_or_throw<PrimType>(), val, os, this);
} else if (op->op.same_as(builtin::bitwise_xor())) {
PrintBinaryIntrinsic(op, " ^ ", os, this);
} else if (op->op.same_as(builtin::bitwise_or())) {
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index 5b8ae8c185..a937a96a99 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -233,9 +233,6 @@ class CodeGenC : public ExprFunctor<void(const Expr&,
std::ostream&)>,
virtual void PrintVecConstructor(const PrimType& t, std::ostream& os);
// Get a cast type from to
virtual std::string CastFromTo(std::string value, const PrimType& from,
const PrimType& target);
- std::string CastFromTo(std::string value, DLDataType from, DLDataType
target) {
- return CastFromTo(std::move(value), PrimType(from), PrimType(target));
- }
// Get load of single element with expression
virtual void PrintVecElemLoadExpr(const PrimType& t, int i, const
std::string& value,
std::ostream& os);
diff --git a/src/target/source/codegen_c_host.cc
b/src/target/source/codegen_c_host.cc
index c37dee2957..709b3ec4a9 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -122,24 +122,23 @@ void CodeGenCHost::PrintFuncPrefix(std::ostream& os) {
// NOLINT(*)
}
void CodeGenCHost::PrintType(const PrimType& type, std::ostream& os) { //
NOLINT(*)
- const DLDataType& t = type->dtype;
- int lanes = static_cast<int16_t>(t.lanes);
- if (t.code == kDLOpaqueHandle && !(t.bits == 0 && lanes == 0)) {
+ int lanes = type.lanes();
+ if (type.MatchesCode(DLDataTypeCode::kDLOpaqueHandle) && !type.IsVoid()) {
TVM_FFI_ICHECK_EQ(lanes, 1) << "does not support vector types";
os << "void*";
return;
}
- if (t.code == kDLOpaqueHandle && t.bits == 0 && lanes == 0) {
+ if (type.IsVoid()) {
os << "void";
return;
}
- if (t.code == kDLBool && lanes == 1) {
+ if (type.MatchesCode(DLDataTypeCode::kDLBool) && lanes == 1) {
os << "bool";
return;
}
bool fail = false;
- if (t.code == kDLFloat) {
- switch (t.bits) {
+ if (type.MatchesCode(DLDataTypeCode::kDLFloat)) {
+ switch (type.bits()) {
case 16:
os << "half";
break;
@@ -158,11 +157,11 @@ void CodeGenCHost::PrintType(const PrimType& type,
std::ostream& os) { // NOLIN
os << lanes;
return;
}
- } else if (t.code == kDLUInt || t.code == kDLInt) {
- if (t.code == kDLUInt) {
+ } else if (type.MatchesCode(DLDataTypeCode::kDLUInt,
DLDataTypeCode::kDLInt)) {
+ if (type.MatchesCode(DLDataTypeCode::kDLUInt)) {
os << 'u';
}
- switch (t.bits) {
+ switch (type.bits()) {
case 8:
os << "int8_t";
break;
@@ -188,14 +187,14 @@ void CodeGenCHost::PrintType(const PrimType& type,
std::ostream& os) { // NOLIN
return;
}
}
- TVM_FFI_THROW(InternalError) << "Cannot convert type " << t << " to C type";
+ TVM_FFI_THROW(InternalError) << "Cannot convert type " << type->dtype << "
to C type";
}
void CodeGenCHost::VisitExpr_(const BroadcastNode* op, std::ostream& os) { //
NOLINT(*)
std::string v = PrintExpr(op->value);
- int lanes =
static_cast<int16_t>(op->ty.as_or_throw<PrimType>()->dtype.lanes);
+ int lanes = op->ty.as_or_throw<PrimType>().lanes();
os << "((";
- PrintType(op->ty.as_or_throw<PrimType>()->dtype, os);
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
os << ")(";
for (int i = 0; i < lanes; ++i) {
if (i != 0) os << ", ";
diff --git a/src/target/source/codegen_source_base.h
b/src/target/source/codegen_source_base.h
index 0dbeb91325..95b6c00e70 100644
--- a/src/target/source/codegen_source_base.h
+++ b/src/target/source/codegen_source_base.h
@@ -59,7 +59,6 @@ class CodeGenSourceBase {
* \param os The stream to print the ctype into
*/
virtual void PrintType(const PrimType& type, std::ostream& os); // NOLINT(*)
- void PrintType(DLDataType type, std::ostream& os) {
PrintType(PrimType(type), os); }
/*!
* Print Type representation of type type.
* \param type The type representation.
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 425c2a979a..d1b137ca3d 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -1714,6 +1714,8 @@ class VectorTypeRewriter : public StmtExprMutator {
: var_remap_(var_remap) {}
private:
+ using StmtExprMutator::VisitExpr_;
+
Expr VisitExpr_(const VarNode* op) final {
if (auto it = var_remap_.find(op); it != var_remap_.end()) {
return it->second;
diff --git a/src/tirx/transform/tile_primitive_dispatch.cc
b/src/tirx/transform/tile_primitive_dispatch.cc
index 1f3cf77feb..37a0dfae6a 100644
--- a/src/tirx/transform/tile_primitive_dispatch.cc
+++ b/src/tirx/transform/tile_primitive_dispatch.cc
@@ -121,6 +121,7 @@ class ElectSyncFinder : public StmtExprVisitor {
}
private:
+ using StmtExprVisitor::VisitExpr_;
using StmtExprVisitor::VisitStmt_;
void VisitExpr_(const CallNode* op) final {
@@ -149,6 +150,7 @@ class ScopeIdVarFinder : public StmtExprVisitor {
private:
explicit ScopeIdVarFinder(const std::vector<Var>& vars) : vars_(vars) {}
+ using StmtExprVisitor::VisitExpr_;
using StmtExprVisitor::VisitStmt_;
void VisitExpr_(const VarNode* op) final {
