[lang] Fix inline ptx pointer outputs and reject read-write operands

Signed-off-by: Asher Mancinelli <amancinelli@nvidia.com>

Author: ashermancinelli

experimental/cuda-lang/src/cuda/lang/__init__.py

@@ -168,6 +168,7 @@
     memory_barrier,
     bitcast,
     assert_,
+    clock,
 )
 from cuda.tile._stub import (
     Constant,
@@ -500,4 +501,5 @@
     "MatrixLoadSourceFormat",
     "MatrixStoreShape",
     "_debug",
+    "clock",
 )

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

@@ -62,14 +62,14 @@ def require_inline_ptx_pair(var: Var) -> tuple[Var, Var]:
     "+": InlinePTX.RMWMode.READ_WRITE,
 }
 
-_INLINE_PTX_TYPECODES = {
+_INLINE_PTX_TYPECODES = (
     "h",
     "r",
     "l",
     "f",
     "d",
-    "C",
-}
+    "p",
+)
 
 _INLINE_PTX_SCALAR_DTYPE_FROM_TYPECODE = {
     "h": datatype.int16,
@@ -97,12 +97,13 @@ def parse_inline_ptx_constraint(var: Var) -> tuple[str, InlinePTX.RMWMode, str]:
             f"Unknown constraint rmw modifier {prefix!r}, expected "
             "'' (meaning readonly), '+' (meaning readwrite), or '=' (meaning writeonly)"
         )
-
     if type_char not in _INLINE_PTX_TYPECODES:
         expected = ", ".join(_INLINE_PTX_TYPECODES)
         raise TypeCheckingError(
             f"Unknown constraint dtype {type_char!r}, expected one of {expected}"
         )
+    if mode is InlinePTX.RMWMode.READ_WRITE:
+        raise TypeCheckingError("Read-write inline_ptx constraints are not supported")
 
     return constraint_str, mode, type_char
 
@@ -111,14 +112,17 @@ def validate_inline_ptx_operand(
     constraint_str: str, mode: InlinePTX.RMWMode, type_char: str, value: Var
 ) -> InlinePTXOperand:
     if mode is InlinePTX.RMWMode.WRITE_ONLY:
