tiled view atomic ops

Signed-off-by: Boyan Li <boyanl@nvidia.com>

Author: blinxt

changelog.d/atomic-add-bf16.md

@@ -0,0 +1,5 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- `ct.atomic_add()` now supports `bfloat16` operands on Hopper (sm_90) and
+  newer architectures.

changelog.d/tiled-view-atomic.md

@@ -0,0 +1,7 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- New `TiledView.atomic_add`, `TiledView.atomic_max`, `TiledView.atomic_min`,
+  `TiledView.atomic_and`, `TiledView.atomic_or`, and `TiledView.atomic_xor`
+  methods for performing element-wise atomic read-modify-write operations on
+  a tiled view at a given tile index.

src/cuda/tile/_ir/op_impl.py

@@ -158,8 +158,8 @@ def _check_version():
             if min_version is not None and cur_version < min_version:
                 raise TileUnsupportedFeatureError(
                     f"{stub.__name__} requires tileiras "
-                    f"{min_version.major()}.{min_version.minor()} or later. "
-                    f"Current version is {cur_version.major()}.{cur_version.minor()}."
+                    f"{min_version.as_string()} or later. "
+                    f"Current version is {cur_version.as_string()}."
                 )
 
         def decorate(func):

src/cuda/tile/_ir/ops.py

