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" diff --git a/experimental/cuda-lang/src/cuda/lang/__init__.py b/experimental/cuda-lang/src/cuda/lang/__init__.py index 22e32365..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, @@ -234,9 +236,12 @@ mbarrier_has_layout, mbarrier_invalidate, mbarrier_arrive, + mbarrier_arrive_nocomplete, + mbarrier_arrive_drop_nocomplete, mbarrier_arrive_expect_transaction, mbarrier_expect_transaction, mbarrier_complete_transaction, + mbarrier_pending_count, mbarrier_test_wait, mbarrier_test_wait_parity, mbarrier_try_wait, @@ -431,9 +436,12 @@ "mbarrier_has_layout", "mbarrier_invalidate", "mbarrier_arrive", + "mbarrier_arrive_nocomplete", + "mbarrier_arrive_drop_nocomplete", "mbarrier_arrive_expect_transaction", "mbarrier_expect_transaction", "mbarrier_complete_transaction", + "mbarrier_pending_count", "mbarrier_test_wait", "mbarrier_test_wait_parity", "mbarrier_try_wait", @@ -466,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/_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/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/_ir/op_impl/mbarrier_impl.py b/experimental/cuda-lang/src/cuda/lang/_ir/op_impl/mbarrier_impl.py index baa0054a..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 @@ -8,7 +8,12 @@ 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_integral_scalar_type, + 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 +38,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 +48,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" @@ -118,6 +125,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, @@ -179,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/_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/_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 32cb126a..58081770 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 @@ -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, @@ -139,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, @@ -205,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 @@ -237,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/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_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(): 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..7b02be27 100644 --- a/experimental/cuda-lang/test/test_mbarrier.py +++ b/experimental/cuda-lang/test/test_mbarrier.py @@ -6,11 +6,16 @@ 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 -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"} @@ -111,6 +116,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): @@ -168,6 +192,155 @@ 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("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]) 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() 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