-        if type_char == "C":
-            # write-only arguments require specifying the output data type, but we don't
-            # expose a dtype for pointers. Disallow this for now.
-            raise TypeCheckingError(
-                "Write-only pointer outputs are not supported for inline_ptx"
+        actual_dtype = require_dtype_spec(value)
+        if type_char == "p":
+            if not is_pointer_dtype(actual_dtype):
+                raise TypeCheckingError(
+                    f"Expected a pointer dtype for constraint {constraint_str}, "
+                    f"got {actual_dtype}"
+                )
+            return InlinePTXOperand(
+                mode=mode, type_code=type_char, value=actual_dtype
             )
 
-        actual_dtype = require_dtype_spec(value)
         expected_dtype = _INLINE_PTX_SCALAR_DTYPE_FROM_TYPECODE[type_char]
         if actual_dtype != expected_dtype:
             raise TypeCheckingError(
@@ -127,7 +131,7 @@ def validate_inline_ptx_operand(
             )
         return InlinePTXOperand(mode=mode, type_code=type_char, value=actual_dtype)
 
-    if type_char == "C":
+    if type_char == "p":
         require_pointer_type(value)
         return InlinePTXOperand(mode=mode, type_code=type_char, value=value)
 
@@ -192,7 +196,15 @@ def rewrite(match: re.Match[str]) -> str:
 
         return ptx_interpolation_replacements[index]
 
-    mlir_ptx_code = _INLINE_PTX_PLACEHOLDER_RE.sub(rewrite, ptx_code)
+    ptx_fragments = []
+    for fragment in ptx_code.split("%%"):
+        fragment = _INLINE_PTX_PLACEHOLDER_RE.sub(rewrite, fragment)
+        if "%" in fragment:
+            raise TypeCheckingError(
+                "Literal percent signs in inline_ptx must be escaped as '%%'"
+            )
+        ptx_fragments.append(fragment)
+    mlir_ptx_code = "%".join(ptx_fragments)
     return (
         mlir_ptx_code,
         tuple(ro_args),

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

@@ -41,22 +41,6 @@ def matrix_impl_registry() -> ImplRegistry:
     return _registry
 
 
-def ldmatrix_intrinsic_name(
-    shape: MatrixLoadShape,
-    count: int,
-    transpose: bool,
-    source_format: MatrixLoadSourceFormat | None,
-) -> str:
-    name = f"llvm.nvvm.ldmatrix.sync.aligned.{shape.value}.x{count}"
-    if transpose:
-        name += ".trans"
-    if source_format is MatrixLoadSourceFormat.B6X16_P32:
-        return name + ".b8x16.b6x16_p32"
-    if source_format is MatrixLoadSourceFormat.B4X16_P64:
-        return name + ".b8x16.b4x16_p64"
-    return name + (".b16" if shape is MatrixLoadShape.M8N8 else ".b8")
-
-
 @impl(load_store_matrix.load_matrix)
 def load_matrix_impl(
     src: Var,

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

@@ -7,10 +7,15 @@
 import cuda.lang as cl
 from cuda.lang._execution import stub, function
 from cuda.lang._exception import TypeCheckingError
-from cuda.tile._stub import Array as TileArray, cdiv as tile_cdiv
+from cuda.tile._stub import (
+    Array as TileArray,
+    cdiv as tile_cdiv,
+    static_assert,
+    static_eval,
+)
 from cuda.lang._enums import MemoryOrder
 from cuda.tile._memory_model import MemoryScope, MemorySpace
-from cuda.lang._datatype import DType
+from cuda.lang._datatype import DType, uint32, uint64
 from .types import Pointer, Scalar, Vector
 
 T = TypeVar("T")
@@ -408,7 +413,7 @@ def vote_ballot_sync(predicate: bool, mask: int = FULL_MASK) -> int:
 def _inline_ptx(ptx_code: str, *constraint_pairs: tuple) -> tuple:
     """Execute inline PTX.
 
-    The API mirrors CUDA C++'s device-side `asm` statement:
+    The API follows CUDA C++'s device-side `asm` statement:
     `cl._inline_ptx(ptx_code, (constraint1, value1), (constraint2, value2), ...)`.
 
     Args:
@@ -419,14 +424,12 @@ def _inline_ptx(ptx_code: str, *constraint_pairs: tuple) -> tuple:
             Constraints must be compile-time constant strings.
 
             Read-only operands use constraints ``"h"``, ``"r"``, ``"l"``,
-            ``"f"``, ``"d"``, or ``"C"`` and are paired with runtime values.
+            ``"f"``, ``"d"``, or ``"p"`` and are paired with runtime values.
 
             Write-only operands use constraints ``"=h"``, ``"=r"``, ``"=l"``,
-            ``"=f"``, or ``"=d"`` and are paired with dtype specs.
-            This determines the type of the output.
-
-            Read-write operands use constraints ``"+h"``, ``"+r"``, ``"+l"``,
-            ``"+f"``, ``"+d"``, or ``"+C"`` and are paired with runtime values.
+            ``"=f"``, ``"=d"``, or ``"=p"`` and are paired with dtype specs.
+            Use a pointer dtype with ``"=p"``. The dtype determines the type
+            of the output.
 
     Returns:
 
@@ -467,7 +470,11 @@ def _inline_ptx(ptx_code: str, *constraint_pairs: tuple) -> tuple:
             - ``l``: ``cl.int64``
             - ``f``: ``cl.float32``
             - ``d``: ``cl.float64``
-            - ``C``: pointer value from ``array.get_base_pointer()``
+            - ``p``: pointer value, or a pointer dtype for an output
+        - ``p`` is not in CUDA C++'s inline ptx. The compiler selects the
+          correct register size for the pointer based on its address space.
+        - Use ``%0``, ``%1``, and so on for operands. Escape literal percent
+          signs with a second percent sign, as in ``%%clock``.
 
     """
 
@@ -477,6 +484,15 @@ def ptx_comment(comment: str):
     _inline_ptx(cl.static_eval("// " + comment))
 
 
+@function()
+def clock(dtype=uint32):
+    static_assert(dtype in (uint32, uint64))
+    fn = static_eval(
+        nvvm.read_ptx_sreg_clock if dtype is uint32 else nvvm.read_ptx_sreg_clock64
+    )
+    return fn()
+
+
 @stub
 def atomic_add(
     ptr: Pointer[T],

experimental/cuda-lang/test/test_inline_ptx.py

@@ -4,7 +4,7 @@
 
 import pytest
 import cuda.lang as cl
-from cuda.lang._exception import TypeCheckingError
+from cuda.lang._exception import TypeCheckingError, StaticAssertionError
 import torch
 
 from .util import compile_kernel
@@ -59,7 +59,7 @@ def kernel(inp, out):
         (value,) = cl._inline_ptx(
             "ld.global.u32 %0, [%1];",
             ("=r", cl.int32),
-            ("C", inp_ptr),
+            ("p", inp_ptr),
         )
         out[0] = value
 
@@ -69,6 +69,74 @@ def kernel(inp, out):
     assert out.cpu().tolist() == [42]
 
 
+def test_inline_ptx_pointer_output():
+    @cl.kernel
+    def kernel(inp, out):
+        inp_ptr = inp.get_base_pointer()
+        dtype = cl.pointer_dtype(cl.int32)
+        (ptr,) = cl._inline_ptx(
+            "mov.u64 %0, %1;",
+            ("=p", dtype),
+            ("p", inp_ptr),
+        )
+        cl.static_assert(cl.dtype_of(ptr) == dtype)
+        out[0] = ptr.load()
+
+    inp = torch.tensor([42], dtype=torch.int32, device="cuda")
+    out = torch.zeros(1, dtype=torch.int32, device="cuda")
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (inp, out))
+    assert out.cpu().tolist() == [42]
+
+
+def test_inline_ptx_shared_pointer_output():
+    @cl.kernel
+    def kernel(out):
+        shared = cl.shared_array(1, cl.int32)
+        shared[0] = 42
+        shared_ptr = shared.get_base_pointer()
+        dtype = cl.pointer_dtype(cl.int32, cl.MemorySpace.SHARED)
+        (result,) = cl._inline_ptx(
+            "mov.u32 %0, %1;",
+            ("=p", dtype),
+            ("p", shared_ptr),
+        )
+        cl.static_assert(cl.dtype_of(result) == dtype)
+        out[0] = result.load()
+
+    out = torch.zeros(1, dtype=torch.int32, device="cuda")
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (out,))
+    assert out.cpu().tolist() == [42]
+
+
+def test_inline_ptx_special_register_operand():
+    @cl.kernel
+    def kernel():
+        clock = cl._nvvm.read_ptx_sreg_clock()
+        cl._inline_ptx(
+            "mov.u32 %0, %1;",
+            ("=r", cl.int32),
+            ("r", clock),
+        )
+
+    compile_kernel(kernel, assert_in_ptx="%clock")
+
+
+@pytest.mark.xfail(
+    strict=True,
+    reason="needs llvm version bump",
+)
+def test_inline_ptx_escaped_special_register():
+    @cl.kernel
+    def kernel():
+        cl._inline_ptx("mov.u32 %0, %%clock;", ("=r", cl.int32))
+
+    compile_kernel(
+        kernel,
+        assert_in_ptx="%clock",
+        assert_not_in_ptx="%%clock",
+    )
+
+
 class TestInlinePTXErrors:
 
     def test_invalid_type_constraint(self):
@@ -82,6 +150,53 @@ def kernel():
             ),
         )
 
+    def test_cuda_c_constraint_is_not_supported(self):
+        def kernel():
+            cl._inline_ptx("// no operation", ("C", 0))
+
+        compile_kernel(
+            kernel,
+            raises=pytest.raises(
+                TypeCheckingError, match="Unknown constraint dtype 'C'"
+            ),
+        )
+
+    def test_pointer_output_requires_pointer_dtype(self):
+        def kernel():
+            cl._inline_ptx("mov.u64 %0, 0;", ("=p", cl.int64))
+
+        compile_kernel(
+            kernel,
+            raises=pytest.raises(
+                TypeCheckingError,
+                match="Expected a pointer dtype for constraint =p, got int64",
+            ),
+        )
+
+    def test_read_write_constraint_is_not_supported(self):
+        def kernel():
+            cl._inline_ptx("add.u32 %0, %0, 1;", ("+r", 2))
+
+        compile_kernel(
+            kernel,
+            raises=pytest.raises(
+                TypeCheckingError,
+                match="Read-write inline_ptx constraints are not supported",
+            ),
+        )
+
+    def test_special_register_is_not_supported(self):
+        def kernel():
+            cl._inline_ptx("mov.u32 %0, %clock;", ("=r", cl.int32))
+
+        compile_kernel(
+            kernel,
+            raises=pytest.raises(
+                TypeCheckingError,
+                match="Literal percent signs in inline_ptx must be escaped",
+            ),
+        )
+
     def test_invalid_rmw_constraint(self):
         def kernel():
             cl._inline_ptx(
@@ -96,3 +211,25 @@ def kernel():
                 TypeCheckingError, match="Unknown constraint rmw modifier '@'"
             ),
         )
+
+
+@pytest.mark.parametrize("dtype", (cl.uint32, cl.uint64))
+def test_clock(dtype):
+    @cl.kernel
+    def kernel():
+        cl.shared_array(1, dtype)[0] = cl.clock(dtype)
+
+    check = "%clock"
+    if dtype is cl.uint64:
+        check += "64"
+
+    compile_kernel(kernel, assert_in_ptx=check)
+
+
+@pytest.mark.parametrize("dtype", (cl.int16, cl.int32, cl.int64, cl.float32, cl.bool_))
+def test_clock_invalid_dtype(dtype):
+    @cl.kernel
+    def kernel():
+        cl.shared_array(1, dtype)[0] = cl.clock(dtype)
+
+    compile_kernel(kernel, raises=pytest.raises(StaticAssertionError))