[Frontend] Improve error when descriptor atomic_{min,max} cannot fallback (#6865)

Currently we fail in the middle of the tensor descriptor rewrite pass,
whereas this pre-empts it and raises the error from the frontend.

Author: peterbell10

python/test/unit/language/test_tensor_descriptor.py

@@ -1563,8 +1563,7 @@ def alloc_fn(size: int, align: int, stream: Optional[int]):
     fallback_supported = dtype in FALLBACK_SUPPORTED_REDUCE_DTYPES[kind]
     supported = native_supported if is_native else fallback_supported
     if not supported:
-        exc_type = CompilationError if not native_supported else RuntimeError
-        with pytest.raises(exc_type):
+        with pytest.raises(CompilationError):
             kernel[(grid_m, grid_n)](out_desc, out, inp, M, N, M_BLOCK, N_BLOCK, kind, num_ctas=num_ctas)
         return
 

python/triton/language/semantic.py

@@ -4,6 +4,8 @@
 from typing import List, Optional, Sequence, Tuple, TypeVar
 import numbers
 
+from triton.runtime import driver
+
 from .._C.libtriton import ir
 from . import core as tl
 
@@ -1180,17 +1182,28 @@ def descriptor_atomic_add(desc: tl.tensor_descriptor_base, value: tl.tensor, off
     return tl.tensor(builder.create_descriptor_reduce(kind, desc.handle, value.handle, offsets), tl.void)
 
 
+def _has_native_tma():
+    target = driver.active.get_current_target()
+    return (target.backend == "cuda" and target.arch >= 90)
+
+
+def _descriptor_atomic_min_max_supported(dtype):
+    assert dtype in {tl.uint32, tl.int32, tl.uint64, tl.int64, tl.float16, tl.bfloat16}, "Unsupported dtype"
+    if dtype in {tl.float16, tl.bfloat16}:
+        assert _has_native_tma(), "16-bit float types require native tma support"
+
+
 def descriptor_atomic_min(desc: tl.tensor_descriptor_base, value: tl.tensor, offsets, builder: ir.builder) -> tl.tensor:
     validate_store_like(desc, value, offsets)
-    assert desc.dtype in {tl.uint32, tl.int32, tl.uint64, tl.int64, tl.float16, tl.bfloat16}, "Unsupported dtype"
+    _descriptor_atomic_min_max_supported(desc.dtype)
     offsets = _convert_to_ir_values(builder, offsets, require_i64=False)
     kind = ir.DESCRIPTOR_REDUCE_KIND.MIN
     return tl.tensor(builder.create_descriptor_reduce(kind, desc.handle, value.handle, offsets), tl.void)
 
 
 def descriptor_atomic_max(desc: tl.tensor_descriptor_base, value: tl.tensor, offsets, builder: ir.builder) -> tl.tensor:
     validate_store_like(desc, value, offsets)
-    assert desc.dtype in {tl.uint32, tl.int32, tl.uint64, tl.int64, tl.float16, tl.bfloat16}, "Unsupported dtype"
+    _descriptor_atomic_min_max_supported(desc.dtype)
     offsets = _convert_to_ir_values(builder, offsets, require_i64=False)
     kind = ir.DESCRIPTOR_REDUCE_KIND.MAX
     return tl.tensor(builder.create_descriptor_reduce(kind, desc.handle, value.handle, offsets), tl.void)