From 8c532eb044948cbe495c3ff39c307ff68e6ea042 Mon Sep 17 00:00:00 2001 From: lezcano Date: Wed, 15 Apr 2026 12:44:27 +0200 Subject: [PATCH 1/3] [Gluon] Expose TMAReduction via an optional enum. tbh, I think this is also the way we should represent this in our IR, but well. --- python/src/gluon_ir.cc | 23 +++++++++-------- python/test/gluon/test_frontend.py | 25 +++++++++++++++++++ .../gluon/language/nvidia/blackwell/tma.py | 2 ++ .../gluon/language/nvidia/hopper/tma.py | 17 +++++++++++-- 4 files changed, 54 insertions(+), 13 deletions(-) diff --git a/python/src/gluon_ir.cc b/python/src/gluon_ir.cc index 1bcf493ddbd3..e5bad9f5cda1 100644 --- a/python/src/gluon_ir.cc +++ b/python/src/gluon_ir.cc @@ -928,17 +928,18 @@ void init_gluon_ir(py::module &&m) { self.create( descPtr, coord, offsetsRange, barrier, result, pred, multicast); }) - .def("create_async_tma_copy_local_to_global", - [](GluonOpBuilder &self, Value descPtr, std::vector &coord, - Value src) { - self.create(descPtr, coord, - src); - }) - .def("create_async_tma_reduce", - [](GluonOpBuilder &self, triton::DescriptorReduceKind kind, - Value descPtr, std::vector &coord, Value src) { - self.create(kind, descPtr, coord, src); - }) + .def( + "create_async_tma_copy_local_to_global", + [](GluonOpBuilder &self, Value descPtr, std::vector &coord, + Value src, std::optional red) { + if (red) { + self.create(*red, descPtr, coord, src); + return; + } + self.create(descPtr, coord, src); + }, + py::arg("descPtr"), py::arg("coord"), py::arg("src"), + py::arg("red") = py::none()) .def("create_async_tma_store_wait", [](GluonOpBuilder &self, int pendings) { self.create(pendings); diff --git a/python/test/gluon/test_frontend.py b/python/test/gluon/test_frontend.py index 0262c1cebcec..73b6c1a00deb 100644 --- a/python/test/gluon/test_frontend.py +++ b/python/test/gluon/test_frontend.py @@ -3884,6 +3884,31 @@ def nv_tma_descriptor_store_kernel(input_ptr): """) +@pytest.mark.parametrize("target", [BLACKWELL_TARGET, HOPPER_TARGET]) +def test_nv_tma_descriptor_reduce_kernel(target): + + @gluon.jit + def nv_tma_descriptor_reduce_kernel(input_ptr): + XBLOCK: ttgl.constexpr = 128 + smem_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=32, rank=2) + input_desc = tma.make_tensor_descriptor( + input_ptr, + shape=[XBLOCK, XBLOCK], + strides=[XBLOCK, 1], + block_shape=[XBLOCK, XBLOCK], + layout=smem_layout, + ) + smem = ttgl.allocate_shared_memory(ttgl.float32, [XBLOCK, XBLOCK], smem_layout) + tma.async_copy_shared_to_global(input_desc, [0, 0], smem, red=tma.ReduceKind.ADD) + tma.store_wait(0) + + ptr = MockTensor(ttgl.float32) + module = run_parser(nv_tma_descriptor_reduce_kernel, *make_args(ptr), target) + ttgir = anonymize_ir(module.str_nodebug()) + assert "ttng.async_tma_reduce add," in ttgir + assert "ttng.async_tma_copy_local_to_global" not in ttgir + + @filecheck_test def tmem_constexpr(): tmem_shape: ttgl.constexpr = (64, 64) diff --git a/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py b/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py index 01adc772005c..148c9de920f7 100644 --- a/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py +++ b/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py @@ -3,6 +3,7 @@ from triton.experimental.gluon.language.nvidia.hopper.tma import ( async_copy_global_to_shared, async_copy_shared_to_global, + ReduceKind, store_wait, tensor_descriptor, tensor_descriptor_type, @@ -15,6 +16,7 @@ "async_scatter", "async_copy_global_to_shared", "async_copy_shared_to_global", + "ReduceKind", "store_wait", "tensor_descriptor", "tensor_descriptor_type", diff --git a/python/triton/experimental/gluon/language/nvidia/hopper/tma.py b/python/triton/experimental/gluon/language/nvidia/hopper/tma.py index 96268df97d2c..6064415be580 100644 --- a/python/triton/experimental/gluon/language/nvidia/hopper/tma.py +++ b/python/triton/experimental/gluon/language/nvidia/hopper/tma.py @@ -13,9 +13,12 @@ "async_copy_global_to_shared", "async_copy_global_to_shared_im2col", "async_copy_shared_to_global", + "ReduceKind", "store_wait", ] +ReduceKind = ttgl.ir.DESCRIPTOR_REDUCE_KIND + @dataclass(eq=True) class _tensor_descriptor_type_base(base_type): @@ -256,12 +259,22 @@ def async_copy_global_to_shared_im2col(tensor_desc, coord, offsets, barrier, res @builtin -def async_copy_shared_to_global(tensor_desc, coord, src, _semantic=None): +def async_copy_shared_to_global(tensor_desc, coord, src, red=None, _semantic=None): + """ + Copy data from shared memory to global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + red (tma.ReduceKind | None): Optional reduction operation. If set, atomically reduces into global memory. + """ if _semantic.builder.options.enable_iisan: _emit_alignment_check(tensor_desc, coord, "async_copy_shared_to_global", "innermost coordinate", _semantic=_semantic) coord = _semantic._convert_to_ir_values(coord, require_i64=False) - _semantic.builder.create_async_tma_copy_local_to_global(tensor_desc.handle, coord, src.handle) + red = _unwrap_if_constexpr(red) + _semantic.builder.create_async_tma_copy_local_to_global(tensor_desc.handle, coord, src.handle, red) @builtin From 287adca4b728cace3ece4c7ef786e54889c495ba Mon Sep 17 00:00:00 2001 From: lezcano Date: Wed, 15 Apr 2026 15:06:23 +0200 Subject: [PATCH 2/3] address review --- python/src/gluon_ir.cc | 23 ++-- python/test/gluon/test_frontend.py | 33 ++++-- .../gluon/language/nvidia/blackwell/tma.py | 14 ++- .../gluon/language/nvidia/hopper/tma.py | 106 ++++++++++++++++-- 4 files changed, 148 insertions(+), 28 deletions(-) diff --git a/python/src/gluon_ir.cc b/python/src/gluon_ir.cc index e5bad9f5cda1..1bcf493ddbd3 100644 --- a/python/src/gluon_ir.cc +++ b/python/src/gluon_ir.cc @@ -928,18 +928,17 @@ void init_gluon_ir(py::module &&m) { self.create( descPtr, coord, offsetsRange, barrier, result, pred, multicast); }) - .def( - "create_async_tma_copy_local_to_global", - [](GluonOpBuilder &self, Value descPtr, std::vector &coord, - Value src, std::optional red) { - if (red) { - self.create(*red, descPtr, coord, src); - return; - } - self.create(descPtr, coord, src); - }, - py::arg("descPtr"), py::arg("coord"), py::arg("src"), - py::arg("red") = py::none()) + .def("create_async_tma_copy_local_to_global", + [](GluonOpBuilder &self, Value descPtr, std::vector &coord, + Value src) { + self.create(descPtr, coord, + src); + }) + .def("create_async_tma_reduce", + [](GluonOpBuilder &self, triton::DescriptorReduceKind kind, + Value descPtr, std::vector &coord, Value src) { + self.create(kind, descPtr, coord, src); + }) .def("create_async_tma_store_wait", [](GluonOpBuilder &self, int pendings) { self.create(pendings); diff --git a/python/test/gluon/test_frontend.py b/python/test/gluon/test_frontend.py index 73b6c1a00deb..8dada655c4b7 100644 --- a/python/test/gluon/test_frontend.py +++ b/python/test/gluon/test_frontend.py @@ -3885,10 +3885,18 @@ def nv_tma_descriptor_store_kernel(input_ptr): @pytest.mark.parametrize("target", [BLACKWELL_TARGET, HOPPER_TARGET]) -def test_nv_tma_descriptor_reduce_kernel(target): +@pytest.mark.parametrize(("op_name", "kind"), [ + ("async_atomic_add", "add"), + ("async_atomic_min", "min"), + ("async_atomic_max", "max"), + ("async_atomic_and", "and"), + ("async_atomic_or", "or"), + ("async_atomic_xor", "xor"), +]) +def test_nv_tma_descriptor_reduce_kernel(target, op_name, kind): @gluon.jit - def nv_tma_descriptor_reduce_kernel(input_ptr): + def nv_tma_descriptor_reduce_kernel(input_ptr, OP_NAME: ttgl.constexpr): XBLOCK: ttgl.constexpr = 128 smem_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=32, rank=2) input_desc = tma.make_tensor_descriptor( @@ -3898,14 +3906,25 @@ def nv_tma_descriptor_reduce_kernel(input_ptr): block_shape=[XBLOCK, XBLOCK], layout=smem_layout, ) - smem = ttgl.allocate_shared_memory(ttgl.float32, [XBLOCK, XBLOCK], smem_layout) - tma.async_copy_shared_to_global(input_desc, [0, 0], smem, red=tma.ReduceKind.ADD) + smem = ttgl.allocate_shared_memory(ttgl.int32, [XBLOCK, XBLOCK], smem_layout) + if OP_NAME == "async_atomic_add": + tma.async_atomic_add(input_desc, [0, 0], smem) + elif OP_NAME == "async_atomic_min": + tma.async_atomic_min(input_desc, [0, 0], smem) + elif OP_NAME == "async_atomic_max": + tma.async_atomic_max(input_desc, [0, 0], smem) + elif OP_NAME == "async_atomic_and": + tma.async_atomic_and(input_desc, [0, 0], smem) + elif OP_NAME == "async_atomic_or": + tma.async_atomic_or(input_desc, [0, 0], smem) + elif OP_NAME == "async_atomic_xor": + tma.async_atomic_xor(input_desc, [0, 0], smem) tma.store_wait(0) - ptr = MockTensor(ttgl.float32) - module = run_parser(nv_tma_descriptor_reduce_kernel, *make_args(ptr), target) + ptr = MockTensor(ttgl.int32) + module = run_parser(nv_tma_descriptor_reduce_kernel, *make_args(ptr, op_name), target) ttgir = anonymize_ir(module.str_nodebug()) - assert "ttng.async_tma_reduce add," in ttgir + assert f"ttng.async_tma_reduce {kind}," in ttgir assert "ttng.async_tma_copy_local_to_global" not in ttgir diff --git a/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py b/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py index 148c9de920f7..8056fb2c2bd5 100644 --- a/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py +++ b/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py @@ -1,9 +1,14 @@ import triton.experimental.gluon.language._core as ttgl from triton.experimental.gluon.language._core import builtin from triton.experimental.gluon.language.nvidia.hopper.tma import ( + async_atomic_add, + async_atomic_and, async_copy_global_to_shared, async_copy_shared_to_global, - ReduceKind, + async_atomic_max, + async_atomic_min, + async_atomic_or, + async_atomic_xor, store_wait, tensor_descriptor, tensor_descriptor_type, @@ -14,9 +19,14 @@ __all__ = [ "async_gather", "async_scatter", + "async_atomic_add", + "async_atomic_and", "async_copy_global_to_shared", "async_copy_shared_to_global", - "ReduceKind", + "async_atomic_max", + "async_atomic_min", + "async_atomic_or", + "async_atomic_xor", "store_wait", "tensor_descriptor", "tensor_descriptor_type", diff --git a/python/triton/experimental/gluon/language/nvidia/hopper/tma.py b/python/triton/experimental/gluon/language/nvidia/hopper/tma.py index 6064415be580..d5d75ceaff72 100644 --- a/python/triton/experimental/gluon/language/nvidia/hopper/tma.py +++ b/python/triton/experimental/gluon/language/nvidia/hopper/tma.py @@ -10,15 +10,18 @@ from triton._C import ir __all__ = [ + "async_atomic_add", + "async_atomic_and", "async_copy_global_to_shared", "async_copy_global_to_shared_im2col", "async_copy_shared_to_global", - "ReduceKind", + "async_atomic_max", + "async_atomic_min", + "async_atomic_or", + "async_atomic_xor", "store_wait", ] -ReduceKind = ttgl.ir.DESCRIPTOR_REDUCE_KIND - @dataclass(eq=True) class _tensor_descriptor_type_base(base_type): @@ -259,7 +262,7 @@ def async_copy_global_to_shared_im2col(tensor_desc, coord, offsets, barrier, res @builtin -def async_copy_shared_to_global(tensor_desc, coord, src, red=None, _semantic=None): +def async_copy_shared_to_global(tensor_desc, coord, src, _semantic=None): """ Copy data from shared memory to global memory using TMA. @@ -267,14 +270,103 @@ def async_copy_shared_to_global(tensor_desc, coord, src, red=None, _semantic=Non tensor_desc (tensor_descriptor): Tensor descriptor (tiled). coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. src (ttgl.shared_memory_descriptor): Source memory descriptor. - red (tma.ReduceKind | None): Optional reduction operation. If set, atomically reduces into global memory. """ if _semantic.builder.options.enable_iisan: _emit_alignment_check(tensor_desc, coord, "async_copy_shared_to_global", "innermost coordinate", _semantic=_semantic) coord = _semantic._convert_to_ir_values(coord, require_i64=False) - red = _unwrap_if_constexpr(red) - _semantic.builder.create_async_tma_copy_local_to_global(tensor_desc.handle, coord, src.handle, red) + _semantic.builder.create_async_tma_copy_local_to_global(tensor_desc.handle, coord, src.handle) + + +def _async_atomic_shared_to_global(kind, tensor_desc, coord, src, fn_name: str, _semantic=None): + if _semantic.builder.options.enable_iisan: + _emit_alignment_check(tensor_desc, coord, fn_name, "innermost coordinate", _semantic=_semantic) + coord = _semantic._convert_to_ir_values(coord, require_i64=False) + _semantic.builder.create_async_tma_reduce(kind, tensor_desc.handle, coord, src.handle) + + +@builtin +def async_atomic_add(tensor_desc, coord, src, _semantic=None): + """ + Atomically add data from shared memory into global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + """ + _async_atomic_shared_to_global(ttgl.ir.DESCRIPTOR_REDUCE_KIND.ADD, tensor_desc, coord, src, "async_atomic_add", + _semantic=_semantic) + + +@builtin +def async_atomic_min(tensor_desc, coord, src, _semantic=None): + """ + Atomically compute the minimum of shared memory data and global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + """ + _async_atomic_shared_to_global(ttgl.ir.DESCRIPTOR_REDUCE_KIND.MIN, tensor_desc, coord, src, "async_atomic_min", + _semantic=_semantic) + + +@builtin +def async_atomic_max(tensor_desc, coord, src, _semantic=None): + """ + Atomically compute the maximum of shared memory data and global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + """ + _async_atomic_shared_to_global(ttgl.ir.DESCRIPTOR_REDUCE_KIND.MAX, tensor_desc, coord, src, "async_atomic_max", + _semantic=_semantic) + + +@builtin +def async_atomic_and(tensor_desc, coord, src, _semantic=None): + """ + Atomically bitwise-and data from shared memory into global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + """ + _async_atomic_shared_to_global(ttgl.ir.DESCRIPTOR_REDUCE_KIND.AND, tensor_desc, coord, src, "async_atomic_and", + _semantic=_semantic) + + +@builtin +def async_atomic_or(tensor_desc, coord, src, _semantic=None): + """ + Atomically bitwise-or data from shared memory into global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + """ + _async_atomic_shared_to_global(ttgl.ir.DESCRIPTOR_REDUCE_KIND.OR, tensor_desc, coord, src, "async_atomic_or", + _semantic=_semantic) + + +@builtin +def async_atomic_xor(tensor_desc, coord, src, _semantic=None): + """ + Atomically bitwise-xor data from shared memory into global memory using TMA. + + Args: + tensor_desc (tensor_descriptor): Tensor descriptor (tiled). + coord (Sequence[int | ttgl.constexpr | ttgl.tensor]): Coordinates in the destination tensor. + src (ttgl.shared_memory_descriptor): Source memory descriptor. + """ + _async_atomic_shared_to_global(ttgl.ir.DESCRIPTOR_REDUCE_KIND.XOR, tensor_desc, coord, src, "async_atomic_xor", + _semantic=_semantic) @builtin From 92ea603d4254f20db19ae4c59b0bb8c8462db212 Mon Sep 17 00:00:00 2001 From: lezcano Date: Wed, 15 Apr 2026 15:36:58 +0200 Subject: [PATCH 3/3] sort --- .../experimental/gluon/language/nvidia/blackwell/tma.py | 8 ++++---- .../experimental/gluon/language/nvidia/hopper/tma.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py b/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py index 8056fb2c2bd5..de88f517a3e7 100644 --- a/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py +++ b/python/triton/experimental/gluon/language/nvidia/blackwell/tma.py @@ -3,12 +3,12 @@ from triton.experimental.gluon.language.nvidia.hopper.tma import ( async_atomic_add, async_atomic_and, - async_copy_global_to_shared, - async_copy_shared_to_global, async_atomic_max, async_atomic_min, async_atomic_or, async_atomic_xor, + async_copy_global_to_shared, + async_copy_shared_to_global, store_wait, tensor_descriptor, tensor_descriptor_type, @@ -21,12 +21,12 @@ "async_scatter", "async_atomic_add", "async_atomic_and", - "async_copy_global_to_shared", - "async_copy_shared_to_global", "async_atomic_max", "async_atomic_min", "async_atomic_or", "async_atomic_xor", + "async_copy_global_to_shared", + "async_copy_shared_to_global", "store_wait", "tensor_descriptor", "tensor_descriptor_type", diff --git a/python/triton/experimental/gluon/language/nvidia/hopper/tma.py b/python/triton/experimental/gluon/language/nvidia/hopper/tma.py index d5d75ceaff72..da6f23c02768 100644 --- a/python/triton/experimental/gluon/language/nvidia/hopper/tma.py +++ b/python/triton/experimental/gluon/language/nvidia/hopper/tma.py @@ -12,13 +12,13 @@ __all__ = [ "async_atomic_add", "async_atomic_and", - "async_copy_global_to_shared", - "async_copy_global_to_shared_im2col", - "async_copy_shared_to_global", "async_atomic_max", "async_atomic_min", "async_atomic_or", "async_atomic_xor", + "async_copy_global_to_shared", + "async_copy_global_to_shared_im2col", + "async_copy_shared_to_global", "store_wait", ]