cuda-core 1.0 — пишем CUDA-ядра на Python без C++ (ну почти)

• Исследователи, прототипирующие нестандартные операции для обучения моделей
  
• Инференс-инженеры, оптимизирующие горячие пути под конкретное железо
  
• Авторы ML-библиотек, которым нужен GPU-доступ без C++ build system

• InMemoryProgramCache — in-process LRU-кэш на время жизни процесса
  
• FileStreamProgramCache — персистентный кэш на диске между запусками

from cuda.core import Device, Program, ProgramOptions, ObjectCode
from cuda.core.utils import InMemoryProgramCache
d = Device(0)
d.set_current()
cuda_src = r"""
extern "C" __global__ void saxpy(
    float alpha, const float* __restrict__ x,
    const float* __restrict__ y, float* out, int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) out[idx] = alpha * x[idx] + y[idx];
}
"""
arch = f'sm_{d.compute_capability.major}{d.compute_capability.minor}'
cache = InMemoryProgramCache()
cache_key = f'saxpy_{arch}'
# Первый запуск — компилируем и кладём в кэш (~15 мс)
if cache.get(cache_key) is None:
    prog = Program(cuda_src, 'c++', options=ProgramOptions(arch=arch))
    cache[cache_key] = prog.compile('cubin')
# Все последующие — мгновенно (~0.05 мс, 300× быстрее)
kernel = ObjectCode.from_cubin(cache.get(cache_key)).get_kernel('saxpy')

