This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new ee2b3a4ce [FIX] Reject unsupported CUDA vector expressions during
codegen (#84)
ee2b3a4ce is described below
commit ee2b3a4ce350ef707b064c88afdc81c9a1295f6f
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 26 09:56:34 2026 -0400
[FIX] Reject unsupported CUDA vector expressions during codegen (#84)
CUDA codegen now reports unsupported vector types, broadcast
constructors and operators before handing source to nvcc. The checks
stay at the existing emitters and preserve supported packed storage,
casts and arithmetic paths. Adds one focused regression for unsupported
FP8 vector addition.
---
src/backend/cuda/codegen/codegen_cuda.cc | 33 +++++++++++++++++++++++++++-----
src/backend/cuda/codegen/codegen_cuda.h | 1 +
tests/python/tirx/codegen/test_cuda.py | 15 +++++++++++++++
3 files changed, 44 insertions(+), 5 deletions(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index f0227c604..366aa0d52 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -575,7 +575,7 @@ void CodeGenCUDA::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*)
} else {
fail = true;
}
- return;
+ if (!fail) return;
} else if (IsFloat4(t)) {
codegen_tags_.insert("fp4");
if (lanes <= 4) {
@@ -583,7 +583,7 @@ void CodeGenCUDA::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*)
} else {
fail = true;
}
- return;
+ if (!fail) return;
} else if (t == PrimType::Bool()) {
os << "bool";
return;
@@ -731,8 +731,11 @@ void CodeGenCUDA::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*)
os << "longlong3";
} else if (lanes == 4) {
os << "longlong4";
+ } else {
+ fail = true;
}
- return;
+ if (!fail) return;
+ break;
}
default:
@@ -747,7 +750,7 @@ void CodeGenCUDA::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*)
return;
}
}
- TVM_FFI_THROW(InternalError) << "Cannot convert type " << t << " to CUDA
type";
+ TVM_FFI_THROW(RuntimeError) << "CUDA codegen does not support type " <<
t->dtype;
}
void CodeGenCUDA::PrintVecConstructor(const PrimType& t, std::ostream& os) {
@@ -845,6 +848,10 @@ void CodeGenCUDA::PrintVecBinaryOpNanPreserving(const
std::string& op, const Pri
void CodeGenCUDA::PrintVecBinaryOp(const std::string& op, const PrimType& t,
PrimExpr lhs,
PrimExpr rhs, std::ostream& os) { //
NOLINT(*)
+ if (IsFloat8(lhs.ty())) {
+ TVM_FFI_THROW(RuntimeError) << "CUDA codegen does not support '" << op <<
"' for "
+ << lhs.ty()->dtype;
+ }
// Declare the result.
std::string sret = name_supply_->FreshName("_");
this->PrintIndent();
@@ -1516,6 +1523,17 @@ void CodeGenCUDA::Dispatch_(const prim::RampObj* op,
std::ostream& os) {
os << sret;
}
+void CodeGenCUDA::Dispatch_(const prim::BitwiseNotObj* op, std::ostream& os) {
+ PrimType t = op->a.ty();
+ // Packed integers represented by a scalar still support the C bitwise
operator.
+ bool packed_scalar = ((t.bits() == 1 || t.bits() == 4) && t.bits() *
t.lanes() <= 32) ||
+ (t.bits() == 8 && t.lanes() == 4);
+ if (t.lanes() > 1 && !t.MatchesCode(DLDataTypeCode::kDLBool) &&
!packed_scalar) {
+ TVM_FFI_THROW(RuntimeError) << "CUDA codegen does not support '~' for " <<
t->dtype;
+ }
+ CodeGenC::Dispatch_(op, os);
+}
+
void CodeGenCUDA::Dispatch_(const prim::BroadcastObj* op, std::ostream& os) {
// NOLINT(*)
PrimType op_ty = op->ty.as_or_throw<PrimType>();
int lanes = op_ty.lanes();
@@ -1574,7 +1592,9 @@ 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);
+ if (lanes != 2 && lanes != 4) {
+ TVM_FFI_THROW(RuntimeError) << "CUDA codegen does not support Broadcast
for " << op_ty->dtype;
+ }
std::string v = PrintExprToSSA(op->value);
// Implicit conversion from float back to fp8
PrintType(op_ty, os);
@@ -1631,6 +1651,9 @@ void CodeGenCUDA::Dispatch_(const prim::BroadcastObj* op,
std::ostream& os) { /
}
}
+ if (lanes > 4) {
+ TVM_FFI_THROW(RuntimeError) << "CUDA codegen does not support Broadcast
for " << op_ty->dtype;
+ }
std::string v = PrintExprToSSA(op->value);
PrintVecConstructor(op_ty, os);
os << '(';
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index e502d3c0f..d09496e22 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -86,6 +86,7 @@ class CodeGenCUDA final : public CodeGenC {
void Dispatch_(const FloatImmObj* op, std::ostream& os) final;
void Dispatch_(const CallObj* op, std::ostream& os) final;
void Dispatch_(const prim::CastObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::BitwiseNotObj* op, std::ostream& os) final;
void Dispatch_(const prim::MinObj* op, std::ostream& os) final; // NOLINT(*)
void Dispatch_(const prim::MaxObj* op, std::ostream& os) final; // NOLINT(*)
template <typename T>
diff --git a/tests/python/tirx/codegen/test_cuda.py
b/tests/python/tirx/codegen/test_cuda.py
index 751a8ff6a..9aabca3d4 100644
--- a/tests/python/tirx/codegen/test_cuda.py
+++ b/tests/python/tirx/codegen/test_cuda.py
@@ -72,6 +72,21 @@ def test_isnan_scalar_and_vector_codegen(dtype):
assert "!=" in source
+def test_unsupported_fp8_vector_arithmetic(monkeypatch):
+ # Source-only codegen must reject the operation before invoking nvcc.
+ monkeypatch.setenv("TVM_TIRX_COMPILE_FORCE_FALLBACK", "1")
+
+ @T.function
+ def func(A: T.Tensor((1,), "float8_e5m2x2"), B: T.Tensor((1,),
"float8_e5m2x2")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ B[0] = A[0] + A[0]
+
+ with pytest.raises(
+ RuntimeError, match=re.escape("CUDA codegen does not support '+' for
float8_e5m2x2")
+ ):
+ _get_source(func)
+
+
def _helper_source(src: str, helper_name: str) -> str:
pattern = rf"__forceinline__ __device__ [^{{;]+ {re.escape(helper_name)}\("
match = re.search(pattern, src)
