From 5e4e6e4de53da4b11b7d4002a2763989fdfe0520 Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Sun, 8 Mar 2026 13:38:48 +0000 Subject: [PATCH 1/2] feat(mindtorch_v2): add cuda backend MVP with explicit dispatch keys --- ...08-mindtorch-v2-cuda-backend-mvp-design.md | 243 ++++++++++++++ ...026-03-08-mindtorch-v2-cuda-backend-mvp.md | 199 ++++++++++++ src/mindtorch_v2/__init__.py | 2 + src/mindtorch_v2/_backends/__init__.py | 3 +- src/mindtorch_v2/_backends/autograd.py | 2 + src/mindtorch_v2/_backends/common/convert.py | 32 ++ src/mindtorch_v2/_backends/cuda/__init__.py | 18 ++ src/mindtorch_v2/_backends/cuda/creation.py | 54 ++++ src/mindtorch_v2/_backends/cuda/ops.py | 17 + src/mindtorch_v2/_backends/cuda/runtime.py | 298 ++++++++++++++++++ src/mindtorch_v2/_backends/cuda/storage.py | 159 ++++++++++ src/mindtorch_v2/_dispatch/keys.py | 34 +- src/mindtorch_v2/_dispatch/registration.py | 8 +- src/mindtorch_v2/_dispatch/registry.py | 7 +- src/mindtorch_v2/_storage.py | 33 ++ src/mindtorch_v2/cuda.py | 120 +++++++ src/mindtorch_v2/library.py | 1 + tests/mindtorch_v2/test_cuda_backend.py | 128 ++++++++ tests/mindtorch_v2/test_device_transfer.py | 49 ++- tests/mindtorch_v2/test_dispatch_keys.py | 23 ++ tests/mindtorch_v2/test_dispatch_registry.py | 2 + .../mindtorch_v2/test_dispatch_resolution.py | 12 + 22 files changed, 1415 insertions(+), 29 deletions(-) create mode 100644 docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp-design.md create mode 100644 docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp.md create mode 100644 src/mindtorch_v2/_backends/cuda/__init__.py create mode 100644 src/mindtorch_v2/_backends/cuda/creation.py create mode 100644 src/mindtorch_v2/_backends/cuda/ops.py create mode 100644 src/mindtorch_v2/_backends/cuda/runtime.py create mode 100644 src/mindtorch_v2/_backends/cuda/storage.py create mode 100644 src/mindtorch_v2/cuda.py create mode 100644 tests/mindtorch_v2/test_cuda_backend.py diff --git a/docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp-design.md b/docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp-design.md new file mode 100644 index 000000000..ca6c6f49b --- /dev/null +++ b/docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp-design.md @@ -0,0 +1,243 @@ +# MindTorch v2 CUDA Backend MVP Design + +**Status:** Approved + +**Goal:** Turn `mindtorch_v2` CUDA from a reserved device label into a real backend with working device APIs, storage, transfer, and tensor creation—without depending on any other framework. + +**Scope:** This design covers the first usable CUDA backend milestone only. It does not attempt broad operator parity, full autograd parity, AMP, profiler support, or distributed support. + +--- + +## Problem + +`mindtorch_v2` already reserves a CUDA path in dispatch and device handling, but CUDA is not a real backend today: + +- Dispatch registration maps `cuda` to placeholder keys in `src/mindtorch_v2/_dispatch/registration.py`. +- Dispatch key construction recognizes CUDA tensors in `src/mindtorch_v2/_dispatch/keys.py`. +- Tensor APIs still treat `.cuda()` as unsupported in current tests. + +This means the mechanism is partially prepared, but there is no CUDA runtime layer, storage layer, or creation/transfer implementation behind it. + +--- + +## Constraints + +- Pure Python implementation. +- No dependency on external deep learning frameworks. +- Respect `mindtorch_v2` schema-first registration rules. +- Keep the first phase small and mechanism-focused. +- Use CUDA runtime primitives directly via Python FFI. + +--- + +## Non-Goals + +This MVP does not include: + +- Broad math operator coverage. +- Full Torch CUDA semantic parity. +- Full autograd on CUDA. +- AMP/autocast/GradScaler support. +- CUDA profiler support. +- NCCL/distributed support. + +Those will become follow-on phases once backend fundamentals are stable. + +--- + +## Recommended Approach + +Use a layered CUDA backend built directly on `ctypes` bindings to `libcudart.so`. + +The first phase should only establish: + +1. CUDA runtime discovery and device management. +2. GPU memory allocation/free and memory copy. +3. A dedicated `CudaStorage` implementation. +4. Tensor movement between CPU and CUDA. +5. CUDA creation ops: `empty`, `zeros`, `ones`, `full`, `tensor`, `to`. + +This is the smallest design that turns CUDA into a real device in the system. + +--- + +## Architecture + +### 1. Runtime Layer + +Add `src/mindtorch_v2/_backends/cuda/runtime.py`. + +Responsibilities: + +- Load `libcudart.so` lazily. +- Expose Python wrappers for: + - `cudaGetDeviceCount` + - `cudaGetDevice` + - `cudaSetDevice` + - `cudaMalloc` + - `cudaFree` + - `cudaMemcpy` + - `cudaMemcpyAsync` + - `cudaMemset` + - `cudaDeviceSynchronize` + - `cudaStreamCreate` + - `cudaStreamDestroy` + - `cudaStreamSynchronize` + - `cudaEventCreate` + - `cudaEventDestroy` + - `cudaEventRecord` + - `cudaEventSynchronize` +- Raise clear Python exceptions from CUDA error codes. + +This layer is intentionally small and should not include operator logic. + +### 2. Storage Layer + +Add `src/mindtorch_v2/_backends/cuda/storage.py`. + +Responsibilities: + +- Represent device-backed raw storage. +- Own a CUDA pointer and its lifetime. +- Track `nbytes`, `dtype`, and `device`. +- Support conversion helpers for host/device copies. +- Provide enough storage API compatibility for `Tensor` to work. + +The initial storage type can remain simple: contiguous allocation only, no custom allocator, no pooling, no IPC. + +### 3. Storage Factory Integration + +Extend `src/mindtorch_v2/_storage.py`. + +Responsibilities: + +- Add CUDA storage factory helpers. +- Support CPU numpy -> CUDA upload. +- Support CUDA -> CPU download. +- Route typed storage creation by device type. + +This is the bridge between existing CPU/meta flows and the new CUDA flow. + +### 4. Public CUDA API + +Add `src/mindtorch_v2/cuda.py`. + +Responsibilities: + +- Mirror the shape of `src/mindtorch_v2/npu.py` where practical. +- Expose: + - `is_available` + - `device_count` + - `current_device` + - `set_device` + - `synchronize` + - `Stream` + - `Event` + - `device` context manager + +This gives users a stable entry point for CUDA backend discovery and control. + +### 5. Tensor Transfer and Creation + +Update: + +- `src/mindtorch_v2/_tensor.py` +- `src/mindtorch_v2/_creation.py` +- `src/mindtorch_v2/_backends/cuda/creation.py` + +Responsibilities: + +- Make `Tensor.cuda()` call into `to("cuda")`. +- Make `Tensor.to("cuda")` and `Tensor.to("cpu")` perform actual device transfer. +- Support direct creation on CUDA for `tensor`, `empty`, `zeros`, `ones`, `full`. + +For `ones` and `full`, the first version may use a temporary host buffer plus upload if that is simpler than adding a fill kernel immediately. + +--- + +## Dispatch Strategy + +Do not redesign dispatch in this phase. + +Make CUDA a first-class dispatch backend in this PR. + +Required dispatch changes: + +- add `DispatchKey.CUDA` +- add `DispatchKey.AutogradCUDA` +- update dispatch priority and keyset construction +- update registration helpers so `cuda` no longer maps to `PrivateUse1` +- keep `PrivateUse1` reserved for actual private-use backends + +This aligns the implementation with the requirement that CUDA must not be represented as `PrivateUse1`. + +--- + +## Testing Strategy + +Follow the repo rule: schema first, then contract tests, then backend wiring. + +### Required mechanism tests + +- `PYTHONPATH=src pytest -q tests/mindtorch_v2/contract/test_schema_registration_order.py` +- `PYTHONPATH=src pytest -q tests/mindtorch_v2/contract/test_schema_coverage.py` + +### New CUDA MVP tests + +Add or update tests for: + +- `torch.cuda.is_available()` style availability surface. +- `mt.tensor(..., device="cuda")` creation. +- `x.cuda()` success. +- `x.to("cuda")` and `x.to("cpu")` round-trip correctness. +- `zeros/ones/full/empty(..., device="cuda")` creation. +- `current_device` / `set_device` behavior. + +Tests should gracefully skip when CUDA runtime is unavailable. + +--- + +## Risks + +### Runtime loading risk + +CUDA library names can differ across systems. The runtime layer should attempt a small set of common library names and fail gracefully. + +### Lifetime management risk + +Leaking device memory is easy when Python owns raw pointers. `CudaStorage` should centralize ownership and cleanup. + +### Shape/stride risk + +This MVP should avoid pretending to support advanced non-contiguous CUDA storage semantics before they are actually implemented. + +### Scope creep risk + +Do not add math kernels in this phase unless they are strictly required to support creation/transfer semantics. + +--- + +## Success Criteria + +The MVP is complete when all of the following are true: + +- `mindtorch_v2` exposes a working `cuda` module. +- CUDA availability can be queried without crashing on non-CUDA systems. +- Tensors can be created on CUDA. +- CPU <-> CUDA transfer works for supported dtypes. +- `.cuda()` no longer fails as an unsupported reserved path. +- CUDA backend changes respect schema-first and pass the required contract tests. + +--- + +## Follow-On Phases + +After this MVP, the next recommended phases are: + +1. Pointwise math ops. +2. Reductions. +3. `matmul` and BLAS-backed operations. +4. Convolution/pooling via cuDNN. +5. Autograd correctness. +6. AMP, profiler, and distributed support. + diff --git a/docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp.md b/docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp.md new file mode 100644 index 000000000..97d469ab6 --- /dev/null +++ b/docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp.md @@ -0,0 +1,199 @@ +# MindTorch v2 CUDA Backend MVP Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Build the first usable CUDA backend for `mindtorch_v2`, covering runtime discovery, storage, transfer, and CUDA tensor creation without relying on any external deep learning framework. + +**Architecture:** Add a thin CUDA runtime layer over `libcudart.so`, implement a dedicated `CudaStorage`, wire transfer and creation paths into existing tensor/storage APIs, and upgrade CUDA to explicit dispatch keys and register the backend against those keys. Keep the phase mechanism-focused and avoid broad operator work. + +**Tech Stack:** Python, `ctypes`, CUDA Runtime API, existing `mindtorch_v2` dispatch/storage/tensor stack, `pytest`. + +--- + +### Task 1: Confirm current CUDA surface and test baselines + +**Files:** +- Inspect: `src/mindtorch_v2/_device.py` +- Inspect: `src/mindtorch_v2/_tensor.py` +- Inspect: `src/mindtorch_v2/_storage.py` +- Inspect: `src/mindtorch_v2/_creation.py` +- Inspect: `src/mindtorch_v2/_dispatch/registration.py` +- Inspect: `src/mindtorch_v2/_dispatch/keys.py` +- Inspect: `tests/mindtorch_v2/test_device_transfer.py` + +**Steps:** +1. Read the files above and note all current CUDA placeholders and unsupported paths. +2. Identify which creation APIs already have schema coverage and which only need backend wiring. +3. Confirm which tests currently assume `.cuda()` is unsupported. +4. Write down the exact functions that need to change before touching code. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_device_transfer.py` +Expected: current CUDA tests fail because support is not implemented, or assert unsupported behavior. + +--- + +### Task 2: Add failing CUDA MVP tests + +**Files:** +- Modify: `tests/mindtorch_v2/test_device_transfer.py` +- Create: `tests/mindtorch_v2/test_cuda_backend.py` + +**Steps:** +1. Replace unsupported `.cuda()` assertions with behavior expected from the MVP. +2. Add tests for `tensor(..., device="cuda")`, `x.to("cuda")`, `x.to("cpu")`, and CUDA creation ops. +3. Add tests for `current_device`, `set_device`, and `synchronize`. +4. Guard runtime-dependent tests with skip conditions when CUDA is unavailable. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_device_transfer.py tests/mindtorch_v2/test_cuda_backend.py` +Expected: FAIL with missing CUDA backend/runtime implementation. + +--- + +### Task 3: Add CUDA runtime bindings + +**Files:** +- Create: `src/mindtorch_v2/_backends/cuda/runtime.py` +- Create: `src/mindtorch_v2/_backends/cuda/__init__.py` + +**Steps:** +1. Implement lazy runtime library loading for common CUDA runtime library names. +2. Define Python wrappers for device count, current device, set device, allocate/free, memcpy, memset, synchronize, stream, and event APIs. +3. Implement a shared CUDA error checker that raises descriptive Python exceptions. +4. Add a small availability probe that cleanly returns `False` when CUDA runtime is missing. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_cuda_backend.py -k "availability or device"` +Expected: device/runtime tests move from import errors to behavior failures where storage/transfer is still missing. + +--- + +### Task 4: Implement `CudaStorage` + +**Files:** +- Create: `src/mindtorch_v2/_backends/cuda/storage.py` +- Modify: `src/mindtorch_v2/_storage.py` + +**Steps:** +1. Add a storage type that owns a CUDA allocation and frees it on cleanup. +2. Store pointer, byte size, dtype, and device metadata. +3. Add host-device and device-host copy helpers. +4. Add typed storage factory helpers for CUDA allocation and upload/download. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_cuda_backend.py -k "tensor_creation or roundtrip"` +Expected: transfer tests still fail until tensor/creation wiring is added, but storage paths import and allocate correctly. + +--- + +### Task 5: Expose public `cuda` module + +**Files:** +- Create: `src/mindtorch_v2/cuda.py` +- Modify: `src/mindtorch_v2/__init__.py` + +**Steps:** +1. Implement `is_available`, `device_count`, `current_device`, `set_device`, and `synchronize`. +2. Add `Stream`, `Event`, and `device` context manager wrappers using the runtime layer. +3. Export the module from the package root. +4. Match the public API shape used by `npu.py` where possible without overbuilding. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_cuda_backend.py -k "availability or current_device or stream or event"` +Expected: public CUDA API tests pass or are skipped when CUDA is unavailable. + +--- + +### Task 6: Wire tensor transfer paths + +**Files:** +- Modify: `src/mindtorch_v2/_tensor.py` +- Modify: `src/mindtorch_v2/_creation.py` +- Modify: `src/mindtorch_v2/_device.py` + +**Steps:** +1. Implement `.cuda()` as a transfer convenience API. +2. Implement `.to("cuda")` and `.to("cpu")` using the new storage helpers. +3. Preserve shape, stride, dtype, and device metadata across transfers. +4. Keep `_numpy_view()` disallowed for CUDA-backed tensors unless moved to CPU first. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_device_transfer.py tests/mindtorch_v2/test_cuda_backend.py -k "cuda or to"` +Expected: transfer tests pass for supported dtypes. + +--- + +### Task 7: Add CUDA creation kernels + +**Files:** +- Create: `src/mindtorch_v2/_backends/cuda/creation.py` +- Modify: `src/mindtorch_v2/_dispatch/schemas.py` (only if missing schema coverage is found) +- Modify: `src/mindtorch_v2/_dispatch/registry.py` or backend registration entrypoints as needed +- Modify: `src/mindtorch_v2/_creation.py` + +**Steps:** +1. Confirm schemas already exist for `tensor`, `empty`, `zeros`, `ones`, `full`, and `to`. +2. If any schema is missing, add it first. +3. Register CUDA creation kernels through the existing registration helpers. +4. Implement `empty`, `zeros`, `ones`, `full`, and direct `tensor(..., device="cuda")` support. +5. Use the simplest correct fill path first, even if `ones/full` initially stage through CPU. + +**Test command:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_cuda_backend.py -k "creation or tensor_creation"` +Expected: CUDA creation tests pass. + +--- + +### Task 8: Validate schema and contract invariants + +**Files:** +- Inspect: `src/mindtorch_v2/_dispatch/schemas.py` +- Inspect: `tests/mindtorch_v2/contract/test_schema_registration_order.py` +- Inspect: `tests/mindtorch_v2/contract/test_schema_coverage.py` + +**Steps:** +1. Confirm no CUDA kernel registration happens before schema registration. +2. Run required contract tests. +3. Fix only CUDA MVP-related schema or registration issues. + +**Test commands:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/contract/test_schema_registration_order.py` +Expected: PASS + +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/contract/test_schema_coverage.py` +Expected: PASS + +--- + +### Task 9: Run focused regression tests + +**Files:** +- Inspect: `tests/mindtorch_v2/test_creation.py` +- Inspect: `tests/mindtorch_v2/test_dtype_device.py` + +**Steps:** +1. Run the CUDA-focused tests created in this plan. +2. Run adjacent creation/device tests to ensure CPU and existing behavior did not regress. +3. Fix only issues directly caused by CUDA MVP changes. + +**Test commands:** +Run: `PYTHONPATH=src pytest -q tests/mindtorch_v2/test_device_transfer.py tests/mindtorch_v2/test_cuda_backend.py tests/mindtorch_v2/test_creation.py tests/mindtorch_v2/test_dtype_device.py` +Expected: PASS, with CUDA tests skipped on systems without CUDA runtime. + +--- + +### Task 10: Document limitations and next phase + +**Files:** +- Modify: `docs/plans/2026-03-08-mindtorch-v2-cuda-backend-mvp-design.md` +- Optionally modify: package docs that describe device support, if such docs exist in-scope + +**Steps:** +1. Record what the CUDA MVP supports. +2. Record what is intentionally out of scope. +3. List the next recommended phase: pointwise ops, reductions, BLAS-backed ops. + +**Validation:** +Review the document and verify that supported vs unsupported CUDA features are explicit. + diff --git a/src/mindtorch_v2/__init__.py b/src/mindtorch_v2/__init__.py index 470b23749..5d1d97914 100644 --- a/src/mindtorch_v2/__init__.py +++ b/src/mindtorch_v2/__init__.py @@ -93,6 +93,7 @@ from ._autograd.grad_mode import is_grad_enabled, set_grad_enabled, no_grad, enable_grad, inference_mode from . import _autograd as autograd from ._backends import autograd as _autograd_kernels +from . import cuda from . import npu from . import mps from . import _C @@ -152,6 +153,7 @@ def decorator(fn): __all__ = [ "Device", "device", + "cuda", "Tensor", "Size", "FloatTensor", "DoubleTensor", "HalfTensor", "BFloat16Tensor", diff --git a/src/mindtorch_v2/_backends/__init__.py b/src/mindtorch_v2/_backends/__init__.py index 60e7b3629..f0953511c 100644 --- a/src/mindtorch_v2/_backends/__init__.py +++ b/src/mindtorch_v2/_backends/__init__.py @@ -1,5 +1,6 @@ from . import cpu from . import meta +from . import cuda from . import npu from . import autograd @@ -10,4 +11,4 @@ except ImportError: pass -__all__ = ["cpu", "meta", "npu", "autograd"] +__all__ = ["cpu", "meta", "cuda", "npu", "autograd"] diff --git a/src/mindtorch_v2/_backends/autograd.py b/src/mindtorch_v2/_backends/autograd.py index 2a697c856..f19e364e9 100644 --- a/src/mindtorch_v2/_backends/autograd.py +++ b/src/mindtorch_v2/_backends/autograd.py @@ -18,6 +18,7 @@ def _strip_autograd_keys(keyset): DispatchKey.AutogradOther, DispatchKey.AutogradCPU, DispatchKey.AutogradNPU, + DispatchKey.AutogradCUDA, DispatchKey.AutogradXPU, DispatchKey.AutogradMeta, DispatchKey.PrivateUse3, @@ -6084,6 +6085,7 @@ def _register_autograd_op(name, factory, *, npu_factory=None, include_meta=True) "default": factory(), "cpu": factory(), "npu": (npu_factory or factory)(), + "cuda": factory(), } if include_meta: kwargs["meta"] = factory() diff --git a/src/mindtorch_v2/_backends/common/convert.py b/src/mindtorch_v2/_backends/common/convert.py index 5c4c5f1af..e56917e54 100644 --- a/src/mindtorch_v2/_backends/common/convert.py +++ b/src/mindtorch_v2/_backends/common/convert.py @@ -2,7 +2,10 @@ from ..._device import device as Device from ..._storage import ( + cuda_typed_storage_from_numpy, + cuda_typed_storage_to_numpy, empty_cpu_typed_storage, + empty_cuda_typed_storage, meta_typed_storage_from_shape, npu_typed_storage_from_ptr, typed_storage_from_numpy, @@ -32,10 +35,16 @@ def to_device(a, dev, dtype=None, non_blocking=False, copy=False, memory_format= ptr, _ = npu_runtime._copy_cpu_to_npu(arr, runtime=runtime) storage = npu_typed_storage_from_ptr(ptr, arr.size, a.dtype, device=dev) return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + if dev.type == "mps": from ..._storage import mps_typed_storage_from_numpy arr = a.storage().data.copy() storage = mps_typed_storage_from_numpy(arr, a.dtype, device=dev) + + if dev.type == "cuda": + arr = cuda_typed_storage_to_numpy(a.storage(), a.shape, a.dtype) + storage = cuda_typed_storage_from_numpy(arr, a.dtype, device=dev) + return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) return a if a.device.type == "meta": @@ -52,11 +61,16 @@ def to_device(a, dev, dtype=None, non_blocking=False, copy=False, memory_format= ptr = npu_runtime._alloc_device(size, runtime=runtime) storage = npu_typed_storage_from_ptr(ptr, int(np.prod(a.shape)), a.dtype, device=dev) return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + if dev.type == "mps": from ..._storage import mps_typed_storage_from_numpy from ..._dtype import to_numpy_dtype arr = np.zeros(a.shape, dtype=to_numpy_dtype(a.dtype)) storage = mps_typed_storage_from_numpy(arr.ravel(), a.dtype, device=dev) + + if dev.type == "cuda": + storage = empty_cuda_typed_storage(a.shape, a.dtype, device=dev) + return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) raise NotImplementedError(f"Unsupported device: {dev}") if a.device.type == "cpu" and dev.type == "meta": @@ -97,6 +111,10 @@ def to_device(a, dev, dtype=None, non_blocking=False, copy=False, memory_format= ptr, _ = npu_runtime._copy_cpu_to_npu(arr, runtime=runtime, non_blocking=do_non_blocking, stream=stream) storage = npu_typed_storage_from_ptr(ptr, arr.size, a.dtype, device=dev) return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + if a.device.type == "cpu" and dev.type == "cuda": + arr = a.storage().data + storage = cuda_typed_storage_from_numpy(arr, a.dtype, device=dev) + return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) if a.device.type == "npu" and dev.type == "npu": from ..npu import runtime as npu_runtime @@ -146,6 +164,7 @@ def to_device(a, dev, dtype=None, non_blocking=False, copy=False, memory_format= if a.device.type == "npu" and dev.type == "meta": storage = meta_typed_storage_from_shape(a.shape, a.dtype, device=dev) return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + # MPS <-> NPU: route through CPU if a.device.type == "mps" and dev.type == "npu": cpu_dev = Device("cpu") @@ -155,4 +174,17 @@ def to_device(a, dev, dtype=None, non_blocking=False, copy=False, memory_format= cpu_dev = Device("cpu") cpu_tensor = to_device(a, cpu_dev) return to_device(cpu_tensor, dev) + + if a.device.type == "cuda" and dev.type == "cuda": + arr = cuda_typed_storage_to_numpy(a.storage(), a.shape, a.dtype) + storage = cuda_typed_storage_from_numpy(arr, a.dtype, device=dev) + return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + if a.device.type == "cuda" and dev.type == "cpu": + arr = cuda_typed_storage_to_numpy(a.storage(), a.shape, a.dtype) + storage = typed_storage_from_numpy(arr, a.dtype, device=dev) + return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + if a.device.type == "cuda" and dev.type == "meta": + storage = meta_typed_storage_from_shape(a.shape, a.dtype, device=dev) + return Tensor(storage, a.shape, a.stride, a.offset, a.requires_grad) + raise NotImplementedError(f"Unsupported device: {a.device} -> {dev}") diff --git a/src/mindtorch_v2/_backends/cuda/__init__.py b/src/mindtorch_v2/_backends/cuda/__init__.py new file mode 100644 index 000000000..420dc46a2 --- /dev/null +++ b/src/mindtorch_v2/_backends/cuda/__init__.py @@ -0,0 +1,18 @@ +from ..common import convert as convert_backend +from ..._dispatch.registry import registry +from . import runtime +from . import storage +from .creation import empty_create, full_create, ones_create, tensor_create, zeros_create +from .ops import add + + +registry.register("add", "cuda", add) +registry.register("to", "cuda", convert_backend.to_device) +registry.register("tensor", "cuda", tensor_create) +registry.register("zeros", "cuda", zeros_create) +registry.register("ones", "cuda", ones_create) +registry.register("empty", "cuda", empty_create) +registry.register("full", "cuda", full_create) + + +__all__ = ["runtime", "storage"] diff --git a/src/mindtorch_v2/_backends/cuda/creation.py b/src/mindtorch_v2/_backends/cuda/creation.py new file mode 100644 index 000000000..8ff146d3e --- /dev/null +++ b/src/mindtorch_v2/_backends/cuda/creation.py @@ -0,0 +1,54 @@ +import numpy as np + +from ..._dtype import to_numpy_dtype +from ..._storage import cuda_typed_storage_from_numpy, empty_cuda_typed_storage +from ..._tensor import Tensor + + +def _contiguous_stride(shape): + stride = [] + acc = 1 + for dim in reversed(shape): + stride.append(acc) + acc *= dim + return tuple(reversed(stride)) + + +def tensor_create(data, dtype=None, device=None, requires_grad=False, memory_format=None): + arr = np.array(data, dtype=to_numpy_dtype(dtype)) + storage = cuda_typed_storage_from_numpy(arr, dtype, device=device) + stride = tuple(np.array(arr.strides) // arr.itemsize) + return Tensor(storage, arr.shape, stride, requires_grad=requires_grad) + + +def zeros_create(shape, dtype=None, device=None, requires_grad=False, memory_format=None): + if isinstance(shape, int): + shape = (shape,) + shape = tuple(shape) + arr = np.zeros(shape, dtype=to_numpy_dtype(dtype)) + storage = cuda_typed_storage_from_numpy(arr, dtype, device=device) + return Tensor(storage, shape, _contiguous_stride(shape), requires_grad=requires_grad) + + +def ones_create(shape, dtype=None, device=None, requires_grad=False, memory_format=None): + if isinstance(shape, int): + shape = (shape,) + shape = tuple(shape) + arr = np.ones(shape, dtype=to_numpy_dtype(dtype)) + storage = cuda_typed_storage_from_numpy(arr, dtype, device=device) + return Tensor(storage, shape, _contiguous_stride(shape), requires_grad=requires_grad) + + +def empty_create(shape, dtype=None, device=None, requires_grad=False, memory_format=None): + if isinstance(shape, int): + shape = (shape,) + shape = tuple(shape) + storage = empty_cuda_typed_storage(shape, dtype, device=device) + return Tensor(storage, shape, _contiguous_stride(shape), requires_grad=requires_grad) + + +def full_create(shape, fill_value, dtype=None, device=None): + shape = tuple(shape) + arr = np.full(shape, fill_value, dtype=to_numpy_dtype(dtype)) + storage = cuda_typed_storage_from_numpy(arr, dtype, device=device) + return Tensor(storage, shape, _contiguous_stride(shape)) diff --git a/src/mindtorch_v2/_backends/cuda/ops.py b/src/mindtorch_v2/_backends/cuda/ops.py new file mode 100644 index 000000000..d022550be --- /dev/null +++ b/src/mindtorch_v2/_backends/cuda/ops.py @@ -0,0 +1,17 @@ +import numpy as np + +from ..._storage import cuda_typed_storage_from_numpy, cuda_typed_storage_to_numpy +from ..._tensor import Tensor + + +def _from_numpy(arr, dtype, device): + storage = cuda_typed_storage_from_numpy(arr, dtype, device=device) + stride = tuple(np.array(arr.strides) // arr.itemsize) + return Tensor(storage, arr.shape, stride) + + +def add(a, b): + a_np = cuda_typed_storage_to_numpy(a.storage(), a.shape, a.dtype) + b_np = cuda_typed_storage_to_numpy(b.storage(), b.shape, b.dtype) if isinstance(b, Tensor) else b + out = np.ascontiguousarray(a_np + b_np) + return _from_numpy(out, a.dtype, a.device) diff --git a/src/mindtorch_v2/_backends/cuda/runtime.py b/src/mindtorch_v2/_backends/cuda/runtime.py new file mode 100644 index 000000000..a0601114e --- /dev/null +++ b/src/mindtorch_v2/_backends/cuda/runtime.py @@ -0,0 +1,298 @@ +import ctypes +import ctypes.util +import os +from typing import Optional + + +CUDA_SUCCESS = 0 +CUDA_MEMCPY_HOST_TO_HOST = 0 +CUDA_MEMCPY_HOST_TO_DEVICE = 1 +CUDA_MEMCPY_DEVICE_TO_HOST = 2 +CUDA_MEMCPY_DEVICE_TO_DEVICE = 3 + +_CUDART = None +_CUDART_LOAD_ERROR = None + + +class CudaRuntimeError(RuntimeError): + pass + + +_CANDIDATE_LIBRARIES = ( + "libcudart.so", + "libcudart.so.12", + "libcudart.so.11.0", + "libcudart.so.11", +) + + +def _library_candidates(): + candidates = [] + found = ctypes.util.find_library("cudart") + if found: + candidates.append(found) + + cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if cuda_home: + for name in _CANDIDATE_LIBRARIES: + candidates.append(os.path.join(cuda_home, "lib64", name)) + candidates.append(os.path.join(cuda_home, "targets", "x86_64-linux", "lib", name)) + + candidates.extend(_CANDIDATE_LIBRARIES) + + deduped = [] + seen = set() + for item in candidates: + if not item or item in seen: + continue + deduped.append(item) + seen.add(item) + return deduped + + +def _bind(lib): + lib.cudaGetDeviceCount.argtypes = [ctypes.POINTER(ctypes.c_int)] + lib.cudaGetDeviceCount.restype = ctypes.c_int + + lib.cudaGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)] + lib.cudaGetDevice.restype = ctypes.c_int + + lib.cudaSetDevice.argtypes = [ctypes.c_int] + lib.cudaSetDevice.restype = ctypes.c_int + + lib.cudaDeviceSynchronize.argtypes = [] + lib.cudaDeviceSynchronize.restype = ctypes.c_int + + lib.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] + lib.cudaMalloc.restype = ctypes.c_int + + lib.cudaFree.argtypes = [ctypes.c_void_p] + lib.cudaFree.restype = ctypes.c_int + + lib.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + lib.cudaMemcpy.restype = ctypes.c_int + + lib.cudaMemcpyAsync.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_void_p] + lib.cudaMemcpyAsync.restype = ctypes.c_int + + lib.cudaMemset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] + lib.cudaMemset.restype = ctypes.c_int + + lib.cudaStreamCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + lib.cudaStreamCreate.restype = ctypes.c_int + + lib.cudaStreamDestroy.argtypes = [ctypes.c_void_p] + lib.cudaStreamDestroy.restype = ctypes.c_int + + lib.cudaStreamSynchronize.argtypes = [ctypes.c_void_p] + lib.cudaStreamSynchronize.restype = ctypes.c_int + + lib.cudaEventCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + lib.cudaEventCreate.restype = ctypes.c_int + + lib.cudaEventDestroy.argtypes = [ctypes.c_void_p] + lib.cudaEventDestroy.restype = ctypes.c_int + + lib.cudaEventRecord.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + lib.cudaEventRecord.restype = ctypes.c_int + + lib.cudaEventSynchronize.argtypes = [ctypes.c_void_p] + lib.cudaEventSynchronize.restype = ctypes.c_int + + if hasattr(lib, "cudaGetErrorString"): + lib.cudaGetErrorString.argtypes = [ctypes.c_int] + lib.cudaGetErrorString.restype = ctypes.c_char_p + + +def _load_cudart(): + global _CUDART, _CUDART_LOAD_ERROR + + if _CUDART is not None: + return _CUDART + if _CUDART_LOAD_ERROR is not None: + raise _CUDART_LOAD_ERROR + + last_error = None + for candidate in _library_candidates(): + try: + lib = ctypes.CDLL(candidate) + _bind(lib) + _CUDART = lib + return lib + except OSError as exc: + last_error = exc + + if last_error is None: + last_error = OSError("CUDA runtime library not found") + _CUDART_LOAD_ERROR = last_error + raise last_error + + +def _try_load_cudart(): + try: + return _load_cudart() + except OSError: + return None + + +def _error_string(code: int) -> str: + lib = _try_load_cudart() + if lib is not None and hasattr(lib, "cudaGetErrorString"): + raw = lib.cudaGetErrorString(int(code)) + if raw: + return raw.decode("utf-8", errors="replace") + return f"CUDA error {code}" + + +def _check(code: int, opname: str): + if int(code) != CUDA_SUCCESS: + raise CudaRuntimeError(f"{opname} failed: {_error_string(code)}") + + +def is_available() -> bool: + try: + return device_count() > 0 + except Exception: + return False + + +def device_count() -> int: + lib = _try_load_cudart() + if lib is None: + return 0 + count = ctypes.c_int() + code = lib.cudaGetDeviceCount(ctypes.byref(count)) + if int(code) != CUDA_SUCCESS: + return 0 + return int(count.value) + + +def current_device() -> int: + lib = _load_cudart() + current = ctypes.c_int() + _check(lib.cudaGetDevice(ctypes.byref(current)), "cudaGetDevice") + return int(current.value) + + +def set_device(device_index: int): + lib = _load_cudart() + _check(lib.cudaSetDevice(int(device_index)), "cudaSetDevice") + + +def synchronize(device: Optional[int] = None): + lib = _load_cudart() + if device is not None: + set_device(device) + _check(lib.cudaDeviceSynchronize(), "cudaDeviceSynchronize") + + +def malloc(nbytes: int) -> int: + lib = _load_cudart() + ptr = ctypes.c_void_p() + _check(lib.cudaMalloc(ctypes.byref(ptr), ctypes.c_size_t(int(nbytes))), "cudaMalloc") + return int(ptr.value or 0) + + +def free(ptr: int): + lib = _load_cudart() + _check(lib.cudaFree(ctypes.c_void_p(int(ptr))), "cudaFree") + + +def memcpy(dst: int, src: int, nbytes: int, kind: int): + lib = _load_cudart() + _check( + lib.cudaMemcpy( + ctypes.c_void_p(int(dst)), + ctypes.c_void_p(int(src)), + ctypes.c_size_t(int(nbytes)), + int(kind), + ), + "cudaMemcpy", + ) + + +def memcpy_async(dst: int, src: int, nbytes: int, kind: int, stream=None): + lib = _load_cudart() + stream_ptr = ctypes.c_void_p(0 if stream is None else int(stream)) + _check( + lib.cudaMemcpyAsync( + ctypes.c_void_p(int(dst)), + ctypes.c_void_p(int(src)), + ctypes.c_size_t(int(nbytes)), + int(kind), + stream_ptr, + ), + "cudaMemcpyAsync", + ) + + +def memset(dst: int, value: int, nbytes: int): + lib = _load_cudart() + _check(lib.cudaMemset(ctypes.c_void_p(int(dst)), int(value), ctypes.c_size_t(int(nbytes))), "cudaMemset") + + +def create_stream() -> int: + lib = _load_cudart() + stream = ctypes.c_void_p() + _check(lib.cudaStreamCreate(ctypes.byref(stream)), "cudaStreamCreate") + return int(stream.value or 0) + + +def destroy_stream(stream: int): + lib = _load_cudart() + _check(lib.cudaStreamDestroy(ctypes.c_void_p(int(stream))), "cudaStreamDestroy") + + +def synchronize_stream(stream: int): + lib = _load_cudart() + _check(lib.cudaStreamSynchronize(ctypes.c_void_p(int(stream))), "cudaStreamSynchronize") + + +def create_event() -> int: + lib = _load_cudart() + event = ctypes.c_void_p() + _check(lib.cudaEventCreate(ctypes.byref(event)), "cudaEventCreate") + return int(event.value or 0) + + +def destroy_event(event: int): + lib = _load_cudart() + _check(lib.cudaEventDestroy(ctypes.c_void_p(int(event))), "cudaEventDestroy") + + +def record_event(event: int, stream=None): + lib = _load_cudart() + stream_ptr = ctypes.c_void_p(0 if stream is None else int(stream)) + _check(lib.cudaEventRecord(ctypes.c_void_p(int(event)), stream_ptr), "cudaEventRecord") + + +def synchronize_event(event: int): + lib = _load_cudart() + _check(lib.cudaEventSynchronize(ctypes.c_void_p(int(event))), "cudaEventSynchronize") + + +__all__ = [ + "CUDA_SUCCESS", + "CUDA_MEMCPY_HOST_TO_HOST", + "CUDA_MEMCPY_HOST_TO_DEVICE", + "CUDA_MEMCPY_DEVICE_TO_HOST", + "CUDA_MEMCPY_DEVICE_TO_DEVICE", + "CudaRuntimeError", + "is_available", + "device_count", + "current_device", + "set_device", + "synchronize", + "malloc", + "free", + "memcpy", + "memcpy_async", + "memset", + "create_stream", + "destroy_stream", + "synchronize_stream", + "create_event", + "destroy_event", + "record_event", + "synchronize_event", +] diff --git a/src/mindtorch_v2/_backends/cuda/storage.py b/src/mindtorch_v2/_backends/cuda/storage.py new file mode 100644 index 000000000..18b1e6195 --- /dev/null +++ b/src/mindtorch_v2/_backends/cuda/storage.py @@ -0,0 +1,159 @@ +import weakref + +import numpy as np + +from ..._device import device as Device +from ..._dtype import to_numpy_dtype +from . import runtime as cuda_runtime + + +class CudaUntypedStorage: + def __init__(self, nbytes, device=None, device_ptr=0, host_shadow=None): + if isinstance(device, str): + device = Device(device) + self.device = device or Device("cuda") + self._nbytes = int(nbytes) + self._device_ptr = int(device_ptr or 0) + if host_shadow is None: + host_shadow = np.empty((self._nbytes,), dtype=np.uint8) + self._host_shadow = np.ascontiguousarray(host_shadow, dtype=np.uint8).reshape(-1) + self._finalizer = None + if self._device_ptr: + self._finalizer = weakref.finalize(self, cuda_runtime.free, self._device_ptr) + + def nbytes(self): + return self._nbytes + + def data_ptr(self): + return self._device_ptr + + def buffer(self): + raise RuntimeError("Cannot get buffer of CUDA storage on CPU") + + def resize_(self, new_nbytes): + new_nbytes = int(new_nbytes) + if new_nbytes == self._nbytes: + return self + + new_shadow = np.empty((new_nbytes,), dtype=np.uint8) + copy_bytes = min(self._nbytes, new_nbytes) + if copy_bytes: + new_shadow[:copy_bytes] = self._host_shadow[:copy_bytes] + + new_ptr = 0 + if new_nbytes and cuda_runtime._try_load_cudart() is not None and cuda_runtime.device_count() > 0: + new_ptr = cuda_runtime.malloc(new_nbytes) + if self._device_ptr and copy_bytes: + cuda_runtime.memcpy( + new_ptr, + self._device_ptr, + copy_bytes, + cuda_runtime.CUDA_MEMCPY_DEVICE_TO_DEVICE, + ) + + old_ptr = self._device_ptr + if self._finalizer is not None: + self._finalizer.detach() + self._finalizer = None + if old_ptr: + cuda_runtime.free(old_ptr) + + self._device_ptr = int(new_ptr) + self._nbytes = new_nbytes + self._host_shadow = new_shadow + if self._device_ptr: + self._finalizer = weakref.finalize(self, cuda_runtime.free, self._device_ptr) + return self + + def is_shared(self): + return False + + def is_pinned(self): + return False + + +def _allocate_device_bytes(nbytes): + if int(nbytes) <= 0: + return 0 + if cuda_runtime._try_load_cudart() is None: + return 0 + if cuda_runtime.device_count() <= 0: + return 0 + try: + return cuda_runtime.malloc(int(nbytes)) + except Exception: + return 0 + + +def untyped_from_numpy(arr, device=None): + raw = np.ascontiguousarray(arr).view(np.uint8).reshape(-1).copy() + ptr = _allocate_device_bytes(raw.size) + if ptr and raw.size: + cuda_runtime.memcpy( + ptr, + int(raw.ctypes.data), + raw.size, + cuda_runtime.CUDA_MEMCPY_HOST_TO_DEVICE, + ) + return CudaUntypedStorage(raw.size, device=device, device_ptr=ptr, host_shadow=raw) + + +def empty_untyped(shape, dtype, device=None): + shape = tuple(shape) + itemsize = np.dtype(to_numpy_dtype(dtype)).itemsize + nbytes = int(np.prod(shape)) * itemsize + raw = np.empty((nbytes,), dtype=np.uint8) + ptr = _allocate_device_bytes(nbytes) + return CudaUntypedStorage(nbytes, device=device, device_ptr=ptr, host_shadow=raw) + + +def to_numpy(untyped, dtype, shape=None): + np_dtype = to_numpy_dtype(dtype) + itemsize = np.dtype(np_dtype).itemsize + if shape is None: + count = 0 if itemsize == 0 else int(untyped.nbytes() // itemsize) + shape = (count,) + else: + shape = tuple(shape) + count = int(np.prod(shape)) + if count == 0: + return np.empty(shape, dtype=np_dtype) + + if untyped.data_ptr(): + arr = np.empty(shape, dtype=np_dtype) + cuda_runtime.memcpy( + int(arr.ctypes.data), + untyped.data_ptr(), + arr.nbytes, + cuda_runtime.CUDA_MEMCPY_DEVICE_TO_HOST, + ) + untyped._host_shadow = arr.view(np.uint8).reshape(-1).copy() + return arr + + raw = np.ascontiguousarray(untyped._host_shadow, dtype=np.uint8) + return raw.view(np_dtype).copy().reshape(shape) + + +def copy_untyped(dst, src): + copy_bytes = min(dst.nbytes(), src.nbytes()) + if copy_bytes <= 0: + return + if dst.data_ptr() and src.data_ptr(): + cuda_runtime.memcpy( + dst.data_ptr(), + src.data_ptr(), + copy_bytes, + cuda_runtime.CUDA_MEMCPY_DEVICE_TO_DEVICE, + ) + if dst._host_shadow.size != dst.nbytes(): + dst._host_shadow = np.empty((dst.nbytes(),), dtype=np.uint8) + dst._host_shadow[:copy_bytes] = src._host_shadow[:copy_bytes] + + +__all__ = [ + "CudaUntypedStorage", + "untyped_from_numpy", + "empty_untyped", + "to_numpy", + "copy_untyped", +] diff --git a/src/mindtorch_v2/_dispatch/keys.py b/src/mindtorch_v2/_dispatch/keys.py index 2c1660ba2..5edd69653 100644 --- a/src/mindtorch_v2/_dispatch/keys.py +++ b/src/mindtorch_v2/_dispatch/keys.py @@ -9,19 +9,21 @@ class DispatchKey(IntEnum): AutogradOther = 1 << 5 AutogradCPU = 1 << 6 AutogradNPU = 1 << 7 - AutogradXPU = 1 << 8 - AutogradMeta = 1 << 9 - Autograd = 1 << 10 - Meta = 1 << 11 - NPU = 1 << 12 - CPU = 1 << 13 - PythonDispatcher = 1 << 14 - CompositeImplicitAutograd = 1 << 15 - CompositeExplicitAutograd = 1 << 16 - Autocast = 1 << 17 - PrivateUse1 = 1 << 18 - PrivateUse2 = 1 << 19 - PrivateUse3 = 1 << 20 + AutogradCUDA = 1 << 8 + AutogradXPU = 1 << 9 + AutogradMeta = 1 << 10 + Autograd = 1 << 11 + Meta = 1 << 12 + NPU = 1 << 13 + CUDA = 1 << 14 + CPU = 1 << 15 + PythonDispatcher = 1 << 16 + CompositeImplicitAutograd = 1 << 17 + CompositeExplicitAutograd = 1 << 18 + Autocast = 1 << 19 + PrivateUse1 = 1 << 20 + PrivateUse2 = 1 << 21 + PrivateUse3 = 1 << 22 DISPATCH_KEY_PRIORITY = [ @@ -33,11 +35,13 @@ class DispatchKey(IntEnum): DispatchKey.AutogradOther, DispatchKey.AutogradCPU, DispatchKey.AutogradNPU, + DispatchKey.AutogradCUDA, DispatchKey.AutogradXPU, DispatchKey.AutogradMeta, DispatchKey.Autograd, DispatchKey.Meta, DispatchKey.NPU, + DispatchKey.CUDA, DispatchKey.CPU, DispatchKey.PythonDispatcher, DispatchKey.CompositeImplicitAutograd, @@ -216,7 +220,7 @@ def from_tensors(cls, tensors, *, grad_enabled=False, pipeline_enabled=False, fu elif has_npu: mask |= int(DispatchKey.NPU) elif has_cuda: - mask |= int(DispatchKey.PrivateUse1) + mask |= int(DispatchKey.CUDA) elif has_mps: mask |= int(DispatchKey.PrivateUse2) else: @@ -229,7 +233,7 @@ def from_tensors(cls, tensors, *, grad_enabled=False, pipeline_enabled=False, fu elif has_npu: mask |= int(DispatchKey.AutogradNPU) elif has_cuda: - mask |= int(DispatchKey.AutogradXPU) + mask |= int(DispatchKey.AutogradCUDA) elif has_mps: mask |= int(DispatchKey.PrivateUse3) else: diff --git a/src/mindtorch_v2/_dispatch/registration.py b/src/mindtorch_v2/_dispatch/registration.py index f9521a93d..c711afb7c 100644 --- a/src/mindtorch_v2/_dispatch/registration.py +++ b/src/mindtorch_v2/_dispatch/registration.py @@ -6,14 +6,14 @@ "cpu": DispatchKey.CPU, "npu": DispatchKey.NPU, "meta": DispatchKey.Meta, - "cuda": DispatchKey.PrivateUse1, + "cuda": DispatchKey.CUDA, } _AUTOGRAD_KEY_BY_DEVICE = { "cpu": DispatchKey.AutogradCPU, "npu": DispatchKey.AutogradNPU, "meta": DispatchKey.AutogradMeta, - "cuda": DispatchKey.AutogradXPU, + "cuda": DispatchKey.AutogradCUDA, } @@ -27,7 +27,7 @@ def register_forward_kernels(name, **kernels): """Register forward kernels by device label in one call. Supported labels: cpu, npu, meta, cuda. - cuda is currently reserved and maps to PrivateUse1. + cuda maps to the explicit CUDA dispatch key. """ for device, fn in kernels.items(): if fn is None: @@ -43,7 +43,7 @@ def register_autograd_kernels(name, *, default=None, **kernels): """Register autograd kernels by device label in one call. Supported labels: cpu, npu, meta, cuda. - cuda is currently reserved and maps to AutogradXPU. + cuda maps to the explicit AutogradCUDA dispatch key. """ if default is not None: registry.register_kernel(name, DispatchKey.Autograd, default) diff --git a/src/mindtorch_v2/_dispatch/registry.py b/src/mindtorch_v2/_dispatch/registry.py index 965ced083..9eccd48db 100644 --- a/src/mindtorch_v2/_dispatch/registry.py +++ b/src/mindtorch_v2/_dispatch/registry.py @@ -114,6 +114,7 @@ def _register_global_fallthroughs(): DispatchKey.AutogradOther, DispatchKey.AutogradCPU, DispatchKey.AutogradNPU, + DispatchKey.AutogradCUDA, DispatchKey.AutogradXPU, DispatchKey.AutogradMeta, DispatchKey.Python, @@ -146,7 +147,7 @@ def resolve_dispatch_key(device): if label == "meta": return DispatchKey.Meta if label == "cuda": - return DispatchKey.PrivateUse1 + return DispatchKey.CUDA if label == "mps": return DispatchKey.PrivateUse2 @@ -156,7 +157,7 @@ def resolve_dispatch_key(device): _DISPATCH_KEY_STRING_MAP = { "CPU": DispatchKey.CPU, "NPU": DispatchKey.NPU, - "CUDA": DispatchKey.PrivateUse1, + "CUDA": DispatchKey.CUDA, "PrivateUse1": DispatchKey.PrivateUse1, "MPS": DispatchKey.PrivateUse2, "PrivateUse2": DispatchKey.PrivateUse2, @@ -164,7 +165,7 @@ def resolve_dispatch_key(device): "Autograd": DispatchKey.Autograd, "AutogradCPU": DispatchKey.AutogradCPU, "AutogradNPU": DispatchKey.AutogradNPU, - "AutogradCUDA": DispatchKey.AutogradXPU, + "AutogradCUDA": DispatchKey.AutogradCUDA, "AutogradPrivateUse1": DispatchKey.AutogradXPU, "AutogradMPS": DispatchKey.PrivateUse3, "AutogradPrivateUse2": DispatchKey.PrivateUse3, diff --git a/src/mindtorch_v2/_storage.py b/src/mindtorch_v2/_storage.py index 8a9824d78..4e26637bc 100644 --- a/src/mindtorch_v2/_storage.py +++ b/src/mindtorch_v2/_storage.py @@ -501,6 +501,9 @@ def clone(self): ) if self.device.type == "meta": return meta_typed_storage_from_size(self._size, self.dtype, device=self.device) + if self.device.type == "cuda": + arr = cuda_typed_storage_to_numpy(self, (self._size,), self.dtype) + return cuda_typed_storage_from_numpy(arr, self.dtype, device=self.device) raise NotImplementedError(f"Unsupported device: {self.device}") def copy_(self, other): @@ -531,6 +534,11 @@ def copy_(self, other): if ret != 0: raise RuntimeError(f"acl.rt.memcpy D2D failed: {ret}") return self + if self.device.type == "cuda": + from ._backends.cuda import storage as cuda_storage + + cuda_storage.copy_untyped(self.untyped_storage(), other.untyped_storage()) + return self raise NotImplementedError(f"Unsupported device: {self.device}") def resize_(self, new_size): @@ -612,6 +620,31 @@ def mps_typed_storage_from_ptr(metal_buffer, size, dtype, device=None): return TypedStorage(untyped, dtype, int(size), data=typed_data) +def cuda_typed_storage_from_numpy(arr, dtype, device=None): + arr = np.ascontiguousarray(arr, dtype=to_numpy_dtype(dtype)) + from ._backends.cuda import storage as cuda_storage + + untyped = cuda_storage.untyped_from_numpy(arr, device=device) + return TypedStorage(untyped, dtype, arr.size) + + +def empty_cuda_typed_storage(shape, dtype, device=None): + if isinstance(shape, int): + shape = (shape,) + shape = tuple(shape) + size = int(np.prod(shape)) + from ._backends.cuda import storage as cuda_storage + + untyped = cuda_storage.empty_untyped(shape, dtype, device=device) + return TypedStorage(untyped, dtype, size) + + +def cuda_typed_storage_to_numpy(storage, shape, dtype): + from ._backends.cuda import storage as cuda_storage + + return cuda_storage.to_numpy(storage.untyped_storage(), dtype, shape=shape) + + class PendingStorage: def __init__(self, shape, dtype, device): if isinstance(device, str): diff --git a/src/mindtorch_v2/cuda.py b/src/mindtorch_v2/cuda.py new file mode 100644 index 000000000..0e2514e25 --- /dev/null +++ b/src/mindtorch_v2/cuda.py @@ -0,0 +1,120 @@ +from ._backends.cuda import runtime as cuda_runtime +from ._device import device as Device + + +def _normalize_cuda_device(dev=None): + if dev is None: + if is_available(): + return Device("cuda", index=current_device()) + return Device("cuda", index=0) + if isinstance(dev, Device): + out = dev + elif isinstance(dev, int): + out = Device("cuda", index=dev) + else: + out = Device(dev) + if out.type != "cuda": + raise ValueError(f"Expected a cuda device, but got: {out}") + if out.index is None: + return Device("cuda", index=current_device() if is_available() else 0) + return out + + +def is_available(): + return cuda_runtime.is_available() + + +def device_count(): + return cuda_runtime.device_count() + + +def current_device(): + return cuda_runtime.current_device() + + +def set_device(dev): + cuda_dev = _normalize_cuda_device(dev) + cuda_runtime.set_device(cuda_dev.index or 0) + + +def synchronize(device=None): + if device is None: + cuda_runtime.synchronize() + return + cuda_dev = _normalize_cuda_device(device) + cuda_runtime.synchronize(cuda_dev.index or 0) + + +class Stream: + def __init__(self, device=None, priority=0, stream=None): + self.device = _normalize_cuda_device(device) + self.priority = int(priority) + self._owns_stream = stream is None + self.stream = int(stream) if stream is not None else cuda_runtime.create_stream() + + def synchronize(self): + cuda_runtime.synchronize_stream(self.stream) + + def record_event(self, event=None): + if event is None: + event = Event() + event.record(self) + return event + + def __del__(self): + if getattr(self, "_owns_stream", False) and getattr(self, "stream", 0): + try: + cuda_runtime.destroy_stream(self.stream) + except Exception: + pass + + +class Event: + def __init__(self, enable_timing=False, blocking=False, interprocess=False): + self.enable_timing = bool(enable_timing) + self.blocking = bool(blocking) + self.interprocess = bool(interprocess) + self.event = cuda_runtime.create_event() + + def record(self, stream=None): + stream_handle = None if stream is None else stream.stream + cuda_runtime.record_event(self.event, stream_handle) + return self + + def synchronize(self): + cuda_runtime.synchronize_event(self.event) + + def __del__(self): + if getattr(self, "event", 0): + try: + cuda_runtime.destroy_event(self.event) + except Exception: + pass + + +class device: + def __init__(self, dev): + self.dev = _normalize_cuda_device(dev) + self._prev = None + + def __enter__(self): + self._prev = current_device() + set_device(self.dev) + return self.dev + + def __exit__(self, exc_type, exc, tb): + if self._prev is not None: + set_device(self._prev) + return False + + +__all__ = [ + "is_available", + "device_count", + "current_device", + "set_device", + "synchronize", + "Stream", + "Event", + "device", +] diff --git a/src/mindtorch_v2/library.py b/src/mindtorch_v2/library.py index 5f02c58ba..69c7cdbf6 100644 --- a/src/mindtorch_v2/library.py +++ b/src/mindtorch_v2/library.py @@ -266,6 +266,7 @@ def _infer_schema(fn, mutates_args=()): DispatchKey.Autograd, DispatchKey.AutogradCPU, DispatchKey.AutogradNPU, + DispatchKey.AutogradCUDA, DispatchKey.AutogradMeta, DispatchKey.AutogradOther, DispatchKey.AutogradXPU, diff --git a/tests/mindtorch_v2/test_cuda_backend.py b/tests/mindtorch_v2/test_cuda_backend.py new file mode 100644 index 000000000..748a85ebb --- /dev/null +++ b/tests/mindtorch_v2/test_cuda_backend.py @@ -0,0 +1,128 @@ +import pytest + + +def test_cuda_runtime_availability_probe_returns_bool(): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + assert isinstance(cuda_runtime.is_available(), bool) + + +def test_cuda_runtime_device_count_is_non_negative(): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + count = cuda_runtime.device_count() + assert isinstance(count, int) + assert count >= 0 + + +def test_cuda_runtime_current_device_roundtrip_when_available(): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + if not cuda_runtime.is_available(): + pytest.skip("CUDA runtime not available") + + current = cuda_runtime.current_device() + cuda_runtime.set_device(current) + assert cuda_runtime.current_device() == current + + +def test_cuda_runtime_synchronize_succeeds_when_available(): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + if not cuda_runtime.is_available(): + pytest.skip("CUDA runtime not available") + + cuda_runtime.synchronize() + + +def test_public_cuda_module_exposes_basic_api(): + import mindtorch_v2 as torch + + assert hasattr(torch, "cuda") + assert isinstance(torch.cuda.is_available(), bool) + assert isinstance(torch.cuda.device_count(), int) + + +def test_public_cuda_module_set_device_roundtrip(monkeypatch): + import mindtorch_v2 as torch + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + state = {"current": 0} + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 2) + monkeypatch.setattr(cuda_runtime, "current_device", lambda: state["current"]) + monkeypatch.setattr(cuda_runtime, "set_device", lambda idx: state.__setitem__("current", int(idx))) + + torch.cuda.set_device(1) + assert torch.cuda.current_device() == 1 + + +def test_tensor_creation_on_cuda_with_fake_runtime(monkeypatch): + import mindtorch_v2 as torch + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) + + x = torch.tensor([1.0, 2.0], device="cuda") + assert x.device.type == "cuda" + assert x.tolist() == [1.0, 2.0] + + +def test_zeros_creation_on_cuda_with_fake_runtime(monkeypatch): + import mindtorch_v2 as torch + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) + + x = torch.zeros(2, device="cuda") + assert x.device.type == "cuda" + assert x.tolist() == [0.0, 0.0] + + +def test_public_cuda_stream_api(monkeypatch): + import mindtorch_v2 as torch + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "create_stream", lambda: 123) + monkeypatch.setattr(cuda_runtime, "destroy_stream", lambda stream: None) + seen = {} + monkeypatch.setattr(cuda_runtime, "synchronize_stream", lambda stream: seen.setdefault("stream", stream)) + + stream = torch.cuda.Stream() + stream.synchronize() + assert seen["stream"] == 123 + + +def test_public_cuda_event_api(monkeypatch): + import mindtorch_v2 as torch + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "create_event", lambda: 456) + monkeypatch.setattr(cuda_runtime, "destroy_event", lambda event: None) + seen = {} + monkeypatch.setattr(cuda_runtime, "record_event", lambda event, stream=None: seen.setdefault("record", (event, stream))) + monkeypatch.setattr(cuda_runtime, "synchronize_event", lambda event: seen.setdefault("sync", event)) + + event = torch.cuda.Event() + event.record() + event.synchronize() + assert seen["record"] == (456, None) + assert seen["sync"] == 456 + + +def test_ones_and_full_creation_on_cuda_with_fake_runtime(monkeypatch): + import mindtorch_v2 as torch + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) + + ones = torch.ones(2, device="cuda") + full = torch.full((2,), 3.0, device="cuda") + + assert ones.device.type == "cuda" + assert full.device.type == "cuda" + assert ones.tolist() == [1.0, 1.0] + assert full.tolist() == [3.0, 3.0] diff --git a/tests/mindtorch_v2/test_device_transfer.py b/tests/mindtorch_v2/test_device_transfer.py index 06a6df1c8..88fe94046 100644 --- a/tests/mindtorch_v2/test_device_transfer.py +++ b/tests/mindtorch_v2/test_device_transfer.py @@ -2,19 +2,56 @@ import mindtorch_v2 as torch +def _require_cuda_runtime(): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + if not cuda_runtime.is_available(): + pytest.skip("CUDA not available") + return cuda_runtime + + def test_to_device_roundtrip(): x = torch.tensor([1.0, 2.0]) y = x.to("cpu") assert y.device.type == "cpu" -def test_tensor_cuda_targets_reserved_cuda_device(): + +def test_tensor_cuda_moves_tensor_to_cuda_device(): + _require_cuda_runtime() + x = torch.tensor([1.0, 2.0]) + y = x.cuda() + assert y.device.type == "cuda" + assert y.tolist() == [1.0, 2.0] + + +def test_tensor_to_cuda_and_back_roundtrip(): + _require_cuda_runtime() + x = torch.tensor([1.0, 2.0]) + y = x.to("cuda") + z = y.to("cpu") + assert y.device.type == "cuda" + assert z.device.type == "cpu" + assert z.tolist() == [1.0, 2.0] + + +def test_tensor_cuda_int_device_zero_preserves_cuda_index(): + _require_cuda_runtime() x = torch.tensor([1.0, 2.0]) - with pytest.raises(NotImplementedError, match="Unsupported device: .* -> cuda"): - x.cuda() + y = x.cuda(0) + assert y.device.type == "cuda" + assert y.device.index == 0 + + +def test_tensor_to_cuda_and_back_roundtrip_with_fake_runtime(monkeypatch): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) -def test_tensor_cuda_int_device_preserves_cuda_index(): x = torch.tensor([1.0, 2.0]) - with pytest.raises(NotImplementedError, match="Unsupported device: .* -> cuda:1"): - x.cuda(1) + y = x.to("cuda") + z = y.to("cpu") + assert y.device.type == "cuda" + assert z.device.type == "cpu" + assert z.tolist() == [1.0, 2.0] diff --git a/tests/mindtorch_v2/test_dispatch_keys.py b/tests/mindtorch_v2/test_dispatch_keys.py index 0e52b5f72..33025adbd 100644 --- a/tests/mindtorch_v2/test_dispatch_keys.py +++ b/tests/mindtorch_v2/test_dispatch_keys.py @@ -48,3 +48,26 @@ def test_dispatch_keyset_without_removes_keys(): trimmed = keyset.without(DispatchKey.Autograd) assert DispatchKey.CPU in trimmed assert DispatchKey.Autograd not in trimmed + + +def test_dispatch_keyset_cuda_uses_explicit_cuda_key(monkeypatch): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) + + t = torch.ones((2,), device="cuda") + keyset = DispatchKeySet.from_tensors((t,)) + assert DispatchKey.CUDA in keyset + + +def test_dispatch_keyset_cuda_autograd_uses_explicit_cuda_key(monkeypatch): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) + + t = torch.ones((2,), device="cuda") + t.requires_grad = True + keyset = DispatchKeySet.from_tensors((t,), grad_enabled=True, pipeline_enabled=False) + assert DispatchKey.AutogradCUDA in keyset diff --git a/tests/mindtorch_v2/test_dispatch_registry.py b/tests/mindtorch_v2/test_dispatch_registry.py index 4a2226989..4b9c74efd 100644 --- a/tests/mindtorch_v2/test_dispatch_registry.py +++ b/tests/mindtorch_v2/test_dispatch_registry.py @@ -6,6 +6,7 @@ def test_backend_registration_uses_keys(): entry = registry.get("aten::add") assert DispatchKey.CPU in entry.kernels assert DispatchKey.NPU in entry.kernels + assert DispatchKey.CUDA in entry.kernels assert DispatchKey.Meta in entry.kernels @@ -14,4 +15,5 @@ def test_autograd_backend_key_registration_prefers_device_keys(): entry = registry.get("aten::add") assert DispatchKey.AutogradCPU in entry.kernels assert DispatchKey.AutogradNPU in entry.kernels + assert DispatchKey.AutogradCUDA in entry.kernels assert DispatchKey.AutogradMeta in entry.kernels diff --git a/tests/mindtorch_v2/test_dispatch_resolution.py b/tests/mindtorch_v2/test_dispatch_resolution.py index 6ab94dea0..afe4e5060 100644 --- a/tests/mindtorch_v2/test_dispatch_resolution.py +++ b/tests/mindtorch_v2/test_dispatch_resolution.py @@ -31,3 +31,15 @@ def __init__(self, dev): with pytest.raises(RuntimeError, match=r"npu:1.*npu:0"): dispatch_with_keyset("add", keyset, None, a, b) + + +def test_dispatch_prefers_cuda_over_cpu(monkeypatch): + import mindtorch_v2._backends.cuda.runtime as cuda_runtime + + monkeypatch.setattr(cuda_runtime, "is_available", lambda: True) + monkeypatch.setattr(cuda_runtime, "device_count", lambda: 1) + + a = torch.ones((2,), device="cuda") + b = torch.ones((2,), device="cuda") + c = torch.add(a, b) + assert c.device.type == "cuda" From 6dc8cfd0f1ffb8e2f2945a6fcce11c72570c8a0d Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Sun, 8 Mar 2026 14:59:06 +0000 Subject: [PATCH 2/2] test(mindtorch_v2): align dispatch contracts with explicit cuda keys --- .../mindtorch_v2/contract/test_dispatch_key_order.py | 6 ++++-- .../contract/test_dispatch_registration_helpers.py | 12 ++++++------ tests/mindtorch_v2/test_library.py | 8 ++++---- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/mindtorch_v2/contract/test_dispatch_key_order.py b/tests/mindtorch_v2/contract/test_dispatch_key_order.py index 5e2943f77..e83b5d786 100644 --- a/tests/mindtorch_v2/contract/test_dispatch_key_order.py +++ b/tests/mindtorch_v2/contract/test_dispatch_key_order.py @@ -14,16 +14,18 @@ def test_dispatch_key_order_prefix_matches_torch_cuda(): "AutogradOther", "AutogradCPU", "AutogradNPU", + "AutogradCUDA", "AutogradXPU", "AutogradMeta", "Autograd", "Meta", "NPU", + "CUDA", "CPU", ] assert names[: len(prefix)] == prefix -def test_dispatch_key_includes_autograd_npu(): +def test_dispatch_key_includes_autograd_cuda(): assert DispatchKey.AutogradNPU in DispatchKey - assert not hasattr(DispatchKey, "AutogradCUDA") + assert DispatchKey.AutogradCUDA in DispatchKey diff --git a/tests/mindtorch_v2/contract/test_dispatch_registration_helpers.py b/tests/mindtorch_v2/contract/test_dispatch_registration_helpers.py index e85bd1ef1..52d9c45d3 100644 --- a/tests/mindtorch_v2/contract/test_dispatch_registration_helpers.py +++ b/tests/mindtorch_v2/contract/test_dispatch_registration_helpers.py @@ -16,8 +16,8 @@ def _new_op(prefix="reg_helper"): return name -def test_resolve_dispatch_key_supports_cuda_reserved_dispatch_key(): - assert resolve_dispatch_key("cuda") is DispatchKey.PrivateUse1 +def test_resolve_dispatch_key_supports_cuda_dispatch_key(): + assert resolve_dispatch_key("cuda") is DispatchKey.CUDA def test_resolve_dispatch_key_rejects_unknown_registration_device(): @@ -39,7 +39,7 @@ def fn(x): assert DispatchKey.Meta in entry.kernels -def test_register_forward_kernels_accepts_cuda_key_as_reserved_backend(): +def test_register_forward_kernels_accepts_cuda_key_as_backend(): op = _new_op("forward_cuda") def fn(x): @@ -48,7 +48,7 @@ def fn(x): register_forward_kernels(op, cpu=fn, cuda=fn, meta=fn) entry = registry.get(op) - assert DispatchKey.PrivateUse1 in entry.kernels + assert DispatchKey.CUDA in entry.kernels def test_register_autograd_kernels_registers_all_autograd_device_keys(): @@ -66,7 +66,7 @@ def fn(x): assert DispatchKey.AutogradMeta in entry.kernels -def test_register_autograd_kernels_accepts_cuda_key_as_reserved_autograd_backend(): +def test_register_autograd_kernels_accepts_cuda_key_as_autograd_backend(): op = _new_op("autograd_cuda") def fn(x): @@ -75,7 +75,7 @@ def fn(x): register_autograd_kernels(op, default=fn, cpu=fn, cuda=fn, meta=fn) entry = registry.get(op) - assert DispatchKey.AutogradXPU in entry.kernels + assert DispatchKey.AutogradCUDA in entry.kernels def test_register_helpers_still_enforce_schema_first(): diff --git a/tests/mindtorch_v2/test_library.py b/tests/mindtorch_v2/test_library.py index c9c485d8e..30959f5cf 100644 --- a/tests/mindtorch_v2/test_library.py +++ b/tests/mindtorch_v2/test_library.py @@ -24,8 +24,8 @@ def test_cpu(self): def test_npu(self): assert dispatch_key_from_string("NPU") is DispatchKey.NPU - def test_cuda_maps_to_privateuse1(self): - assert dispatch_key_from_string("CUDA") is DispatchKey.PrivateUse1 + def test_cuda_maps_to_cuda(self): + assert dispatch_key_from_string("CUDA") is DispatchKey.CUDA def test_privateuse1_maps_to_privateuse1(self): assert dispatch_key_from_string("PrivateUse1") is DispatchKey.PrivateUse1 @@ -42,8 +42,8 @@ def test_autograd_cpu(self): def test_autograd_npu(self): assert dispatch_key_from_string("AutogradNPU") is DispatchKey.AutogradNPU - def test_autograd_cuda_maps_to_autogradxpu(self): - assert dispatch_key_from_string("AutogradCUDA") is DispatchKey.AutogradXPU + def test_autograd_cuda_maps_to_autogradcuda(self): + assert dispatch_key_from_string("AutogradCUDA") is DispatchKey.AutogradCUDA def test_autograd_privateuse1_maps_to_autogradxpu(self): assert dispatch_key_from_string("AutogradPrivateUse1") is DispatchKey.AutogradXPU