From 76e1db7ff57f177353a75dc6af1285f4e9721699 Mon Sep 17 00:00:00 2001 From: Asher Mancinelli Date: Mon, 5 Oct 2026 09:42:05 -0700 Subject: [PATCH 1/6] [lang] Make target info observable in hir2ir Signed-off-by: Asher Mancinelli --- .../cuda-lang/src/cuda/lang/_compile.py | 36 ++++++++++++------- .../cuda-lang/src/cuda/lang/_ir/ir.py | 9 +++++ .../cuda/lang/_ir/op_impl/mbarrier_impl.py | 8 ++++- .../cuda/lang/_ir/type_checking_helpers.py | 21 +++++++++++ .../lang/_passes/ir2mlir/pass_definition.py | 8 +++-- .../cuda-lang/src/cuda/lang/_target.py | 27 ++++++-------- experimental/cuda-lang/test/test_math.py | 20 ++++++----- experimental/cuda-lang/test/test_mbarrier.py | 21 ++++++++++- experimental/cuda-lang/test/test_target.py | 30 ++++++++++++++++ 9 files changed, 138 insertions(+), 42 deletions(-) diff --git a/experimental/cuda-lang/src/cuda/lang/_compile.py b/experimental/cuda-lang/src/cuda/lang/_compile.py index c4afaf46..51b07eb6 100644 --- a/experimental/cuda-lang/src/cuda/lang/_compile.py +++ b/experimental/cuda-lang/src/cuda/lang/_compile.py @@ -2,6 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 +from cuda.tile._exception import InternalError from copy import deepcopy from functools import total_ordering import sys @@ -326,7 +327,27 @@ def _dump(phase: str, contents: object) -> None: if signature.symbol is None: signature = signature.with_mangled_symbol(function.pyfunc.__name__) - ctx = ctx or ir.IRContext(log_ir_on_error=log_flags.log_hir or log_flags.log_ir) + with timer.phase("ir.target_resolution"): + if gpu_name is None or arch is None: + cc = compute_capability or get_compute_capability() + suffix = "a" if cc >= (9, 0) else "" + gpu_name = gpu_name or cc.gpu_name + suffix + arch = arch or cc.arch + suffix + + target_info = TargetInfo.from_arch(arch) + + if ctx is None: + ctx = ir.IRContext( + log_ir_on_error=log_flags.log_hir or log_flags.log_ir, + target_info=target_info, + ) + if ctx.execution_space != "device": + raise ValueError("Cannot compile device code with a host IR context") + if ctx.target_info != target_info: + raise InternalError( + f"Target info on supplied context ({ctx.target_info}) " + f"does not match target info inferred from the architecture ({target_info})" + ) with timer.phase("ir.hir2ir"): func_ir = get_function_ir( @@ -347,15 +368,6 @@ def _dump(phase: str, contents: object) -> None: if log_flags.log_flattened_ir: _dump("Flattened IR", flattened_ir) - with timer.phase("ir.target_resolution"): - if gpu_name is None or arch is None: - cc = compute_capability or get_compute_capability() - suffix = "a" if cc >= (9, 0) else "" - gpu_name = gpu_name or cc.gpu_name + suffix - arch = arch or cc.arch + suffix - - target_info = TargetInfo.from_arch(arch) - need_nvvm = log_flags.log_nvvm or keep_nvvm need_ptx = log_flags.log_ptx or keep_ptx @@ -367,10 +379,10 @@ def _dump(phase: str, contents: object) -> None: program = nvvm.create_program() program.add_module(bitcode, "main") program.add_module(libdevice, "libdevice") - ptx = program.compile(["-arch=" + cc.arch]) + ptx = program.compile(["-arch=" + arch]) ptx_compiler: PtxCompiler = PtxCompiler.get() - cubin = ptx_compiler.compile(ptx, cc.gpu_name) + cubin = ptx_compiler.compile(ptx, gpu_name) if need_nvvm: # FIXME: actually find the llvm-dis instead of assuming it is in the PATH diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/ir.py b/experimental/cuda-lang/src/cuda/lang/_ir/ir.py index ccc16887..be79ac67 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/ir.py +++ b/experimental/cuda-lang/src/cuda/lang/_ir/ir.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field import itertools from collections import defaultdict +from cuda.lang._target import TargetInfo from cuda.tile._ir.ir import ( Block as TileBlock, Builder as TileBuilder, @@ -92,6 +93,7 @@ def __init__( log_ir_on_error: bool = True, *, execution_space: ExecutionSpace = "device", + target_info: TargetInfo | None = None, ): from cuda.lang._ir.type import LangTypingHooks self._block_names: dict[int, str] = {} @@ -99,6 +101,13 @@ def __init__( super().__init__(log_ir_on_error, tileiras_version=None, typing_hooks=LangTypingHooks(), execution_space=execution_space) + if execution_space == "host" and target_info is not None: + raise ValueError("Cannot set a GPU target on a host IR context") + self._target_info = target_info + + @property + def target_info(self) -> TargetInfo | None: + return self._target_info def make_block(self, name: str, loc: Loc, params: tuple[Var, ...] = ()) -> Block: block = Block(self, loc) diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py index baa0054a..babf12e2 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py +++ b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py @@ -8,7 +8,11 @@ from cuda.lang._ir.ir import Var, add_operation from cuda.lang._ir.op_defs import RawLLVMIntrinsic from cuda.lang._ir.type import MemorySpace, ScalarTy -from cuda.lang._ir.type_checking_helpers import is_none, require_mbarrier_ptr +from cuda.lang._ir.type_checking_helpers import ( + is_none, + require_mbarrier_ptr, + require_target_greater_equal, +) from cuda.lang._stub import mbarrier from .inline_ptx_impl import inline_ptx from cuda.tile._ir.arithmetic_ops import astype @@ -33,6 +37,7 @@ def mbarrier_initialize_impl(mbar: Var, participants: Var, layout: Var) -> None: require_mbarrier_ptr(mbar, (MemorySpace.SHARED,)) layout_suffix = "" if not is_none(layout): + require_target_greater_equal(9, 0, "mbarrier_initialize with an explicit layout") layout = require_constant_enum(layout, MbarrierLayout) layout_suffix = f".layout::v{layout.value}" participants = astype(participants, datatype.int32) @@ -42,6 +47,7 @@ def mbarrier_initialize_impl(mbar: Var, participants: Var, layout: Var) -> None: @impl(mbarrier.mbarrier_has_layout) def mbarrier_has_layout_impl(mbar: Var, layout: Var) -> Var: + require_target_greater_equal(9, 0, "mbarrier_has_layout") require_mbarrier_ptr(mbar, (MemorySpace.SHARED,)) layout = require_constant_enum(layout, MbarrierLayout) instruction = f"mbarrier.check_layout.layout::v{layout.value}.shared::cta.b64" diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/type_checking_helpers.py b/experimental/cuda-lang/src/cuda/lang/_ir/type_checking_helpers.py index fdbd42b8..27cdd7af 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/type_checking_helpers.py +++ b/experimental/cuda-lang/src/cuda/lang/_ir/type_checking_helpers.py @@ -5,7 +5,9 @@ from typing import Any, Callable import cuda.lang._datatype as datatype +from cuda.lang._exception import UnsupportedFeatureError from cuda.lang._ir.type import MemorySpace, ScalarTy, VectorTy, PointerTy +from cuda.lang._target import TargetInfo from cuda.tile import DType from cuda.tile._ir.ir import Var from cuda.tile._ir.op_impl import ( # noqa: F401 @@ -27,6 +29,25 @@ ) +def require_target_exact(major: int, minor: int, arch_flag: str, message: str) -> None: + target = TargetInfo.get_current() + required_target = TargetInfo(major, minor, arch_flag or None) + if target != required_target: + raise UnsupportedFeatureError( + f"GPU target {target} does not match required target {required_target}:" + + message + ) + + +def require_target_greater_equal(major: int, minor: int, message: str) -> None: + target = TargetInfo.get_current() + if (target.major, target.minor) < (major, minor): + raise UnsupportedFeatureError( + f"GPU target {target} is below required compute capability {major}.{minor}:" + + message + ) + + def is_none(var: Var): return var.is_constant() and var.get_constant() is None diff --git a/experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py b/experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py index b1e1469e..8c478b36 100644 --- a/experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py +++ b/experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py @@ -22,7 +22,7 @@ import cuda.lang._mlir.extras.types as T import cuda.lang._ir.type as ir_type from cuda.lang.compilation import KernelSignature -from cuda.lang._target import TargetFeature, TargetInfo +from cuda.lang._target import TargetInfo import cuda.lang._datatype as datatype from cuda.tile import _cext from cuda.tile._datatype import PointerInfo, is_integral @@ -520,7 +520,8 @@ def lower_fma( and not operation.relu and not operation.oob ) - if packed_f32x2 and context.target_info.supports(TargetFeature.PACKED_F32X2): + target = context.target_info + if packed_f32x2 and (target.major, target.minor) >= (10, 0): result = mlir.add_operation( name="nvvm.fma.packed.f32x2", result_type=mlir_result_type, @@ -685,7 +686,8 @@ def _lower_binary_arith_with_nvvm_modifiers( and res_type.length % 2 == 0 and res_dtype == datatype.float32 ) - if packed_f32x2 and context.target_info.supports(TargetFeature.PACKED_F32X2): + target = context.target_info + if packed_f32x2 and (target.major, target.minor) >= (10, 0): result = mlir.add_operation( name=f"nvvm.{operation.fn}.packed.f32x2", result_type=result_mlir_type, diff --git a/experimental/cuda-lang/src/cuda/lang/_target.py b/experimental/cuda-lang/src/cuda/lang/_target.py index b8f0b241..0a8552c3 100644 --- a/experimental/cuda-lang/src/cuda/lang/_target.py +++ b/experimental/cuda-lang/src/cuda/lang/_target.py @@ -4,12 +4,7 @@ import re from dataclasses import dataclass -from enum import Enum, auto -from functools import cache - - -class TargetFeature(Enum): - PACKED_F32X2 = auto() +from ._exception import InternalError @dataclass(frozen=True) @@ -22,8 +17,14 @@ class TargetInfo: def from_arch(cls, arch: str) -> "TargetInfo": return _parse_target_name(arch, "compute") - def supports(self, feature: TargetFeature) -> bool: - return feature in _features_for_target(self) + @staticmethod + def get_current() -> "TargetInfo": + """Return the GPU target of the active IR builder's context.""" + from cuda.lang._ir.ir import TileBuilder + ctx = TileBuilder.get_current().ir_ctx + if ctx.target_info is None: + raise InternalError("The current IR context has no GPU target") + return ctx.target_info def _parse_target_name(name: str, prefix: str) -> TargetInfo: @@ -40,12 +41,4 @@ def _parse_target_name(name: str, prefix: str) -> TargetInfo: ) -@cache -def _features_for_target(target: TargetInfo) -> frozenset[TargetFeature]: - features = set() - if (target.major, target.minor) >= (10, 0): - features.add(TargetFeature.PACKED_F32X2) - return frozenset(features) - - -__all__ = ("TargetFeature", "TargetInfo") +__all__ = ("TargetInfo",) diff --git a/experimental/cuda-lang/test/test_math.py b/experimental/cuda-lang/test/test_math.py index dbca45c4..ade44b98 100644 --- a/experimental/cuda-lang/test/test_math.py +++ b/experimental/cuda-lang/test/test_math.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from test.util import compile_kernel +from dataclasses import replace import cuda.lang as cl import cuda.lang._datatype as datatype import builtins @@ -18,7 +19,8 @@ from cuda.lang.compilation import KernelSignature from cuda.lang._exception import CompilerExecutionError, TypeCheckingError from cuda.lang._fp_utils import _FLOAT_SMALLEST_NORMAL, isnormal -from cuda.lang._target import TargetFeature, TargetInfo +from cuda.lang._target import TargetInfo +from cuda.lang._passes.ir2mlir import pass_definition from .util import make_symbolic_tensor @@ -1683,14 +1685,16 @@ def kernel(): def test_arith_f32x2_nvvm_toolchain_packing( monkeypatch, op_name, device_op, rounding_mode, flush_to_zero ): - original_supports = TargetInfo.supports + original_lowering = pass_definition._lower_binary_arith_with_nvvm_modifiers - def supports_without_packed_f32x2(self, feature): - if feature is TargetFeature.PACKED_F32X2: - return False - return original_supports(self, feature) + def lower_without_packed_f32x2(context, operation, res_type): + return original_lowering( + replace(context, target_info=TargetInfo(9, 0)), operation, res_type + ) - monkeypatch.setattr(TargetInfo, "supports", supports_without_packed_f32x2) + monkeypatch.setattr( + pass_definition, "_lower_binary_arith_with_nvvm_modifiers", lower_without_packed_f32x2 + ) def kernel(): values = cl.shared_array(2, cl.float32) @@ -1705,7 +1709,7 @@ def kernel(): rnd = rounding_mode.name.lower() ftz = ".ftz" if flush_to_zero else "" - # If this starts emitting packed PTX, remove the special PACKED_F32X2 + # If this starts emitting packed PTX, remove the special packed f32x2 # target lowering and rely on the toolchain to pack nvvm.addf/subf. compile_kernel( kernel, diff --git a/experimental/cuda-lang/test/test_mbarrier.py b/experimental/cuda-lang/test/test_mbarrier.py index 2726fd6e..fdbff6b2 100644 --- a/experimental/cuda-lang/test/test_mbarrier.py +++ b/experimental/cuda-lang/test/test_mbarrier.py @@ -6,7 +6,7 @@ import pytest import cuda.lang as cl -from cuda.lang._exception import TypeCheckingError +from cuda.lang._exception import TypeCheckingError, UnsupportedFeatureError from cuda.lang._ir.ops import RawLLVMIntrinsic from cuda.lang.compilation import KernelSignature @@ -111,6 +111,25 @@ def kernel(): ) +@pytest.mark.parametrize("layout", tuple(cl.MbarrierLayout)) +@pytest.mark.parametrize("operation", (cl.mbarrier_initialize, cl.mbarrier_has_layout)) +def test_mbarrier_layout_rejects_unsupported_target(operation, layout): + initialize = operation is cl.mbarrier_initialize + + def kernel(): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + if initialize: + operation(mbar, 1, layout=layout) + else: + operation(mbar, layout) + + with pytest.raises(UnsupportedFeatureError) as error: + cl.compile_simt(kernel, [KernelSignature([])], arch="compute_80", gpu_name="sm_100a") + assert "TargetInfo(major=8, minor=0, suffix=None)" in error.value.message + assert "9.0" in error.value.message + assert operation.__name__ in error.value.message + + @pytest.mark.parametrize("layout", (*cl.MbarrierLayout, "V0", "V1")) @require_hopper_or_newer() def test_mbarrier_layout(layout): diff --git a/experimental/cuda-lang/test/test_target.py b/experimental/cuda-lang/test/test_target.py index c666aea3..f573290f 100644 --- a/experimental/cuda-lang/test/test_target.py +++ b/experimental/cuda-lang/test/test_target.py @@ -5,6 +5,14 @@ import pytest from cuda.lang._target import TargetInfo +from cuda.lang._exception import InternalError +from cuda.lang._ir.ir import ( + Builder, + IRContext, + Loc, + Region, + TileBuilder, +) @pytest.mark.parametrize( @@ -30,3 +38,25 @@ def test_target_info_from_arch(arch, expected): def test_target_info_from_invalid_arch(arch): with pytest.raises(ValueError, match="invalid CUDA target name"): TargetInfo.from_arch(arch) + + +def test_host_context_rejects_target_info(): + target = TargetInfo.from_arch("compute_100a") + with pytest.raises(ValueError, match="host IR context"): + IRContext(execution_space="host", target_info=target) + + +def test_target_info_get_current(): + target = TargetInfo.from_arch("compute_100f") + ctx = IRContext(target_info=target) + loc = Loc.unknown() + with Builder(Region(ctx), loc): + assert TargetInfo.get_current() is target + + +@pytest.mark.parametrize("execution_space", ("device", "host")) +def test_target_info_get_current_without_target(execution_space): + ctx = IRContext(execution_space=execution_space) + with TileBuilder(ctx, Loc.unknown()): + with pytest.raises(InternalError, match="current IR context has no GPU target"): + TargetInfo.get_current() From 8589a7866787b496969813a01cc8f4e781a93220 Mon Sep 17 00:00:00 2001 From: Asher Mancinelli Date: Tue, 6 Oct 2026 08:13:07 -0700 Subject: [PATCH 2/6] [lang] Support mbarrier_{drop_}nocomplete intrinsics Signed-off-by: Asher Mancinelli --- .../cuda-lang/src/cuda/lang/__init__.py | 4 + .../cuda/lang/_ir/op_impl/mbarrier_impl.py | 14 ++++ .../cuda-lang/src/cuda/lang/_stub/mbarrier.py | 28 +++++++ experimental/cuda-lang/test/test_mbarrier.py | 84 ++++++++++++++++++- 4 files changed, 129 insertions(+), 1 deletion(-) diff --git a/experimental/cuda-lang/src/cuda/lang/__init__.py b/experimental/cuda-lang/src/cuda/lang/__init__.py index 22e32365..2bb5e270 100644 --- a/experimental/cuda-lang/src/cuda/lang/__init__.py +++ b/experimental/cuda-lang/src/cuda/lang/__init__.py @@ -234,6 +234,8 @@ mbarrier_has_layout, mbarrier_invalidate, mbarrier_arrive, + mbarrier_arrive_nocomplete, + mbarrier_arrive_drop_nocomplete, mbarrier_arrive_expect_transaction, mbarrier_expect_transaction, mbarrier_complete_transaction, @@ -431,6 +433,8 @@ "mbarrier_has_layout", "mbarrier_invalidate", "mbarrier_arrive", + "mbarrier_arrive_nocomplete", + "mbarrier_arrive_drop_nocomplete", "mbarrier_arrive_expect_transaction", "mbarrier_expect_transaction", "mbarrier_complete_transaction", diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py index babf12e2..da2c400d 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py +++ b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py @@ -124,6 +124,20 @@ def mbarrier_arrive_impl( return results[0] if return_type else None +@impl(mbarrier.mbarrier_arrive_nocomplete, fixed_args=['noComplete']) +@impl(mbarrier.mbarrier_arrive_drop_nocomplete, fixed_args=['drop.noComplete']) +def mbarrier_arrive_nocomplete_impl(kind: str, mbar: Var, count: Var) -> Var: + require_mbarrier_ptr(mbar, (MemorySpace.SHARED,)) + count = astype(count, datatype.int32) + intrinsic = f"llvm.nvvm.mbarrier.arrive.{kind}.shared" + return add_operation( + RawLLVMIntrinsic, + ScalarTy(datatype.uint64), + intrinsic=intrinsic, + operands_=(mbar, count), + ) + + @impl(mbarrier.mbarrier_arrive_expect_transaction) def mbarrier_arrive_expect_transaction_impl( mbar: Var, diff --git a/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py b/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py index 32cb126a..1f185a7d 100644 --- a/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py +++ b/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py @@ -80,6 +80,34 @@ def mbarrier_arrive( """ +@stub +def mbarrier_arrive_nocomplete(mbar, count: int) -> "uint64": + """Arrive without completing the current phase. + + Args: + mbar: Pointer to mbarrier in block-local shared memory. + count: Amount to subtract from the pending arrival count. + + Returns: + Opaque 64-bit token for the phase before the arrival. + """ + ... + + +@stub +def mbarrier_arrive_drop_nocomplete(mbar, count: int) -> "uint64": + """Arrive without completing the phase and drop expected arrivals. + + Args: + mbar: Pointer to mbarrier in block-local shared memory. + count: Amount to subtract from the pending and expected arrival counts. + + Returns: + Opaque 64-bit token for the phase before the arrival. + """ + ... + + @stub def mbarrier_arrive_expect_transaction( mbar, diff --git a/experimental/cuda-lang/test/test_mbarrier.py b/experimental/cuda-lang/test/test_mbarrier.py index fdbff6b2..baf1bca8 100644 --- a/experimental/cuda-lang/test/test_mbarrier.py +++ b/experimental/cuda-lang/test/test_mbarrier.py @@ -10,7 +10,12 @@ from cuda.lang._ir.ops import RawLLVMIntrinsic from cuda.lang.compilation import KernelSignature -from .util import compile_kernel, make_symbolic_tensor, require_hopper_or_newer +from .util import ( + compile_kernel, + make_symbolic_scalar, + make_symbolic_tensor, + require_hopper_or_newer, +) HOPPER_TARGET = {"gpu_name": "sm_90", "arch": "compute_90"} @@ -187,6 +192,83 @@ def kernel(): WAIT_MEMORY_ORDERS = [cl.MemoryOrder.ACQUIRE, cl.MemoryOrder.RELAXED] +NO_COMPLETE_ARRIVALS = { + cl.mbarrier_arrive_nocomplete: "mbarrier.arrive.noComplete.shared.b64", + cl.mbarrier_arrive_drop_nocomplete: "mbarrier.arrive_drop.noComplete.shared.b64", +} + + +@pytest.mark.parametrize(("operation", "instruction"), NO_COMPLETE_ARRIVALS.items()) +@pytest.mark.parametrize("count_dtype", (cl.int32, cl.uint32)) +def test_arrive_nocomplete_ptx(operation, instruction, count_dtype): + @cl.kernel + def kernel(count): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + operation(mbar, count) + + compile_kernel( + kernel, + signature=KernelSignature([make_symbolic_scalar(count_dtype)]), + assert_in_ptx=instruction, + gpu_name="sm_80", + arch="compute_80", + ) + + +@pytest.mark.parametrize("operation", NO_COMPLETE_ARRIVALS.keys()) +@pytest.mark.parametrize("count", (1, 2)) +@require_hopper_or_newer() +def test_arrive_nocomplete_phases(operation, count): + next_count = 1 if operation is cl.mbarrier_arrive_drop_nocomplete else count + 1 + + @cl.kernel + def kernel(out, count): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + cl.mbarrier_initialize(mbar, count + 1) + state = operation(mbar, count) + out[0] = cl.mbarrier_test_wait(mbar, state) + cl.mbarrier_arrive(mbar) + out[1] = cl.mbarrier_test_wait(mbar, state) + state = cl.mbarrier_arrive(mbar, next_count) + out[2] = cl.mbarrier_test_wait(mbar, state) + cl.mbarrier_invalidate(mbar) + + out = torch.zeros(3, dtype=torch.bool).cuda(0) + cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (out, count)) + assert out.cpu().tolist() == [False, True, True] + + +class TestArriveNocompleteRequirePointer: + + @staticmethod + def shared_cluster_kernel(op): + @cl.kernel + def kernel(): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + mbar = cl.map_shared_to_cluster(mbar, 0) + op(mbar, 1) + return kernel + + @staticmethod + def int_kernel(op): + @cl.kernel + def kernel(): + i64_ptr = cl.shared_array(1, cl.int64, alignment=8).pointer() + op(i64_ptr, 1) + return kernel + + @pytest.mark.parametrize("operation", NO_COMPLETE_ARRIVALS.keys()) + @pytest.mark.parametrize("factory", (shared_cluster_kernel, int_kernel)) + def test_requires_pointer(self, operation, factory): + kernel = factory(operation) + + compile_kernel( + kernel, + raises=pytest.raises(TypeCheckingError), + **HOPPER_TARGET, + ) + + @pytest.mark.parametrize("scope", SCOPES) @pytest.mark.parametrize("memory_order", ARRIVE_MEMORY_ORDERS) @pytest.mark.parametrize("drop", [False, True]) From 8234d331f79379c672bede77b2bf70b3a98c74cd Mon Sep 17 00:00:00 2001 From: Asher Mancinelli Date: Tue, 6 Oct 2026 09:28:43 -0700 Subject: [PATCH 3/6] [lang] Support mbarrier_pending_count intrinsic Signed-off-by: Asher Mancinelli --- .../cuda-lang/src/cuda/lang/__init__.py | 2 + .../cuda/lang/_ir/op_impl/mbarrier_impl.py | 14 ++++ .../cuda-lang/src/cuda/lang/_stub/mbarrier.py | 16 ++++- experimental/cuda-lang/test/test_mbarrier.py | 72 +++++++++++++++++++ 4 files changed, 103 insertions(+), 1 deletion(-) diff --git a/experimental/cuda-lang/src/cuda/lang/__init__.py b/experimental/cuda-lang/src/cuda/lang/__init__.py index 2bb5e270..9b4e808b 100644 --- a/experimental/cuda-lang/src/cuda/lang/__init__.py +++ b/experimental/cuda-lang/src/cuda/lang/__init__.py @@ -239,6 +239,7 @@ mbarrier_arrive_expect_transaction, mbarrier_expect_transaction, mbarrier_complete_transaction, + mbarrier_pending_count, mbarrier_test_wait, mbarrier_test_wait_parity, mbarrier_try_wait, @@ -438,6 +439,7 @@ "mbarrier_arrive_expect_transaction", "mbarrier_expect_transaction", "mbarrier_complete_transaction", + "mbarrier_pending_count", "mbarrier_test_wait", "mbarrier_test_wait_parity", "mbarrier_try_wait", diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py index da2c400d..3309a1ff 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py +++ b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py @@ -10,6 +10,7 @@ from cuda.lang._ir.type import MemorySpace, ScalarTy from cuda.lang._ir.type_checking_helpers import ( is_none, + require_integral_scalar_type, require_mbarrier_ptr, require_target_greater_equal, ) @@ -199,6 +200,19 @@ def mbarrier_complete_transaction_impl(mbar: Var, bytes: Var, scope: Var) -> Var ) +@impl(mbarrier.mbarrier_pending_count) +def mbarrier_pending_count_impl(state: Var) -> Var: + require_target_greater_equal(8, 0, "mbarrier_pending_count") + require_integral_scalar_type(state, bitwidth=64) + state = astype(state, datatype.int64) + return add_operation( + RawLLVMIntrinsic, + ScalarTy(datatype.uint32), + intrinsic="llvm.nvvm.mbarrier.pending.count", + operands_=(state,), + ) + + @impl(mbarrier.mbarrier_test_wait) def mbarrier_test_wait_impl( mbar: Var, state: Var, scope: Var, memory_order: Var diff --git a/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py b/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py index 1f185a7d..18dbb665 100644 --- a/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py +++ b/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py @@ -6,7 +6,7 @@ from cuda.lang._enums import MbarrierLayout, MbarrierScope from cuda.lang._execution import stub -from cuda.lang._datatype import uint64, bool_ +from cuda.lang._datatype import uint32, uint64, bool_ from cuda.lang._enums import MemoryOrder @@ -167,6 +167,20 @@ def mbarrier_complete_transaction( ... +@stub +def mbarrier_pending_count(state) -> "uint32": + """Query the pending arrival count from the opaque mbarrier state. + + Args: + state: 64-bit token returned by ``mbarrier_arrive_drop_nocomplete`` + or ``mbarrier_arrive_nocomplete``. + + Returns: + Pending arrival count before the arrival that produced ``state``. + """ + ... + + @stub def mbarrier_test_wait( mbar, diff --git a/experimental/cuda-lang/test/test_mbarrier.py b/experimental/cuda-lang/test/test_mbarrier.py index baf1bca8..7b02be27 100644 --- a/experimental/cuda-lang/test/test_mbarrier.py +++ b/experimental/cuda-lang/test/test_mbarrier.py @@ -269,6 +269,78 @@ def test_requires_pointer(self, operation, factory): ) +@pytest.mark.parametrize("dtype", (cl.int64, cl.uint64)) +def test_pending_count_ptx(dtype): + @cl.kernel + def kernel(out): + out[0] = cl.mbarrier_pending_count(dtype(0)) + + compile_kernel( + kernel, + signature=KernelSignature( + [ + make_symbolic_tensor(1, cl.uint32), + ] + ), + assert_in_nvvm="llvm.nvvm.mbarrier.pending.count", + assert_in_ptx="mbarrier.pending_count", + gpu_name="sm_80", + arch="compute_80", + ) + + +@pytest.mark.parametrize("operation", ( + cl._nvvm.mbarrier_arrive_noComplete_shared, + cl._nvvm.mbarrier_arrive_drop_noComplete_shared, +)) +@pytest.mark.parametrize("count", (1, 2)) +@require_hopper_or_newer() +def test_pending_count_before_arrival(operation, count): + @cl.kernel + def kernel(out, count): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + cl.mbarrier_initialize(mbar, count + 3) + state = operation(mbar, cl.int32(count)) + out[0] = cl.mbarrier_pending_count(state) + next_state = operation(mbar, cl.int32(1)) + out[1] = cl.mbarrier_pending_count(next_state) + out[2] = cl.mbarrier_pending_count(cl.uint64(state)) + cl.mbarrier_arrive(mbar, 2) + cl.mbarrier_invalidate(mbar) + + out = torch.zeros(3, dtype=torch.uint32).cuda(0) + cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (out, count)) + assert out.cpu().tolist() == [count + 3, 3, count + 3] + + +@pytest.mark.parametrize("state_dtype", (cl.int32, cl.uint32, cl.float64)) +def test_pending_count_rejects_invalid_state_type(state_dtype): + @cl.kernel + def kernel(state): + cl.mbarrier_pending_count(state) + + compile_kernel( + kernel, + signature=KernelSignature([make_symbolic_scalar(state_dtype)]), + raises=pytest.raises(TypeCheckingError, match="scalar integral"), + **HOPPER_TARGET, + ) + + +def test_pending_count_rejects_unsupported_target(): + @cl.kernel + def kernel(state): + cl.mbarrier_pending_count(state) + + compile_kernel( + kernel, + signature=KernelSignature([make_symbolic_scalar(cl.uint64)]), + raises=pytest.raises(UnsupportedFeatureError, match="mbarrier_pending_count"), + gpu_name="sm_75", + arch="compute_75", + ) + + @pytest.mark.parametrize("scope", SCOPES) @pytest.mark.parametrize("memory_order", ARRIVE_MEMORY_ORDERS) @pytest.mark.parametrize("drop", [False, True]) From baffa50249abe5f69ea4a2758454f82f95ee0479 Mon Sep 17 00:00:00 2001 From: Asher Mancinelli Date: Tue, 6 Oct 2026 11:05:30 -0700 Subject: [PATCH 4/6] [lang] Support intrinsics for signaling arrival of async copies on mbarriers Also updated the mbarrier wait spinloop utilities to use the signature of the wrapped intrinsics with the default time hint so users get better argument information in their LSPs. Signed-off-by: Asher Mancinelli --- .../cuda-lang/src/cuda/lang/__init__.py | 4 + .../cuda/lang/_ir/op_impl/copy_async_impl.py | 15 ++++ .../src/cuda/lang/_stub/_nvvm_support.py | 2 + .../src/cuda/lang/_stub/copy_async.py | 20 +++++ .../cuda-lang/src/cuda/lang/_stub/mbarrier.py | 26 ++++++- .../cuda-lang/test/test_copy_async.py | 76 ++++++++++++++++++- 6 files changed, 138 insertions(+), 5 deletions(-) diff --git a/experimental/cuda-lang/src/cuda/lang/__init__.py b/experimental/cuda-lang/src/cuda/lang/__init__.py index 9b4e808b..d525ed45 100644 --- a/experimental/cuda-lang/src/cuda/lang/__init__.py +++ b/experimental/cuda-lang/src/cuda/lang/__init__.py @@ -113,6 +113,8 @@ from ._stub.copy_async import ( TMALoadMode, TMAStoreMode, + copy_async_mbarrier_arrive, + copy_async_mbarrier_arrive_no_increment, copy_async_bulk_tensor_shared_to_global, copy_async_bulk_tensor_global_to_shared, copy_async_bulk_commit_group, @@ -472,6 +474,8 @@ "grid_dependency_control_wait", "grid_dependency_control_launch_dependents", "bitcast", + "copy_async_mbarrier_arrive", + "copy_async_mbarrier_arrive_no_increment", "copy_async_bulk_tensor_shared_to_global", "copy_async_bulk_tensor_global_to_shared", "copy_async_bulk_commit_group", diff --git a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/copy_async_impl.py b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/copy_async_impl.py index b8438c8d..3488c263 100644 --- a/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/copy_async_impl.py +++ b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/copy_async_impl.py @@ -25,6 +25,7 @@ require_none, require_optional, require_pointer_in_memory_space, + require_target_greater_equal, require_uniform_int_tuple_type, tensor_map_descriptor_pointer_like, ) @@ -43,6 +44,20 @@ def _optional_operand(value): return None if value is None or is_none(value) else value +@impl(copy_async.copy_async_mbarrier_arrive, fixed_args=["arrive"]) +@impl(copy_async.copy_async_mbarrier_arrive_no_increment, fixed_args=["arrive.noinc"]) +def copy_async_mbarrier_arrive_impl(kind: str, mbar: Var) -> None: + require_target_greater_equal(8, 0, "copy_async_mbarrier_arrive") + require_mbarrier_ptr(mbar, (MemorySpace.SHARED,)) + intrinsic = f"llvm.nvvm.cp.async.mbarrier.{kind}.shared" + add_operation_variadic( + RawLLVMIntrinsic, + (), + intrinsic=intrinsic, + operands_=(mbar,), + ) + + def validate_g2s_mode(mode: copy_async.TMALoadMode, im2col_count: int) -> None: match mode: case copy_async.TMALoadMode.TILE | copy_async.TMALoadMode.TILE_GATHER4: diff --git a/experimental/cuda-lang/src/cuda/lang/_stub/_nvvm_support.py b/experimental/cuda-lang/src/cuda/lang/_stub/_nvvm_support.py index ca0630b3..17671472 100644 --- a/experimental/cuda-lang/src/cuda/lang/_stub/_nvvm_support.py +++ b/experimental/cuda-lang/src/cuda/lang/_stub/_nvvm_support.py @@ -259,6 +259,8 @@ def _implicit_cast_with_fallback(src: Var, target_dtype: DType, error_context: s try: return implicit_cast(src, target_dtype, error_context) except (TypeCheckingError, InvalidValueError): + # TODO: should we bitcast pointer of the same bitwidth to the expected + # type? if not (datatype.is_integral(src.get_type().dtype) and datatype.is_integral(target_dtype)): raise diff --git a/experimental/cuda-lang/src/cuda/lang/_stub/copy_async.py b/experimental/cuda-lang/src/cuda/lang/_stub/copy_async.py index 009d9224..0027751e 100644 --- a/experimental/cuda-lang/src/cuda/lang/_stub/copy_async.py +++ b/experimental/cuda-lang/src/cuda/lang/_stub/copy_async.py @@ -7,6 +7,26 @@ from . import nvvm as _nvvm +@stub +def copy_async_mbarrier_arrive(mbar) -> None: + """Signal an arrival after earlier ``cp.async`` copies complete. + + Args: + mbar: Pointer to an mbarrier in shared memory. + """ + ... + + +@stub +def copy_async_mbarrier_arrive_no_increment(mbar) -> None: + """Signal an asynchronous-copy arrival without incrementing the count. + + Args: + mbar: Pointer to an mbarrier in shared memory. + """ + ... + + @stub def copy_async_bulk_tensor_global_to_shared( src_tensor_map_descriptor, diff --git a/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py b/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py index 18dbb665..58081770 100644 --- a/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py +++ b/experimental/cuda-lang/src/cuda/lang/_stub/mbarrier.py @@ -247,11 +247,20 @@ def mbarrier_try_wait( """ -def mbarrier_wait(*args, time_hint=10_000, **kwargs): +def mbarrier_wait( + mbar, + state, + *, + time_hint: int = 10_000, + scope: MbarrierScope = MbarrierScope.BLOCK, + memory_order: WaitMemoryOrder = MemoryOrder.ACQUIRE, +): """Synchronously wait for an mbarrier to complete. Accepts the same arguments as :func:``mbarrier_try_wait``. """ - while not mbarrier_try_wait(*args, time_hint=time_hint, **kwargs): + while not mbarrier_try_wait( + mbar, state, time_hint=time_hint, scope=scope, memory_order=memory_order + ): pass @@ -279,9 +288,18 @@ def mbarrier_try_wait_parity( """ -def mbarrier_wait_parity(*args, time_hint=10_000, **kwargs): +def mbarrier_wait_parity( + mbar, + parity: int, + *, + time_hint: int = 10_000, + scope: MbarrierScope = MbarrierScope.BLOCK, + memory_order: WaitMemoryOrder = MemoryOrder.ACQUIRE, +): """Synchronously wait for an mbarrier to complete for a given phase parity. Accepts the same arguments as :func:``mbarrier_try_wait_parity``. """ - while not mbarrier_try_wait_parity(*args, time_hint=time_hint, **kwargs): + while not mbarrier_try_wait_parity( + mbar, parity, time_hint=time_hint, scope=scope, memory_order=memory_order + ): pass diff --git a/experimental/cuda-lang/test/test_copy_async.py b/experimental/cuda-lang/test/test_copy_async.py index 10615e7a..9dd378a5 100644 --- a/experimental/cuda-lang/test/test_copy_async.py +++ b/experimental/cuda-lang/test/test_copy_async.py @@ -3,9 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +import torch import cuda.lang as cl -from cuda.lang._exception import TypeCheckingError +from cuda.lang._exception import TypeCheckingError, UnsupportedFeatureError from cuda.lang.compilation import CallingConvention, KernelSignature, TensorMapConstraint from cuda.tile._cext import cconv_v3_enabled @@ -13,6 +14,7 @@ compile_kernel, make_symbolic_scalar, make_symbolic_tensor, + require_hopper_or_newer, ) @@ -20,6 +22,78 @@ SM100_TARGET = {"gpu_name": "sm_100a", "arch": "compute_100a"} +COPY_ASYNC_ARRIVALS = { + cl.copy_async_mbarrier_arrive: "cp.async.mbarrier.arrive", + cl.copy_async_mbarrier_arrive_no_increment: "cp.async.mbarrier.arrive.noinc", +} + + +@pytest.mark.parametrize(("operation", "instruction"), COPY_ASYNC_ARRIVALS.items()) +def test_copy_async_mbarrier_arrive_ptx(operation, instruction): + @cl.kernel + def kernel(): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + operation(mbar) + + compile_kernel( + kernel, + assert_in_ptx=instruction, + gpu_name="sm_80", + arch="compute_80", + ) + + +@pytest.mark.parametrize("operation", COPY_ASYNC_ARRIVALS.keys()) +def test_copy_async_mbarrier_arrive_unsupported_target(operation): + @cl.kernel + def kernel(): + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + operation(mbar) + + compile_kernel( + kernel, + raises=pytest.raises(UnsupportedFeatureError, match="copy_async_mbarrier_arrive"), + gpu_name="sm_75", + arch="compute_75", + ) + + +@pytest.mark.parametrize("arrival_operation", COPY_ASYNC_ARRIVALS.keys()) +@pytest.mark.parametrize("arrivals", (1, 2)) +@require_hopper_or_newer() +def test_copy_async_mbarrier_arrive_completion(arrival_operation, arrivals): + participants = 32 + if arrival_operation is cl.copy_async_mbarrier_arrive_no_increment: + participants += arrivals + size = arrivals * 4 + + @cl.kernel + def kernel(src, out): + tid = cl.thread_index(0) + memory = cl.shared_array(size, cl.int32, alignment=16) + mbar = cl.shared_array(1, cl.mbarrier, alignment=8).pointer() + if tid == 0: + cl.mbarrier_initialize(mbar, participants) + cl.barrier_sync_block_aligned() + + if tid == 0: + for i in range(arrivals): + ptr = src.pointer(4 * i) + ptr = cl.bitcast(ptr, cl.opaque_pointer_dtype('GLOBAL')) + cl._nvvm.cp_async_ca_shared_global_16(memory.pointer(4 * i), ptr) + arrival_operation(mbar) + + state = cl.mbarrier_arrive(mbar) + cl.mbarrier_wait(mbar, state) + out[tid] = memory[tid % size] + + src = torch.arange(1, size + 1, dtype=torch.int32).cuda(0) + out = torch.zeros(32, dtype=torch.int32).cuda(0) + cl.launch(torch.cuda.current_stream(), (1,), (32,), kernel, (src, out)) + expected = src.cpu()[torch.arange(32) % size] + assert torch.equal(out.cpu(), expected) + + class CopyAsyncPtxTestBase: @staticmethod def signature(): From be52df31e40b3858d2b400fa4b9e37242f4e5c21 Mon Sep 17 00:00:00 2001 From: Camille Dunning Date: Tue, 6 Oct 2026 11:17:22 -0700 Subject: [PATCH 5/6] upgrade pytest Addresses https://nspect.nvidia.com/vulnerability?id=BDSA-2025-122508 Signed-off-by: Camille Dunning --- test/requirements-ft.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/requirements-ft.txt b/test/requirements-ft.txt index f0e554c2..982d2b43 100644 --- a/test/requirements-ft.txt +++ b/test/requirements-ft.txt @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 setuptools==84.0.0 -pytest==9.0.1 +pytest==9.0.3 pytest-benchmark==5.1.0 pytest-cov==6.2.1 pytest-env==1.2.0 From 3ebeb4eca3ce0c2d2a47395cfd4999c44730390f Mon Sep 17 00:00:00 2001 From: Camille Dunning Date: Tue, 6 Oct 2026 18:59:24 -0700 Subject: [PATCH 6/6] update GitHub workflow image tags Signed-off-by: Camille Dunning --- .github/workflows/ci.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index eb21b6d3..282d56e2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,12 +32,12 @@ jobs: build_py313: ghcr.io/nvidia/cutile-python/build_py_3.13_x86_64:2026-08-17-e66dd3d7ef7a build_py314: ghcr.io/nvidia/cutile-python/build_py_3.14_x86_64:2026-08-17-931831a8d099 build_py314t: ghcr.io/nvidia/cutile-python/build_py_3.14t_x86_64:2026-08-17-28bcee962449 - test_py310: ghcr.io/nvidia/cutile-python/test_py_3.10_x86_64:2026-08-17-8549280618e3 - test_py311: ghcr.io/nvidia/cutile-python/test_py_3.11_x86_64:2026-08-17-a437217a163c - test_py312: ghcr.io/nvidia/cutile-python/test_py_3.12_x86_64:2026-08-17-85d044574dad - test_py313: ghcr.io/nvidia/cutile-python/test_py_3.13_x86_64:2026-08-17-7f71fb21a2a9 - test_py314: ghcr.io/nvidia/cutile-python/test_py_3.14_x86_64:2026-08-17-4e620388d66f - test_py314t: ghcr.io/nvidia/cutile-python/test_py_3.14t_x86_64:2026-08-17-68e4e9c53f83 + test_py310: ghcr.io/nvidia/cutile-python/test_py_3.10_x86_64:2026-10-07-930ad3efefbb + test_py311: ghcr.io/nvidia/cutile-python/test_py_3.11_x86_64:2026-10-07-6c9a98aa0164 + test_py312: ghcr.io/nvidia/cutile-python/test_py_3.12_x86_64:2026-10-07-e9e9091b2165 + test_py313: ghcr.io/nvidia/cutile-python/test_py_3.13_x86_64:2026-10-07-747bd6e13da8 + test_py314: ghcr.io/nvidia/cutile-python/test_py_3.14_x86_64:2026-10-07-014972b99f54 + test_py314t: ghcr.io/nvidia/cutile-python/test_py_3.14t_x86_64:2026-10-07-6fc242eed61d steps: - run: echo "Defining image tags"