@@ -160,8 +160,8 @@ def check_rd_and_ftz(fn: str, rounding_mode: Optional[RoundingMode], flush_to_ze
             if cur_version < min_version:
                 raise TileUnsupportedFeatureError(
                     f'{fn} rounding_mode={rounding_mode.value} requires tileiras '
-                    f'{min_version.major()}.{min_version.minor()} or later. '
-                    f'Current version is {cur_version.major()}.{cur_version.minor()}.')
+                    f'{min_version.as_string()} or later. '
+                    f'Current version is {cur_version.as_string()}.')
         if not datatype.is_unrestricted_float(dtype):
             raise TileTypeError(
                 f'Rounding mode can only be used for unrestricted float types, '
@@ -362,10 +362,21 @@ def validate_memory_order_and_scope(
             f"Invalid memory order for {opcode}. "
             f"Got {memory_order}, expected one of {formatted_expected}"
         )
+
+    if memory_scope not in operation_type.VALID_MEMORY_SCOPES:
+        formatted_expected = ", ".join(
+            str(scope) for scope in operation_type.VALID_MEMORY_SCOPES
+        )
+        raise TileTypeError(
+            f"Invalid memory scope for {opcode}. "
+            f"Got {memory_scope}, expected one of {formatted_expected}"
+        )
+
     if memory_order == MemoryOrder.WEAK and memory_scope != MemoryScope.NONE:
         raise TileTypeError(
             f"{opcode} with WEAK memory ordering cannot specify a memory scope"
         )
+
     if memory_order != MemoryOrder.WEAK and memory_scope == MemoryScope.NONE:
         raise TileTypeError(
             f"{opcode} with {memory_order.name} memory ordering requires a memory scope"

src/cuda/tile/_ir/ops_utils.py

@@ -451,15 +451,12 @@ def __str__(self):
 
 # ============== TiledView Type ===============
 
-
+@dataclass(frozen=True)
 class TiledViewTy(Type):
-    def __init__(self, array_ty: ArrayTy, tile_shape: tuple[int, ...],
-                 padding_mode: PaddingMode,
-                 traversal_steps: Optional[tuple[int, ...]] = None):
-        self.array_ty = array_ty
-        self.tile_shape = tile_shape
-        self.padding_mode = padding_mode
-        self.traversal_steps = traversal_steps
+    array_ty: ArrayTy
+    tile_shape: tuple[int, ...]
+    padding_mode: PaddingMode
+    traversal_steps: Optional[tuple[int, ...]] = None
 
     def is_aggregate(self) -> bool:
         return True
@@ -480,17 +477,6 @@ def ndim(self):
     def dtype(self):
         return self.array_ty.dtype
 
-    def __eq__(self, other: "Type"):
-        return (isinstance(other, TiledViewTy)
-                and self.array_ty == other.array_ty
-                and self.tile_shape == other.tile_shape
-                and self.padding_mode == other.padding_mode
-                and self.traversal_steps == other.traversal_steps)
-
-    def __hash__(self):
-        return hash(("TiledViewTy", self.array_ty, self.tile_shape,
-                     self.padding_mode, self.traversal_steps))
-
     def __str__(self):
         return (f"TiledView[{self.array_ty},tile_shape={self.tile_shape},"
                 f"padding_mode={self.padding_mode},traversal_steps={self.traversal_steps}]")

src/cuda/tile/_ir/type.py

@@ -5,9 +5,11 @@
 import math
 from typing import Callable
 from cuda.tile._ir.ir import Block
-from cuda.tile._ir.ops import TypedConst
+from cuda.tile._ir.ops import TileAtomicRMW, TileAtomicRedView, TypedConst, AtomicRMWMode
 from cuda.tile._ir.type import TileTy, PointerTy, Type
-from cuda.tile._datatype import DType, float4_e2m1fn, float8_e4m3fn, float8_e5m2, float8_e8m0fnu
+from cuda.tile._datatype import (
+    DType, float4_e2m1fn, float8_e4m3fn, float8_e5m2, float8_e8m0fnu, bfloat16
+)
 from cuda.tile._bytecode.version import BytecodeVersion
 from cuda.tile._exception import TileUnsupportedFeatureError, TileValueError
 
@@ -65,8 +67,33 @@ def _check_const_value(op: TypedConst):
             raise TileValueError(msg, loc=op.loc)
 
 
-def _check_dtype(dtype: DType, sm_arch: str, version: BytecodeVersion, loc):
-    sm_number = int(sm_arch.removeprefix("sm_"))
+def _check_atomic_rmw_dtype(op: TileAtomicRedView | TileAtomicRMW,
+                            sm_arch: str,
+                            sm_number: int,
+                            version: BytecodeVersion):
+    dtypes = (_extract_dtypes(op.view.try_get_type())
+              if isinstance(op, TileAtomicRedView) else
+              _extract_dtypes(op.result_vars[0].try_get_type()))
+    if not (op.mode == AtomicRMWMode.ADD_FLOAT and bfloat16 in dtypes):
+        return
+
+    if sm_number < 90:
+        raise TileUnsupportedFeatureError(
+            f"{bfloat16} is not supported by atomic add on {sm_arch}",
+            loc=op.loc
+        )
+
+    min_version = BytecodeVersion.V_13_3
+    if version < min_version:
+        raise TileUnsupportedFeatureError(
+            f"{bfloat16} on atomic add requires tileiras"
+            f" {min_version.as_string()} or later."
+            f" Current version is {version.as_string()}.",
+            loc=op.loc
+        )
+
+
+def _check_dtype(dtype: DType, sm_arch: str, sm_number: int, version: BytecodeVersion, loc):
     min_sm = _DTYPE_MIN_SM.get(dtype)
     if min_sm is not None and sm_number < min_sm:
         raise TileUnsupportedFeatureError(
@@ -78,17 +105,21 @@ def _check_dtype(dtype: DType, sm_arch: str, version: BytecodeVersion, loc):
     if min_version is not None and version < min_version:
         raise TileUnsupportedFeatureError(
             f"{dtype} requires tileiras"
-            f" {min_version.major()}.{min_version.minor()} or later."
-            f" Current version is {version.major()}.{version.minor()}.",
+            f" {min_version.as_string()} or later."
+            f" Current version is {version.as_string()}.",
             loc=loc,
         )
 
 
 def check_dtype_support(root_block: Block, sm_arch: str, version: BytecodeVersion) -> None:
+    sm_number = int(sm_arch.removeprefix("sm_"))
     for op in root_block.traverse():
         if isinstance(op, TypedConst):
             _check_const_value(op)
 
+        if isinstance(op, (TileAtomicRedView, TileAtomicRMW)):
+            _check_atomic_rmw_dtype(op, sm_arch, sm_number, version)
+
         all_dtypes = set().union(*(_extract_dtypes(v.try_get_type()) for v in op.all_inputs()))
         for dtype in all_dtypes:
-            _check_dtype(dtype, sm_arch, version, op.loc)
+            _check_dtype(dtype, sm_arch, sm_number, version, op.loc)

src/cuda/tile/_passes/check_dtype_support.py

@@ -15,7 +15,7 @@
 from cuda.tile._ir.ops import (
     Break, Continue, EndBranch, IfElse,
     JoinTokens, Loop, MakeToken,
-    TileAtomicCAS, TileAtomicRMW, LoadPointer,
+    TileAtomicCAS, TileAtomicRMW, LoadPointer, TileAtomicRedView,
     TileLoad, StorePointer,
     TileStore, TileAssert, TilePrintf,
 )
@@ -132,7 +132,8 @@ def get_memory_effects(cur_op):
             return EMPTY_MEMORY_EFFECTS
 
         has_acquire_order = False
-        if isinstance(cur_op, (TileAtomicCAS, TileAtomicRMW, TileLoad, TileStore)):
+        if isinstance(cur_op, (TileAtomicCAS, TileAtomicRMW, TileAtomicRedView,
+                               TileLoad, TileStore)):
             has_acquire_order = memory_order_has_acquire(cur_op.memory_order)
 
         return MemoryEffects({dataflow_result[_get_input_var(cur_op).name].alias_set: effect},
@@ -214,7 +215,7 @@ def _to_token_order_in_block(block: Block,
             token_map[last_op_key] = result_tok
             token_map[last_store_key] = result_tok
 
-        elif isinstance(op, (TileAtomicCAS, TileAtomicRMW)):
+        elif isinstance(op, (TileAtomicCAS, TileAtomicRMW, TileAtomicRedView)):
             alias_set = context.dataflow_result[_get_input_var(op).name].alias_set
             last_op_key = _last_op_key(alias_set)
             last_store_key = _last_store_key(alias_set)
@@ -224,7 +225,11 @@ def _to_token_order_in_block(block: Block,
             if maybe_input_tok_join_op:
                 operations.append(maybe_input_tok_join_op)
 
-            _, result_tok = op.result_vars
+            if isinstance(op, TileAtomicRedView):
+                [result_tok] = op.result_vars
+            else:
+                _, result_tok = op.result_vars
+
             operations.append(dataclasses.replace(op, token=input_tok))
 
             token_map[last_op_key] = result_tok
@@ -460,7 +465,7 @@ def _get_parallel_stores(
     alias_set_to_mem_ops = defaultdict(list)
     for op in loop_op.body.operations:
         if isinstance(op, (TileLoad, StorePointer, LoadPointer, TileStore,
-                           TileAtomicCAS, TileAtomicRMW)):
+                           TileAtomicCAS, TileAtomicRMW, TileAtomicRedView)):
             alias_set = context.dataflow_result[_get_input_var(op).name].alias_set
             alias_set_to_mem_ops[alias_set].append(op)
 

src/cuda/tile/_passes/token_order.py

@@ -707,6 +707,57 @@ def __rmatmul__(self, other) -> "Tile":
 TileOrScalar = Union[Tile, Scalar]
 
 
+def _doc_tv_atomic_rmw_op(*, testoutput: int):
+    def decorator(f):
+        op_name = f.__name__
+        f.__doc__ += f"""\
+
+        This method does not return a value.
+
+        For each individual element, the operation is performed atomically,
+        but the operation as a whole is not atomic, and the order of individual
+        writes is unspecified.
+
+        `update`'s shape must be broadcastable to :attr:`tile_shape`.
+
+        For a tile that partially extends beyond the tiled view boundaries,
+        out-of-bound elements are ignored.
+        If the tile lies entirely outside the tiled view, the behavior is
+        undefined.
+
+        Use this operation instead of ct.{op_name} for better performance
+        when modified value is not needed.
+
+        Args:
+            index (tuple[int,...]): An index in the |tiled view|'s tile space.
+            update (Tile): The update values.
+
+        Returns:
+            None:
+
+        Examples:
+
+            .. testcode::
+                :template: setup_only.py
+
+                @ct.kernel
+                def kernel(x):
+                    tv = x.tiled_view(4)
+                    update = ct.full((4,), 1, dtype=ct.int32)
+                    tv.{op_name}(0, update)
+
+                x = torch.zeros(4, dtype=torch.int32, device='cuda')
+                ct.launch(stream, (1,), kernel, (x,))
+                print(x.tolist())
+
+            .. testoutput::
+
+                {[testoutput] * 4}
+    """
+        return f
+    return decorator
+
+
 class TiledView:
     """Class for |tiled view| objects."""
 
@@ -836,6 +887,62 @@ def kernel(x):
                 [99, 99, 99, 99, 0, 0, 0, 0]
         """
 
+    @_doc_tv_atomic_rmw_op(testoutput=(0 + 1))
+    @stub
+    def atomic_add(self, index: Shape, update: Tile) -> None:
+        """Atomically adds `update` to the |tiled view| at the given tile `index`.
+
+        If `update`'s dtype differs from the view's dtype, an implicit cast is
+        performed.
+        """
+
+    @_doc_tv_atomic_rmw_op(testoutput=max(0, 1))
+    @stub
+    def atomic_max(self, index: Shape, update: Tile) -> None:
+        """Atomically applies element-wise maximum with update to the |tiled view|
+        at the given tile `index`.
+
+        If `update`'s dtype differs from the view's dtype, an implicit cast is
+        performed.
+        """
+
+    @_doc_tv_atomic_rmw_op(testoutput=min(0, 1))
+    @stub
+    def atomic_min(self, index: Shape, update: Tile) -> None:
+        """Atomically applies element-wise minimum with update to the |tiled view|
+        at the given tile `index`.
+
+        If `update`'s dtype differs from the view's dtype, an implicit cast is
+        performed.
+        """
+
+    @_doc_tv_atomic_rmw_op(testoutput=(0 & 1))
+    @stub
+    def atomic_and(self, index: Shape, update: Tile) -> None:
+        """Atomically applies bitwise AND of `update` to the |tiled view| at the given tile `index`.
+
+        `update`'s dtype must exactly match the view's dtype; no implicit cast is
+        performed.
+        """
+
+    @_doc_tv_atomic_rmw_op(testoutput=(0 | 1))
+    @stub
+    def atomic_or(self, index: Shape, update: Tile) -> None:
+        """Atomically applies bitwise OR of `update` to the |tiled view| at the given tile `index`.
+
+        `update`'s dtype must exactly match the view's dtype; no implicit cast is
+        performed.
+        """
+
+    @_doc_tv_atomic_rmw_op(testoutput=(0 ^ 1))
+    @stub
+    def atomic_xor(self, index: Shape, update: Tile) -> None:
+        """Atomically applies bitwise XOR of `update` to the |tiled view| at the given tile `index`.
+
+        `update`'s dtype must exactly match the view's dtype; no implicit cast is
+        performed.
+        """
+
 
 ###############################################################################
 # Constantness Hints