[lang] Add load/store matrix intrinsics

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

Author: ashermancinelli

experimental/cuda-lang/docs/source/operations.rst

@@ -171,6 +171,25 @@ Warp shuffle
     shfl_xor_sync
 
 
+Warp Matrix Load and Store
+--------------------------
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+
+    load_matrix
+    store_matrix
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: autosummary/class_no_init.rst
+
+    MatrixLoadShape
+    MatrixStoreShape
+    MatrixLoadSourceFormat
+
+
 .. _operations-tensor-map:
 
 TensorMap

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

@@ -34,6 +34,15 @@
     prefetch_uniform,
     prefetch_tensor_map,
 )
+from ._stub.load_store_matrix import (
+    load_matrix,
+    store_matrix,
+)
+from ._enums import (
+    MatrixLoadShape,
+    MatrixLoadSourceFormat,
+    MatrixStoreShape,
+)
 from ._stub.math import (
     add,
     sub,
@@ -464,4 +473,9 @@
     "prefetch",
     "prefetch_uniform",
     "prefetch_tensor_map",
+    "load_matrix",
+    "store_matrix",
+    "MatrixLoadShape",
+    "MatrixLoadSourceFormat",
+    "MatrixStoreShape",
 )

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

@@ -141,6 +141,22 @@ class PrefetchLevel(Enum):
     L2 = auto()
 
 
+class MatrixLoadShape(Enum):
+    M8N8 = "m8n8"
+    M8N16 = "m8n16"
+    M16N16 = "m16n16"
+
+
+class MatrixStoreShape(Enum):
+    M8N8 = "m8n8"
+    M16N8 = "m16n8"
+
+
+class MatrixLoadSourceFormat(Enum):
+    B6X16_P32 = "b6x16_p32"
+    B4X16_P64 = "b4x16_p64"
+
+
 __all__ = (
     "MemorySpace",
     "MemoryScope",
@@ -164,4 +180,7 @@ class PrefetchLevel(Enum):
     "BarrierReductionKind",
     "CachePolicy",
     "PrefetchLevel",
+    "MatrixStoreShape",
+    "MatrixLoadShape",
+    "MatrixLoadSourceFormat",
 )

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

@@ -0,0 +1,170 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import cuda.lang._datatype as datatype
+from cuda.lang._enums import (
+    MatrixLoadShape,
+    MatrixLoadSourceFormat,
+    MatrixStoreShape,
+)
+from cuda.lang._exception import InvalidValueError, TypeCheckingError
+from cuda.lang._ir.ir import Var
+from cuda.lang._ir.op_defs import RawNVVMIntrinsic
+from cuda.lang._ir.op_impl.vector_impl import (
+    vector_getitem,
+    vector_undef,
+    vector_with_item,
+)
+from cuda.lang._ir.type import MemorySpace, ScalarTy, VectorTy
+from cuda.lang._ir.type_checking_helpers import (
+    is_none,
+    require_pointer_in_memory_space,
+)
+from cuda.lang._stub import load_store_matrix
+from cuda.tile._datatype import is_integral
+from cuda.tile._ir.op_impl import (
+    ImplRegistry,
+    require_constant_bool,
+    require_constant_enum,
+    require_constant_int,
+)
+from cuda.tile._ir.core_ops import strictly_typed_const
+from cuda.tile._ir.ir import add_operation_variadic
+
+
+_registry = ImplRegistry()
+impl = _registry.impl
+
+
+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,
+    shape: Var,
+    count: Var,
+    transpose: Var,
+    source_format: Var,
+) -> Var:
+    require_pointer_in_memory_space(src, (MemorySpace.SHARED,))
+    shape_value = require_constant_enum(shape, MatrixLoadShape)
+    count_value = require_constant_int(count)
+    transpose_value = require_constant_bool(transpose)
+    source_format_value = (
+        None
+        if is_none(source_format)
+        else require_constant_enum(source_format, MatrixLoadSourceFormat)
+    )
+    if count_value not in (1, 2, 4):
+        raise InvalidValueError("count must be 1, 2, or 4")
+    register_count = count_value * (
+        2 if shape_value is MatrixLoadShape.M16N16 else 1
+    )
+    register_type = ScalarTy(datatype.int32)
+
+    name = f"llvm.nvvm.ldmatrix.sync.aligned.{shape_value.value}.x{count_value}"
+
+    if transpose_value:
+        name += ".trans"
+
+    if source_format_value is MatrixLoadSourceFormat.B6X16_P32:
+        name += ".b8x16.b6x16_p32"
+    elif source_format_value is MatrixLoadSourceFormat.B4X16_P64:
+        name += ".b8x16.b4x16_p64"
+    else:
+        name += ".b16" if shape_value is MatrixLoadShape.M8N8 else ".b8"
+
+    registers = add_operation_variadic(
+        RawNVVMIntrinsic,
+        (register_type,) * register_count,
+        intrinsic=name,
+        operands_=(src,),
+    )
+    if register_count == 1:
+        return registers[0]
+
+    result = vector_undef(VectorTy(datatype.int32, register_count))
+    for index, register in enumerate(registers):
+        result = vector_with_item(result, index, register)
+    return result
+
+
+def _store_register_count(values: Var) -> int:
+    value_type = values.get_type()
+    match value_type:
+        case ScalarTy() as st:
+            dtype = st.dtype
+            count = 1
+        case VectorTy() as vt:
+            dtype = vt.element_dtype
+            count = vt.length
+        case _:
+            raise TypeCheckingError(
+                "Expected a scalar or vector of 32-bit integers"
+            )
+
+    if not is_integral(dtype) or dtype.bitwidth != 32:
+        # TODO: is this too restrictive? should we bitcast if the operand is
+        # not integral but is 32 bits wide?
+        raise TypeCheckingError(
+            "Expected a scalar or vector of 32-bit integers, "
+            f"but got {value_type}"
+        )
+    if count not in (1, 2, 4):
+        raise InvalidValueError("Matrix store register count must be 1, 2, or 4")
+    return count
+
+
+@impl(load_store_matrix.store_matrix)
+def store_matrix_impl(
+    dst: Var,
+    values: Var,
+    shape: Var,
+    transpose: Var,
+) -> None:
+    require_pointer_in_memory_space(dst, (MemorySpace.SHARED,))
+    shape_value = require_constant_enum(shape, MatrixStoreShape)
+    transpose_value = require_constant_bool(transpose)
+    register_count = _store_register_count(values)
+    if shape_value is MatrixStoreShape.M16N8 and not transpose_value:
+        raise InvalidValueError("M16N8 requires transpose=True")
+
+    if register_count == 1:
+        registers = (values,)
+    else:
+        index_type = ScalarTy(datatype.int32)
+        registers = tuple(
+            vector_getitem(values, strictly_typed_const(index, index_type))
+            for index in range(register_count)
+        )
+
+    name = f"llvm.nvvm.stmatrix.sync.aligned.{shape_value.value}.x{register_count}"
+    if transpose_value:
+        name += ".trans"
+    name += ".b16" if shape_value is MatrixStoreShape.M8N8 else ".b8"
+
+    add_operation_variadic(
+        RawNVVMIntrinsic,
+        (),
+        intrinsic=name,
+        operands_=(dst, *registers),
+    )

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

