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 7e316ee47 Reject unsupported scalar FP8 CUDA expressions (#88)
7e316ee47 is described below
commit 7e316ee471981df8545550d943d0af4c2a44d2f3
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 26 12:52:00 2026 -0400
Reject unsupported scalar FP8 CUDA expressions (#88)
CUDA scalar arithmetic, comparisons and min/max now reject unsupported
FP8 operands before emitting source. A small operator-templated check
covers scalar loops produced from vectorized input while preserving
existing FP16 arithmetic and FP8 memory/conversion paths.
Extend the existing FP8 regression to cover the scalar-loop case.
---
src/backend/cuda/codegen/codegen_cuda.cc | 65 ++++++++++++++++++++++++++++++++
src/backend/cuda/codegen/codegen_cuda.h | 11 ++++++
tests/python/tirx/codegen/test_cuda.py | 13 ++++++-
3 files changed, 88 insertions(+), 1 deletion(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 366aa0d52..6d4f6b228 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -61,6 +61,14 @@ TVM_FFI_INLINE bool IsFloat8(const PrimType& ty) {
DLDataTypeCode::kDLFloat8_e5m2fnuz,
DLDataTypeCode::kDLFloat8_e8m0fnu);
}
+template <typename Op>
+void CheckBinaryOpSupported(const Op* op) {
+ if (IsFloat8(op->a.ty())) {
+ TVM_FFI_THROW(RuntimeError) << "CUDA codegen does not support " <<
Op::_type_key << " for "
+ << op->a.ty()->dtype;
+ }
+}
+
TVM_FFI_INLINE bool IsFloat6(const PrimType& ty) {
return ty.MatchesCode(DLDataTypeCode::kDLFloat6_e2m3fn,
DLDataTypeCode::kDLFloat6_e3m2fn);
}
@@ -806,11 +814,68 @@ void CodeGenCUDA::PrintMinMaxNanPreservingImpl(const T*
op, const char* opstr,
EndScope(ssa_scope);
}
+void CodeGenCUDA::Dispatch_(const prim::AddObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::SubObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::MulObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::DivObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::ModObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::EQObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::NEObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::LTObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::LEObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::GTObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::GEObj* op, std::ostream& os) {
+ CheckBinaryOpSupported(op);
+ CodeGenC::Dispatch_(op, os);
+}
+
void CodeGenCUDA::Dispatch_(const prim::MinObj* op, std::ostream& os) { //
NOLINT(*)
+ CheckBinaryOpSupported(op);
this->PrintMinMaxNanPreservingImpl(op, "min", os);
}
void CodeGenCUDA::Dispatch_(const prim::MaxObj* op, std::ostream& os) { //
NOLINT(*)
+ CheckBinaryOpSupported(op);
this->PrintMinMaxNanPreservingImpl(op, "max", os);
}
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index d09496e22..5cd88137b 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -87,6 +87,17 @@ class CodeGenCUDA final : public CodeGenC {
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::AddObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::SubObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::MulObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::DivObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::ModObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::EQObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::NEObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::LTObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::LEObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::GTObj* op, std::ostream& os) final;
+ void Dispatch_(const prim::GEObj* 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 9aabca3d4..1d25915c0 100644
--- a/tests/python/tirx/codegen/test_cuda.py
+++ b/tests/python/tirx/codegen/test_cuda.py
@@ -82,10 +82,21 @@ def test_unsupported_fp8_vector_arithmetic(monkeypatch):
B[0] = A[0] + A[0]
with pytest.raises(
- RuntimeError, match=re.escape("CUDA codegen does not support '+' for
float8_e5m2x2")
+ RuntimeError, match=re.escape("CUDA codegen does not support prim.Add
for float8_e5m2x2")
):
_get_source(func)
+ @T.function
+ def vectorized_loop(A: T.Tensor((2,), "float8_e5m2"), B: T.Tensor((2,),
"float8_e5m2")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
+ for i in T.vectorized(2):
+ B[i] = A[i] + A[i]
+
+ with pytest.raises(
+ RuntimeError, match=re.escape("CUDA codegen does not support prim.Add
for float8_e5m2")
+ ):
+ _get_source(vectorized_loop)
+
def _helper_source(src: str, helper_name: str) -> str:
pattern = rf"__forceinline__ __device__ [^{{;]+ {re.escape(helper_name)}\("