import numpy as np
import ctypes
import threading
from cuda.core import (
    Device, SMResourceOptions, ContextOptions,
    Program, ProgramOptions, LaunchConfig, launch, ObjectCode,
    DeviceMemoryResource, DeviceMemoryResourceOptions
)
from cuda.core.utils import InMemoryProgramCache
from cuda.bindings import runtime as cudart
N = 5_000_000
d = Device(0)
d.set_current()
total_sm = d.resources.sm.sm_count
print(f"GPU: {d.name}, Total SMs: {total_sm}")  # Total SMs: 128
# Делим 128 SM на два раздела по 64
sm_opts = SMResourceOptions(count=[total_sm // 2, total_sm // 2])
groups, _ = d.resources.sm.split(sm_opts)
ctx_a = d.create_context(ContextOptions([groups[0]]))
ctx_b = d.create_context(ContextOptions([groups[1]]))
print(f"Context A: {ctx_a.resources.sm.sm_count} SMs, is_green={ctx_a.is_green}")
print(f"Context B: {ctx_b.resources.sm.sm_count} SMs, is_green={ctx_b.is_green}")
# Context A: 64 SMs, is_green=True
# Context B: 64 SMs, is_green=True
cuda_src = r"""
extern "C" __global__ void saxpy(
    float alpha, const float* __restrict__ x,
    const float* __restrict__ y, float* out, int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) out[idx] = alpha * x[idx] + y[idx];
}
"""
arch = f"sm_{d.compute_capability.major}{d.compute_capability.minor}"
cache = InMemoryProgramCache()
cache_key = f"saxpy_{arch}"
if cache.get(cache_key) is None:
    prog = Program(cuda_src, "c++", options=ProgramOptions(arch=arch))
    cache[cache_key] = prog.compile("cubin")
def run_task(ctx, task_name, alpha, fill_x, fill_y):
    d.set_current()  # нужен primary context в каждом потоке
    # stream, созданный через green context, привязан к его SM-партиции
    stream = ctx.create_stream()
    mr = DeviceMemoryResource(0, DeviceMemoryResourceOptions())
    kernel = ObjectCode.from_cubin(cache.get(cache_key)).get_kernel("saxpy")
    size = N * 4  # float32
    buf_x = mr.allocate(size, stream=stream)
    buf_y = mr.allocate(size, stream=stream)
    buf_o = mr.allocate(size, stream=stream)
    x_h = np.full(N, fill_x, dtype=np.float32)
    y_h = np.full(N, fill_y, dtype=np.float32)
    cudart.cudaMemcpy(buf_x.handle, x_h.ctypes.data, size, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
    cudart.cudaMemcpy(buf_y.handle, y_h.ctypes.data, size, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
    cfg = LaunchConfig(grid=(N + 255) // 256, block=256)
    launch(stream, cfg, kernel,
           np.float32(alpha), buf_x.handle, buf_y.handle, buf_o.handle, ctypes.c_int(N))
    stream.sync()
    out_h = np.empty(N, dtype=np.float32)
    cudart.cudaMemcpy(out_h.ctypes.data, buf_o.handle, size, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
    print(f"  {task_name}: result={out_h[0]:.1f}")
    buf_x.close(); buf_y.close(); buf_o.close()
    stream.close()
# Запускаем два ядра параллельно — каждое на своих 64 SM
t_a = threading.Thread(target=run_task, args=(ctx_a, "Task-A (64 SM)", 3.0, 2.0, 1.0))
t_b = threading.Thread(target=run_task, args=(ctx_b, "Task-B (64 SM)", 5.0, 4.0, 2.0))
t_a.start(); t_b.start()
t_a.join(); t_b.join()
# Task-A (64 SM): result=7.0
# Task-B (64 SM): result=22.0
ctx_a.close(); ctx_b.close()

from cuda.core import system, Device
print(f"Devices: {system.get_num_devices()}")
d = Device(0)
d.set_current()
sd = d.to_system_device()
util = sd.utilization
mem  = sd.memory_info
print(f"{sd.name}: GPU {util.gpu}%, MEM {util.memory}%")
print(f"Memory: {mem.used // 10242} MB / {mem.total // 10242} MB")

import numpy as np
import ctypes
from cuda.core import (
    Device, DeviceMemoryResource, DeviceMemoryResourceOptions,
    Program, ProgramOptions, LaunchConfig, launch, ObjectCode
)
from cuda.core.utils import InMemoryProgramCache
from cuda.bindings import runtime as cudart
N = 10_000_000
FLOAT_BYTES = 4
# 1. Устройство и стрим
d = Device(0)
d.set_current()
stream = d.create_stream()
print(f"GPU: {d.name}, SM: {d.properties.multiprocessor_count}")
# 2. Memory pool
mr = DeviceMemoryResource(0, DeviceMemoryResourceOptions())
# 3. Kernel: SAXPY = alpha * X + Y
cuda_src = r"""
extern "C" __global__ void saxpy(
    float alpha, const float* __restrict__ x,
    const float* __restrict__ y, float* out, int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) out[idx] = alpha * x[idx] + y[idx];
}
"""
# 4. Компиляция с кэшем
arch = f'sm_{d.compute_capability.major}{d.compute_capability.minor}'
cache = InMemoryProgramCache()
cache_key = f'saxpy_{arch}'
if cache.get(cache_key) is None:
    prog = Program(cuda_src, 'c++', options=ProgramOptions(arch=arch))
    cache[cache_key] = prog.compile('cubin')
kernel = ObjectCode.from_cubin(cache.get(cache_key)).get_kernel('saxpy')
# 5. GPU-память
size_bytes = N * FLOAT_BYTES
buf_x   = mr.allocate(size_bytes, stream=stream)
buf_y   = mr.allocate(size_bytes, stream=stream)
buf_out = mr.allocate(size_bytes, stream=stream)
# 6. H→D
x_host = np.ones(N, dtype=np.float32) * 2.0
y_host = np.ones(N, dtype=np.float32) * 1.0
alpha  = np.float32(3.0)
cudart.cudaMemcpy(buf_x.handle, x_host.ctypes.data, size_bytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
cudart.cudaMemcpy(buf_y.handle, y_host.ctypes.data, size_bytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
# 7. Запуск
cfg = LaunchConfig(grid=(N + 255) // 256, block=256)
launch(stream, cfg, kernel, alpha, buf_x.handle, buf_y.handle, buf_out.handle, ctypes.c_int(N))
# 8. D→H + проверка
stream.sync()
out_host = np.empty(N, dtype=np.float32)
cudart.cudaMemcpy(out_host.ctypes.data, buf_out.handle, size_bytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
assert np.allclose(out_host, 3.0 * x_host + y_host)
print(f"SAXPY OK: {out_host[0]:.1f}")  # 7.0
buf_x.close(); buf_y.close(); buf_out.close()
stream.close()

• CUDA 12+ обязателен
  
• Green Contexts — Ada Lovelace (sm_89) и новее, CUDA 12.4+
  
• Process checkpointing — только Linux

# CUDA 12 (драйвер 525+)
pip install "cuda-core[cu12]"
# CUDA 13 (драйвер 570+)
pip install "cuda-core[cu13]"