This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 30c842daa [Fix] Preserve effectful operand evaluation in C codegen
(#50)
30c842daa is described below
commit 30c842daa6385ca3c1a1f01daa83d454e40ec334
Author: Shushi Hong <[email protected]>
AuthorDate: Sat Oct 24 20:28:12 2026 -0400
[Fix] Preserve effectful operand evaluation in C codegen (#50)
This PR fixes operand evaluation in C-based codegen by preventing
distinct calls with side effects from being merged, evaluating each
operand once when expanding vector operations, and avoiding stale reads
after writes. It introduces `PrintExprToSSA` to make expression printing
and binding explicit, with minimal C host and CUDA regression tests
checking results and evaluation counts.
---
src/backend/cuda/codegen/codegen_cuda.cc | 70 +++++++++++++------------
src/backend/metal/codegen/codegen_metal.cc | 2 +-
src/backend/opencl/codegen/codegen_opencl.cc | 13 +++--
src/backend/webgpu/codegen/codegen_webgpu.cc | 14 ++---
src/codegen/c/codegen_c.cc | 34 +++++++-----
src/codegen/c/codegen_c.h | 10 ++++
src/codegen/c/codegen_c_host.cc | 10 ++--
src/codegen/c/codegen_source_base.cc | 10 ++--
src/codegen/c/codegen_source_base.h | 8 +--
tests/python/backend/cuda/codegen/test_basic.py | 28 ++++++++++
tests/python/codegen/c/test_basic.py | 23 ++++++++
11 files changed, 145 insertions(+), 77 deletions(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 32a91899b..bed6b7d66 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -774,21 +774,11 @@ void CodeGenCUDA::PrintMinMaxNanPreservingImpl(const T*
op, const char* opstr,
this->PrintVecBinaryOpNanPreserving(std::string(opstr), op_ty, op->a,
op->b, cmp, os);
return;
}
- // Bind both operands once (SSAGetID, in the expression's own scope: the
- // same binding the host side and PrintVecBinaryOp use), then reference the
- // temporaries: the expressions must not be re-evaluated per clause. Both
- // are bound before the constant classification below, so a constant NaN
- // lhs still evaluates the rhs and only the select is dropped. A
- // float32/float64 constant comes back from SSAGetID as itself (PrintConst
- // registers its text with MarkConst); a float16/bfloat16 constant is
- // wrapped in a conversion call that is not registered and gets a temporary
- // like any other operand. The bindings live in their own scope so a later
- // statement that prints the same text (e.g.
- // `red_buf[0] = max(red_buf[0], shuffle_down(...))` repeated by warp
- // reduction) does not hit the cache and read the pre-write value.
+ // Evaluate both operands before simplifying the result. Even a constant NaN
+ // lhs must preserve effects in the rhs; repeated clauses use the bound
values.
int ssa_scope = BeginScope();
- std::string va = SSAGetID(PrintExpr(op->a), op->a.ty());
- std::string vb = SSAGetID(PrintExpr(op->b), op->b.ty());
+ std::string va = PrintExprToSSA(op->a);
+ std::string vb = PrintExprToSSA(op->b);
const FloatImmObj* ca = AsFloatImm(op->a);
const FloatImmObj* cb = AsFloatImm(op->b);
if (ca != nullptr && std::isnan(ca->value)) {
@@ -830,8 +820,8 @@ void CodeGenCUDA::PrintVecBinaryOpNanPreserving(const
std::string& op, const Pri
stream << ' ' << sret << ";\n";
int ssa_scope = BeginScope();
{
- std::string vlhs = SSAGetID(PrintExpr(lhs), lhs.ty());
- std::string vrhs = SSAGetID(PrintExpr(rhs), rhs.ty());
+ std::string vlhs = PrintExprToSSA(lhs);
+ std::string vrhs = PrintExprToSSA(rhs);
for (int i = 0, lanes = t.lanes(); i < lanes; ++i) {
std::ostringstream value_temp;
value_temp << "((";
@@ -863,8 +853,8 @@ void CodeGenCUDA::PrintVecBinaryOp(const std::string& op,
const PrimType& t, Pri
int ssa_scope = BeginScope();
{
// Unpack into individual ops.
- std::string vlhs = SSAGetID(PrintExpr(lhs), lhs.ty());
- std::string vrhs = SSAGetID(PrintExpr(rhs), rhs.ty());
+ std::string vlhs = PrintExprToSSA(lhs);
+ std::string vrhs = PrintExprToSSA(rhs);
for (int i = 0, lanes = t.lanes(); i < lanes; ++i) {
std::ostringstream value_temp;
@@ -1087,7 +1077,7 @@ void CodeGenCUDA::Dispatch_(const prim::CastObj* op,
std::ostream& os) {
this->PrintType(target_ty, stream);
stream << ' ' << sret << ";\n";
{
- std::string src = SSAGetID(PrintExpr(op->value), from_ty);
+ std::string src = PrintExprToSSA(op->value);
for (int i = 0, lanes = from_ty.lanes(); i < lanes; ++i) {
std::ostringstream val;
val << "(";
@@ -1136,7 +1126,7 @@ void CodeGenCUDA::PrintCallExtern(Type ret_type,
ffi::String global_symbol,
size_t arg_begin = static_cast<size_t>(skip_first_arg);
for (size_t i = arg_begin; i < args.size(); ++i) {
Expr arg = args[i];
- std::string val = SSAGetID(PrintExpr(arg), arg->ty);
+ std::string val = PrintExprToSSA(arg);
sargs.push_back(std::move(val));
}
@@ -1180,9 +1170,20 @@ void CodeGenCUDA::Dispatch_(const CallObj* op,
std::ostream& os) {
if (op->op.same_as(prim::isnan_op()) &&
op->ty.as_or_throw<PrimType>().lanes() > 1) {
// CUDA vector comparisons do not produce a lane-wise boolean vector.
- // Reuse the scalar comparison path for every lane, including packed half.
+ // Evaluate the input once, then compare each lane with itself.
PrimExpr value = op->args[0].as_or_throw<PrimExpr>();
- PrintVecBinaryOp("!=", op->ty.as_or_throw<PrimType>(), value, value, os);
+ std::string input = PrintExprToSSA(value);
+ PrimType result_ty = op->ty.as_or_throw<PrimType>();
+ std::string result = name_supply_->FreshName("_");
+ PrintIndent();
+ PrintType(result_ty, stream);
+ stream << ' ' << result << ";\n";
+ for (int i = 0; i < result_ty.lanes(); ++i) {
+ std::ostringstream lane;
+ PrintVecElemLoad(input, value.ty(), i, lane);
+ PrintVecElemStore(result, result_ty, i, "(" + lane.str() + " != " +
lane.str() + ")");
+ }
+ os << result;
return;
}
if (auto opt_call_opt = op->op.as<Op>()) {
@@ -1293,7 +1294,7 @@ void CodeGenCUDA::Dispatch_(const CallObj* op,
std::ostream& os) {
if (lanes == 1) {
// The case of lane=1 is same as the normal reinterpret,
// except that we allow the src and dst dtype to have different number
of bits.
- std::string rhs = SSAGetID(PrintExpr(value), src_ty);
+ std::string rhs = PrintExprToSSA(value);
os << "(*(";
this->PrintType(tgt_ty, os);
os << " *)(&(" << rhs << ")))";
@@ -1420,11 +1421,12 @@ void CodeGenCUDA::Dispatch_(const prim::RampObj* op,
std::ostream& os) {
PrimType op_ty = op->ty.as_or_throw<PrimType>();
int lanes = op_ty.lanes();
if (lanes <= 4) {
+ std::string base = PrintExprToSSA(op->base);
+ std::string stride = PrintExprToSSA(op->stride);
PrintVecConstructor(op_ty, os);
os << "(";
for (int i = 0; i < lanes; i++) {
- os << "(" << PrintExpr(op->base) << ")"
- << "+(" << PrintExpr(op->stride) << "*" << i << ")";
+ os << "(" << base << ")+(" << stride << "*" << i << ")";
if (i != lanes - 1) os << ", ";
}
os << ")";
@@ -1439,8 +1441,8 @@ void CodeGenCUDA::Dispatch_(const prim::RampObj* op,
std::ostream& os) {
stream << ' ' << sret << ";\n";
int ssa_scope = BeginScope();
{
- std::string vbase = SSAGetID(PrintExpr(op->base), op->base.ty());
- std::string vstride = SSAGetID(PrintExpr(op->stride), op->stride.ty());
+ std::string vbase = PrintExprToSSA(op->base);
+ std::string vstride = PrintExprToSSA(op->stride);
for (int i = 0; i < lanes; ++i) {
std::ostringstream value_temp;
value_temp << "(" << vbase << ")+(" << vstride << "*" << i << ")";
@@ -1471,7 +1473,7 @@ void CodeGenCUDA::Dispatch_(const prim::BroadcastObj* op,
std::ostream& os) { /
}
if (op_ty.MatchesElementType(DLDataTypeCode::kDLFloat, 16)) {
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
PrintVecConstructor(op_ty, os);
os << '(';
if (lanes <= 4) {
@@ -1490,7 +1492,7 @@ void CodeGenCUDA::Dispatch_(const prim::BroadcastObj* op,
std::ostream& os) { /
}
if (op_ty.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
PrintVecConstructor(op_ty, os);
os << '(';
if (lanes > 4) {
@@ -1510,7 +1512,7 @@ void CodeGenCUDA::Dispatch_(const prim::BroadcastObj* op,
std::ostream& os) { /
if (IsFloat8(op_ty) || IsFloat4(op_ty)) {
TVM_FFI_ICHECK(lanes == 1 || lanes == 2 || lanes == 4);
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
// Implicit conversion from float back to fp8
PrintType(op_ty, os);
os << "(make_float" << lanes << "(";
@@ -1566,7 +1568,7 @@ void CodeGenCUDA::Dispatch_(const prim::BroadcastObj* op,
std::ostream& os) { /
}
}
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
PrintVecConstructor(op_ty, os);
os << '(';
for (int i = 0; i < lanes; ++i) {
@@ -1593,9 +1595,9 @@ void CodeGenCUDA::Dispatch_(const prim::SelectObj* op,
std::ostream& os) {
this->PrintType(op_ty, stream);
stream << ' ' << r_var << ";\n";
{
- std::string c_var = SSAGetID(PrintExpr(op->condition), op_ty);
- std::string t_var = SSAGetID(PrintExpr(op->true_value), op_ty);
- std::string f_var = SSAGetID(PrintExpr(op->false_value), op_ty);
+ std::string c_var = PrintExprToSSA(op->condition);
+ std::string t_var = PrintExprToSSA(op->true_value);
+ std::string f_var = PrintExprToSSA(op->false_value);
// The condition is stored as an ushort vector.
int lanes = op_ty.lanes();
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 07ca9bd7a..43b8297a9 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -449,7 +449,7 @@ void CodeGenMetal::Dispatch_(const prim::SelectObj* op,
std::ostream& os) { //
}
void CodeGenMetal::Dispatch_(const prim::BroadcastObj* op, std::ostream& os) {
// NOLINT(*)
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
PrintType(op->ty.as_or_throw<PrimType>(), os);
os << "(";
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc
b/src/backend/opencl/codegen/codegen_opencl.cc
index eb23638b2..9b3a577bd 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -516,7 +516,8 @@ void CodeGenOpenCL::Dispatch_(const CallObj* op,
std::ostream& os) {
std::stringstream ss;
this->PrintExpr(op->args[5].as_or_throw<PrimExpr>(), ss);
std::string value;
- value = this->SSAGetID(ss.str(), buffer_type.WithLanes(channel_size /
buffer_type.bits()));
+ value = this->SSAGetID(ss.str(), buffer_type.WithLanes(channel_size /
buffer_type.bits()),
+ SideEffect(op->args[5]) == CallEffectKind::kPure);
if (channel_size == 64) {
os << "write_imageh(";
} else if (channel_size == 128) {
@@ -571,7 +572,8 @@ void CodeGenOpenCL::Dispatch_(const CallObj* op,
std::ostream& os) {
this->PrintExpr(IntImm::Int32(0), ss);
ss << "))))";
- std::string rhs = SSAGetID(ss.str(), op_ty.WithLanes(data_lanes));
+ // An earlier image read may precede a write to the same texture.
+ std::string rhs = SSAGetID(ss.str(), op_ty.WithLanes(data_lanes),
/*can_reuse=*/false);
if (auto ramp = op->args.back().as<prim::RampObj>()) {
const auto* base = ramp->base.as<IntImmObj>();
const auto* lanes_imm = ramp->lanes.as<IntImmObj>();
@@ -622,7 +624,7 @@ void CodeGenOpenCL::Dispatch_(const CallObj* op,
std::ostream& os) {
}
void CodeGenOpenCL::Dispatch_(const prim::BroadcastObj* op, std::ostream& os)
{ // NOLINT(*)
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
os << "((";
PrintType(op->ty.as_or_throw<PrimType>(), os);
@@ -635,13 +637,14 @@ void CodeGenOpenCL::Dispatch_(const prim::BroadcastObj*
op, std::ostream& os) {
}
void CodeGenOpenCL::Dispatch_(const prim::RampObj* op, std::ostream& os) { //
NOLINT(*)
+ std::string base = PrintExprToSSA(op->base);
+ std::string stride = PrintExprToSSA(op->stride);
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++) {
- os << "(" << PrintExpr(op->base) << ")"
- << "+(" << PrintExpr(op->stride) << "*" << i << ")";
+ os << "(" << base << ")+(" << stride << "*" << i << ")";
if (i != lanes - 1) os << ", ";
}
os << "))";
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index ffb60267b..e18fb9ecc 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -384,7 +384,7 @@ void CodeGenWebGPU::BindThreadIndex(const PrimVar& var,
const ffi::String& threa
// We allow runtime to spread the load out to blockIdx.z so it can be a
large number.
os << "(blockIdx.z * gridDim.x + blockIdx.x)";
std::string tidx = os.str();
- std::string aggregated_bidx = SSAGetID(os.str(), var.ty());
+ std::string aggregated_bidx = SSAGetID(os.str(), var.ty(),
/*can_reuse=*/true);
var_idmap_[var.get()] = aggregated_bidx;
} else {
os << "(" << thread_tag << ")";
@@ -472,7 +472,7 @@ void CodeGenWebGPU::PrintVecElemStore(const std::string&
vec, const PrimType& t,
}
void CodeGenWebGPU::Dispatch_(const prim::BroadcastObj* op, std::ostream& os)
{ // NOLINT(*)
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
PrintType(op->ty.as_or_throw<PrimType>(), os);
os << "(";
@@ -660,7 +660,7 @@ void CodeGenWebGPU::Dispatch_(const TensorLoadObj* op,
std::ostream& os) { // N
tvm::PVar<PrimExpr> base;
if (tvm::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());
+ std::string base_vid = PrintExprToSSA(base.Eval());
PrintType(element_ty.WithLanes(value_ty.lanes()), os);
os << "(";
for (int i = 0; i < lanes; ++i) {
@@ -670,7 +670,7 @@ void CodeGenWebGPU::Dispatch_(const TensorLoadObj* op,
std::ostream& os) { // N
os << ")";
} else {
// vec3<f32>(buf[index[0]], buf[index[1]], buf[index[2]]);
- std::string index_vid = SSAGetID(PrintExpr(index), index.ty());
+ std::string index_vid = PrintExprToSSA(index);
PrintType(element_ty.WithLanes(value_ty.lanes()), os);
os << "(";
for (int i = 0; i < lanes; ++i) {
@@ -742,12 +742,12 @@ void CodeGenWebGPU::Dispatch_(const TensorStoreObj* op) {
TVM_FFI_ICHECK_EQ(element_ty.lanes(), 1) << "Can only vector load scalar
array";
TVM_FFI_ICHECK(value_ty.WithLanes(1) == element_ty)
<< "WebGPU vector stire requires base type to match";
- std::string value_vid = PrintExpr(op->value);
+ std::string value_vid = PrintExprToSSA(op->value);
tvm::PVar<PrimExpr> base;
if (tvm::ramp(base, 1, value_ty.lanes()).Match(index)) {
// buf[base + 0] = value[0]
// buf[base + 1] = value[1]
- std::string base_vid = SSAGetID(PrintExpr(base.Eval()),
base.Eval().ty());
+ std::string base_vid = PrintExprToSSA(base.Eval());
for (int i = 0; i < value_ty.lanes(); ++i) {
this->PrintIndent();
stream << buffer_vid << "[" << base_vid << " + " << i << "] = " <<
value_vid << "[" << i
@@ -756,7 +756,7 @@ void CodeGenWebGPU::Dispatch_(const TensorStoreObj* op) {
} else {
// buf[index[0]] = value[0]
// buf[index[1]] = value[1]
- std::string index_vid = SSAGetID(PrintExpr(index), index.ty());
+ std::string index_vid = PrintExprToSSA(index);
for (int i = 0; i < value_ty.lanes(); ++i) {
this->PrintIndent();
stream << buffer_vid << "[" << index_vid << "[" << i << "]] = " <<
value_vid << "[" << i
diff --git a/src/codegen/c/codegen_c.cc b/src/codegen/c/codegen_c.cc
index 430dff054..bbe009df6 100644
--- a/src/codegen/c/codegen_c.cc
+++ b/src/codegen/c/codegen_c.cc
@@ -42,6 +42,7 @@
#include <cctype>
#include <iomanip>
#include <limits>
+#include <utility>
#include "codegen_params.h"
@@ -225,12 +226,19 @@ void CodeGenC::PrintExpr(const PrimExpr& n, std::ostream&
os) { // NOLINT(*)
if (print_ssa_form_) {
std::ostringstream temp;
Dispatch(n, temp);
- os << SSAGetID(temp.str(), n.ty());
+ os << SSAGetID(temp.str(), n.ty(), SideEffect(n) == CallEffectKind::kPure);
} else {
Dispatch(n, os);
}
}
+std::string CodeGenC::PrintExprToSSA(const Expr& expr) {
+ std::string source = PrintExpr(expr);
+ bool can_reuse = SideEffect(expr) == CallEffectKind::kPure;
+ // Emit a declaration if needed; the caller prints the returned variable or
literal.
+ return CodeGenSourceBase::SSAGetID(std::move(source), expr->ty, can_reuse);
+}
+
void CodeGenC::PrintExpr(const Expr& n, std::ostream& os) { // NOLINT(*)
if (auto prim = n.as<PrimExpr>()) {
PrintExpr(prim.value(), os);
@@ -896,17 +904,14 @@ void CodeGenC::Dispatch_(const CallObj* op, std::ostream&
os) { // NOLINT(*)
source_dtype.lanes() * source_dtype.bits())
<< "reinterpret expects source and target to have the same number of
bits";
int ssa_scope = BeginScope();
- std::string rhs = SSAGetID(PrintExpr(op->args[0]), source_dtype);
+ std::string rhs = PrintExprToSSA(op->args[0]);
os << "(*(";
this->PrintType(target_dtype, os);
os << " *)(&(" << rhs << ")))";
EndScope(ssa_scope);
} else if (op->op.same_as(prim::isnan_op())) {
- os << "(";
- this->PrintExpr(op->args[0], os);
- os << " != ";
- this->PrintExpr(op->args[0], os);
- os << ")";
+ std::string value = PrintExprToSSA(op->args[0]);
+ os << "(" << value << " != " << value << ")";
} else if (op->op.same_as(tirx::gpu_thread_invariant_op())) {
os << "(";
this->PrintExpr(op->args[0], os);
@@ -1040,7 +1045,7 @@ void CodeGenC::Dispatch_(const TensorLoadObj* op,
std::ostream& os) { // NOLINT
HandleVolatileLoads(ref, op, os);
} else {
std::ostringstream svalue_expr;
- std::string sindex = SSAGetID(PrintExpr(index), index.ty());
+ std::string sindex = PrintExprToSSA(index);
std::string vid = GetVarID(buffer_var.get());
PrimType elem_type = value_ty.WithLanes(1);
for (int i = 0; i < lanes; ++i) {
@@ -1095,8 +1100,8 @@ void CodeGenC::Dispatch_(const TensorStoreObj* op) {
int vec_scope = BeginScope();
// store elements separately
- std::string index = SSAGetID(PrintExpr(index_expr), index_expr.ty());
- std::string value = SSAGetID(PrintExpr(op->value), op->value.ty());
+ std::string index = PrintExprToSSA(index_expr);
+ std::string value = PrintExprToSSA(op->value);
std::string vid = GetVarID(buffer_var.get());
for (int i = 0; i < value_ty.lanes(); ++i) {
this->PrintIndent();
@@ -1160,14 +1165,15 @@ void CodeGenC::Dispatch_(const prim::LetObj* op,
std::ostream& os) { // NOLINT(
}
void CodeGenC::Dispatch_(const prim::RampObj* op, std::ostream& os) { //
NOLINT(*)
+ std::string base = PrintExprToSSA(op->base);
+ std::string stride = PrintExprToSSA(op->stride);
// NOTE: C have comma expression so cannot use (int2)(v0, v1)
// instead should use int2(v0, v1)
PrintType(op->ty.as_or_throw<PrimType>(), os);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
os << "(";
for (int i = 0; i < lanes; i++) {
- os << "(" << PrintExpr(op->base) << ")"
- << "+(" << PrintExpr(op->stride) << "*" << i << ")";
+ os << "(" << base << ")+(" << stride << "*" << i << ")";
if (i != lanes - 1) os << ", ";
}
os << ")";
@@ -1188,7 +1194,7 @@ void CodeGenC::Dispatch_(const prim::ShuffleObj* op,
std::ostream& os) { // NOL
// print each elements
if (op->vectors.size() > 1) {
for (const PrimExpr& vec : op->vectors) {
- std::string vec_value = this->PrintExpr(vec);
+ std::string vec_value = PrintExprToSSA(vec);
if (vec.ty().lanes() == 1) {
concat_vec.push_back(vec_value);
} else {
@@ -1203,7 +1209,7 @@ void CodeGenC::Dispatch_(const prim::ShuffleObj* op,
std::ostream& os) { // NOL
}
} else {
// Extract elements from a single vector-type value.
- std::string vec_value = "(" + this->PrintExpr(op->vectors[0]) + ")";
+ std::string vec_value = "(" + PrintExprToSSA(op->vectors[0]) + ")";
if (op->vectors[0].ty().lanes() == 1) {
concat_vec.push_back(vec_value);
} else {
diff --git a/src/codegen/c/codegen_c.h b/src/codegen/c/codegen_c.h
index d3b8212b5..a44ed8755 100644
--- a/src/codegen/c/codegen_c.h
+++ b/src/codegen/c/codegen_c.h
@@ -266,6 +266,16 @@ class CodeGenC : public tvm::ExprFunctor<void(const Expr&,
std::ostream&)>,
}
protected:
+ /*!
+ * \brief Print an expression, bind its value when needed, and return the
variable or literal.
+ * Only pure expressions may reuse an earlier binding. Reads and calls with
+ * effects are evaluated separately on each invocation; repeat the returned
+ * variable when one operand is needed more than once.
+ * \param expr The expression to print.
+ * \return The variable or literal holding the expression's value.
+ */
+ std::string PrintExprToSSA(const Expr& expr);
+
/*! \brief Print a C string literal with proper escaping of special chars. */
void PrintEscapedCString(const std::string& str, std::ostream& os);
// Print reference to struct location
diff --git a/src/codegen/c/codegen_c_host.cc b/src/codegen/c/codegen_c_host.cc
index 92b367fd5..55a917f0a 100644
--- a/src/codegen/c/codegen_c_host.cc
+++ b/src/codegen/c/codegen_c_host.cc
@@ -197,7 +197,7 @@ void CodeGenCHost::PrintType(const PrimType& type,
std::ostream& os) { // NOLIN
}
void CodeGenCHost::Dispatch_(const prim::BroadcastObj* op, std::ostream& os) {
// NOLINT(*)
- std::string v = PrintExpr(op->value);
+ std::string v = PrintExprToSSA(op->value);
int lanes = op->ty.as_or_throw<PrimType>().lanes();
os << "((";
PrintType(op->ty.as_or_throw<PrimType>(), os);
@@ -382,12 +382,8 @@ void CodeGenCHost::Dispatch_(const prim::MaxObj* op,
std::ostream& os) { // NOL
void CodeGenCHost::PrintMinMax(const PrimExpr& lhs, const PrimExpr& rhs, bool
is_min,
std::ostream& os) { // NOLINT(*)
- std::ostringstream temp_lhs;
- Dispatch(lhs, temp_lhs);
- std::string lhs_value = "(" + SSAGetID(temp_lhs.str(), lhs.ty()) + ")";
- std::ostringstream temp_rhs;
- Dispatch(rhs, temp_rhs);
- std::string rhs_value = "(" + SSAGetID(temp_rhs.str(), rhs.ty()) + ")";
+ std::string lhs_value = "(" + PrintExprToSSA(lhs) + ")";
+ std::string rhs_value = "(" + PrintExprToSSA(rhs) + ")";
PrimType dtype = lhs.ty().as_or_throw<PrimType>();
FloatNaNStatus lhs_nan_status = GetFloatNaNStatus(lhs);
FloatNaNStatus rhs_nan_status = GetFloatNaNStatus(rhs);
diff --git a/src/codegen/c/codegen_source_base.cc
b/src/codegen/c/codegen_source_base.cc
index 47096b0a8..1714f51c9 100644
--- a/src/codegen/c/codegen_source_base.cc
+++ b/src/codegen/c/codegen_source_base.cc
@@ -34,13 +34,13 @@ void CodeGenSourceBase::ClearFuncState() {
scope_mark_.clear();
}
-std::string CodeGenSourceBase::SSAGetID(std::string src, const Type& t) {
+std::string CodeGenSourceBase::SSAGetID(std::string src, const Type& t, bool
can_reuse) {
TVM_FFI_ICHECK(t.as<PrimTypeObj>() || t.as<PtrTypeObj>())
<< "Cannot assign an SSA value of type " << t;
if (name_supply_->ContainsName(src)) return src;
- auto it = ssa_assign_map_.find(src);
- if (it != ssa_assign_map_.end()) {
- if (scope_mark_.at(it->second.scope_id)) {
+ if (can_reuse) {
+ auto it = ssa_assign_map_.find(src);
+ if (it != ssa_assign_map_.end() && scope_mark_.at(it->second.scope_id)) {
return it->second.vid;
}
}
@@ -48,7 +48,7 @@ std::string CodeGenSourceBase::SSAGetID(std::string src,
const Type& t) {
// use v_ prefix so it works for most systems
e.vid = name_supply_->FreshName("v_");
e.scope_id = static_cast<int>(scope_mark_.size() - 1);
- ssa_assign_map_[src] = e;
+ if (can_reuse) ssa_assign_map_[src] = e;
this->PrintIndent();
PrintSSAAssign(e.vid, src, t);
return e.vid;
diff --git a/src/codegen/c/codegen_source_base.h
b/src/codegen/c/codegen_source_base.h
index 3f09f8150..4ed1767ab 100644
--- a/src/codegen/c/codegen_source_base.h
+++ b/src/codegen/c/codegen_source_base.h
@@ -99,12 +99,12 @@ class CodeGenSourceBase {
*/
std::string GetVarID(const tvm::VarObj* v) const;
/*!
- * \brief Get the SSA ID corresponds to src
- * If necessary, generate new assignment
- * \param src The source expression
+ * \brief Bind source text to an SSA variable, emitting an assignment if
needed.
+ * \param src The source expression.
* \param t The type of the expression.
+ * \param can_reuse Whether an earlier binding of this text may be reused.
*/
- std::string SSAGetID(std::string src, const Type& t);
+ std::string SSAGetID(std::string src, const Type& t, bool can_reuse);
/*!
* \brief mark the beginning of a new scope
* \return The scope id.
diff --git a/tests/python/backend/cuda/codegen/test_basic.py
b/tests/python/backend/cuda/codegen/test_basic.py
index cc6517ec0..5fa70ba54 100644
--- a/tests/python/backend/cuda/codegen/test_basic.py
+++ b/tests/python/backend/cuda/codegen/test_basic.py
@@ -924,6 +924,34 @@ def test_min_max_nan_preserving_cuda(dtype, form, *,
backend_config):
tvm_tirx.testing.run_with_gpu_lock(run_and_check)
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
+def test_atomic_operand_evaluation():
+ @T.function
+ def main(P: T.Tensor((1,), "float32"), O: T.Tensor((1,), "float32x4")):
+ with T.launch_thread("threadIdx.x", 1):
+ # Two independent atomics run once each; broadcast copies their
max.
+ O[0] = T.Broadcast(
+ T.max(
+ T.cuda.atomic_add(T.address_of(P[0]), T.float32(1)),
+ T.cuda.atomic_add(T.address_of(P[0]), T.float32(1)),
+ ),
+ 4,
+ )
+
+ executable = tvm_tirx.tirx.compile(main, target="cuda")
+
+ def run_and_check():
+ dev = tvm_tirx.cuda(0)
+ counter = tvm_tirx.runtime.tensor(np.zeros(1, dtype="float32"), dev)
+ output = tvm_tirx.runtime.empty((1,), "float32x4", dev)
+ executable(counter, output)
+ np.testing.assert_array_equal(output.numpy(), [[1] * 4])
+ np.testing.assert_array_equal(counter.numpy(), [2])
+
+ tvm_tirx.testing.run_with_gpu_lock(run_and_check)
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_round_ties_to_even(*, backend_config):
diff --git a/tests/python/codegen/c/test_basic.py
b/tests/python/codegen/c/test_basic.py
index 2ab26b80f..e0aec80d3 100644
--- a/tests/python/codegen/c/test_basic.py
+++ b/tests/python/codegen/c/test_basic.py
@@ -190,6 +190,29 @@ def test_round():
check_c()
+def test_min_max_reads_after_store():
+ @I.ir_module
+ class Module:
+ # Keep both reads in one SSA cache scope, without public ABI checks.
+ @T.function(private=True)
+ def check(p: T.handle("float32"), o: T.handle("float32")):
+ P = T.decl_tensor(p, ty_args=[T.Tensor((1,), "float32", data=p)])
+ Out = T.decl_tensor(o, ty_args=[T.Tensor((2,), "float32", data=o)])
+ Out[0] = T.max(P[0], T.float32(0))
+ P[0] = 8.0
+ Out[1] = T.max(P[0], T.float32(0))
+
+ @T.function
+ def main(P: T.Tensor((1,), "float32"), Out: T.Tensor((2,), "float32")):
+ Module.check(P.data, Out.data)
+
+ executable = tvm_tirx.tirx.compile(Module, target="c")
+ data = tvm_tirx.runtime.tensor(np.ones(1, dtype="float32"))
+ output = tvm_tirx.runtime.empty((2,), "float32")
+ executable(data, output)
+ np.testing.assert_array_equal(output.numpy(), [1, 8])
+
+
def test_subroutine_call():
@I.ir_module
class Module:
