ricci 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
ricci-0.0.1/.gitignore ADDED
@@ -0,0 +1,10 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ build/
5
+ dist/
6
+ .venv/
7
+ venv/
8
+ .pytest_cache/
9
+ .ruff_cache/
10
+ .mypy_cache/
ricci-0.0.1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 The Ricci authors
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
ricci-0.0.1/PKG-INFO ADDED
@@ -0,0 +1,44 @@
1
+ Metadata-Version: 2.5
2
+ Name: ricci
3
+ Version: 0.0.1
4
+ Summary: A shape-typed tensor language whose compiler emits NVIDIA PTX directly.
5
+ Project-URL: Homepage, https://github.com/bsowlx/ricci
6
+ Project-URL: Issues, https://github.com/bsowlx/ricci/issues
7
+ Author: The Ricci authors
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: compiler,cuda,deep-learning,gpu,ptx,tensor,type-system
11
+ Classifier: Development Status :: 1 - Planning
12
+ Classifier: Environment :: GPU :: NVIDIA CUDA
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
16
+ Classifier: Topic :: Software Development :: Compilers
17
+ Requires-Python: >=3.10
18
+ Description-Content-Type: text/markdown
19
+
20
+ # Ricci
21
+
22
+ **Ricci** is a tensor language where shape bugs are compile errors, gradients are a compiler pass, and GPU kernels are emitted straight as NVIDIA PTX, with no LLVM and no nvcc.
23
+
24
+ It is named after Gregorio Ricci-Curbastro, who invented tensor calculus. Ricci programs end in `.ri`.
25
+
26
+ ```rust
27
+ fn forward<B>(p: Params, x: f32[B, 784]) -> f32[B, 10] {
28
+ let h = relu(x @ p.w1 + p.b1); // shapes are checked at compile time
29
+ h @ p.w2 + p.b2
30
+ }
31
+ ```
32
+
33
+ ## Status
34
+
35
+ Early development. This first release ships the CUDA driver runtime and one command:
36
+
37
+ ```
38
+ pip install ricci
39
+ ricci doctor
40
+ ```
41
+
42
+ `ricci doctor` finds the NVIDIA driver and lists your GPUs. It reports which PTX versions the driver accepts, then JIT-compiles and runs a small PTX kernel to check that your machine can run Ricci's output. It needs only the NVIDIA driver, not the CUDA toolkit. On machines without an NVIDIA GPU it says so; Ricci will use a CPU reference backend there.
43
+
44
+ The lexer, parser, shape checker, autodiff and PTX code generator are in progress at [github.com/bsowlx/ricci](https://github.com/bsowlx/ricci).
@@ -0,0 +1,25 @@
1
+ # Ricci
2
+
3
+ **Ricci** is a tensor language where shape bugs are compile errors, gradients are a compiler pass, and GPU kernels are emitted straight as NVIDIA PTX, with no LLVM and no nvcc.
4
+
5
+ It is named after Gregorio Ricci-Curbastro, who invented tensor calculus. Ricci programs end in `.ri`.
6
+
7
+ ```rust
8
+ fn forward<B>(p: Params, x: f32[B, 784]) -> f32[B, 10] {
9
+ let h = relu(x @ p.w1 + p.b1); // shapes are checked at compile time
10
+ h @ p.w2 + p.b2
11
+ }
12
+ ```
13
+
14
+ ## Status
15
+
16
+ Early development. This first release ships the CUDA driver runtime and one command:
17
+
18
+ ```
19
+ pip install ricci
20
+ ricci doctor
21
+ ```
22
+
23
+ `ricci doctor` finds the NVIDIA driver and lists your GPUs. It reports which PTX versions the driver accepts, then JIT-compiles and runs a small PTX kernel to check that your machine can run Ricci's output. It needs only the NVIDIA driver, not the CUDA toolkit. On machines without an NVIDIA GPU it says so; Ricci will use a CPU reference backend there.
24
+
25
+ The lexer, parser, shape checker, autodiff and PTX code generator are in progress at [github.com/bsowlx/ricci](https://github.com/bsowlx/ricci).
@@ -0,0 +1,36 @@
1
+ [build-system]
2
+ requires = ["hatchling>=1.27"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "ricci"
7
+ version = "0.0.1"
8
+ description = "A shape-typed tensor language whose compiler emits NVIDIA PTX directly."
9
+ readme = "docs/pypi.md"
10
+ requires-python = ">=3.10"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ authors = [{ name = "The Ricci authors" }]
14
+ keywords = ["compiler", "tensor", "gpu", "cuda", "ptx", "deep-learning", "type-system"]
15
+ classifiers = [
16
+ "Development Status :: 1 - Planning",
17
+ "Environment :: GPU :: NVIDIA CUDA",
18
+ "Intended Audience :: Science/Research",
19
+ "Programming Language :: Python :: 3",
20
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
21
+ "Topic :: Software Development :: Compilers",
22
+ ]
23
+ dependencies = []
24
+
25
+ [project.urls]
26
+ Homepage = "https://github.com/bsowlx/ricci"
27
+ Issues = "https://github.com/bsowlx/ricci/issues"
28
+
29
+ [project.scripts]
30
+ ricci = "ricci.cli:main"
31
+
32
+ [tool.hatch.build.targets.wheel]
33
+ packages = ["src/ricci"]
34
+
35
+ [tool.hatch.build.targets.sdist]
36
+ include = ["src/ricci", "docs/pypi.md", "LICENSE", "pyproject.toml"]
@@ -0,0 +1,3 @@
1
+ """Ricci: a shape-typed tensor language whose compiler emits NVIDIA PTX directly."""
2
+
3
+ __version__ = "0.0.1"
@@ -0,0 +1,5 @@
1
+ import sys
2
+
3
+ from ricci.cli import main
4
+
5
+ sys.exit(main())
@@ -0,0 +1,29 @@
1
+ """Command-line entry point: `ricci <command>`."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+
7
+ from ricci import __version__
8
+
9
+
10
+ def main(argv: list[str] | None = None) -> int:
11
+ parser = argparse.ArgumentParser(
12
+ prog="ricci",
13
+ description="Ricci: a shape-typed tensor language that compiles to NVIDIA PTX.",
14
+ )
15
+ parser.add_argument("--version", action="version", version=f"ricci {__version__}")
16
+ commands = parser.add_subparsers(dest="command", metavar="<command>")
17
+ commands.add_parser(
18
+ "doctor",
19
+ help="check the NVIDIA driver and GPU, then run a PTX smoke test",
20
+ )
21
+
22
+ args = parser.parse_args(argv)
23
+ if args.command == "doctor":
24
+ from ricci.doctor import run
25
+
26
+ return run()
27
+
28
+ parser.print_help()
29
+ return 0
@@ -0,0 +1,129 @@
1
+ """`ricci doctor`: can this machine run the PTX that Ricci generates?
2
+
3
+ 1. Find the NVIDIA driver and list GPUs.
4
+ 2. Ask the driver's JIT which PTX ISA versions it accepts for this GPU.
5
+ 3. Load a hand-written vector-add kernel, run it on 16M floats, check the
6
+ result and report the achieved memory bandwidth.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import array
12
+ import ctypes
13
+ import random
14
+ import sys
15
+
16
+ from ricci import __version__
17
+ from ricci.runtime.cuda import CudaError, Driver, DriverNotFound, PtxError
18
+
19
+ # PTX ISA versions to probe, oldest first. 7.8 is the first with sm_89 (RTX 40);
20
+ # 8.7 is the first with sm_120 (RTX 50).
21
+ PTX_VERSIONS = ["7.8", "8.0", "8.1", "8.2", "8.3", "8.4", "8.5", "8.6", "8.7", "8.8", "9.0", "9.1", "9.2"]
22
+
23
+ # c[i] = a[i] + b[i] for i < n. Written by hand; Ricci's code generator will emit
24
+ # kernels of exactly this shape.
25
+ VADD_PTX = """\
26
+ .version {version}
27
+ .target {target}
28
+ .address_size 64
29
+
30
+ .visible .entry vadd(.param .u64 pa, .param .u64 pb, .param .u64 pc, .param .u32 pn)
31
+ {{
32
+ .reg .pred %p<2>;
33
+ .reg .b32 %r<6>;
34
+ .reg .f32 %f<4>;
35
+ .reg .b64 %rd<8>;
36
+
37
+ ld.param.u64 %rd1, [pa];
38
+ ld.param.u64 %rd2, [pb];
39
+ ld.param.u64 %rd3, [pc];
40
+ ld.param.u32 %r1, [pn];
41
+ mov.u32 %r2, %ctaid.x;
42
+ mov.u32 %r3, %ntid.x;
43
+ mov.u32 %r4, %tid.x;
44
+ mad.lo.s32 %r5, %r2, %r3, %r4; // i = blockIdx.x * blockDim.x + threadIdx.x
45
+ setp.ge.u32 %p1, %r5, %r1;
46
+ @%p1 bra DONE; // threads past the end do nothing
47
+ mul.wide.u32 %rd4, %r5, 4; // byte offset of element i
48
+ add.s64 %rd5, %rd1, %rd4;
49
+ add.s64 %rd6, %rd2, %rd4;
50
+ add.s64 %rd7, %rd3, %rd4;
51
+ ld.global.f32 %f1, [%rd5];
52
+ ld.global.f32 %f2, [%rd6];
53
+ add.f32 %f3, %f1, %f2;
54
+ st.global.f32 [%rd7], %f3;
55
+ DONE:
56
+ ret;
57
+ }}
58
+ """
59
+
60
+
61
+ def _say(status: str, message: str) -> None:
62
+ print(f" [{status:^4}] {message}")
63
+
64
+
65
+ def run() -> int:
66
+ print(f"ricci {__version__} doctor (Python {sys.version.split()[0]}, {sys.platform})\n")
67
+
68
+ try:
69
+ driver = Driver()
70
+ except DriverNotFound as exc:
71
+ _say("info", str(exc))
72
+ _say("info", "Ricci will run programs on its CPU reference backend on this machine.")
73
+ return 0
74
+ except CudaError as exc:
75
+ _say("fail", f"the NVIDIA driver is present but could not start: {exc}")
76
+ return 1
77
+
78
+ major, minor = driver.version()
79
+ _say("ok", f"NVIDIA driver found, supports CUDA {major}.{minor}")
80
+ devices = driver.devices()
81
+ if not devices:
82
+ _say("fail", "the driver reports no GPUs")
83
+ return 1
84
+ for d in devices:
85
+ _say("ok", f"GPU {d.index}: {d.name}, {d.target}, {d.sm_count} SMs, {d.total_memory / 2**30:.1f} GiB")
86
+
87
+ gpu = devices[0]
88
+ ctx = driver.context(gpu.index)
89
+
90
+ accepted = []
91
+ for version in PTX_VERSIONS:
92
+ try:
93
+ ctx.load_ptx(VADD_PTX.format(version=version, target=gpu.target)).unload()
94
+ accepted.append(version)
95
+ except PtxError:
96
+ pass
97
+ if not accepted:
98
+ _say("fail", f"the driver accepted none of PTX ISA {PTX_VERSIONS[0]}..{PTX_VERSIONS[-1]} for {gpu.target}")
99
+ return 1
100
+ _say("ok", f"driver JIT accepts PTX ISA {accepted[0]} to {accepted[-1]} for {gpu.target}")
101
+
102
+ module = ctx.load_ptx(VADD_PTX.format(version=accepted[0], target=gpu.target))
103
+ kernel = module.function("vadd")
104
+ n = 1 << 24
105
+ nbytes = 4 * n
106
+ a = array.array("f", (float(i % 1000) for i in range(n)))
107
+ b = array.array("f", (float((7 * i) % 1000) for i in range(n)))
108
+ c = array.array("f", bytes(nbytes))
109
+ da, db, dc = ctx.alloc(nbytes), ctx.alloc(nbytes), ctx.alloc(nbytes)
110
+ try:
111
+ ctx.to_device(da, a.buffer_info()[0], nbytes)
112
+ ctx.to_device(db, b.buffer_info()[0], nbytes)
113
+ block = 256
114
+ args = [da, db, dc, ctypes.c_uint32(n)]
115
+ ms = ctx.time_ms(lambda: ctx.launch(kernel, ((n + block - 1) // block, 1, 1), (block, 1, 1), args))
116
+ ctx.to_host(c.buffer_info()[0], dc, nbytes)
117
+ finally:
118
+ for ptr in (da, db, dc):
119
+ ctx.free(ptr)
120
+ module.unload()
121
+
122
+ wrong = sum(1 for i in random.sample(range(n), 10_000) if c[i] != a[i] + b[i])
123
+ if wrong:
124
+ _say("fail", f"vector add gave {wrong} wrong values out of 10,000 checked")
125
+ return 1
126
+ gbps = 3 * nbytes / (ms / 1e3) / 1e9
127
+ _say("ok", f"PTX vector add on 16M floats: correct, {ms:.2f} ms, {gbps:.0f} GB/s")
128
+ print("\nThis machine can run Ricci's GPU backend.")
129
+ return 0
@@ -0,0 +1 @@
1
+ """Runtimes that execute compiled Ricci kernels."""
@@ -0,0 +1,236 @@
1
+ """A minimal binding to the CUDA *driver* API through ctypes.
2
+
3
+ Ricci never needs the CUDA toolkit (nvcc, NVRTC) at run time. The compiler
4
+ writes PTX text; the NVIDIA driver, which every machine with an NVIDIA GPU
5
+ already has, JIT-compiles that PTX to machine code inside `cuModuleLoadDataEx`.
6
+
7
+ Only the handful of driver calls Ricci uses are wrapped here: init, device
8
+ queries, PTX loading, memory, kernel launch and event timing.
9
+
10
+ Pitfall worth remembering: device pointers are 64-bit, so they must be passed
11
+ as `c_uint64` / `c_void_p`, never as bare Python ints (ctypes would truncate
12
+ them to 32 bits).
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import ctypes
18
+ import sys
19
+ from dataclasses import dataclass
20
+ from typing import Any
21
+
22
+ # CUdevice_attribute values from cuda.h
23
+ _ATTR_MULTIPROCESSOR_COUNT = 16
24
+ _ATTR_COMPUTE_CAPABILITY_MAJOR = 75
25
+ _ATTR_COMPUTE_CAPABILITY_MINOR = 76
26
+
27
+ # CUjit_option values from cuda.h
28
+ _JIT_INFO_LOG_BUFFER = 3
29
+ _JIT_INFO_LOG_BUFFER_SIZE_BYTES = 4
30
+ _JIT_ERROR_LOG_BUFFER = 5
31
+ _JIT_ERROR_LOG_BUFFER_SIZE_BYTES = 6
32
+
33
+ _LOG_SIZE = 8192
34
+
35
+
36
+ class CudaError(RuntimeError):
37
+ """A CUDA driver call returned an error code."""
38
+
39
+
40
+ class DriverNotFound(CudaError):
41
+ """No NVIDIA driver library on this machine (e.g. macOS, or no NVIDIA GPU)."""
42
+
43
+
44
+ class PtxError(CudaError):
45
+ """The driver's JIT rejected a PTX module; `log` holds ptxas's message."""
46
+
47
+ def __init__(self, message: str, log: str) -> None:
48
+ super().__init__(f"{message}\n{log}" if log else message)
49
+ self.log = log
50
+
51
+
52
+ def _load_library() -> Any:
53
+ if sys.platform == "win32":
54
+ candidates, loader = ["nvcuda.dll"], getattr(ctypes, "WinDLL")
55
+ elif sys.platform.startswith("linux"):
56
+ candidates, loader = ["libcuda.so.1", "libcuda.so"], ctypes.CDLL
57
+ else:
58
+ raise DriverNotFound(f"CUDA is not available on {sys.platform}")
59
+ for name in candidates:
60
+ try:
61
+ return loader(name)
62
+ except OSError:
63
+ continue
64
+ raise DriverNotFound("no NVIDIA driver found (looked for " + ", ".join(candidates) + ")")
65
+
66
+
67
+ @dataclass(frozen=True)
68
+ class Device:
69
+ index: int
70
+ name: str
71
+ compute_capability: tuple[int, int]
72
+ sm_count: int
73
+ total_memory: int # bytes
74
+
75
+ @property
76
+ def target(self) -> str:
77
+ """PTX target name, e.g. 'sm_89' for an RTX 40-series GPU."""
78
+ major, minor = self.compute_capability
79
+ return f"sm_{major}{minor}"
80
+
81
+
82
+ class Driver:
83
+ """The process-wide CUDA driver. Create one, then ask it for a `Context`."""
84
+
85
+ def __init__(self) -> None:
86
+ self.lib = _load_library()
87
+ self.check(self.lib.cuInit(0), "cuInit")
88
+
89
+ def check(self, result: int, what: str) -> None:
90
+ if result != 0:
91
+ name = ctypes.c_char_p()
92
+ self.lib.cuGetErrorName(result, ctypes.byref(name))
93
+ label = name.value.decode() if name.value else f"error {result}"
94
+ raise CudaError(f"{what} failed: {label}")
95
+
96
+ def version(self) -> tuple[int, int]:
97
+ """Highest CUDA version the driver supports, e.g. (13, 0)."""
98
+ v = ctypes.c_int()
99
+ self.check(self.lib.cuDriverGetVersion(ctypes.byref(v)), "cuDriverGetVersion")
100
+ return v.value // 1000, (v.value % 1000) // 10
101
+
102
+ def devices(self) -> list[Device]:
103
+ count = ctypes.c_int()
104
+ self.check(self.lib.cuDeviceGetCount(ctypes.byref(count)), "cuDeviceGetCount")
105
+ return [self._device(i) for i in range(count.value)]
106
+
107
+ def _device(self, index: int) -> Device:
108
+ dev = ctypes.c_int()
109
+ self.check(self.lib.cuDeviceGet(ctypes.byref(dev), index), "cuDeviceGet")
110
+ name = ctypes.create_string_buffer(256)
111
+ self.check(self.lib.cuDeviceGetName(name, 256, dev), "cuDeviceGetName")
112
+
113
+ def attr(code: int) -> int:
114
+ out = ctypes.c_int()
115
+ self.check(self.lib.cuDeviceGetAttribute(ctypes.byref(out), code, dev), "cuDeviceGetAttribute")
116
+ return out.value
117
+
118
+ mem = ctypes.c_size_t()
119
+ self.check(self.lib.cuDeviceTotalMem_v2(ctypes.byref(mem), dev), "cuDeviceTotalMem")
120
+ return Device(
121
+ index=index,
122
+ name=name.value.decode(errors="replace"),
123
+ compute_capability=(attr(_ATTR_COMPUTE_CAPABILITY_MAJOR), attr(_ATTR_COMPUTE_CAPABILITY_MINOR)),
124
+ sm_count=attr(_ATTR_MULTIPROCESSOR_COUNT),
125
+ total_memory=mem.value,
126
+ )
127
+
128
+ def context(self, device_index: int = 0) -> Context:
129
+ return Context(self, device_index)
130
+
131
+
132
+ class Module:
133
+ """A loaded PTX module; look kernels up by name."""
134
+
135
+ def __init__(self, ctx: Context, handle: ctypes.c_void_p) -> None:
136
+ self._ctx = ctx
137
+ self.handle = handle
138
+
139
+ def function(self, name: str) -> ctypes.c_void_p:
140
+ fn = ctypes.c_void_p()
141
+ self._ctx.driver.check(
142
+ self._ctx.lib.cuModuleGetFunction(ctypes.byref(fn), self.handle, name.encode()),
143
+ f"cuModuleGetFunction({name})",
144
+ )
145
+ return fn
146
+
147
+ def unload(self) -> None:
148
+ self._ctx.lib.cuModuleUnload(self.handle)
149
+
150
+
151
+ class Context:
152
+ """A CUDA context on one device: memory, modules, launches and timing."""
153
+
154
+ def __init__(self, driver: Driver, device_index: int) -> None:
155
+ self.driver = driver
156
+ self.lib = driver.lib
157
+ dev = ctypes.c_int()
158
+ driver.check(self.lib.cuDeviceGet(ctypes.byref(dev), device_index), "cuDeviceGet")
159
+ self.handle = ctypes.c_void_p()
160
+ driver.check(self.lib.cuDevicePrimaryCtxRetain(ctypes.byref(self.handle), dev), "cuDevicePrimaryCtxRetain")
161
+ driver.check(self.lib.cuCtxSetCurrent(self.handle), "cuCtxSetCurrent")
162
+
163
+ # -- modules -----------------------------------------------------------------
164
+
165
+ def load_ptx(self, ptx: str) -> Module:
166
+ """JIT-compile PTX text with the driver's built-in ptxas."""
167
+ error_log = ctypes.create_string_buffer(_LOG_SIZE)
168
+ info_log = ctypes.create_string_buffer(_LOG_SIZE)
169
+ options = (ctypes.c_int * 4)(
170
+ _JIT_ERROR_LOG_BUFFER, _JIT_ERROR_LOG_BUFFER_SIZE_BYTES,
171
+ _JIT_INFO_LOG_BUFFER, _JIT_INFO_LOG_BUFFER_SIZE_BYTES,
172
+ )
173
+ values = (ctypes.c_void_p * 4)(
174
+ ctypes.cast(error_log, ctypes.c_void_p), ctypes.c_void_p(_LOG_SIZE),
175
+ ctypes.cast(info_log, ctypes.c_void_p), ctypes.c_void_p(_LOG_SIZE),
176
+ )
177
+ handle = ctypes.c_void_p()
178
+ result = self.lib.cuModuleLoadDataEx(ctypes.byref(handle), ptx.encode() + b"\0", 4, options, values)
179
+ if result != 0:
180
+ raise PtxError(f"the driver rejected the PTX module (error {result})",
181
+ error_log.value.decode(errors="replace").strip())
182
+ return Module(self, handle)
183
+
184
+ # -- memory --------------------------------------------------------------------
185
+
186
+ def alloc(self, nbytes: int) -> ctypes.c_uint64:
187
+ ptr = ctypes.c_uint64()
188
+ self.driver.check(self.lib.cuMemAlloc_v2(ctypes.byref(ptr), ctypes.c_size_t(nbytes)), "cuMemAlloc")
189
+ return ptr
190
+
191
+ def free(self, ptr: ctypes.c_uint64) -> None:
192
+ self.driver.check(self.lib.cuMemFree_v2(ptr), "cuMemFree")
193
+
194
+ def to_device(self, dst: ctypes.c_uint64, host_address: int, nbytes: int) -> None:
195
+ self.driver.check(
196
+ self.lib.cuMemcpyHtoD_v2(dst, ctypes.c_void_p(host_address), ctypes.c_size_t(nbytes)), "cuMemcpyHtoD")
197
+
198
+ def to_host(self, host_address: int, src: ctypes.c_uint64, nbytes: int) -> None:
199
+ self.driver.check(
200
+ self.lib.cuMemcpyDtoH_v2(ctypes.c_void_p(host_address), src, ctypes.c_size_t(nbytes)), "cuMemcpyDtoH")
201
+
202
+ # -- launch and timing -------------------------------------------------------------
203
+
204
+ def launch(self, fn: ctypes.c_void_p, grid: tuple[int, int, int], block: tuple[int, int, int],
205
+ args: list[Any], shared_bytes: int = 0) -> None:
206
+ """Launch a kernel. `args` are ctypes values (c_uint64 pointers, c_uint32, c_float...).
207
+
208
+ The driver wants an array of *pointers to* each argument, hence the byref casts.
209
+ """
210
+ params = (ctypes.c_void_p * len(args))(*(ctypes.cast(ctypes.byref(a), ctypes.c_void_p) for a in args))
211
+ self.driver.check(
212
+ self.lib.cuLaunchKernel(fn, *grid, *block, shared_bytes, None, params, None), "cuLaunchKernel")
213
+
214
+ def synchronize(self) -> None:
215
+ self.driver.check(self.lib.cuCtxSynchronize(), "cuCtxSynchronize")
216
+
217
+ def time_ms(self, body: Any, repeat: int = 20, warmup: int = 3) -> float:
218
+ """Median-free quick timing: average milliseconds per call of `body()` over `repeat` calls."""
219
+ for _ in range(warmup):
220
+ body()
221
+ self.synchronize()
222
+ start, stop = ctypes.c_void_p(), ctypes.c_void_p()
223
+ self.driver.check(self.lib.cuEventCreate(ctypes.byref(start), 0), "cuEventCreate")
224
+ self.driver.check(self.lib.cuEventCreate(ctypes.byref(stop), 0), "cuEventCreate")
225
+ try:
226
+ self.driver.check(self.lib.cuEventRecord(start, None), "cuEventRecord")
227
+ for _ in range(repeat):
228
+ body()
229
+ self.driver.check(self.lib.cuEventRecord(stop, None), "cuEventRecord")
230
+ self.driver.check(self.lib.cuEventSynchronize(stop), "cuEventSynchronize")
231
+ elapsed = ctypes.c_float()
232
+ self.driver.check(self.lib.cuEventElapsedTime(ctypes.byref(elapsed), start, stop), "cuEventElapsedTime")
233
+ return elapsed.value / repeat
234
+ finally:
235
+ self.lib.cuEventDestroy_v2(start)
236
+ self.lib.cuEventDestroy_v2(stop)