From 58468f87ca73e3c6819270c9dce346320c4952bf Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Thu, 10 Sep 2026 14:15:43 +0100
Subject: [PATCH] Fix dtype validation for conditional SM targets

Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com>
---
 src/cuda/tile/_passes/check_dtype_support.py |  7 ++-
 test/test_conditional_arch.py                | 60 ++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 test/test_conditional_arch.py

diff --git a/src/cuda/tile/_passes/check_dtype_support.py b/src/cuda/tile/_passes/check_dtype_support.py
index 0d40a0d3..80ae59ef 100644
--- a/src/cuda/tile/_passes/check_dtype_support.py
+++ b/src/cuda/tile/_passes/check_dtype_support.py
@@ -134,7 +134,12 @@ def _check_dtype(dtype: DType, sm_arch: str | None, sm_number: int | None,
 def check_dtype_support(root_block: Block, sm_arch: str | None,
                         version: BytecodeVersion) -> None:
     # Skip arch check by setting sm_number to None when sm_arch is not provided
-    sm_number = int(sm_arch.removeprefix("sm_")) if sm_arch is not None else None
+    sm_number = None
+    if sm_arch is not None:
+        arch_number = sm_arch.removeprefix("sm_")
+        if arch_number.endswith(("a", "f")):
+            arch_number = arch_number[:-1]
+        sm_number = int(arch_number)
     for op in root_block.traverse():
         if isinstance(op, TypedConst):
             _check_const_value(op)
diff --git a/test/test_conditional_arch.py b/test/test_conditional_arch.py
new file mode 100644
index 00000000..e7230def
--- /dev/null
+++ b/test/test_conditional_arch.py
@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from io import BytesIO
+
+import pytest
+
+import cuda.tile as ct
+from cuda.tile._cext import CallingConvention
+from cuda.tile._exception import TileUnsupportedFeatureError
+
+
+def compile_bytecode(pyfunc, arch, version):
+    kernel = ct.kernel(pyfunc)
+    sig = ct.compilation.KernelSignature(
+        [], CallingConvention.cutile_python_v1(), symbol="kernel")
+    output = BytesIO()
+    ct.compilation.export_kernel(kernel, [sig], output_file=output, gpu_code=arch,
+                                 output_format="tileir_bytecode", bytecode_version=version)
+    assert output.getvalue()
+
+
+@pytest.mark.parametrize("arch", ["sm_90", "sm_90a", "sm_100a", "sm_100f", "sm_120a", "sm_121a"])
+def test_conditional_arch_dtype_check(arch):
+    def kernel():
+        t = ct.full((2,), 1.5, dtype=ct.float32)
+        ct.printf("%f", t)
+
+    compile_bytecode(kernel, arch, "13.2")
+
+
+@pytest.mark.parametrize("arch", ["sm_90a", "sm_100f"])
+def test_conditional_arch_preserves_dtype_limits(arch):
+    def kernel():
+        t = ct.full((2,), 1.5, dtype=ct.float8_e5m3fnu)
+        ct.printf("%f", t)
+
+    with pytest.raises(TileUnsupportedFeatureError, match=f"is not supported on {arch}"):
+        compile_bytecode(kernel, arch, "13.4")
+
+
+@pytest.mark.parametrize("arch", ["sm_100a", "sm_100f"])
+def test_conditional_arch_preserves_bytecode_limits(arch):
+    def kernel():
+        t = ct.full((2,), 1.5, dtype=ct.float4_e2m1fn)
+        ct.printf("%f", t)
+
+    with pytest.raises(TileUnsupportedFeatureError,
+                       match=r"float4_e2m1fn requires tileiras 13\.3"):
+        compile_bytecode(kernel, arch, "13.2")
+
+
+@pytest.mark.parametrize("arch", ["sm_100aa", "sm_100af", "sm_100ff", "sm_100x"])
+def test_conditional_arch_rejects_invalid_suffix(arch):
+    def kernel():
+        ct.printf("%d", ct.bid(0))
+
+    with pytest.raises(ValueError, match="invalid literal for int"):
+        compile_bytecode(kernel, arch, "13.2")