@@ -157,6 +157,7 @@
 from .op_impl.barrier_impl import barrier_impl_registry
 from .op_impl.mbarrier_impl import mbarrier_impl_registry
 from .op_impl.inline_ptx_impl import inline_ptx_impl_registry
+from .op_impl.matrix_impl import matrix_impl_registry
 
 cuda_lang_impl_registry = ImplRegistry()
 cuda_lang_impl_registry.update(core_impl_registry())
@@ -174,6 +175,7 @@
 cuda_lang_impl_registry.update(copy_async_impl_registry())
 cuda_lang_impl_registry.update(barrier_impl_registry())
 cuda_lang_impl_registry.update(mbarrier_impl_registry())
+cuda_lang_impl_registry.update(matrix_impl_registry())
 
 impl = cuda_lang_impl_registry.impl
 

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

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .._enums import (
+    MatrixStoreShape,
+    MatrixLoadShape,
+    MatrixLoadSourceFormat,
+)
+from cuda.lang._execution import stub
+from typing import Literal
+
+
+@stub()
+def load_matrix(
+    src,
+    /,
+    *,
+    shape: MatrixLoadShape,
+    count: Literal[1, 2, 4] = 1,
+    transpose: bool = False,
+    source_format: MatrixLoadSourceFormat | None = None,
+):
+    """
+    Collectively load one or more matrices from shared memory for mma instruction
+
+    Args:
+        src: Pointer to a matrix row in shared memory.
+        shape: Shape and element size.
+        count: Number of matrices to load.
+        transpose: False loads the matrix in row-major order, True loads in
+            column-major order.
+        source_format: Packed source format.
+
+    Returns:
+        Scalar or vector of 32 bit integers depending on ``count``.
+    """
+    ...
+
+
+@stub()
+def store_matrix(
+    dst,
+    values,
+    /,
+    *,
+    shape: MatrixStoreShape,
+    transpose: bool = False,
+):
+    """
+    Collectively store one or more matrices to shared memory.
+
+    Args:
+        dst: Pointer to a matrix row in shared memory.
+        values: One 32-bit scalar or a ``Vector`` of 1, 2, or 4 32-bit values.
+            The number of values selects the number of matrices.
+        shape: Matrix shape.
+        transpose: False stores the matrix in row-major order, True stores in
+            column-major order.
+    """
+    ...