Complete ordering to memory_order rename in pointer load/store

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

experimental/cuda-lang/src/cuda/lang/_ir/op_defs.py

@@ -71,9 +71,9 @@ class StorePointer(Operation, opcode="store_pointer", memory_effect=MemoryEffect
     value: Var = operand()
     alignment: Optional[int] = attribute()
     volatile: bool = attribute(default=False)
-    ordering: Optional[MemoryOrder] = attribute(default=None)
+    memory_order: Optional[MemoryOrder] = attribute(default=None)
 
-    valid_orderings = (
+    valid_memory_orders = (
         None,
         MemoryOrder.WEAK,
         MemoryOrder.RELAXED,
@@ -86,9 +86,9 @@ class LoadPointer(Operation, opcode="load_pointer", memory_effect=MemoryEffect.L
     pointer: Var = operand()
     alignment: Optional[int] = attribute()
     volatile: bool = attribute(default=False)
-    ordering: Optional[MemoryOrder] = attribute(default=None)
+    memory_order: Optional[MemoryOrder] = attribute(default=None)
 
-    valid_orderings = (
+    valid_memory_orders = (
         None,
         MemoryOrder.WEAK,
         MemoryOrder.RELAXED,

experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py

@@ -61,21 +61,21 @@ def _mbar_space_scope_suffix(scope: MbarrierScope, space: MemorySpace) -> str:
     return ".scope." + scope.value + ".space." + space_str
 
 
-def require_mbarrier_ordering(
-    ordering_var: Var,
-    valid_orderings: tuple[MemoryOrder, ...],
+def require_mbarrier_memory_order(
+    memory_order_var: Var,
+    valid_memory_orders: tuple[MemoryOrder, ...],
 ) -> MemoryOrder:
-    ordering = require_constant_enum(ordering_var, MemoryOrder)
-    if ordering not in valid_orderings:
-        formatted = ", ".join(str(o) for o in valid_orderings)
+    memory_order = require_constant_enum(memory_order_var, MemoryOrder)
+    if memory_order not in valid_memory_orders:
+        formatted = ", ".join(str(o) for o in valid_memory_orders)
         raise TypeCheckingError(
-            f"Invalid mbarrier memory order {ordering}, expected one of {formatted}"
+            f"Invalid mbarrier memory order {memory_order}, expected one of {formatted}"
         )
-    return ordering
+    return memory_order
 
 
-ARRIVE_ORDERINGS = (MemoryOrder.RELEASE, MemoryOrder.RELAXED)
-WAIT_ORDERINGS = (MemoryOrder.ACQUIRE, MemoryOrder.RELAXED)
+ARRIVE_MEMORY_ORDERS = (MemoryOrder.RELEASE, MemoryOrder.RELAXED)
+WAIT_MEMORY_ORDERS = (MemoryOrder.ACQUIRE, MemoryOrder.RELAXED)
 
 
 @impl(mbarrier.mbarrier_arrive)
@@ -89,7 +89,7 @@ def mbarrier_arrive_impl(
     count = astype(count, datatype.int32)
     drop = require_constant_bool(drop)
     scope = require_constant_enum(scope, MbarrierScope)
-    memory_order = require_mbarrier_ordering(memory_order, ARRIVE_ORDERINGS)
+    memory_order = require_mbarrier_memory_order(memory_order, ARRIVE_MEMORY_ORDERS)
     space = require_mbarrier_ptr(mbar).memory_space
     intrinsic = "llvm.nvvm.mbarrier.arrive"
     if drop:
@@ -119,7 +119,7 @@ def mbarrier_arrive_expect_transaction_impl(
     bytes = astype(bytes, datatype.int32)
     drop = require_constant_bool(drop)
     scope = require_constant_enum(scope, MbarrierScope)
-    memory_order = require_mbarrier_ordering(memory_order, ARRIVE_ORDERINGS)
+    memory_order = require_mbarrier_memory_order(memory_order, ARRIVE_MEMORY_ORDERS)
     space = require_mbarrier_ptr(mbar).memory_space
     intrinsic = "llvm.nvvm.mbarrier.arrive"
     if drop:
@@ -176,7 +176,7 @@ def mbarrier_test_wait_impl(
     scope = require_constant_enum(scope, MbarrierScope)
     state = astype(state, datatype.int64)
     require_mbarrier_ptr(mbar, (MemorySpace.SHARED,))
-    memory_order = require_mbarrier_ordering(memory_order, WAIT_ORDERINGS)
+    memory_order = require_mbarrier_memory_order(memory_order, WAIT_MEMORY_ORDERS)
     intrinsic = "llvm.nvvm.mbarrier.test.wait"
     if memory_order is MemoryOrder.RELAXED:
         intrinsic += ".relaxed"
@@ -196,7 +196,7 @@ def mbarrier_test_wait_parity_impl(
     require_mbarrier_ptr(mbar, (MemorySpace.SHARED,))
     parity = astype(parity, datatype.int32)
     scope = require_constant_enum(scope, MbarrierScope)
-    memory_order = require_mbarrier_ordering(memory_order, WAIT_ORDERINGS)
+    memory_order = require_mbarrier_memory_order(memory_order, WAIT_MEMORY_ORDERS)
     intrinsic = "llvm.nvvm.mbarrier.test.wait.parity"
     if memory_order is MemoryOrder.RELAXED:
         intrinsic += ".relaxed"
@@ -220,7 +220,7 @@ def mbarrier_try_wait_impl(
     require_mbarrier_ptr(mbar, (MemorySpace.SHARED,))
     state = astype(state, datatype.int64)
     scope = require_constant_enum(scope, MbarrierScope)
-    memory_order = require_mbarrier_ordering(memory_order, WAIT_ORDERINGS)
+    memory_order = require_mbarrier_memory_order(memory_order, WAIT_MEMORY_ORDERS)
     intrinsic = "llvm.nvvm.mbarrier.try.wait"
     args = (mbar, state)
     if not is_none(time_hint):
@@ -249,7 +249,7 @@ def mbarrier_try_wait_parity_impl(
     require_mbarrier_ptr(mbar, (MemorySpace.SHARED,))
     parity = astype(parity, datatype.int32)
     scope = require_constant_enum(scope, MbarrierScope)
-    memory_order = require_mbarrier_ordering(memory_order, WAIT_ORDERINGS)
+    memory_order = require_mbarrier_memory_order(memory_order, WAIT_MEMORY_ORDERS)
     intrinsic = "llvm.nvvm.mbarrier.try.wait.parity"
     args = (mbar, parity)
     if not is_none(time_hint):

experimental/cuda-lang/src/cuda/lang/_ir/op_impl/pointer_impl.py

@@ -210,7 +210,7 @@ def pointer_getitem(object: Var[PointerTy], key: Var[Type]):
         pointer=pointer,
         volatile=False,
         alignment=None,
-        ordering=None,
+        memory_order=None,
     )
 
 
@@ -226,7 +226,7 @@ def pointer_setitem(object: Var[PointerTy], key: Var[Type], value: Var[Type]):
         value=value,
         alignment=None,
         volatile=False,
-        ordering=None,
+        memory_order=None,
     )
 
 
@@ -258,7 +258,7 @@ def array_setitem(object: Var, key: Var, value: Var):
         value=value,
         alignment=None,
         volatile=False,
-        ordering=None,
+        memory_order=None,
     )
 
 
@@ -277,14 +277,14 @@ def pointer_load(
     count: Var,
     alignment: Var,
     volatile: Var,
-    ordering: Var,
+    memory_order: Var,
 ) -> Var:
     pointee_dtype = require_pointer_type(pointer).pointee_dtype
     count = require_optional_constant_int(count)
     volatile = require_constant_bool(volatile)
     alignment = require_optional_alignment(alignment)
-    ordering = require_pointer_memory_order(LoadPointer, ordering)
-    is_atomic = ordering not in (None, MemoryOrder.WEAK)
+    memory_order = require_pointer_memory_order(LoadPointer, memory_order)
+    is_atomic = memory_order not in (None, MemoryOrder.WEAK)
     if count is None or count == 1:
         if is_atomic:
             _require_atomic_scalar_dtype(pointee_dtype, "load")
@@ -303,7 +303,7 @@ def pointer_load(
         pointer=pointer,
         volatile=volatile,
         alignment=alignment,
-        ordering=ordering,
+        memory_order=memory_order,
     )
 
 
@@ -312,18 +312,18 @@ def pointer_store(
     value: Var,
     alignment: Var,
     volatile: Var,
-    ordering: Var,
+    memory_order: Var,
 ) -> None:
     pointer_ty = require_pointer_type(pointer)
     volatile = require_constant_bool(volatile)
     alignment = require_optional_alignment(alignment)
-    ordering = require_pointer_memory_order(StorePointer, ordering)
+    memory_order = require_pointer_memory_order(StorePointer, memory_order)
 
     pointee_dtype = pointer_ty.pointee_dtype
     value = implicit_cast(value, pointee_dtype,
                           "Stored value type is incompatible with pointer type")
 
-    is_atomic = ordering not in (None, MemoryOrder.WEAK)
+    is_atomic = memory_order not in (None, MemoryOrder.WEAK)
     if is_atomic:
         if isinstance(value.get_type(), VectorTy):
             raise TypeCheckingError(
@@ -340,7 +340,7 @@ def pointer_store(
         value=value,
         volatile=volatile,
         alignment=alignment,
-        ordering=ordering,
+        memory_order=memory_order,
     )
 
 

experimental/cuda-lang/src/cuda/lang/_ir/type_checking_helpers.py

@@ -235,19 +235,19 @@ def require_optional_alignment(alignment: Var) -> int | None:
 
 def require_pointer_memory_order(
     operation: type[LoadPointer] | type[StorePointer],
-    ordering_var: Var,
+    memory_order_var: Var,
 ):
-    ordering = require_optional_constant_enum(ordering_var, MemoryOrder)
-    if ordering in operation.valid_orderings:
-        return ordering
+    memory_order = require_optional_constant_enum(memory_order_var, MemoryOrder)
+    if memory_order in operation.valid_memory_orders:
+        return memory_order
 
     formatted_expected = ", ".join(
-        "None" if order is None else str(order) for order in operation.valid_orderings
+        "None" if order is None else str(order) for order in operation.valid_memory_orders
     )
     operation_name = "load" if operation is LoadPointer else "store"
     raise make_type_checking_error(
         f"Invalid memory order for Pointer.{operation_name}. "
-        f"Got {ordering}, expected one of {formatted_expected}"
+        f"Got {memory_order}, expected one of {formatted_expected}"
     )
 
 

experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py

@@ -834,7 +834,7 @@ def lower_vector_getitem(self, operation: ops.VectorGetItem) -> Sequence[mlir.Va
     @lower_operation.register
     def lower_load_pointer(self, operation: ops.LoadPointer) -> Sequence[mlir.Value]:
         ptr_dtype = operation.pointer.get_type().pointer_dtype
-        ordering = _get_llvm_memory_ordering(operation.ordering)
+        ordering = _get_llvm_memory_ordering(operation.memory_order)
         info = PointerInfo(ptr_dtype)
         assert not info.opaque, f"Expected a typed pointer, got {ptr_dtype}"
         result_type = ir_type_to_mlir_type(operation.result_var.get_type())
@@ -851,7 +851,7 @@ def lower_load_pointer(self, operation: ops.LoadPointer) -> Sequence[mlir.Value]
     @lower_operation.register
     def lower_store_pointer(self, operation: ops.StorePointer) -> Sequence[mlir.Value]:
         pointer = self.get_var(operation.pointer)
-        ordering = _get_llvm_memory_ordering(operation.ordering)
+        ordering = _get_llvm_memory_ordering(operation.memory_order)
         value = self.get_var(operation.value)
         mlir.llvm.add_StoreOp(
             value=value,

experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py

@@ -10,8 +10,8 @@
 from cuda.tile._memory_model import MemoryOrder
 
 
-ArriveOrdering = Literal[MemoryOrder.RELAXED, MemoryOrder.RELEASE]
-WaitOrdering = Literal[MemoryOrder.RELAXED, MemoryOrder.ACQUIRE]
+ArriveMemoryOrder = Literal[MemoryOrder.RELAXED, MemoryOrder.RELEASE]
+WaitMemoryOrder = Literal[MemoryOrder.RELAXED, MemoryOrder.ACQUIRE]
 
 
 @stub
@@ -42,7 +42,7 @@ def mbarrier_arrive(
     *,
     drop: bool = False,
     scope: MbarrierScope = MbarrierScope.BLOCK,
-    memory_order: ArriveOrdering = MemoryOrder.RELEASE,
+    memory_order: ArriveMemoryOrder = MemoryOrder.RELEASE,
 ) -> "uint64 | None":
     """Arrive at ``mbar``. When the mbarrier resides in ``MemorySpace.SHARED``,
     an opaque 64-bit value capturing the phase of the mbarrier object _prior_
@@ -69,7 +69,7 @@ def mbarrier_arrive_expect_transaction(
     *,
     drop: bool = False,
     scope: MbarrierScope = MbarrierScope.BLOCK,
-    memory_order: ArriveOrdering = MemoryOrder.RELEASE,
+    memory_order: ArriveMemoryOrder = MemoryOrder.RELEASE,
 ) -> "uint64 | None":
     """Arrive at ``mbar`` and add expected transaction bytes.
 
@@ -127,7 +127,7 @@ def mbarrier_test_wait(
     state,
     *,
     scope: MbarrierScope = MbarrierScope.BLOCK,
-    memory_order: WaitOrdering = MemoryOrder.ACQUIRE,
+    memory_order: WaitMemoryOrder = MemoryOrder.ACQUIRE,
 ) -> "bool_":
     """Non-blocking test whether ``mbar`` has completed.
 
@@ -148,7 +148,7 @@ def mbarrier_test_wait_parity(
     parity: int,
     *,
     scope: MbarrierScope = MbarrierScope.BLOCK,
-    memory_order: WaitOrdering = MemoryOrder.ACQUIRE,
+    memory_order: WaitMemoryOrder = MemoryOrder.ACQUIRE,
 ) -> "bool_":
     """Phase-parity variant of ``mbarrier_test_wait``.
     ``parity`` is the 0/1 integer parity of the phase to test for.
@@ -171,7 +171,7 @@ def mbarrier_try_wait(
     *,
     time_hint: int | None = None,
     scope: MbarrierScope = MbarrierScope.BLOCK,
-    memory_order: WaitOrdering = MemoryOrder.ACQUIRE,
+    memory_order: WaitMemoryOrder = MemoryOrder.ACQUIRE,
 ) -> "bool_":
     """Bounded-wait test whether ``mbar`` has completed.
 
@@ -202,7 +202,7 @@ def mbarrier_try_wait_parity(
     *,
     time_hint: int | None = None,
     scope: MbarrierScope = MbarrierScope.BLOCK,
-    memory_order: WaitOrdering = MemoryOrder.ACQUIRE,
+    memory_order: WaitMemoryOrder = MemoryOrder.ACQUIRE,
 ) -> "bool_":
     """Phase-parity variant of ``mbarrier_try_wait``.
     ``parity`` is the 0/1 integer parity of the phase to test for.