From 043c838db7968fb931ed353a0965d994946f3a31 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Tue, 1 Sep 2026 00:01:39 -0400 Subject: [PATCH 1/2] Version 17: Add PyTorch tensor interoperability --- README.md | 103 +++++++++++--- src/mytriton/compiler.py | 2 +- src/mytriton/cuda_utils.py | 231 +++++++++++++++++++++++++------ src/mytriton/runtime_args.py | 79 +++++++++++ src/mytriton/trace.py | 12 +- tests/test_runtime_args.py | 154 +++++++++++++++++++++ tests/test_torch_interop.py | 255 +++++++++++++++++++++++++++++++++++ 7 files changed, 774 insertions(+), 62 deletions(-) create mode 100644 src/mytriton/runtime_args.py create mode 100644 tests/test_runtime_args.py create mode 100644 tests/test_torch_interop.py diff --git a/README.md b/README.md index e4ef03b..38a5049 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,10 @@ MLIR's GPU/NVVM stack to a cubin. mapping, broadcast-aware register arithmetic and pointer construction, register-valued loop-carried accumulators, masked multi-result stores, and CUDA execution tests in which each thread computes several C elements. +- [ver17](https://github.com/pbelevich/mytriton/tree/ver17): optional PyTorch + tensor arguments, framework-independent runtime array metadata, zero-copy + DLPack conversion for CUDA tensors, same-device validation, execution on the + current PyTorch CUDA stream, and Torch-backed CUDA and MLIR execution tests. ## AST frontend @@ -362,6 +366,61 @@ K-tiles. Version 16 separates the physical thread tile from the logical dot output, carries register-tile accumulators through runtime K-loops, and emits a masked store for every result owned by a thread. +## PyTorch tensor interoperability + +Runtime pointer arguments may be NumPy arrays, CuPy arrays, or PyTorch tensors. +All three become the same `ptr` parameter in typed SSA, so the frontend, +optimizer, and backend source are independent of the Python array framework. + +CPU NumPy arrays and CPU PyTorch tensors are compilation-only inputs. A CUDA +PyTorch tensor compiles and executes the kernel directly: + +```python +import torch + +n = 1_000 +block = 256 +x = torch.ones(n, device="cuda", dtype=torch.float32) +y = torch.ones(n, device="cuda", dtype=torch.float32) +out = torch.empty_like(x) + +add_kernel[ + lambda meta: (triton.cdiv(n, meta["BLOCK"]),) +]( + x, + y, + out, + n, + BLOCK=block, +) +``` + +CuPy remains the internal CUDA compiler and launcher. At the runtime boundary, +a Torch CUDA tensor is detached from autograd metadata and exported through +DLPack: + +```text +torch.Tensor -- detach -- DLPack -- zero-copy CuPy view + | + v + RawKernel/cubin launch +``` + +`detach()` does not copy storage. It only makes the raw tensor memory +exportable through DLPack, matching a low-level Triton launch: tensors with +`requires_grad=True` are accepted as pointers, but the launch does not create +an autograd graph or provide a backward operation. + +Torch launches run in `torch.cuda.current_stream()` by wrapping it with +`cupy.cuda.Stream.from_external()`. DLPack conversion and kernel execution +happen inside the same stream context, so work queued before and after the +kernel remains correctly ordered without a global device synchronization. + +One launch must use either CuPy CUDA arrays or Torch CUDA tensors, not a mixture +of the two frameworks. All array arguments must be on the same CUDA device. +Mixing CPU and CUDA arrays is also rejected. As elsewhere in the current MVP, +runtime arrays must be C-contiguous and have `float32` elements. + ## Example ```python @@ -405,10 +464,10 @@ print(src) The first result contains the expression-tree operations built by the AST frontend. The second contains optimized typed SSA operations, and the third contains generated source for the selected backend. The default backend is -CUDA, so `src` is CUDA C++. With NumPy arguments, compilation stops there. If -the arguments are CuPy arrays and a CUDA GPU is available, the generated kernel -is also compiled and launched. Shared expressions such as `offsets` and `mask` -are lowered once and referenced by their SSA values wherever they are reused. +CUDA, so `src` is CUDA C++. With NumPy or CPU Torch arguments, compilation stops +there. With CuPy arrays or CUDA Torch tensors, the generated kernel is also +compiled and launched. Shared expressions such as `offsets` and `mask` are +lowered once and referenced by their SSA values wherever they are reused. For example, part of the resulting SSA looks like this: @@ -505,13 +564,14 @@ module attributes {gpu.container_module} { } ``` -For NumPy arguments, the MLIR backend stops after source generation, so MLIR -Python bindings are not required just to inspect the emitted MLIR. For CuPy -arguments, the backend runs a small pass pipeline that attaches an NVVM target, -converts GPU operations to NVVM, emits a GPU binary, extracts the cubin, loads -it through CuPy, and launches it with the same grid and thread-block size used -by the CUDA backend. CuPy arrays are passed using the ranked-memref ABI: -allocated pointer, aligned pointer, offset, size, and stride. +For NumPy or CPU Torch arguments, the MLIR backend stops after source +generation, so MLIR Python bindings are not required just to inspect the +emitted MLIR. For CuPy arrays or CUDA Torch tensors, the backend runs a small +pass pipeline that attaches an NVVM target, converts GPU operations to NVVM, +emits a GPU binary, extracts the cubin, loads it through CuPy, and launches it +with the same grid and thread-block size used by the CUDA backend. CUDA arrays +are passed using the ranked-memref ABI: allocated pointer, aligned pointer, +offset, size, and stride. The test kernels also include a copy, 2D matrix add, ReLU through `tl.maximum`, leaky ReLU through `tl.where`, sigmoid through negation, @@ -553,8 +613,10 @@ these rewrite passes because they are not region-aware yet. ## Current limitations - Generated backend source is returned as a string. Execution requires CuPy - built for the installed CUDA version and an available CUDA GPU; NumPy inputs - remain compilation-only. + built for the installed CUDA version and an available CUDA GPU. CUDA launch + arguments may be homogeneous CuPy arrays or PyTorch CUDA tensors; PyTorch is + imported only for Torch execution. NumPy arrays and CPU Torch tensors remain + compilation-only. - `MYTRITON_BACKEND` can be `cuda` or `mlir`. The CUDA backend is the default and supports the full current mytriton test language. The MLIR backend is an experimental MVP for 1D elementwise kernels. MLIR source generation does not @@ -568,7 +630,10 @@ these rewrite passes because they are not region-aware yet. assigning to the induction variable is rejected. `if`/`while`, `break`/`continue`, `for/else`, and other symbolic Python control flow are not supported. -- Runtime array arguments must be C-contiguous `float32` arrays. +- Runtime array arguments must be C-contiguous `float32` arrays. One execution + cannot mix CPU and CUDA arrays, CuPy and Torch CUDA arrays, or arrays from + different CUDA devices. Raw launches accept Torch tensors with + `requires_grad=True`, but do not participate in PyTorch autograd. - The launch grid is evaluated and used for CUDA execution, but it is not represented in the IR. - The CUDA kernel layout is inferred from block-shaped operands of observable @@ -624,8 +689,9 @@ these rewrite passes because they are not region-aware yet. code. It does not yet support 2D program IDs, reductions, `expand_dims`, Boolean `&`, `tl.maximum`, `tl.minimum`, `tl.where`, negation, `tl.exp`, `tl.static_range`, runtime `range`, or matrix multiplication. -- MLIR execution currently supports only 1D C-contiguous CuPy arrays because it - builds one-dimensional memref descriptors. +- MLIR execution currently supports only 1D C-contiguous CUDA arrays because it + builds one-dimensional memref descriptors. Torch CUDA tensors are normalized + to zero-copy CuPy views before those descriptors are constructed. - The SSA IR has structured `for` regions and loop-carried `iter_args`/`yield` values, but it has no general basic blocks, conditional control flow, or phi nodes outside this loop representation. @@ -648,6 +714,11 @@ To enable CUDA execution with CUDA 12, install the matching CuPy wheel: python -m pip install -e ".[cuda12]" ``` +PyTorch is an optional runtime integration rather than a project dependency. +Install a PyTorch build matching the local CUDA environment separately. When +PyTorch is available, CUDA tensors can be passed directly to kernels; CuPy is +still required internally for CUDA source compilation and kernel launch. + MLIR cubin execution requires Python bindings importable as `mlir.ir` and `mlir.passmanager`, plus an MLIR build that includes the GPU/NVVM passes needed by `gpu-module-to-binary`. These bindings are intentionally not listed as a diff --git a/src/mytriton/compiler.py b/src/mytriton/compiler.py index b9fe067..b93d9e8 100644 --- a/src/mytriton/compiler.py +++ b/src/mytriton/compiler.py @@ -118,7 +118,7 @@ def launch(*args: P.args, **kwargs: P.kwargs) -> CompilationResult: chip = None if backend == "mlir": try: - chip = cuda_chip() + chip = cuda_chip(runtime_args) except CudaUnavailableError: chip = "sm_80" diff --git a/src/mytriton/cuda_utils.py b/src/mytriton/cuda_utils.py index 4a8edd4..af727ee 100644 --- a/src/mytriton/cuda_utils.py +++ b/src/mytriton/cuda_utils.py @@ -1,9 +1,13 @@ import importlib +from collections.abc import Iterator +from contextlib import contextmanager from typing import Any, Protocol, TypeGuard import numpy as np -CudaKernelCache = dict[tuple[object, str], Any] +from .runtime_args import array_arg_info + +CudaKernelCache = dict[tuple[object, str, int], Any] class _ArrayFlagsLike(Protocol): @@ -48,6 +52,15 @@ def cuda_module(): return cp +def _torch_module() -> Any: + try: + return importlib.import_module("torch") + except (ImportError, OSError) as error: + raise CudaUnavailableError( + "PyTorch is required for Torch CUDA tensor execution" + ) from error + + def _is_cupy_array(value: object) -> TypeGuard[_CupyArrayLike]: module = type(value).__module__ return module == "cupy" or module.startswith("cupy.") @@ -61,27 +74,118 @@ def _convert_runtime_arg(value: object) -> object: return value +def _normalize_cuda_array_args( + cp, + runtime_args: tuple[object, ...], +) -> tuple[object, ...]: + normalized = [] + + for value in runtime_args: + info = array_arg_info(value) + + if info is None: + normalized.append(value) + continue + + if not info.is_cuda: + raise TypeError(f"cannot use {info.framework} CPU array as a CUDA argument") + + if info.framework == "cupy": + normalized.append(value) + continue + + if info.framework == "torch": + detach = getattr(value, "detach", None) + if not callable(detach): + raise TypeError("Torch CUDA array does not support detach()") + + # A raw kernel launch is not an autograd operation. Detaching makes + # the tensor exportable through DLPack while preserving its storage. + normalized.append(cp.from_dlpack(detach())) + continue + + raise TypeError(f"unsupported CUDA array framework: {info.framework}") + + return tuple(normalized) + + def cuda_execution_required( runtime_args: tuple[object, ...], *, backend_name: str ) -> bool: - array_args = [ - value - for value in runtime_args - if hasattr(value, "dtype") and hasattr(value, "flags") + array_infos = [ + info for value in runtime_args if (info := array_arg_info(value)) is not None ] - cupy_array_args = [value for value in array_args if _is_cupy_array(value)] - if not cupy_array_args: + if not array_infos: + return False + + cuda_infos = [info for info in array_infos if info.is_cuda] + + if not cuda_infos: return False - if len(cupy_array_args) != len(array_args): + if len(cuda_infos) != len(array_infos): + raise TypeError( + f"{backend_name} execution does not support mixed CPU and CUDA arrays" + ) + + frameworks = {info.framework for info in cuda_infos} + if len(frameworks) != 1: + rendered = ", ".join(sorted(frameworks)) raise TypeError( - f"{backend_name} execution does not support mixed NumPy and CuPy arrays" + f"{backend_name} execution does not support mixed CUDA " + f"array frameworks: {rendered}" + ) + + device_indices = {info.device_index for info in cuda_infos} + if len(device_indices) != 1: + rendered = ", ".join( + "unknown" if index is None else str(index) + for index in sorted( + device_indices, + key=lambda index: -1 if index is None else index, + ) + ) + raise TypeError( + f"{backend_name} execution requires one CUDA device, got: {rendered}" ) return True +@contextmanager +def _cuda_launch_context( + cp, + runtime_args: tuple[object, ...], +) -> Iterator[None]: + cuda_infos = [ + info + for value in runtime_args + if ((info := array_arg_info(value)) is not None and info.is_cuda) + ] + + if not cuda_infos: + raise RuntimeError("CUDA launch requires at least one CUDA array") + + launch_info = cuda_infos[0] + device_index = launch_info.device_index + + if device_index is None: + raise TypeError("CUDA array device index is unavailable") + + with cp.cuda.Device(device_index): + if launch_info.framework == "torch": + torch = _torch_module() + torch_stream = torch.cuda.current_stream( + device=device_index, + ) + + with cp.cuda.Stream.from_external(torch_stream): + yield + else: + yield + + def execute_cuda_if_needed( *, kernel_cache: CudaKernelCache, @@ -91,35 +195,66 @@ def execute_cuda_if_needed( threads_per_block: int, runtime_args: tuple[object, ...], ) -> None: - # NumPy calls are compilation-only, including on CUDA machines. + # CPU arrays are compilation-only, including on CUDA machines. if not cuda_execution_required(runtime_args, backend_name="CUDA"): return cp = cuda_module() - max_threads = cp.cuda.Device().attributes["MaxThreadsPerBlock"] - if threads_per_block > max_threads: - raise ValueError( - f"CUDA block size {threads_per_block} exceeds device limit {max_threads}" - ) - cache_key = (cuda_src, kernel_name) - if cache_key not in kernel_cache: - kernel_cache[cache_key] = cp.RawKernel( + with _cuda_launch_context(cp, runtime_args): + max_threads = cp.cuda.Device().attributes["MaxThreadsPerBlock"] + if threads_per_block > max_threads: + raise ValueError( + f"CUDA block size {threads_per_block} " + f"exceeds device limit {max_threads}" + ) + + cache_key = ( cuda_src, kernel_name, - options=("--std=c++14",), + cp.cuda.Device().id, + ) + if cache_key not in kernel_cache: + kernel_cache[cache_key] = cp.RawKernel( + cuda_src, + kernel_name, + options=("--std=c++14",), + ) + + normalized_args = _normalize_cuda_array_args( + cp, + runtime_args, + ) + + kernel_cache[cache_key]( + launch_grid, + (threads_per_block,), + tuple(_convert_runtime_arg(value) for value in normalized_args), ) - kernel_cache[cache_key]( - launch_grid, - (threads_per_block,), - tuple(_convert_runtime_arg(value) for value in runtime_args), - ) +def cuda_chip(runtime_args: tuple[object, ...] = ()) -> str: + cuda_infos = [ + info + for value in runtime_args + if ((info := array_arg_info(value)) is not None and info.is_cuda) + ] + + device_index = None + if cuda_infos: + cuda_execution_required(runtime_args, backend_name="MLIR") + device_index = cuda_infos[0].device_index + + if device_index is None: + raise TypeError("CUDA array device index is unavailable") -def cuda_chip() -> str: cp = cuda_module() - return f"sm_{cp.cuda.Device().compute_capability}" + + if device_index is None: + return f"sm_{cp.cuda.Device().compute_capability}" + + with cp.cuda.Device(device_index): + return f"sm_{cp.cuda.Device().compute_capability}" def _convert_mlir_memref_args(runtime_args: tuple[object, ...]) -> tuple[object, ...]: @@ -157,25 +292,37 @@ def execute_mlir_cubin_if_needed( threads_per_block: int, runtime_args: tuple[object, ...], ) -> None: - # NumPy calls are compile-only, same behavior as CUDA backend. + # CPU arrays are compilation-only, same behavior as the CUDA backend. if not cuda_execution_required(runtime_args, backend_name="MLIR"): return cp = cuda_module() - max_threads = cp.cuda.Device().attributes["MaxThreadsPerBlock"] - if threads_per_block > max_threads: - raise ValueError( - f"CUDA block size {threads_per_block} exceeds device limit {max_threads}" + + with _cuda_launch_context(cp, runtime_args): + max_threads = cp.cuda.Device().attributes["MaxThreadsPerBlock"] + if threads_per_block > max_threads: + raise ValueError( + f"CUDA block size {threads_per_block} " + f"exceeds device limit {max_threads}" + ) + + cache_key = ( + cubin, + kernel_name, + cp.cuda.Device().id, + ) + if cache_key not in kernel_cache: + module = cp.cuda.function.Module() + module.load(cubin) + kernel_cache[cache_key] = module.get_function(kernel_name) + + normalized_args = _normalize_cuda_array_args( + cp, + runtime_args, ) - cache_key = (cubin, kernel_name) - if cache_key not in kernel_cache: - module = cp.cuda.function.Module() - module.load(cubin) - kernel_cache[cache_key] = module.get_function(kernel_name) - - kernel_cache[cache_key]( - launch_grid, - (threads_per_block,), - _convert_mlir_memref_args(runtime_args), - ) + kernel_cache[cache_key]( + launch_grid, + (threads_per_block,), + _convert_mlir_memref_args(normalized_args), + ) diff --git a/src/mytriton/runtime_args.py b/src/mytriton/runtime_args.py new file mode 100644 index 0000000..c3fe4fa --- /dev/null +++ b/src/mytriton/runtime_args.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass +from typing import Literal, cast + +import numpy as np + +ArrayFramework = Literal["numpy", "cupy", "torch"] +ArrayDevice = Literal["cpu", "cuda"] + + +@dataclass(frozen=True) +class RuntimeArrayInfo: + framework: ArrayFramework + device: ArrayDevice + device_index: int | None + dtype_name: str + c_contiguous: bool + + @property + def is_cuda(self) -> bool: + return self.device == "cuda" + + +def _module_matches(module: str, root: str) -> bool: + return module == root or module.startswith(f"{root}.") + + +def _device_index(value: object) -> int | None: + return value if type(value) is int else None + + +def array_arg_info(value: object) -> RuntimeArrayInfo | None: + if isinstance(value, np.ndarray): + return RuntimeArrayInfo( + framework="numpy", + device="cpu", + device_index=None, + dtype_name=str(value.dtype), + c_contiguous=bool(value.flags.c_contiguous), + ) + + module = type(value).__module__ + + if _module_matches(module, "cupy"): + dtype = getattr(value, "dtype", None) + flags = getattr(value, "flags", None) + device = getattr(value, "device", None) + + if dtype is None or flags is None or device is None: + return None + + return RuntimeArrayInfo( + framework="cupy", + device="cuda", + device_index=_device_index(getattr(device, "id", None)), + dtype_name=str(dtype), + c_contiguous=bool(getattr(flags, "c_contiguous", False)), + ) + + if _module_matches(module, "torch"): + dtype = getattr(value, "dtype", None) + device = getattr(value, "device", None) + is_contiguous = getattr(value, "is_contiguous", None) + + if dtype is None or device is None or not callable(is_contiguous): + return None + + device_type = getattr(device, "type", None) + if device_type not in ("cpu", "cuda"): + return None + + return RuntimeArrayInfo( + framework="torch", + device=cast(ArrayDevice, device_type), + device_index=_device_index(getattr(device, "index", None)), + dtype_name=str(dtype).removeprefix("torch."), + c_contiguous=bool(is_contiguous()), + ) + + return None diff --git a/src/mytriton/trace.py b/src/mytriton/trace.py index dde9d93..82afd03 100644 --- a/src/mytriton/trace.py +++ b/src/mytriton/trace.py @@ -5,6 +5,8 @@ import numpy as np +from .runtime_args import array_arg_info + # ---------------------------- # Language API # ---------------------------- @@ -527,11 +529,15 @@ def current(): def _make_param(name, value) -> Param: - if hasattr(value, "dtype") and hasattr(value, "flags"): - if str(value.dtype) != "float32": + array_info = array_arg_info(value) + + if array_info is not None: + if array_info.dtype_name != "float32": raise TypeError(f"{name}: only float32 arrays are supported") - if not value.flags.c_contiguous: + + if not array_info.c_contiguous: raise TypeError(f"{name}: only C-contiguous arrays are supported") + return Param(name, PTR_F32) if isinstance(value, (int, np.integer)): diff --git a/tests/test_runtime_args.py b/tests/test_runtime_args.py new file mode 100644 index 0000000..6a8b705 --- /dev/null +++ b/tests/test_runtime_args.py @@ -0,0 +1,154 @@ +import pytest + +from mytriton.cuda_utils import cuda_execution_required +from mytriton.runtime_args import RuntimeArrayInfo, array_arg_info + + +class FakeDType: + def __init__(self, name: str) -> None: + self.name = name + + def __str__(self) -> str: + return self.name + + +class FakeDevice: + def __init__(self, device_type: str, index: int | None) -> None: + self.type = device_type + self.index = index + self.id = index + + +class FakeFlags: + def __init__(self, *, c_contiguous: bool) -> None: + self.c_contiguous = c_contiguous + + +class FakeTorchTensor: + __module__ = "torch" + + def __init__( + self, + *, + device: str, + device_index: int | None, + c_contiguous: bool = True, + ) -> None: + self.dtype = FakeDType("torch.float32") + self.device = FakeDevice(device, device_index) + self._c_contiguous = c_contiguous + + def is_contiguous(self) -> bool: + return self._c_contiguous + + +class FakeCupyArray: + __module__ = "cupy" + + def __init__(self, *, device_index: int) -> None: + self.dtype = FakeDType("float32") + self.device = FakeDevice("cuda", device_index) + self.flags = FakeFlags(c_contiguous=True) + + +def test_array_arg_info_recognizes_torch_cpu_tensor() -> None: + tensor = FakeTorchTensor( + device="cpu", + device_index=None, + c_contiguous=False, + ) + + assert array_arg_info(tensor) == RuntimeArrayInfo( + framework="torch", + device="cpu", + device_index=None, + dtype_name="float32", + c_contiguous=False, + ) + + +def test_array_arg_info_recognizes_torch_cuda_tensor() -> None: + tensor = FakeTorchTensor( + device="cuda", + device_index=2, + ) + + assert array_arg_info(tensor) == RuntimeArrayInfo( + framework="torch", + device="cuda", + device_index=2, + dtype_name="float32", + c_contiguous=True, + ) + + +def test_torch_cpu_arrays_are_compilation_only() -> None: + tensor = FakeTorchTensor( + device="cpu", + device_index=None, + ) + + assert not cuda_execution_required((tensor,), backend_name="CUDA") + + +def test_torch_cuda_arrays_require_execution() -> None: + tensor = FakeTorchTensor( + device="cuda", + device_index=0, + ) + + assert cuda_execution_required((tensor,), backend_name="CUDA") + + +def test_execution_rejects_mixed_cpu_and_cuda_arrays() -> None: + cpu = FakeTorchTensor( + device="cpu", + device_index=None, + ) + cuda = FakeTorchTensor( + device="cuda", + device_index=0, + ) + + with pytest.raises( + TypeError, + match="does not support mixed CPU and CUDA arrays", + ): + cuda_execution_required((cpu, cuda), backend_name="CUDA") + + +def test_execution_rejects_mixed_cuda_frameworks() -> None: + torch_tensor = FakeTorchTensor( + device="cuda", + device_index=0, + ) + cupy_array = FakeCupyArray(device_index=0) + + with pytest.raises( + TypeError, + match="mixed CUDA array frameworks: cupy, torch", + ): + cuda_execution_required( + (torch_tensor, cupy_array), + backend_name="CUDA", + ) + + +def test_execution_rejects_multiple_cuda_devices() -> None: + first = FakeTorchTensor( + device="cuda", + device_index=0, + ) + second = FakeTorchTensor( + device="cuda", + device_index=1, + ) + + with pytest.raises( + TypeError, + match="requires one CUDA device, got: 0, 1", + ): + cuda_execution_required( + (first, second), + backend_name="CUDA", + ) diff --git a/tests/test_torch_interop.py b/tests/test_torch_interop.py new file mode 100644 index 0000000..00b4ec7 --- /dev/null +++ b/tests/test_torch_interop.py @@ -0,0 +1,255 @@ +from typing import Any + +import pytest + +import mytriton as triton +import mytriton.language as tl +from mytriton.cuda_utils import ( + CudaKernelCache, + cuda_module, + execute_cuda_if_needed, +) + + +@triton.jit +def torch_add_kernel( + x, + y, + out, + n, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n + + x_values = tl.load(x + offsets, mask=mask) + y_values = tl.load(y + offsets, mask=mask) + + tl.store( + out + offsets, + x_values + y_values, + mask=mask, + ) + + +def torch_module() -> Any: + return pytest.importorskip("torch") + + +def test_torch_cpu_tensors_are_compile_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + torch = torch_module() + + import mytriton.cuda_utils as cuda_utils + + def fail(): + raise AssertionError("CuPy must not be loaded for CPU Torch tensors") + + monkeypatch.setattr(cuda_utils, "_cupy", fail) + + n = 16 + x = torch.ones(n, dtype=torch.float32) + y = torch.ones(n, dtype=torch.float32) + out = torch.empty(n, dtype=torch.float32) + + torch_add_kernel.clear_cache() + + _, _, cuda_src = torch_add_kernel[(1,)]( + x, + y, + out, + n, + BLOCK_SIZE=16, + ) + + assert "void torch_add_kernel" in cuda_src + + +@pytest.mark.execution +def test_torch_cuda_tensors_execute( + backend: str, +) -> None: + torch = torch_module() + + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is not available") + + assert backend in {"cuda", "mlir"} + + n = 37 + block_size = 32 + + x = torch.arange( + n, + device="cuda", + dtype=torch.float32, + ) + y = ( + torch.arange( + n, + device="cuda", + dtype=torch.float32, + ) + * 0.5 + ) + out = torch.full( + (n,), + float("nan"), + device="cuda", + dtype=torch.float32, + ) + + grid = ((n + block_size - 1) // block_size,) + + torch_add_kernel.clear_cache() + torch_add_kernel[grid]( + x, + y, + out, + n, + BLOCK_SIZE=block_size, + ) + torch.cuda.synchronize() + + torch.testing.assert_close( + out, + x + y, + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.execution +def test_torch_cuda_tensor_requiring_grad_executes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + torch = torch_module() + + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is not available") + + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + n = 32 + x = torch.arange( + n, + device="cuda", + dtype=torch.float32, + requires_grad=True, + ) + y = torch.ones_like(x) + out = torch.empty_like(x, requires_grad=False) + + torch_add_kernel.clear_cache() + torch_add_kernel[(1,)]( + x, + y, + out, + n, + BLOCK_SIZE=n, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(out, x.detach() + y) + assert not out.requires_grad + + +@pytest.mark.execution +def test_torch_launch_uses_current_torch_stream() -> None: + torch = torch_module() + + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is not available") + + cp = cuda_module() + tensor = torch.zeros( + 1, + device="cuda", + dtype=torch.float32, + ) + torch_stream = torch.cuda.Stream(device=tensor.device) + + observed_streams: list[int] = [] + + def probe_kernel( + launch_grid: tuple[int, ...], + block: tuple[int, ...], + args: tuple[object, ...], + ) -> None: + del launch_grid, block, args + observed_streams.append(cp.cuda.get_current_stream().ptr) + + cuda_src = "unused" + kernel_name = "stream_probe" + device_index = tensor.device.index + assert device_index is not None + cache_key = (cuda_src, kernel_name, device_index) + kernel_cache: CudaKernelCache = { + cache_key: probe_kernel, + } + + with torch.cuda.stream(torch_stream): + execute_cuda_if_needed( + kernel_cache=kernel_cache, + cuda_src=cuda_src, + kernel_name=kernel_name, + launch_grid=(1,), + threads_per_block=1, + runtime_args=(tensor,), + ) + + assert observed_streams == [torch_stream.cuda_stream] + + +@pytest.mark.execution +def test_torch_cuda_tensors_execute_on_non_default_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + torch = torch_module() + + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is not available") + + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + n = 4097 + block_size = 256 + torch_stream = torch.cuda.Stream() + + with torch.cuda.stream(torch_stream): + x = torch.arange( + n, + device="cuda", + dtype=torch.float32, + ) + y = x * 0.25 + out = torch.full( + (n,), + float("nan"), + device="cuda", + dtype=torch.float32, + ) + + grid = ((n + block_size - 1) // block_size,) + + torch_add_kernel.clear_cache() + torch_add_kernel[grid]( + x, + y, + out, + n, + BLOCK_SIZE=block_size, + ) + + actual = out.clone() + expected = x + y + + torch_stream.synchronize() + + torch.testing.assert_close( + actual, + expected, + rtol=1e-5, + atol=1e-5, + ) From ea9e47ace3fa5bf1aa5e97063c26cd7dc663de07 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Tue, 1 Sep 2026 00:04:55 -0400 Subject: [PATCH 2/2] Fix README formatting for CI --- README.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/README.md b/README.md index 38a5049..98bccaa 100644 --- a/README.md +++ b/README.md @@ -384,9 +384,7 @@ x = torch.ones(n, device="cuda", dtype=torch.float32) y = torch.ones(n, device="cuda", dtype=torch.float32) out = torch.empty_like(x) -add_kernel[ - lambda meta: (triton.cdiv(n, meta["BLOCK"]),) -]( +add_kernel[lambda meta: (triton.cdiv(n, meta["BLOCK"]),)]( x, y, out,