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 642ff3fc7 [FIX][CUDA] Correct vector storage, alignment, and FP8 casts
(#68)
642ff3fc7 is described below
commit 642ff3fc7d8aa4ad4bd66ddc28b2d1e886f47b5d
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 25 14:30:57 2026 -0400
[FIX][CUDA] Correct vector storage, alignment, and FP8 casts (#68)
---
src/backend/cuda/codegen/codegen_cuda.cc | 71 ++++++++++++++++++++--
src/backend/cuda/codegen/codegen_cuda.h | 3 +
tests/python/backend/cuda/codegen/test_fp8.py | 20 ++++++
.../backend/cuda/codegen/test_vector_storage.py | 62 +++++++++++++++++++
tests/python/codegen/gpu_common/test_intrinsics.py | 5 +-
5 files changed, 155 insertions(+), 6 deletions(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index bed6b7d66..f0227c604 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -589,10 +589,10 @@ void CodeGenCUDA::PrintType(const PrimType& t,
std::ostream& os) { // NOLINT(*)
return;
} else if (t.MatchesCode(DLDataTypeCode::kDLBool) && lanes > 1) {
// CUDA does not support bool vectors.
- // Use ushort vectors to represent instead.
+ // Use one byte per element, matching scalar bool tensor storage.
int n = lanes;
if (n <= 4) {
- os << "ushort" << n;
+ os << "uchar" << n;
return;
}
} else if (t.MatchesCode(DLDataTypeCode::kDLUInt, DLDataTypeCode::kDLInt)) {
@@ -878,6 +878,36 @@ void CodeGenCUDA::PrintVecBinaryOp(const std::string& op,
const PrimType& t, Pri
os << sret;
}
+std::string CodeGenCUDA::GetBufferRef(const PrimType& t, const VarObj* buffer,
PrimExpr index) {
+ if (!t.IsScalar()) {
+ // CUDA vector representations align to their power-of-two storage
component,
+ // capped at 16 bytes. Bool and FP6 wrappers use one byte per scalar
element.
+ int bits = t.MatchesCode(DLDataTypeCode::kDLBool) || IsFloat6(t) ? 8 :
t.bits();
+ int bytes = (bits * t.lanes() + 7) / 8;
+ int required_alignment = std::min(16, bytes & -bytes);
+ int buffer_alignment =
ffi::GetRef<Var>(buffer).as_or_throw<TensorVar>().ty()->data_alignment;
+ if (buffer_alignment % required_alignment != 0) {
+ TVM_FFI_THROW(ValueError) << "Cannot prove alignment for CUDA vector
access of " << t
+ << ": buffer alignment " << buffer_alignment
+ << " is not a multiple of " <<
required_alignment << " bytes";
+ }
+ }
+ return CodeGenC::GetBufferRef(t, buffer, index);
+}
+
+void CodeGenCUDA::PrintVecStore(const VarObj* buffer, const PrimType& t,
PrimExpr base,
+ const std::string& value) {
+ // Match the conservative offset-alignment proof used for vector loads.
+ // A typed CUDA vector dereference must not reach an unaligned address.
+ sym::ModularSet alignment = sym::Analyzer()->modular_set(base);
+ if (alignment->coeff % t.lanes() != 0 || alignment->base % t.lanes() != 0) {
+ TVM_FFI_THROW(ValueError) << "Cannot prove alignment for CUDA vector store
of " << t
+ << ": base index must be a multiple of " <<
t.lanes() << ", but got "
+ << base;
+ }
+ CodeGenC::PrintVecStore(buffer, t, base, value);
+}
+
void CodeGenCUDA::PrintVecElemLoad(const std::string& vec, const PrimType& t,
int i,
std::ostream& os) { // NOLINT(*)
int lanes = t.lanes();
@@ -1054,6 +1084,39 @@ void CodeGenCUDA::Dispatch_(const prim::CastObj* op,
std::ostream& os) {
PrimType target_ty = op->ty.as_or_throw<PrimType>();
TVM_FFI_ICHECK_EQ(target_ty.lanes(), from_ty.lanes());
+ // Distinct FP8 formats require a numeric conversion even when their CUDA
+ // vector storage types coincide. CUDA's FP8 wrappers do not provide direct
+ // cross-format constructors; float is an exact intermediate for FP8 values.
+ if (IsFloat8(from_ty) && IsFloat8(target_ty) && from_ty.code() !=
target_ty.code()) {
+ std::string source = PrintExprToSSA(op->value);
+ PrimType from_scalar = from_ty.WithLanes(1);
+ PrimType target_scalar = target_ty.WithLanes(1);
+ if (from_ty.IsScalar()) {
+ PrintType(target_scalar, os);
+ os << "(float(" << source << "))";
+ return;
+ }
+ std::string result = name_supply_->FreshName("_");
+ PrintIndent();
+ PrintType(target_ty, stream);
+ stream << " " << result << ";\n";
+ for (int i = 0; i < from_ty.lanes(); ++i) {
+ std::string lane = name_supply_->FreshName("_");
+ PrintIndent();
+ PrintType(from_scalar, stream);
+ stream << " " << lane << ";\n";
+ PrintIndent();
+ stream << lane << ".__x = reinterpret_cast<const unsigned char*>(&" <<
source << ")[" << i
+ << "];\n";
+ PrintIndent();
+ stream << "reinterpret_cast<unsigned char*>(&" << result << ")[" << i <<
"] = ";
+ PrintType(target_scalar, stream);
+ stream << "(float(" << lane << ")).__x;\n";
+ }
+ os << result;
+ return;
+ }
+
// Emit simple C-style type conversion.
if (from_ty.IsScalar()) return CodeGenC::Dispatch_(op, os);
@@ -1599,9 +1662,9 @@ void CodeGenCUDA::Dispatch_(const prim::SelectObj* op,
std::ostream& os) {
std::string t_var = PrintExprToSSA(op->true_value);
std::string f_var = PrintExprToSSA(op->false_value);
- // The condition is stored as an ushort vector.
+ // Use the condition's bool-vector representation.
int lanes = op_ty.lanes();
- PrimType memory_ty = PrimType::UInt(16, lanes);
+ PrimType memory_ty = op->condition.ty();
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 f8292bab0..e502d3c0f 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -71,6 +71,9 @@ class CodeGenCUDA final : public CodeGenC {
std::ostream& os) final; // NOLINT(*)
void PrintVecElemStore(const std::string& vec, const PrimType& t, int i,
const std::string& value) final;
+ void PrintVecStore(const VarObj* buffer, const PrimType& t, PrimExpr base,
+ const std::string& value) final;
+ std::string GetBufferRef(const PrimType& t, const VarObj* buffer, PrimExpr
index) final;
void BindThreadIndex(const PrimVar& var, const ffi::String& thread_tag)
final; // NOLINT(*)
void PrintVecElemLoadExpr(const PrimType& t, int i, const std::string& value,
std::ostream& os) final;
diff --git a/tests/python/backend/cuda/codegen/test_fp8.py
b/tests/python/backend/cuda/codegen/test_fp8.py
index 59b52b60c..473e75e2f 100644
--- a/tests/python/backend/cuda/codegen/test_fp8.py
+++ b/tests/python/backend/cuda/codegen/test_fp8.py
@@ -609,5 +609,25 @@ def test_fp8_fp16_bf16_vectorize_arith(vec_length, dtype):
tvm_tirx.testing.run_with_gpu_lock(run_and_check)
[email protected]
[email protected](not env.has_cuda_compute(9), reason="need CUDA compute >=
9.0")
+def test_cross_format_fp8_conversion():
+ @T.function
+ def main(A: T.Tensor((8,), "float8_e4m3fn"), B: T.Tensor((8,),
"float8_e5m2")):
+ T.func_attr({"tirx.noalias": True})
+ with T.launch_thread("blockIdx.x", 1):
+ with T.launch_thread("threadIdx.x", 1):
+ B[T.Ramp(0, 1, 8)] = T.Cast("float8_e5m2x8", A[T.Ramp(0, 1,
8)])
+
+ compiled = tvm_tirx.tirx.compile(main, target="cuda")
+ values = np.array([0.0, 1.0, -1.0, 2.5, -2.5, 0.125, -0.125, 240.0],
dtype="float8_e4m3fn")
+ device = tvm_tirx.cuda()
+ output = tvm_tirx.runtime.empty((8,), "float8_e5m2", device)
+ compiled(tvm_tirx.runtime.tensor(values, device), output)
+ np.testing.assert_array_equal(
+ output.numpy().astype("float32"),
values.astype("float8_e5m2").astype("float32")
+ )
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
diff --git a/tests/python/backend/cuda/codegen/test_vector_storage.py
b/tests/python/backend/cuda/codegen/test_vector_storage.py
new file mode 100644
index 000000000..fad03231f
--- /dev/null
+++ b/tests/python/backend/cuda/codegen/test_vector_storage.py
@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import numpy as np
+import pytest
+
+import tvm_tirx as tx
+from tvm_tirx.script import tirx as T
+from tvm_tirx.testing import env
+
+
[email protected]
[email protected](not env.has_cuda(), reason="need CUDA")
+def test_bool_storage_uses_one_byte_per_lane():
+ @T.function
+ def compare(
+ A: T.Tensor((32,), "int32"),
+ B: T.Tensor((32,), "int32"),
+ C: T.Tensor((40,), "bool"),
+ ):
+ T.func_attr({"tirx.noalias": True})
+ with T.launch_thread("blockIdx.x", 1):
+ with T.launch_thread("threadIdx.x", 8) as tx:
+ C[T.Ramp(4 + tx * 4, 1, 4)] = A[T.Ramp(tx * 4, 1, 4)] <
B[T.Ramp(tx * 4, 1, 4)]
+
+ device = tx.cuda()
+ a = np.arange(32, dtype="int32") % 3
+ b = np.ones(32, dtype="int32")
+ expected = np.ones(40, dtype="bool")
+ expected[4:-4] = a < b
+ result = tx.runtime.tensor(np.ones_like(expected), device)
+ compiled = tx.tirx.compile(compare, target="cuda")
+ compiled(tx.runtime.tensor(a, device), tx.runtime.tensor(b, device),
result)
+ np.testing.assert_array_equal(result.numpy(), expected)
+
+
[email protected]
[email protected](not env.has_cuda(), reason="need CUDA")
+def test_unaligned_vector_store_is_rejected():
+ @T.function
+ def copy(A: T.Tensor((32,), "int32"), B: T.Tensor((33,), "int32")):
+ T.func_attr({"tirx.noalias": True})
+ with T.launch_thread("blockIdx.x", 1):
+ with T.launch_thread("threadIdx.x", 8) as tx:
+ B[T.Ramp(tx * 4 + 1, 1, 4)] = A[T.Ramp(tx * 4, 1, 4)]
+
+ with pytest.raises(ValueError, match="Cannot prove alignment for CUDA
vector store"):
+ tx.tirx.compile(copy, target="cuda")
diff --git a/tests/python/codegen/gpu_common/test_intrinsics.py
b/tests/python/codegen/gpu_common/test_intrinsics.py
index 1726c5380..7802ba1e4 100644
--- a/tests/python/codegen/gpu_common/test_intrinsics.py
+++ b/tests/python/codegen/gpu_common/test_intrinsics.py
@@ -44,7 +44,7 @@ def test_int_intrin(target, dtype):
if not tvm_tirx.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
test_funcs = [
- (T.clz, lambda x, dtype: int(dtype[-2:]) - (len(bin(x)) - 2)),
+ (T.clz, lambda x, dtype: int(dtype[-2:]) - int(x).bit_length()),
]
for tvm_intrin, np_func in test_funcs:
@@ -65,7 +65,8 @@ def test_int_intrin(target, dtype):
def run_and_check():
dev = tvm_tirx.device_from_target(target)
- a = tvm_tirx.runtime.tensor(np.random.randint(0, 100000,
size=n).astype(dtype), dev)
+ # clz is undefined for zero.
+ a = tvm_tirx.runtime.tensor(np.random.randint(1, 100000,
size=n).astype(dtype), dev)
b = tvm_tirx.runtime.tensor(np.zeros(shape=(n,)).astype(dtype),
dev)
f(a, b)
ref = np.vectorize(partial(np_func, dtype=dtype))(a.numpy())
