[Gluon] Expose TMA atomic ops

python/src/gluon_ir.cc

@@ -928,17 +928,18 @@ void init_gluon_ir(py::module &&m) {
             self.create<ttng::AsyncTMACopyGlobalToLocalOp>(
                 descPtr, coord, offsetsRange, barrier, result, pred, multicast);
           })
-      .def("create_async_tma_copy_local_to_global",
-           [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &coord,
-              Value src) {
-             self.create<ttng::AsyncTMACopyLocalToGlobalOp>(descPtr, coord,
-                                                            src);
-           })
-      .def("create_async_tma_reduce",
-           [](GluonOpBuilder &self, triton::DescriptorReduceKind kind,
-              Value descPtr, std::vector<Value> &coord, Value src) {
-             self.create<ttng::AsyncTMAReduceOp>(kind, descPtr, coord, src);
-           })
+      .def(
+          "create_async_tma_copy_local_to_global",
+          [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &coord,
+             Value src, std::optional<triton::DescriptorReduceKind> red) {
+            if (red) {
+              self.create<ttng::AsyncTMAReduceOp>(*red, descPtr, coord, src);
+              return;
+            }
+            self.create<ttng::AsyncTMACopyLocalToGlobalOp>(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<ttng::TMAStoreWaitOp>(pendings);

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)

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",

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