// SafeC Standard Library — NVIDIA GPU backend (CUDA). See gpu_cuda.h.
//
// Kernels are CUDA C++ source embedded below and compiled on first use with
// NVRTC to a cubin for the selected GPU's own architecture (sm_80, sm_90,
// ...), then loaded with cuModuleLoadData — so one build of this library
// runs on any CUDA GPU the installed NVRTC knows. Compilation takes
// ~0.3 s once per process. Matmuls go to cuBLAS (and cuBLASLt for the
// fused GEMM + ReLU), on the same stream as the kernels.
// safeguard: link linux -L${CUDA_HOME:-/usr/local/cuda}/lib64 -lcuda -lcublas -lcublasLt -lnvrtc
// safeguard: link windows -L${CUDA_PATH}/lib/x64 -lcuda -lcublas -lcublasLt -lnvrtc
#pragma once
#include <std/stddef.h>
#include <std/ml/gpu_cuda.h>
#include <std/ml/cuda_ffi.h>
#include <std/ml/float16.h>
#include <std/math.h>
#include <std/mem.h>
#include <std/thread.h>
namespace std {
// ── Driver API / NVRTC / cuBLAS / cuBLASLt C ABI (cuda.h, nvrtc.h, cublas_v2.h, cublasLt.h) ──
// The driver API, NVRTC, cuBLAS and cuBLASLt are declared in cuda_ffi.h.
extern unsigned char* memcpy(unsigned char* dst, const unsigned char* src, size_t n);
extern int printf(const char* fmt, ...);
#define CUDA_OK 0
#define CUDA_ATTR_SM_COUNT 16 // CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT
#define CUDA_ATTR_CC_MAJOR 75 // CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR
#define CUDA_ATTR_CC_MINOR 76 // CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR
#define CUDA_STREAM_NON_BLOCKING 1U // CU_STREAM_NON_BLOCKING
#define CUDA_CAPTURE_GLOBAL 0 // CU_STREAM_CAPTURE_MODE_GLOBAL
#define CUDA_BLAS_N 0 // CUBLAS_OP_N
#define CUDA_BLAS_T 1 // CUBLAS_OP_T
#define CUDA_R_32F_ 0 // CUDA_R_32F
#define CUDA_R_16F_ 2 // CUDA_R_16F
#define CUDA_R_16BF_ 14 // CUDA_R_16BF
#define CUDA_GEMM_DEFAULT_ -1 // CUBLAS_GEMM_DEFAULT
#define CUDA_COMPUTE_32F 68 // CUBLAS_COMPUTE_32F
#define CUDA_COMPUTE_32F_TF32 77 // CUBLAS_COMPUTE_32F_FAST_TF32
#define CUDA_MATH_DEFAULT 0 // CUBLAS_DEFAULT_MATH
#define CUDA_MATH_TF32 3 // CUBLAS_TF32_TENSOR_OP_MATH
#define LT_DESC_TRANSA 3 // CUBLASLT_MATMUL_DESC_TRANSA
#define LT_DESC_TRANSB 4 // CUBLASLT_MATMUL_DESC_TRANSB
#define LT_DESC_EPILOGUE 7 // CUBLASLT_MATMUL_DESC_EPILOGUE
#define LT_EPILOGUE_RELU 2 // CUBLASLT_EPILOGUE_RELU
#define LT_PREF_MAX_WORKSPACE 1 // CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES
#define LT_RESULT_WORDS 12 // sizeof(cublasLtMatmulHeuristicResult_t) / 8 (algo is the first 64 bytes)
#define LT_CACHE 32
#define CUDA_BLOCK 256U
#define CUDA_MAX_BLOCKS 4096UL
#define CUDA_WORKSPACE_BYTES 33554432UL // 32 MiB: cuBLAS's recommendation for graph capture on Hopper
#define CUDA_STAGING_BYTES 65536UL // pinned bounce buffer for small reads
#define CUDA_PIPE_MIN 1048576UL // transfers at least this big go through the lanes
#define CUDA_PIPE_LANES 4 // parallel copy lanes (more only contend for bandwidth)
#define CUDA_PIPE_CHUNK 2097152UL // bytes per pinned buffer (two per lane)
#define CUDA_EVENT_NO_TIMING 2U // CU_EVENT_DISABLE_TIMING
// ── Kernel source ────────────────────────────────────────────────────────────
// Elementwise kernels use grid-stride loops; reduce_sum and mse_grad reduce
// per warp with shuffles, then per block through shared memory, then one
// atomic per block (or a plain store when the grid is one block, which
// also makes the zeroing memset unnecessary). sgemm_nn is a 64x64-tile
// shared-memory kernel with 4x4 register tiles. The fp16/bf16 kernels
// keep 2-byte storage and do the arithmetic in fp32.
static const char* cuda_kernel_source_() {
return
"\n"
"#define EW_UNARY(name, expr) \\\n"
"extern \"C\" __global__ void name(const float* __restrict__ a, \\\n"
" float* __restrict__ o, unsigned n, float k) { \\\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x; \\\n"
" unsigned stride = gridDim.x * blockDim.x; \\\n"
" for (; i < n; i += stride) { float x = a[i]; o[i] = (expr); } \\\n"
"}\n"
"#define EW_BINARY(name, expr) \\\n"
"extern \"C\" __global__ void name(const float* __restrict__ a, \\\n"
" const float* __restrict__ b, \\\n"
" float* __restrict__ o, unsigned n) { \\\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x; \\\n"
" unsigned stride = gridDim.x * blockDim.x; \\\n"
" for (; i < n; i += stride) { float x = a[i], y = b[i]; o[i] = (expr); } \\\n"
"}\n"
"EW_BINARY(ew_add, x + y)\n"
"EW_BINARY(ew_sub, x - y)\n"
"EW_BINARY(ew_mul, x * y)\n"
"EW_BINARY(ew_div, x / y)\n"
"EW_BINARY(ew_pow, __powf(x, y))\n"
"EW_BINARY(ew_relu_grad, x > 0.0f ? y : 0.0f)\n"
"EW_UNARY(ew_relu, x > 0.0f ? x : 0.0f)\n"
"EW_UNARY(ew_log, __logf(x))\n"
"EW_UNARY(ew_exp, __expf(x))\n"
"EW_UNARY(ew_sqrt, __fsqrt_rn(x))\n"
"EW_UNARY(ew_scale, x * k)\n"
"\n"
"// w -= lr * g, in place\n"
"extern \"C\" __global__ void sgd_update(float* __restrict__ w, const float* __restrict__ g,\n"
" float lr, unsigned n) {\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" unsigned stride = gridDim.x * blockDim.x;\n"
" for (; i < n; i += stride) w[i] -= lr * g[i];\n"
"}\n"
"\n"
"// two SGD updates in one launch: index space [0, n1) is w1, [n1, n1+n2) is w2\n"
"extern \"C\" __global__ void sgd_update2(float* __restrict__ w1, const float* __restrict__ g1, unsigned n1,\n"
" float* __restrict__ w2, const float* __restrict__ g2, unsigned n2,\n"
" float lr) {\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" unsigned stride = gridDim.x * blockDim.x;\n"
" for (; i < n1 + n2; i += stride) {\n"
" if (i < n1) w1[i] -= lr * g1[i]; else w2[i - n1] -= lr * g2[i - n1];\n"
" }\n"
"}\n"
"\n"
"// Block-wide sum of 's'; thread 0 gets the result.\n"
"__device__ float block_sum(float s) {\n"
" for (int off = 16; off > 0; off >>= 1) s += __shfl_down_sync(0xffffffffu, s, off);\n"
" __shared__ float part[32];\n"
" int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;\n"
" if (lane == 0) part[wid] = s;\n"
" __syncthreads();\n"
" s = 0.0f;\n"
" if (wid == 0) {\n"
" s = lane < (int)((blockDim.x + 31) >> 5) ? part[lane] : 0.0f;\n"
" for (int off = 16; off > 0; off >>= 1) s += __shfl_down_sync(0xffffffffu, s, off);\n"
" }\n"
" return s;\n"
"}\n"
"\n"
"// Sum: one partial per block, then one atomicAdd per block (*out zeroed\n"
"// first) or, for a one-block grid, a plain store.\n"
"extern \"C\" __global__ void reduce_sum(const float* __restrict__ a, float* __restrict__ out, unsigned n) {\n"
" float s = 0.0f;\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" unsigned stride = gridDim.x * blockDim.x;\n"
" for (; i < n; i += stride) s += a[i];\n"
" s = block_sum(s);\n"
" if (threadIdx.x == 0) { if (gridDim.x == 1) *out = s; else atomicAdd(out, s); }\n"
"}\n"
"\n"
"// dY = 2 (y - t), loss = sum (y - t)^2 (fused: one pass instead of sub/scale/square/sum)\n"
"extern \"C\" __global__ void mse_grad(const float* __restrict__ y, const float* __restrict__ t,\n"
" float* __restrict__ dy, float* __restrict__ loss, unsigned n) {\n"
" float s = 0.0f;\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" unsigned stride = gridDim.x * blockDim.x;\n"
" for (; i < n; i += stride) { float d = y[i] - t[i]; dy[i] = 2.0f * d; s += d * d; }\n"
" s = block_sum(s);\n"
" if (threadIdx.x == 0) { if (gridDim.x == 1) *loss = s; else atomicAdd(loss, s); }\n"
"}\n"
"\n"
"// ── SGEMM: C[M,N] = A[M,K] * B[K,N], row-major ─────────────────────────────\n"
"// 64x64 output tile per 256-thread block, BK=16, each thread a 4x4 micro-tile.\n"
"extern \"C\" __global__ __launch_bounds__(256) void sgemm_nn(const float* __restrict__ A,\n"
" const float* __restrict__ B, float* __restrict__ C, int M, int N, int K) {\n"
" const int BM = 64, BN = 64, BK = 16;\n"
" __shared__ float As[BK][BM + 4];\n"
" __shared__ float Bs[BK][BN + 4];\n"
" int tid = threadIdx.x;\n"
" int tx = tid % 16, ty = tid / 16;\n"
" int row0 = blockIdx.y * BM, col0 = blockIdx.x * BN;\n"
" float acc[4][4] = {};\n"
" for (int k0 = 0; k0 < K; k0 += BK) {\n"
" for (int l = tid; l < BM * BK; l += 256) {\n"
" int m = l / BK, k = l % BK;\n"
" int gm = row0 + m, gk = k0 + k;\n"
" As[k][m] = (gm < M && gk < K) ? A[(size_t)gm * K + gk] : 0.0f;\n"
" }\n"
" for (int l = tid; l < BK * BN; l += 256) {\n"
" int k = l / BN, n = l % BN;\n"
" int gk = k0 + k, gn = col0 + n;\n"
" Bs[k][n] = (gk < K && gn < N) ? B[(size_t)gk * N + gn] : 0.0f;\n"
" }\n"
" __syncthreads();\n"
"#pragma unroll\n"
" for (int k = 0; k < BK; ++k) {\n"
" float a[4], b[4];\n"
"#pragma unroll\n"
" for (int i = 0; i < 4; ++i) a[i] = As[k][ty * 4 + i];\n"
"#pragma unroll\n"
" for (int j = 0; j < 4; ++j) b[j] = Bs[k][tx * 4 + j];\n"
"#pragma unroll\n"
" for (int i = 0; i < 4; ++i)\n"
"#pragma unroll\n"
" for (int j = 0; j < 4; ++j) acc[i][j] += a[i] * b[j];\n"
" }\n"
" __syncthreads();\n"
" }\n"
"#pragma unroll\n"
" for (int i = 0; i < 4; ++i) {\n"
" int gm = row0 + ty * 4 + i;\n"
" if (gm >= M) continue;\n"
"#pragma unroll\n"
" for (int j = 0; j < 4; ++j) {\n"
" int gn = col0 + tx * 4 + j;\n"
" if (gn < N) C[(size_t)gm * N + gn] = acc[i][j];\n"
" }\n"
" }\n"
"}\n"
"\n"
"// Nearest-neighbour 2x upsample of 'planes' h x w planes: out[p][y][x] = in[p][y/2][x/2].\n"
"// One thread per input value, writing its 2x2 block as two float2 stores.\n"
"extern \"C\" __global__ void upsample2x(const float* __restrict__ in, float* __restrict__ out,\n"
" unsigned planes, unsigned h, unsigned w) {\n"
" unsigned hw = h * w, n = planes * hw;\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" unsigned stride = gridDim.x * blockDim.x;\n"
" for (; i < n; i += stride) {\n"
" unsigned p = i / hw, r = i - p * hw, y = r / w, x = r - y * w;\n"
" float v = in[i];\n"
" float2* o = (float2*)(out + (size_t)p * 4u * hw + (size_t)(2u * y) * (2u * w) + 2u * x);\n"
" o[0] = make_float2(v, v);\n"
" o[w] = make_float2(v, v); // next output row: 2w floats = w float2s on\n"
" }\n"
"}\n"
"\n"
"// dst[dOff + r*dStride + c] = src[sOff + r*sStride + c] for r < rows, c < len; T is\n"
"// float or float4 (all offsets, strides and len then counted in T).\n"
"template <typename T> __device__ void copy_rows_t(const T* __restrict__ src, T* __restrict__ dst,\n"
" unsigned long long sOff, unsigned long long sStride, unsigned long long dOff,\n"
" unsigned long long dStride, unsigned len, unsigned rows) {\n"
" unsigned long long n = (unsigned long long)len * rows;\n"
" for (unsigned long long i = blockIdx.x * (unsigned long long)blockDim.x + threadIdx.x; i < n;\n"
" i += (unsigned long long)gridDim.x * blockDim.x) {\n"
" unsigned r = (unsigned)(i / len), c = (unsigned)(i - (unsigned long long)r * len);\n"
" dst[dOff + r * dStride + c] = src[sOff + r * sStride + c];\n"
" }\n"
"}\n"
"extern \"C\" __global__ void copy_rows(const float* src, float* dst, unsigned long long sOff,\n"
" unsigned long long sStride, unsigned long long dOff, unsigned long long dStride,\n"
" unsigned len, unsigned rows) {\n"
" copy_rows_t<float>(src, dst, sOff, sStride, dOff, dStride, len, rows);\n"
"}\n"
"extern \"C\" __global__ void copy_rows4(const float4* src, float4* dst, unsigned long long sOff,\n"
" unsigned long long sStride, unsigned long long dOff, unsigned long long dStride,\n"
" unsigned len, unsigned rows) {\n"
" copy_rows_t<float4>(src, dst, sOff, sStride, dOff, dStride, len, rows);\n"
"}\n"
"\n"
"// ── fp16 / bf16 storage, fp32 arithmetic ───────────────────────────────────\n"
"__device__ float h2f(unsigned short h) { float f; asm(\"cvt.f32.f16 %0, %1;\" : \"=f\"(f) : \"h\"(h)); return f; }\n"
"__device__ unsigned short f2h(float f) { unsigned short h; asm(\"cvt.rn.f16.f32 %0, %1;\" : \"=h\"(h) : \"f\"(f)); return h; }\n"
"__device__ float b2f(unsigned short b) { return __uint_as_float(((unsigned)b) << 16); }\n"
"__device__ unsigned short f2b(float f) { // round to nearest even; NaN stays NaN\n"
" unsigned u = __float_as_uint(f);\n"
" if ((u & 0x7fffffffu) > 0x7f800000u) return (unsigned short)((u >> 16) | 0x40u);\n"
" u += 0x7fffu + ((u >> 16) & 1u);\n"
" return (unsigned short)(u >> 16);\n"
"}\n"
"template <int BF> __device__ float ld16(unsigned short v) { return BF ? b2f(v) : h2f(v); }\n"
"template <int BF> __device__ unsigned short st16(float v) { return BF ? f2b(v) : f2h(v); }\n"
"template <int BF> __device__ void matmul16(const unsigned short* A, const unsigned short* B,\n"
" unsigned short* C, int M, int N, int K) {\n"
" int col = blockIdx.x * blockDim.x + threadIdx.x;\n"
" int row = blockIdx.y * blockDim.y + threadIdx.y;\n"
" if (row >= M || col >= N) return;\n"
" float s = 0.0f;\n"
" for (int k = 0; k < K; ++k) s += ld16<BF>(A[(size_t)row * K + k]) * ld16<BF>(B[(size_t)k * N + col]);\n"
" C[(size_t)row * N + col] = st16<BF>(s);\n"
"}\n"
"template <int BF> __device__ void sgd16(unsigned short* w, const unsigned short* g, float lr, unsigned n) {\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" unsigned stride = gridDim.x * blockDim.x;\n"
" for (; i < n; i += stride) w[i] = st16<BF>(ld16<BF>(w[i]) - lr * ld16<BF>(g[i]));\n"
"}\n"
"extern \"C\" __global__ void matmul_f16(const unsigned short* A, const unsigned short* B, unsigned short* C, int M, int N, int K) { matmul16<0>(A, B, C, M, N, K); }\n"
"extern \"C\" __global__ void matmul_bf16(const unsigned short* A, const unsigned short* B, unsigned short* C, int M, int N, int K) { matmul16<1>(A, B, C, M, N, K); }\n"
"extern \"C\" __global__ void sgd_f16(unsigned short* w, const unsigned short* g, float lr, unsigned n) { sgd16<0>(w, g, lr, n); }\n"
"extern \"C\" __global__ void sgd_bf16(unsigned short* w, const unsigned short* g, float lr, unsigned n) { sgd16<1>(w, g, lr, n); }\n"
"// CudaHalfBuf ops (grid-stride, fp32 arithmetic, one rounding per stored value)\n"
"#define GS for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
"template <int BF> __device__ void to16(const float* a, unsigned short* o, unsigned n) { GS o[i] = st16<BF>(a[i]); }\n"
"template <int BF> __device__ void from16(const unsigned short* a, float* o, unsigned n) { GS o[i] = ld16<BF>(a[i]); }\n"
"template <int BF> __device__ void relu16(const unsigned short* a, unsigned short* o, unsigned n) {\n"
" GS { float v = ld16<BF>(a[i]); o[i] = v > 0.0f ? a[i] : (unsigned short)0; }\n"
"}\n"
"template <int BF> __device__ void relug16(const unsigned short* z, const unsigned short* g, unsigned short* o, unsigned n) {\n"
" GS o[i] = ld16<BF>(z[i]) > 0.0f ? g[i] : (unsigned short)0;\n"
"}\n"
"template <int BF> __device__ void mse16(const unsigned short* y, const unsigned short* t, unsigned short* dy,\n"
" float* loss, unsigned n) {\n"
" float s = 0.0f;\n"
" GS { float d = ld16<BF>(y[i]) - ld16<BF>(t[i]); dy[i] = st16<BF>(2.0f * d); s += d * d; }\n"
" s = block_sum(s);\n"
" if (threadIdx.x == 0) { if (gridDim.x == 1) *loss = s; else atomicAdd(loss, s); }\n"
"}\n"
"template <int BF> __device__ void sgd2_16(unsigned short* w1, const unsigned short* g1, unsigned n1,\n"
" unsigned short* w2, const unsigned short* g2, unsigned n2, float lr) {\n"
" unsigned n = n1 + n2;\n"
" GS {\n"
" if (i < n1) w1[i] = st16<BF>(ld16<BF>(w1[i]) - lr * ld16<BF>(g1[i]));\n"
" else w2[i - n1] = st16<BF>(ld16<BF>(w2[i - n1]) - lr * ld16<BF>(g2[i - n1]));\n"
" }\n"
"}\n"
"extern \"C\" __global__ void to_f16(const float* a, unsigned short* o, unsigned n) { to16<0>(a, o, n); }\n"
"extern \"C\" __global__ void to_bf16(const float* a, unsigned short* o, unsigned n) { to16<1>(a, o, n); }\n"
"extern \"C\" __global__ void from_f16(const unsigned short* a, float* o, unsigned n) { from16<0>(a, o, n); }\n"
"extern \"C\" __global__ void from_bf16(const unsigned short* a, float* o, unsigned n) { from16<1>(a, o, n); }\n"
"extern \"C\" __global__ void relu_f16(const unsigned short* a, unsigned short* o, unsigned n) { relu16<0>(a, o, n); }\n"
"extern \"C\" __global__ void relu_bf16(const unsigned short* a, unsigned short* o, unsigned n) { relu16<1>(a, o, n); }\n"
"extern \"C\" __global__ void relug_f16(const unsigned short* z, const unsigned short* g, unsigned short* o, unsigned n) { relug16<0>(z, g, o, n); }\n"
"extern \"C\" __global__ void relug_bf16(const unsigned short* z, const unsigned short* g, unsigned short* o, unsigned n) { relug16<1>(z, g, o, n); }\n"
"extern \"C\" __global__ void mse_f16(const unsigned short* y, const unsigned short* t, unsigned short* dy, float* l, unsigned n) { mse16<0>(y, t, dy, l, n); }\n"
"extern \"C\" __global__ void mse_bf16(const unsigned short* y, const unsigned short* t, unsigned short* dy, float* l, unsigned n) { mse16<1>(y, t, dy, l, n); }\n"
"extern \"C\" __global__ void sgd2_f16(unsigned short* w1, const unsigned short* g1, unsigned n1, unsigned short* w2, const unsigned short* g2, unsigned n2, float lr) { sgd2_16<0>(w1, g1, n1, w2, g2, n2, lr); }\n"
"extern \"C\" __global__ void sgd2_bf16(unsigned short* w1, const unsigned short* g1, unsigned n1, unsigned short* w2, const unsigned short* g2, unsigned n2, float lr) { sgd2_16<1>(w1, g1, n1, w2, g2, n2, lr); }\n"
"// Decode ops: tanh GELU, RMSNorm, and a 16-bit-weight GEMV with fp32 x / y\n"
"EW_UNARY(ew_gelu, 0.5f * x * (1.0f + tanhf(0.7978845608028654f * (x + 0.044715f * x * x * x))))\n"
"// One block per row; y may be x (every x value is read before y is written).\n"
"extern \"C\" __global__ void rmsnorm(const float* x, const float* w, float* y, unsigned n, float eps) {\n"
" const float* xr = x + (size_t)blockIdx.x * n;\n"
" float* yr = y + (size_t)blockIdx.x * n;\n"
" float s = 0.0f;\n"
" for (unsigned i = threadIdx.x; i < n; i += blockDim.x) { float v = xr[i]; s += v * v; }\n"
" s = block_sum(s);\n"
" __shared__ float inv;\n"
" if (threadIdx.x == 0) inv = rsqrtf(s / (float)n + eps);\n"
" __syncthreads();\n"
" const float r = inv;\n"
" for (unsigned i = threadIdx.x; i < n; i += blockDim.x) yr[i] = xr[i] * r * (w ? w[i] : 1.0f);\n"
"}\n"
"// 16 bytes of weights that are read once: past L1, so x stays cached there.\n"
"__device__ __forceinline__ uint4 ld_stream(const unsigned short* p) {\n"
" uint4 v;\n"
" asm volatile(\"ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];\"\n"
" : \"=r\"(v.x), \"=r\"(v.y), \"=r\"(v.z), \"=r\"(v.w) : \"l\"(p));\n"
" return v;\n"
"}\n"
"template <int BF> __device__ __forceinline__ float lo16(unsigned u) { return BF ? __uint_as_float(u << 16) : h2f((unsigned short)(u & 0xffffu)); }\n"
"template <int BF> __device__ __forceinline__ float hi16(unsigned u) { return BF ? __uint_as_float(u & 0xffff0000u) : h2f((unsigned short)(u >> 16)); }\n"
"template <int BF> __device__ __forceinline__ float dot8(uint4 w, const float* x) {\n"
" const float4 a = *(const float4*)x, b = *(const float4*)(x + 4);\n"
" return lo16<BF>(w.x) * a.x + hi16<BF>(w.x) * a.y + lo16<BF>(w.y) * a.z + hi16<BF>(w.y) * a.w +\n"
" lo16<BF>(w.z) * b.x + hi16<BF>(w.z) * b.y + lo16<BF>(w.w) * b.z + hi16<BF>(w.w) * b.w;\n"
"}\n"
"// y = W x with W [n][k]. 128 threads = 4 warps; 'wpr' (1, 2 or 4) warps share\n"
"// a row, rows grid-strided. VEC (k % 8 == 0, 16-byte aligned W and x): each\n"
"// lane reads 8 weights per 16-byte load, four loads in flight per warp.\n"
"template <int BF, int VEC> __device__ void gemv16(const unsigned short* __restrict__ W, const float* __restrict__ x,\n"
" float* __restrict__ y, unsigned n, unsigned k, int wpr) {\n"
" __shared__ float part[4];\n"
" const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, rpb = 4 / wpr, sl = warp % wpr;\n"
" const size_t kk = k;\n"
" for (unsigned r0 = blockIdx.x * rpb; r0 < n; r0 += gridDim.x * rpb) {\n"
" const unsigned row = r0 + warp / wpr;\n"
" float s = 0.0f;\n"
" if (row < n) {\n"
" const unsigned short* wr = W + (size_t)row * kk;\n"
" if (VEC) {\n"
" const size_t step = 256u * wpr;\n"
" for (size_t c = sl * 256u + lane * 8u; c < kk; c += 4 * step) {\n"
" uint4 wv[4];\n"
"#pragma unroll\n"
" for (int u = 0; u < 4; ++u) if (c + u * step < kk) wv[u] = ld_stream(wr + c + u * step);\n"
"#pragma unroll\n"
" for (int u = 0; u < 4; ++u) if (c + u * step < kk) s += dot8<BF>(wv[u], x + c + u * step);\n"
" }\n"
" } else {\n"
" for (size_t c = sl * 32u + lane; c < kk; c += 32u * wpr) s += ld16<BF>(wr[c]) * __ldg(x + c);\n"
" }\n"
" }\n"
" for (int off = 16; off > 0; off >>= 1) s += __shfl_down_sync(0xffffffffu, s, off);\n"
" if (wpr == 1) {\n"
" if (lane == 0 && row < n) y[row] = s;\n"
" } else {\n"
" if (lane == 0) part[warp] = s;\n"
" __syncthreads();\n"
" if (sl == 0 && lane == 0 && row < n) { float t = 0.0f; for (int i = 0; i < wpr; ++i) t += part[warp + i]; y[row] = t; }\n"
" __syncthreads();\n"
" }\n"
" }\n"
"}\n"
"extern \"C\" __global__ void __launch_bounds__(128) gemv_f16(const unsigned short* W, const float* x, float* y, unsigned n, unsigned k, int wpr) { gemv16<0, 1>(W, x, y, n, k, wpr); }\n"
"extern \"C\" __global__ void __launch_bounds__(128) gemv_bf16(const unsigned short* W, const float* x, float* y, unsigned n, unsigned k, int wpr) { gemv16<1, 1>(W, x, y, n, k, wpr); }\n"
"extern \"C\" __global__ void __launch_bounds__(128) gemv1_f16(const unsigned short* W, const float* x, float* y, unsigned n, unsigned k, int wpr) { gemv16<0, 0>(W, x, y, n, k, wpr); }\n"
"extern \"C\" __global__ void __launch_bounds__(128) gemv1_bf16(const unsigned short* W, const float* x, float* y, unsigned n, unsigned k, int wpr) { gemv16<1, 0>(W, x, y, n, k, wpr); }\n"
"// ── Mixed precision (the dtype as a runtime flag; _rn intrinsics so nothing\n"
"// is contracted into an FMA: the same arithmetic as the CPU) ───────────────\n"
"__device__ __forceinline__ float ldx(unsigned short v, int bf) { return bf ? b2f(v) : h2f(v); }\n"
"__device__ __forceinline__ unsigned short stx(float v, int bf) { return bf ? f2b(v) : f2h(v); }\n"
"extern \"C\" __global__ void half_to_half(const unsigned short* a, unsigned short* o, unsigned n, int toBf) {\n"
" GS o[i] = stx(ldx(a[i], !toBf), toBf);\n"
"}\n"
"// Up to 32 buffers per launch, blockIdx.y picking one. mode 0: fp32 -> 16-bit;\n"
"// 1: 16-bit -> fp32 times scale; 2: fp32 -> fp32 times scale (in place allowed).\n"
"// Modes 1 and 2 set *flag = 1 when a result is not finite.\n"
"struct CvtTab { const unsigned char* src[32]; unsigned char* dst[32]; unsigned n[32]; };\n"
"#define GTID (blockIdx.x * blockDim.x + threadIdx.x)\n"
"#define GSTRIDE (gridDim.x * blockDim.x)\n"
"__device__ __forceinline__ unsigned pack2(float lo, float hi, int bf) { return (unsigned)stx(lo, bf) | ((unsigned)stx(hi, bf) << 16); }\n"
"__device__ __forceinline__ float lo16x(unsigned u, int bf) { return ldx((unsigned short)(u & 0xffffu), bf); }\n"
"__device__ __forceinline__ float hi16x(unsigned u, int bf) { return ldx((unsigned short)(u >> 16), bf); }\n"
"__device__ __forceinline__ float cvt1(float v, float scale, int& bad) {\n"
" v = __fmul_rn(v, scale);\n"
" if (!isfinite(v)) bad = 1;\n"
" return v;\n"
"}\n"
"// vec: every pointer 16-byte aligned, so each thread moves 4 values at a time\n"
"// (float4 / 8 bytes of 16-bit) and the last n % 4 go one by one.\n"
"extern \"C\" __global__ void cvt_many(CvtTab t, int mode, int bf, float scale, float* flag, int vec) {\n"
" const unsigned b = blockIdx.y, n = t.n[b];\n"
" int bad = 0;\n"
" unsigned head = 0;\n"
" if (vec) {\n"
" const unsigned nv = n / 4;\n"
" head = nv * 4;\n"
" for (unsigned j = GTID; j < nv; j += GSTRIDE) {\n"
" if (mode == 0) {\n"
" const float4 x = ((const float4*)t.src[b])[j];\n"
" ((uint2*)t.dst[b])[j] = make_uint2(pack2(x.x, x.y, bf), pack2(x.z, x.w, bf));\n"
" } else if (mode == 1) {\n"
" const uint2 u = ((const uint2*)t.src[b])[j];\n"
" ((float4*)t.dst[b])[j] = make_float4(cvt1(lo16x(u.x, bf), scale, bad), cvt1(hi16x(u.x, bf), scale, bad),\n"
" cvt1(lo16x(u.y, bf), scale, bad), cvt1(hi16x(u.y, bf), scale, bad));\n"
" } else {\n"
" const float4 x = ((const float4*)t.src[b])[j];\n"
" ((float4*)t.dst[b])[j] = make_float4(cvt1(x.x, scale, bad), cvt1(x.y, scale, bad),\n"
" cvt1(x.z, scale, bad), cvt1(x.w, scale, bad));\n"
" }\n"
" }\n"
" }\n"
" for (unsigned i = head + GTID; i < n; i += GSTRIDE) {\n"
" if (mode == 0) { ((unsigned short*)t.dst[b])[i] = stx(((const float*)t.src[b])[i], bf); continue; }\n"
" float v = mode == 1 ? ldx(((const unsigned short*)t.src[b])[i], bf) : ((const float*)t.src[b])[i];\n"
" ((float*)t.dst[b])[i] = cvt1(v, scale, bad);\n"
" }\n"
" if (bad && flag) *flag = 1.0f;\n"
"}\n"
"// fp32 master-weight steps (nn.h's Optimizer::step), writing the 16-bit copy\n"
"// wh when not null; nothing happens when *skip != 0.\n"
"extern \"C\" __global__ void mp_sgd(float* w, const float* g, float* vel, unsigned short* wh, float lr,\n"
" float mom, float wd, unsigned n, int bf, const float* skip) {\n"
" if (skip && *skip != 0.0f) return;\n"
" GS {\n"
" float gi = __fadd_rn(g[i], __fmul_rn(wd, w[i]));\n"
" if (vel) { float v = __fadd_rn(__fmul_rn(mom, vel[i]), gi); vel[i] = v; gi = v; }\n"
" float nw = __fsub_rn(w[i], __fmul_rn(lr, gi));\n"
" w[i] = nw;\n"
" if (wh) wh[i] = stx(nw, bf);\n"
" }\n"
"}\n"
"struct AdamArgs { float lr, b1, b2, eps, wd, bc1, bc2; int decoupled; };\n"
"__device__ __forceinline__ float adam1(float w, float& m, float& v, float g, const AdamArgs& a) {\n"
" if (!a.decoupled) g = __fadd_rn(g, __fmul_rn(a.wd, w));\n"
" m = __fadd_rn(__fmul_rn(a.b1, m), __fmul_rn(__fsub_rn(1.0f, a.b1), g));\n"
" v = __fadd_rn(__fmul_rn(a.b2, v), __fmul_rn(__fmul_rn(__fsub_rn(1.0f, a.b2), g), g));\n"
" float upd = __fdiv_rn(__fdiv_rn(m, a.bc1), __fadd_rn(__fsqrt_rn(__fdiv_rn(v, a.bc2)), a.eps));\n"
" if (a.decoupled) upd = __fadd_rn(upd, __fmul_rn(a.wd, w));\n"
" return __fsub_rn(w, __fmul_rn(a.lr, upd));\n"
"}\n"
"// One tensor's Adam step over [GTID, n): 4 values per thread when vec (w, m,\n"
"// v, g 16-byte and wh 8-byte aligned), then the tail one by one.\n"
"__device__ __forceinline__ void adam_span(float* w, float* m, float* v, const float* g, unsigned short* wh,\n"
" unsigned n, int bf, int vec, const AdamArgs& a) {\n"
" unsigned head = 0;\n"
" if (vec) {\n"
" const unsigned nv = n / 4;\n"
" head = nv * 4;\n"
" for (unsigned j = GTID; j < nv; j += GSTRIDE) {\n"
" float4 W = ((float4*)w)[j], M = ((float4*)m)[j], V = ((float4*)v)[j];\n"
" const float4 G = ((const float4*)g)[j];\n"
" W.x = adam1(W.x, M.x, V.x, G.x, a); W.y = adam1(W.y, M.y, V.y, G.y, a);\n"
" W.z = adam1(W.z, M.z, V.z, G.z, a); W.w = adam1(W.w, M.w, V.w, G.w, a);\n"
" ((float4*)w)[j] = W; ((float4*)m)[j] = M; ((float4*)v)[j] = V;\n"
" if (wh) ((uint2*)wh)[j] = make_uint2(pack2(W.x, W.y, bf), pack2(W.z, W.w, bf));\n"
" }\n"
" }\n"
" for (unsigned i = head + GTID; i < n; i += GSTRIDE) {\n"
" float mi = m[i], vi = v[i];\n"
" const float nw = adam1(w[i], mi, vi, g[i], a);\n"
" m[i] = mi; v[i] = vi; w[i] = nw;\n"
" if (wh) wh[i] = stx(nw, bf);\n"
" }\n"
"}\n"
"extern \"C\" __global__ void mp_adam(float* w, float* m, float* v, const float* g, unsigned short* wh,\n"
" float lr, float b1, float b2, float eps, float wd, float bc1, float bc2,\n"
" int decoupled, unsigned n, int bf, const float* skip, int vec) {\n"
" if (skip && *skip != 0.0f) return;\n"
" const AdamArgs a = {lr, b1, b2, eps, wd, bc1, bc2, decoupled};\n"
" adam_span(w, m, v, g, wh, n, bf, vec, a);\n"
"}\n"
"// Up to 24 tensors per launch (blockIdx.y picks one): a model's step in a few launches.\n"
"struct AdamTab { float* w[24]; float* m[24]; float* v[24]; const float* g[24]; unsigned short* wh[24]; unsigned n[24]; };\n"
"extern \"C\" __global__ void mp_adam_many(AdamTab t, AdamArgs a, int bf, const float* skip, int vec) {\n"
" if (skip && *skip != 0.0f) return;\n"
" const unsigned b = blockIdx.y;\n"
" if (blockIdx.x * blockDim.x * (vec ? 4u : 1u) >= t.n[b]) return;\n"
" adam_span(t.w[b], t.m[b], t.v[b], t.g[b], t.wh[b], t.n[b], bf, vec, a);\n"
"}\n"
"// 16-bit elementwise: 0 a+b, 1 a-b, 2 a*b, 3 a*k, 4 a + k*b.\n"
"__device__ __forceinline__ float ew1(float x, float y, float k, int op) {\n"
" return op == 0 ? __fadd_rn(x, y) : op == 1 ? __fsub_rn(x, y) : op == 2 ? __fmul_rn(x, y)\n"
" : op == 3 ? __fmul_rn(x, k) : __fadd_rn(x, __fmul_rn(k, y));\n"
"}\n"
"__device__ __forceinline__ unsigned ew2(unsigned x, unsigned y, float k, int op, int bf) {\n"
" return pack2(ew1(lo16x(x, bf), lo16x(y, bf), k, op), ew1(hi16x(x, bf), hi16x(y, bf), k, op), bf);\n"
"}\n"
"// vec: a, b, o 16-byte aligned; 8 values per thread, then the tail.\n"
"extern \"C\" __global__ void half_ew(const unsigned short* a, const unsigned short* b, unsigned short* o,\n"
" float k, unsigned n, int op, int bf, int vec) {\n"
" unsigned head = 0;\n"
" if (vec) {\n"
" const unsigned nv = n / 8;\n"
" head = nv * 8;\n"
" for (unsigned j = GTID; j < nv; j += GSTRIDE) {\n"
" const uint4 x = ((const uint4*)a)[j];\n"
" const uint4 y = op == 3 ? x : ((const uint4*)b)[j];\n"
" ((uint4*)o)[j] = make_uint4(ew2(x.x, y.x, k, op, bf), ew2(x.y, y.y, k, op, bf),\n"
" ew2(x.z, y.z, k, op, bf), ew2(x.w, y.w, k, op, bf));\n"
" }\n"
" }\n"
" for (unsigned i = head + GTID; i < n; i += GSTRIDE)\n"
" o[i] = stx(ew1(ldx(a[i], bf), op == 3 ? 0.0f : ldx(b[i], bf), k, op), bf);\n"
"}\n"
"// y[i] += bias[(i / plane) % channels], then ReLU if relu.\n"
"extern \"C\" __global__ void half_bias(unsigned short* y, const unsigned short* bias, unsigned n,\n"
" unsigned channels, unsigned plane, int relu, int bf) {\n"
" GS {\n"
" float r = __fadd_rn(ldx(y[i], bf), ldx(bias[(i / plane) % channels], bf));\n"
" if (relu && r < 0.0f) r = 0.0f;\n"
" y[i] = stx(r, bf);\n"
" }\n"
"}\n"
"// 16-bit max / avg pooling over planes [rows][h][w], no padding.\n"
"extern \"C\" __global__ void half_pool(const unsigned short* x, unsigned short* y, unsigned rows,\n"
" unsigned h, unsigned w, unsigned oh, unsigned ow, unsigned k,\n"
" unsigned s, int avg, int bf) {\n"
" unsigned plane = oh * ow, n = rows * plane;\n"
" GS {\n"
" unsigned p = i / plane, r = i % plane, oy = r / ow, ox = r % ow;\n"
" const unsigned short* src = x + (size_t)p * h * w + (size_t)(oy * s) * w + ox * s;\n"
" float acc = 0.0f;\n"
" unsigned short best = src[0];\n"
" float bv = ldx(best, bf);\n"
" for (unsigned ky = 0; ky < k; ++ky)\n"
" for (unsigned kx = 0; kx < k; ++kx) {\n"
" unsigned short b = src[ky * w + kx];\n"
" float v = ldx(b, bf);\n"
" acc = __fadd_rn(acc, v);\n"
" if (v > bv) { bv = v; best = b; }\n"
" }\n"
" y[i] = avg ? stx(__fdiv_rn(acc, (float)(k * k)), bf) : best;\n"
" }\n"
"}\n";
}
// Kernel table, in this order.
#define K_ADD 0
#define K_SUB 1
#define K_MUL 2
#define K_DIV 3
#define K_POW 4
#define K_RELU 5
#define K_LOG 6
#define K_EXP 7
#define K_SQRT 8
#define K_SCALE 9
#define K_RELU_GRAD 10
#define K_SGD 11
#define K_SUM 12
#define K_SGEMM 13
#define K_MSE 14
#define K_SGD2 15
#define K_MATMUL_F16 16
#define K_MATMUL_BF16 17
#define K_SGD_F16 18
#define K_SGD_BF16 19
#define K_UPSAMPLE2X 20
#define K_COPY_ROWS 21
#define K_COPY_ROWS4 22
// CudaHalfBuf kernels: each fp16 index is followed by its bf16 twin
#define K_TO_F16 23
#define K_FROM_F16 25
#define K_RELU_F16 27
#define K_RELUG_F16 29
#define K_MSE_F16 31
#define K_SGD2_F16 33
#define K_GELU 35
#define K_RMSNORM 36
#define K_GEMV_F16 37
#define K_GEMV1_F16 39 // scalar loads (k % 8 != 0 or unaligned)
#define K_HALF_CVT 41
#define K_CVT_MANY 42
#define K_MP_SGD 43
#define K_MP_ADAM 44
#define K_HALF_EW 45
#define K_HALF_BIAS 46
#define K_HALF_POOL 47
#define K_MP_ADAM_MANY 48
#define K_COUNT 49
static const char* cuda_kernel_name_(int k) {
if (k == K_ADD) return "ew_add";
if (k == K_SUB) return "ew_sub";
if (k == K_MUL) return "ew_mul";
if (k == K_DIV) return "ew_div";
if (k == K_POW) return "ew_pow";
if (k == K_RELU) return "ew_relu";
if (k == K_LOG) return "ew_log";
if (k == K_EXP) return "ew_exp";
if (k == K_SQRT) return "ew_sqrt";
if (k == K_SCALE) return "ew_scale";
if (k == K_RELU_GRAD) return "ew_relu_grad";
if (k == K_SGD) return "sgd_update";
if (k == K_SUM) return "reduce_sum";
if (k == K_SGEMM) return "sgemm_nn";
if (k == K_MSE) return "mse_grad";
if (k == K_SGD2) return "sgd_update2";
if (k == K_MATMUL_F16) return "matmul_f16";
if (k == K_MATMUL_BF16) return "matmul_bf16";
if (k == K_SGD_F16) return "sgd_f16";
if (k == K_SGD_BF16) return "sgd_bf16";
if (k == K_COPY_ROWS) return "copy_rows";
if (k == K_COPY_ROWS4) return "copy_rows4";
if (k == K_TO_F16) return "to_f16";
if (k == K_TO_F16 + 1) return "to_bf16";
if (k == K_FROM_F16) return "from_f16";
if (k == K_FROM_F16 + 1) return "from_bf16";
if (k == K_RELU_F16) return "relu_f16";
if (k == K_RELU_F16 + 1) return "relu_bf16";
if (k == K_RELUG_F16) return "relug_f16";
if (k == K_RELUG_F16 + 1) return "relug_bf16";
if (k == K_MSE_F16) return "mse_f16";
if (k == K_MSE_F16 + 1) return "mse_bf16";
if (k == K_SGD2_F16) return "sgd2_f16";
if (k == K_SGD2_F16 + 1) return "sgd2_bf16";
if (k == K_GELU) return "ew_gelu";
if (k == K_RMSNORM) return "rmsnorm";
if (k == K_GEMV_F16) return "gemv_f16";
if (k == K_GEMV_F16 + 1) return "gemv_bf16";
if (k == K_GEMV1_F16) return "gemv1_f16";
if (k == K_GEMV1_F16 + 1) return "gemv1_bf16";
if (k == K_HALF_CVT) return "half_to_half";
if (k == K_CVT_MANY) return "cvt_many";
if (k == K_MP_SGD) return "mp_sgd";
if (k == K_MP_ADAM) return "mp_adam";
if (k == K_MP_ADAM_MANY) return "mp_adam_many";
if (k == K_HALF_EW) return "half_ew";
if (k == K_HALF_BIAS) return "half_bias";
if (k == K_HALF_POOL) return "half_pool";
return "upsample2x";
}
// ── Backend state (one per process) ──────────────────────────────────────────
static int g_cuda_state = 0; // 0 = not tried, 1 = ready, -1 = unavailable
static CuContext_* g_cuda_ctx = (CuContext_*)0;
static CudaNativeStream* g_cuda_stream = (CudaNativeStream*)0;
static CublasHandle_* g_cuda_blas = (CublasHandle_*)0;
static CublasLtHandle_* g_cuda_lt = (CublasLtHandle_*)0;
static unsigned long long g_cuda_ws = 0ULL;
static unsigned char* g_cuda_staging = (unsigned char*)0;
// Pinned transfer lanes (allocated on the first large transfer). Lane k
// has its own stream, two pinned buffers g_cuda_pipe[2k + b], each with an
// event recorded after the last copy that used it (g_cuda_pipe_busy = 1
// while that copy may still be in flight), and an end event.
static int g_cuda_pipe_state = 0; // 0 = not tried, 1 = ready, -1 = unavailable
static unsigned char* g_cuda_pipe[CUDA_PIPE_LANES * 2];
static CuEvent_* g_cuda_pipe_ev[CUDA_PIPE_LANES * 2];
static int g_cuda_pipe_busy[CUDA_PIPE_LANES * 2];
static CudaNativeStream* g_cuda_lane_stream[CUDA_PIPE_LANES];
static CuEvent_* g_cuda_lane_end[CUDA_PIPE_LANES];
static int g_cuda_lane_ok[CUDA_PIPE_LANES];
static CuEvent_* g_cuda_pipe_start = (CuEvent_*)0;
// Lanes 1.. run on their own threads (lane 0 on the caller), parked on
// g_cuda_lane_wake between copies: waking three threads costs much less
// than a whole-machine pool would, and they never contend with it.
static unsigned long long g_cuda_lane_mu = 0ULL;
static unsigned long long g_cuda_lane_wake = 0ULL;
static unsigned long long g_cuda_lane_done = 0ULL;
static unsigned long g_cuda_lane_gen = 0UL;
static unsigned long g_cuda_lane_pending = 0UL;
static unsigned long g_cuda_lane_count = 0UL; // lanes in the current copy
static unsigned long g_cuda_lane_threads = 0UL; // helper threads started
static unsigned long g_cuda_lane_ids[CUDA_PIPE_LANES];
static unsigned long g_cuda_lane_born[CUDA_PIPE_LANES]; // g_cuda_lane_gen when each was started
static ?&CudaLaneJob g_cuda_lane_job = null;
static CuFunction_* g_cuda_fn[K_COUNT];
static char g_cuda_name[256];
static int g_cuda_blocks_cap = 0; // 0 = use CUDA_MAX_BLOCKS
static int g_cuda_sms = 0; // the GPU's SM count
static int g_cuda_tf32 = 0;
// cuBLASLt plans, one per (ta, tb, m, n, k, epilogue): descriptor, three
// layouts and the heuristic's algo. -1 in g_lt_key0 = free slot.
static long g_lt_key0[LT_CACHE]; // ta | tb << 1 | tf32 << 2 | epilogue << 3
static long g_lt_m[LT_CACHE];
static long g_lt_n[LT_CACHE];
static long g_lt_k[LT_CACHE];
static CublasLtDesc_* g_lt_desc[LT_CACHE];
static CublasLtLayout_* g_lt_la[LT_CACHE];
static CublasLtLayout_* g_lt_lb[LT_CACHE];
static CublasLtLayout_* g_lt_lc[LT_CACHE];
static unsigned long g_lt_algo[LT_CACHE * LT_RESULT_WORDS];
// Compiles the kernel source for sm_<major><minor> and loads every kernel.
static int g_cuda_cc = 0; // compute capability, major * 10 + minor (90 on an H100/H200)
// Compiles CUDA C++ 'src' with NVRTC to a cubin for the selected GPU
// (-arch=sm_<cc>) and loads it; null on failure (the compile log is printed).
static ?&CudaNativeModule cuda_compile_module_(const char* src, const char* name) {
unsafe {
NvrtcProgram_* prog = (NvrtcProgram_*)0;
if (nvrtcCreateProgram(&prog, src, name, 0, (const char**)0, (const char**)0) != 0) return null;
char arch[32];
int v = g_cuda_cc;
// "-arch=sm_" + two or three digits
const char* prefix = "-arch=sm_";
int p = 0;
while (prefix[p] != (char)0) { arch[p] = prefix[p]; p = p + 1; }
if (v >= 100) { arch[p] = (char)(48 + v / 100); p = p + 1; }
arch[p] = (char)(48 + (v / 10) % 10); p = p + 1;
arch[p] = (char)(48 + v % 10); p = p + 1;
arch[p] = (char)0;
const char* opts[3];
opts[0] = (const char*)&arch[0];
opts[1] = "-std=c++17";
opts[2] = "-default-device";
if (nvrtcCompileProgram(prog, 3, &opts[0]) != 0) {
size_t logSize = 0UL;
nvrtcGetProgramLogSize(prog, &logSize);
char* log = (char*)alloc_raw_((unsigned long)logSize + 1UL);
if (log != (char*)0) {
nvrtcGetProgramLog(prog, log);
printf("std::cuda: kernel compile failed (%s, %s):\n%s\n", name, (const char*)&arch[0], (const char*)log);
dealloc(log);
}
nvrtcDestroyProgram(&prog);
return null;
}
size_t size = 0UL;
nvrtcGetCUBINSize(prog, &size);
char* cubin = (char*)alloc_raw_(size);
if (cubin == (char*)0) { nvrtcDestroyProgram(&prog); return null; }
nvrtcGetCUBIN(prog, cubin);
nvrtcDestroyProgram(&prog);
CudaNativeModule* module = (CudaNativeModule*)0;
int loaded = cuModuleLoadData(&module, (const char*)cubin);
dealloc(cubin);
if (loaded != CUDA_OK) return null;
return module;
}
}
// Compiles + loads this file's kernels for the selected GPU.
static int cuda_build_kernels_() {
unsafe {
CudaNativeModule* module = (CudaNativeModule*)cuda_compile_module_(cuda_kernel_source_(), "safec_cuda.cu");
if (module == (CudaNativeModule*)0) return 0;
int k = 0;
while (k < K_COUNT) {
if (cuModuleGetFunction(&g_cuda_fn[k], module, cuda_kernel_name_(k)) != CUDA_OK) return 0;
k = k + 1;
}
}
return 1;
}
// Initializes everything once: picks the GPU with the most SMs, makes its
// primary context current, creates the stream, the cuBLAS / cuBLASLt
// handles and their shared workspace, the pinned read buffer, and compiles
// + loads the kernels. Later calls only make sure the context is current
// on the calling thread.
static int cuda_ensure_() {
if (g_cuda_state == 1) {
unsafe {
CuContext_* cur = (CuContext_*)0;
cuCtxGetCurrent(&cur);
if (cur != g_cuda_ctx) cuCtxSetCurrent(g_cuda_ctx);
}
return 1;
}
if (g_cuda_state != 0) return 0;
g_cuda_state = -1;
unsafe {
g_cuda_name[0] = (char)0;
if (cuInit(0U) != CUDA_OK) return 0;
int count = 0;
if (cuDeviceGetCount(&count) != CUDA_OK || count <= 0) return 0;
int best = -1;
int bestSms = -1;
int i = 0;
while (i < count) {
int dev = 0;
int sms = 0;
if (cuDeviceGet(&dev, i) == CUDA_OK &&
cuDeviceGetAttribute(&sms, CUDA_ATTR_SM_COUNT, dev) == CUDA_OK && sms > bestSms) {
bestSms = sms;
best = dev;
}
i = i + 1;
}
if (best < 0) return 0;
cuDeviceGetName(&g_cuda_name[0], 256, best);
int major = 0;
int minor = 0;
cuDeviceGetAttribute(&major, CUDA_ATTR_CC_MAJOR, best);
cuDeviceGetAttribute(&minor, CUDA_ATTR_CC_MINOR, best);
g_cuda_blocks_cap = bestSms * 8;
g_cuda_sms = bestSms;
if (cuDevicePrimaryCtxRetain(&g_cuda_ctx, best) != CUDA_OK) return 0;
if (cuCtxSetCurrent(g_cuda_ctx) != CUDA_OK) return 0;
if (cuStreamCreate(&g_cuda_stream, CUDA_STREAM_NON_BLOCKING) != CUDA_OK) return 0;
if (cuMemAlloc_v2(&g_cuda_ws, CUDA_WORKSPACE_BYTES) != CUDA_OK) return 0;
if (cuMemAllocHost_v2(&g_cuda_staging, CUDA_STAGING_BYTES) != CUDA_OK) return 0;
if (cublasCreate_v2(&g_cuda_blas) != 0) return 0;
cublasSetStream_v2(g_cuda_blas, g_cuda_stream);
cublasSetWorkspace_v2(g_cuda_blas, (CuDevPtr_*)(unsigned long)g_cuda_ws, CUDA_WORKSPACE_BYTES);
if (cublasLtCreate(&g_cuda_lt) != 0) g_cuda_lt = (CublasLtHandle_*)0; // optional: gemm_relu falls back
int s = 0;
while (s < LT_CACHE) { g_lt_key0[s] = -1L; s = s + 1; }
g_cuda_cc = major * 10 + minor;
if (!cuda_build_kernels_()) return 0;
}
g_cuda_state = 1;
return 1;
}
int cuda_available() { return cuda_ensure_(); }
const &char cuda_device_name() {
cuda_ensure_();
unsafe { return (const char*)&g_cuda_name[0]; }
}
// ── Buffers ──────────────────────────────────────────────────────────────────
static unsigned long long cuda_alloc_bytes_(unsigned long bytes) {
if (bytes == 0UL || !cuda_ensure_()) return 0ULL;
unsafe {
unsigned long long p = 0ULL;
if (cuMemAlloc_v2(&p, bytes) != CUDA_OK) return 0ULL;
return p;
}
}
// ── Host <-> device copies ───────────────────────────────────────────────────
// A copy from pageable memory makes the driver bounce it through its own
// pinned buffers on one thread, which bounds it at one core's memcpy rate
// (~16 GB/s here, a third of what PCIe 5 moves). Large copies instead split
// into CUDA_PIPE_LANES slices that the thread pool runs side by side: each
// lane copies its slice chunk by chunk between host memory and its two
// pinned buffers in turn, on its own stream, so the DMA of one chunk
// overlaps the copy of the next. Events order the lanes after the work
// already queued on the main stream, and (for writes) the main stream
// after the lanes.
static int cuda_pipe_ensure_() {
if (g_cuda_pipe_state != 0) return g_cuda_pipe_state == 1;
g_cuda_pipe_state = -1;
unsafe {
if (cuEventCreate(&g_cuda_pipe_start, CUDA_EVENT_NO_TIMING) != CUDA_OK) return 0;
int i = 0;
while (i < CUDA_PIPE_LANES * 2) {
g_cuda_pipe_busy[i] = 0;
if (cuMemAllocHost_v2(&g_cuda_pipe[i], CUDA_PIPE_CHUNK) != CUDA_OK ||
cuEventCreate(&g_cuda_pipe_ev[i], CUDA_EVENT_NO_TIMING) != CUDA_OK) return 0;
i = i + 1;
}
i = 0;
while (i < CUDA_PIPE_LANES) {
if (cuStreamCreate(&g_cuda_lane_stream[i], CUDA_STREAM_NON_BLOCKING) != CUDA_OK ||
cuEventCreate(&g_cuda_lane_end[i], CUDA_EVENT_NO_TIMING) != CUDA_OK) return 0;
i = i + 1;
}
}
g_cuda_pipe_state = 1;
return 1;
}
// Waits until no queued copy uses pinned buffer i.
static int cuda_pipe_wait_(int i) {
if (!g_cuda_pipe_busy[i]) return 1;
g_cuda_pipe_busy[i] = 0;
unsafe { return cuEventSynchronize(g_cuda_pipe_ev[i]) == CUDA_OK; }
}
struct CudaLaneJob {
char* host;
unsigned long long dev;
unsigned long bytes;
unsigned long part; // bytes per lane
int toDevice;
};
// Lane k's slice: host -> device (queued; returns once the host bytes are
// staged) or device -> host (returns once they've landed).
static int cuda_lane_run_(&CudaLaneJob j, int k) {
unsigned long base = (unsigned long)k * j.part;
if (base >= j.bytes) return 1;
unsigned long len = j.bytes - base < j.part ? j.bytes - base : j.part;
unsafe {
CuContext_* cur = (CuContext_*)0;
cuCtxGetCurrent(&cur);
if (cur != g_cuda_ctx && cuCtxSetCurrent(g_cuda_ctx) != CUDA_OK) return 0; // pool threads
CudaNativeStream* st = g_cuda_lane_stream[k];
if (cuStreamWaitEvent(st, g_cuda_pipe_start, 0U) != CUDA_OK) return 0;
char* host = j.host + base;
unsigned long long dev = j.dev + (unsigned long long)base;
unsigned long off = 0UL;
unsigned long prevOff = 0UL;
unsigned long prevLen = 0UL;
int b = 0;
while (off < len) {
unsigned long c = len - off < CUDA_PIPE_CHUNK ? len - off : CUDA_PIPE_CHUNK;
int i = 2 * k + b;
if (!cuda_pipe_wait_(i)) return 0;
if (j.toDevice) {
memcpy(g_cuda_pipe[i], (const unsigned char*)(host + off), c);
if (cuMemcpyHtoDAsync_v2(dev + (unsigned long long)off, (const unsigned char*)g_cuda_pipe[i], c, st) != CUDA_OK)
return 0;
} else {
if (cuMemcpyDtoHAsync_v2(g_cuda_pipe[i], dev + (unsigned long long)off, c, st) != CUDA_OK) return 0;
}
if (cuEventRecord(g_cuda_pipe_ev[i], st) != CUDA_OK) return 0;
g_cuda_pipe_busy[i] = 1;
if (!j.toDevice && prevLen > 0UL) {
int p = 2 * k + (1 - b);
if (!cuda_pipe_wait_(p)) return 0;
memcpy((unsigned char*)(host + prevOff), (const unsigned char*)g_cuda_pipe[p], prevLen);
}
prevOff = off;
prevLen = c;
off = off + c;
b = 1 - b;
}
if (!j.toDevice) {
int p = 2 * k + (1 - b);
if (!cuda_pipe_wait_(p)) return 0;
memcpy((unsigned char*)(host + prevOff), (const unsigned char*)g_cuda_pipe[p], prevLen);
}
return cuEventRecord(g_cuda_lane_end[k], st) == CUDA_OK;
}
}
static ?&unsigned long cuda_lane_worker_(?&unsigned long arg) {
unsigned long k = 0UL;
match (arg) { case some(p): { k = *p; } case null: { } }
mutex_lock(&g_cuda_lane_mu);
unsigned long seen = g_cuda_lane_born[k]; // not g_cuda_lane_gen: a copy may already be posted
while (1) {
while (g_cuda_lane_gen == seen) cond_wait(&g_cuda_lane_wake, &g_cuda_lane_mu);
seen = g_cuda_lane_gen;
int take = k < g_cuda_lane_count;
?&CudaLaneJob job = g_cuda_lane_job;
mutex_unlock(&g_cuda_lane_mu);
if (take) {
match (job) {
case some(jb): { g_cuda_lane_ok[k] = cuda_lane_run_(jb, (int)k); }
case null: { g_cuda_lane_ok[k] = 0; }
}
}
mutex_lock(&g_cuda_lane_mu);
if (take) {
g_cuda_lane_pending = g_cuda_lane_pending - 1UL;
if (g_cuda_lane_pending == 0UL) cond_signal(&g_cuda_lane_done);
}
}
return null;
}
// Starts the helper threads on first use; returns how many lanes can run
// (1 = only the caller's, if no thread could be started).
static int g_cuda_lane_inited = 0;
static unsigned long cuda_lanes_start_() {
if (!g_cuda_lane_inited) {
g_cuda_lane_inited = 1;
if (mutex_init(&g_cuda_lane_mu) != 0 || cond_init(&g_cuda_lane_wake) != 0 ||
cond_init(&g_cuda_lane_done) != 0) return 1UL;
unsigned long k = 1UL;
while (k < (unsigned long)CUDA_PIPE_LANES) {
g_cuda_lane_ids[k] = k;
g_cuda_lane_born[k] = g_cuda_lane_gen;
unsigned long long tid = 0ULL;
if (thread_create(&tid, cuda_lane_worker_, (?&unsigned long)&g_cuda_lane_ids[k]) != 0) break;
thread_detach(tid);
g_cuda_lane_threads = k;
k = k + 1UL;
}
}
return g_cuda_lane_threads + 1UL;
}
// Runs a large copy over the lanes; 1 on success.
static int cuda_lanes_copy_(unsigned char* host, unsigned long long dev, unsigned long bytes, int toDevice) {
struct CudaLaneJob j;
unsafe { j.host = (char*)host; }
j.dev = dev;
j.bytes = bytes;
j.toDevice = toDevice;
unsigned long lanes = (bytes + CUDA_PIPE_MIN - 1UL) / CUDA_PIPE_MIN; // >= 1 MiB per lane
unsigned long avail = cuda_lanes_start_();
if (lanes > avail) lanes = avail;
j.part = (bytes / lanes + 4095UL) & ~4095UL;
unsafe {
if (cuEventRecord(g_cuda_pipe_start, g_cuda_stream) != CUDA_OK) return 0;
if (lanes > 1UL) {
mutex_lock(&g_cuda_lane_mu);
g_cuda_lane_job = (?&CudaLaneJob)&j;
g_cuda_lane_count = lanes;
g_cuda_lane_pending = lanes - 1UL;
g_cuda_lane_gen = g_cuda_lane_gen + 1UL;
cond_broadcast(&g_cuda_lane_wake);
mutex_unlock(&g_cuda_lane_mu);
}
g_cuda_lane_ok[0] = cuda_lane_run_(&j, 0);
if (lanes > 1UL) {
mutex_lock(&g_cuda_lane_mu);
while (g_cuda_lane_pending != 0UL) cond_wait(&g_cuda_lane_done, &g_cuda_lane_mu);
mutex_unlock(&g_cuda_lane_mu);
}
}
int ok = 1;
unsigned long k = 0UL;
while (k < lanes) {
if (!g_cuda_lane_ok[k]) ok = 0;
// later work on the main stream waits for this lane's copies
unsafe { if (ok && toDevice && cuStreamWaitEvent(g_cuda_stream, g_cuda_lane_end[k], 0U) != CUDA_OK) ok = 0; }
k = k + 1UL;
}
return ok;
}
// Queues host -> device; returns once 'host' may be reused (the last chunks
// may still be on their way from pinned memory).
static int cuda_write_bytes_(unsigned long long dst, const unsigned char* host, unsigned long bytes) {
unsafe {
if (bytes < CUDA_PIPE_MIN || !cuda_pipe_ensure_())
return cuMemcpyHtoDAsync_v2(dst, host, bytes, g_cuda_stream) == CUDA_OK;
return cuda_lanes_copy_((unsigned char*)host, dst, bytes, 1);
}
}
// Waits for queued work, then copies 'bytes' to 'host'. Small reads (a
// loss, one output value) bounce through the pinned staging buffer: a
// pinned D2H copy is one DMA, where a pageable one adds a driver-side copy.
static int cuda_read_bytes_(unsigned char* host, unsigned long long src, unsigned long bytes) {
unsafe {
if (bytes <= CUDA_STAGING_BYTES) {
if (cuMemcpyDtoHAsync_v2(g_cuda_staging, src, bytes, g_cuda_stream) != CUDA_OK) return 0;
if (cuStreamSynchronize(g_cuda_stream) != CUDA_OK) return 0;
memcpy(host, (const unsigned char*)g_cuda_staging, bytes);
return 1;
}
if (bytes < CUDA_PIPE_MIN || !cuda_pipe_ensure_()) {
if (cuMemcpyDtoHAsync_v2(host, src, bytes, g_cuda_stream) != CUDA_OK) return 0;
return cuStreamSynchronize(g_cuda_stream) == CUDA_OK;
}
return cuda_lanes_copy_(host, src, bytes, 0);
}
}
int cuda_copy_to_device(unsigned long dev, []const unsigned char host) {
if (dev == 0UL || host.len == 0UL || !cuda_ensure_()) return 0;
unsafe { return cuda_write_bytes_((unsigned long long)dev, (const unsigned char*)host.ptr, host.len); }
}
int cuda_copy_to_host([]unsigned char host, unsigned long dev) {
if (dev == 0UL || host.len == 0UL || !cuda_ensure_()) return 0;
unsafe { return cuda_read_bytes_((unsigned char*)host.ptr, (unsigned long long)dev, host.len); }
}
struct CudaBuf CudaBuf::alloc(unsigned long n) {
struct CudaBuf b;
b.dev = 0UL;
b.len = 0UL;
if (n == 0UL) return b;
unsigned long long p = cuda_alloc_bytes_(checked_mul_size(n, sizeof(float)));
if (p == 0ULL) return b;
b.dev = (unsigned long)p;
b.len = n;
return b;
}
int cuda_write_n_(struct CudaBuf dst, const &float host, unsigned long n) {
if (dst.dev == 0UL || n > dst.len || !cuda_ensure_()) return 0;
unsafe { return cuda_write_bytes_((unsigned long long)dst.dev, (const unsigned char*)host, n * sizeof(float)); }
}
int cuda_write(struct CudaBuf dst, []const float host_s) {
unsigned long n = host_s.len;
float host_d_[1];
const &float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (const &float)host_s.ptr; } }
return cuda_write_n_(dst, host, n);
}
struct CudaBuf CudaBuf::upload_n_(const &float host, unsigned long n) {
struct CudaBuf b = CudaBuf::alloc(n);
if (b.dev != 0UL && !cuda_write_n_(b, host, n)) b.free();
return b;
}
struct CudaBuf CudaBuf::upload([]const float host_s) {
unsigned long n = host_s.len;
float host_d_[1];
const &float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (const &float)host_s.ptr; } }
return CudaBuf::upload_n_(host, n);
}
int cuda_read_n_(struct CudaBuf src, &float host, unsigned long n) {
if (src.dev == 0UL || n > src.len || !cuda_ensure_()) return 0;
unsafe { return cuda_read_bytes_((unsigned char*)host, (unsigned long long)src.dev, n * sizeof(float)); }
}
int cuda_read(struct CudaBuf src, []float host_s) {
unsigned long n = host_s.len;
float host_d_[1];
&float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (&float)host_s.ptr; } }
return cuda_read_n_(src, host, n);
}
void CudaBuf::free() {
if (self.dev != 0UL) { unsafe { cuMemFree_v2((unsigned long long)self.dev); } }
self.dev = 0UL;
self.len = 0UL;
}
int cuda_sync() {
if (!cuda_ensure_()) return 0;
unsafe { return cuStreamSynchronize(g_cuda_stream) == CUDA_OK; }
}
int cuda_set_tf32(int on) {
if (!cuda_ensure_()) return 0;
unsafe {
if (cublasSetMathMode(g_cuda_blas, on ? CUDA_MATH_TF32 : CUDA_MATH_DEFAULT) != 0) return 0;
}
g_cuda_tf32 = on ? 1 : 0; // cuBLASLt plans are keyed on this
return 1;
}
// ── Launch helpers ───────────────────────────────────────────────────────────
static unsigned int cuda_blocks_(unsigned long n) {
unsigned long blocks = (n + (unsigned long)CUDA_BLOCK - 1UL) / (unsigned long)CUDA_BLOCK;
unsigned long cap = g_cuda_blocks_cap > 0 ? (unsigned long)g_cuda_blocks_cap : CUDA_MAX_BLOCKS;
if (blocks > cap) blocks = cap;
if (blocks == 0UL) blocks = 1UL;
return (unsigned int)blocks;
}
static int cuda_launch_(int kernel, unsigned long n, CuKernelArg_** params) {
unsafe {
return cuLaunchKernel(g_cuda_fn[kernel], cuda_blocks_(n), 1U, 1U,
CUDA_BLOCK, 1U, 1U, 0U, g_cuda_stream,
params, (CuKernelArg_**)0) == CUDA_OK;
}
}
static int cuda_fits_(struct CudaBuf b, unsigned long n) { return b.dev != 0UL && b.len >= n; }
static int cuda_binary_raw_(int kernel, unsigned long a, unsigned long b, unsigned long out, unsigned long n) {
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)a;
CuDevPtr_* pb = (CuDevPtr_*)b;
CuDevPtr_* po = (CuDevPtr_*)out;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[4];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&pb; params[2] = (CuKernelArg_*)&po; params[3] = (CuKernelArg_*)&nn;
return cuda_launch_(kernel, n, ¶ms[0]);
}
}
static int cuda_unary_raw_(int kernel, unsigned long a, float k, unsigned long out, unsigned long n) {
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)a;
CuDevPtr_* po = (CuDevPtr_*)out;
unsigned int nn = (unsigned int)n;
float kk = k;
CuKernelArg_* params[4];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&po; params[2] = (CuKernelArg_*)&nn; params[3] = (CuKernelArg_*)&kk;
return cuda_launch_(kernel, n, ¶ms[0]);
}
}
static int cuda_binary_(int kernel, struct CudaBuf a, struct CudaBuf b, struct CudaBuf out) {
unsigned long n = out.len;
if (!cuda_ensure_() || !cuda_fits_(a, n) || !cuda_fits_(b, n) || !cuda_fits_(out, n) ||
n > 4294967295UL) return 0;
return cuda_binary_raw_(kernel, a.dev, b.dev, out.dev, n);
}
static int cuda_unary_(int kernel, struct CudaBuf a, float k, struct CudaBuf out) {
unsigned long n = out.len;
if (!cuda_ensure_() || !cuda_fits_(a, n) || !cuda_fits_(out, n) || n > 4294967295UL) return 0;
return cuda_unary_raw_(kernel, a.dev, k, out.dev, n);
}
// out1[0] = reduction over n elements; zeroes out1 first unless the grid is one block.
static int cuda_reduce_launch_(int kernel, unsigned long n, unsigned long out1, CuKernelArg_** params) {
unsafe {
if (cuda_blocks_(n) > 1U &&
cuMemsetD32Async((unsigned long long)out1, 0U, 1UL, g_cuda_stream) != CUDA_OK) return 0;
}
return cuda_launch_(kernel, n, params);
}
// Row-major sgemm on raw device addresses (sizes already checked).
static int cuda_gemm_raw_(unsigned long a, int ta, unsigned long b, int tb, unsigned long c,
unsigned long m, unsigned long n, unsigned long k) {
float one = 1.0f;
float zero = 0.0f;
// Row-major C = op(A) op(B) is column-major C^T = op(B)^T op(A)^T, so
// hand cuBLAS the operands swapped; each operand's leading dimension
// is its stored row length.
unsafe {
return cublasSgemm_v2(g_cuda_blas,
tb ? CUDA_BLAS_T : CUDA_BLAS_N, ta ? CUDA_BLAS_T : CUDA_BLAS_N,
(int)n, (int)m, (int)k, &one,
(const float*)b, tb ? (int)k : (int)n,
(const float*)a, ta ? (int)m : (int)k,
&zero, (float*)c, (int)n) == 0;
}
}
static int cuda_gemm_dims_ok_(unsigned long m, unsigned long n, unsigned long k) {
return m > 0UL && n > 0UL && k > 0UL && m <= 2147483647UL && n <= 2147483647UL && k <= 2147483647UL;
}
// ── Queued ops ───────────────────────────────────────────────────────────────
int cuda_gemm(struct CudaBuf a, int ta, struct CudaBuf b, int tb, struct CudaBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!cuda_ensure_() || !cuda_gemm_dims_ok_(m, n, k) ||
!cuda_fits_(a, m * k) || !cuda_fits_(b, k * n) || !cuda_fits_(c, m * n)) return 0;
return cuda_gemm_raw_(a.dev, ta, b.dev, tb, c.dev, m, n, k);
}
// Finds or builds the cuBLASLt plan for this shape + epilogue; -1 if none.
// 'dtype' is the element type of a, b and c (CUDA_R_32F_/16F_/16BF_); 16-bit
// types accumulate in fp32 (TF32 doesn't apply).
static int cuda_lt_plan_(int ta, int tb, unsigned long m, unsigned long n, unsigned long k, int epilogue, int dtype) {
long key0 = (long)(ta | (tb << 1) | (g_cuda_tf32 << 2) | (epilogue << 3) | (dtype << 8));
int slot = -1;
int s = 0;
while (s < LT_CACHE) {
if (g_lt_key0[s] == key0 && g_lt_m[s] == (long)m && g_lt_n[s] == (long)n && g_lt_k[s] == (long)k)
return s;
if (slot < 0 && g_lt_key0[s] == -1L) slot = s;
s = s + 1;
}
if (slot < 0 || g_cuda_lt == (CublasLtHandle_*)0) return -1;
unsafe {
// Column-major view (see cuda_gemm_raw_): D[n,m] = op(B') op(A'),
// with B' = row-major b (stored n x k if tb, else k x n, read
// column-major) and A' = row-major a.
int opA = tb ? CUDA_BLAS_T : CUDA_BLAS_N;
int opB = ta ? CUDA_BLAS_T : CUDA_BLAS_N;
CublasLtDesc_* desc = (CublasLtDesc_*)0;
int compute = g_cuda_tf32 && dtype == CUDA_R_32F_ ? CUDA_COMPUTE_32F_TF32 : CUDA_COMPUTE_32F;
if (cublasLtMatmulDescCreate(&desc, compute, CUDA_R_32F_) != 0) return -1;
cublasLtMatmulDescSetAttribute(desc, LT_DESC_TRANSA, (const unsigned char*)&opA, sizeof(int));
cublasLtMatmulDescSetAttribute(desc, LT_DESC_TRANSB, (const unsigned char*)&opB, sizeof(int));
cublasLtMatmulDescSetAttribute(desc, LT_DESC_EPILOGUE, (const unsigned char*)&epilogue, sizeof(int));
CublasLtLayout_* la = (CublasLtLayout_*)0;
CublasLtLayout_* lb = (CublasLtLayout_*)0;
CublasLtLayout_* lc = (CublasLtLayout_*)0;
unsigned long aRows = tb ? k : n;
unsigned long aCols = tb ? n : k;
unsigned long bRows = ta ? m : k;
unsigned long bCols = ta ? k : m;
if (cublasLtMatrixLayoutCreate(&la, dtype, (unsigned long long)aRows, (unsigned long long)aCols, (long long)aRows) != 0 ||
cublasLtMatrixLayoutCreate(&lb, dtype, (unsigned long long)bRows, (unsigned long long)bCols, (long long)bRows) != 0 ||
cublasLtMatrixLayoutCreate(&lc, dtype, (unsigned long long)n, (unsigned long long)m, (long long)n) != 0)
return -1;
CublasLtPref_* pref = (CublasLtPref_*)0;
if (cublasLtMatmulPreferenceCreate(&pref) != 0) return -1;
unsigned long long wsBytes = (unsigned long long)CUDA_WORKSPACE_BYTES;
cublasLtMatmulPreferenceSetAttribute(pref, LT_PREF_MAX_WORKSPACE, (const unsigned char*)&wsBytes, sizeof(unsigned long long));
int returned = 0;
int hs = cublasLtMatmulAlgoGetHeuristic(g_cuda_lt, desc, la, lb, lc, lc, pref, 1,
&g_lt_algo[slot * LT_RESULT_WORDS], &returned);
cublasLtMatmulPreferenceDestroy(pref);
if (hs != 0 || returned < 1) return -1;
g_lt_desc[slot] = desc;
g_lt_la[slot] = la;
g_lt_lb[slot] = lb;
g_lt_lc[slot] = lc;
g_lt_m[slot] = (long)m;
g_lt_n[slot] = (long)n;
g_lt_k[slot] = (long)k;
g_lt_key0[slot] = key0;
}
return slot;
}
int cuda_gemm_relu(struct CudaBuf a, int ta, struct CudaBuf b, int tb, struct CudaBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!cuda_ensure_() || !cuda_gemm_dims_ok_(m, n, k) ||
!cuda_fits_(a, m * k) || !cuda_fits_(b, k * n) || !cuda_fits_(c, m * n)) return 0;
int plan = cuda_lt_plan_(ta, tb, m, n, k, LT_EPILOGUE_RELU, CUDA_R_32F_);
if (plan >= 0) {
float one = 1.0f;
float zero = 0.0f;
unsafe {
if (cublasLtMatmul(g_cuda_lt, g_lt_desc[plan], (const float*)&one,
(const CuDevPtr_*)b.dev, g_lt_la[plan], (const CuDevPtr_*)a.dev, g_lt_lb[plan],
(const float*)&zero, (const CuDevPtr_*)c.dev, g_lt_lc[plan],
(CuDevPtr_*)c.dev, g_lt_lc[plan],
(const unsigned long*)&g_lt_algo[plan * LT_RESULT_WORDS],
(CuDevPtr_*)(unsigned long)g_cuda_ws, CUDA_WORKSPACE_BYTES, g_cuda_stream) == 0)
return 1;
}
}
// No cuBLASLt (or no algo for this shape): same result in two launches.
return cuda_gemm_raw_(a.dev, ta, b.dev, tb, c.dev, m, n, k) &&
cuda_unary_raw_(K_RELU, c.dev, 0.0f, c.dev, m * n);
}
int cuda_add(struct CudaBuf a, struct CudaBuf b, struct CudaBuf out) { return cuda_binary_(K_ADD, a, b, out); }
int cuda_sub(struct CudaBuf a, struct CudaBuf b, struct CudaBuf out) { return cuda_binary_(K_SUB, a, b, out); }
int cuda_mul(struct CudaBuf a, struct CudaBuf b, struct CudaBuf out) { return cuda_binary_(K_MUL, a, b, out); }
int cuda_div(struct CudaBuf a, struct CudaBuf b, struct CudaBuf out) { return cuda_binary_(K_DIV, a, b, out); }
int cuda_pow(struct CudaBuf a, struct CudaBuf b, struct CudaBuf out) { return cuda_binary_(K_POW, a, b, out); }
int cuda_relu_grad(struct CudaBuf z, struct CudaBuf g, struct CudaBuf out) { return cuda_binary_(K_RELU_GRAD, z, g, out); }
int cuda_scale(struct CudaBuf a, float k, struct CudaBuf out) { return cuda_unary_(K_SCALE, a, k, out); }
int cuda_relu(struct CudaBuf a, struct CudaBuf out) { return cuda_unary_(K_RELU, a, 0.0f, out); }
int cuda_gelu(struct CudaBuf a, struct CudaBuf out) { return cuda_unary_(K_GELU, a, 0.0f, out); }
int cuda_rmsnorm(struct CudaBuf x, struct CudaBuf w, struct CudaBuf y, unsigned long rows, unsigned long n, float eps) {
if (!cuda_ensure_() || rows == 0UL || n == 0UL || rows > 2147483647UL || n > 4294967295UL) return 0;
if (!cuda_fits_(x, rows * n) || !cuda_fits_(y, rows * n)) return 0;
if (w.dev != 0UL && w.len < n) return 0;
unsafe {
CuDevPtr_* px = (CuDevPtr_*)x.dev;
CuDevPtr_* pw = (CuDevPtr_*)w.dev;
CuDevPtr_* py = (CuDevPtr_*)y.dev;
unsigned int nn = (unsigned int)n;
float e = eps;
CuKernelArg_* params[5];
params[0] = (CuKernelArg_*)&px; params[1] = (CuKernelArg_*)&pw; params[2] = (CuKernelArg_*)&py; params[3] = (CuKernelArg_*)&nn; params[4] = (CuKernelArg_*)&e;
return cuLaunchKernel(g_cuda_fn[K_RMSNORM], (unsigned int)rows, 1U, 1U, CUDA_BLOCK, 1U, 1U, 0U,
g_cuda_stream, ¶ms[0], (CuKernelArg_**)0) == CUDA_OK;
}
}
int cuda_exp(struct CudaBuf a, struct CudaBuf out) { return cuda_unary_(K_EXP, a, 0.0f, out); }
int cuda_log(struct CudaBuf a, struct CudaBuf out) { return cuda_unary_(K_LOG, a, 0.0f, out); }
int cuda_sqrt(struct CudaBuf a, struct CudaBuf out) { return cuda_unary_(K_SQRT, a, 0.0f, out); }
static int cuda_sgd_raw_(int kernel, unsigned long w, unsigned long g, float lr, unsigned long n) {
unsafe {
CuDevPtr_* pw = (CuDevPtr_*)w;
CuDevPtr_* pg = (CuDevPtr_*)g;
float rate = lr;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[4];
params[0] = (CuKernelArg_*)&pw; params[1] = (CuKernelArg_*)&pg; params[2] = (CuKernelArg_*)&rate; params[3] = (CuKernelArg_*)&nn;
return cuda_launch_(kernel, n, ¶ms[0]);
}
}
int cuda_sgd(struct CudaBuf w, struct CudaBuf g, float lr) {
unsigned long n = w.len;
if (!cuda_ensure_() || !cuda_fits_(w, n) || !cuda_fits_(g, n) || n > 4294967295UL) return 0;
return cuda_sgd_raw_(K_SGD, w.dev, g.dev, lr, n);
}
int cuda_sgd2(struct CudaBuf w1, struct CudaBuf g1, struct CudaBuf w2, struct CudaBuf g2, float lr) {
unsigned long n1 = w1.len;
unsigned long n2 = w2.len;
if (!cuda_ensure_() || !cuda_fits_(w1, n1) || !cuda_fits_(g1, n1) ||
!cuda_fits_(w2, n2) || !cuda_fits_(g2, n2) || n1 + n2 > 4294967295UL) return 0;
unsafe {
CuDevPtr_* pw1 = (CuDevPtr_*)w1.dev;
CuDevPtr_* pg1 = (CuDevPtr_*)g1.dev;
CuDevPtr_* pw2 = (CuDevPtr_*)w2.dev;
CuDevPtr_* pg2 = (CuDevPtr_*)g2.dev;
unsigned int nn1 = (unsigned int)n1;
unsigned int nn2 = (unsigned int)n2;
float rate = lr;
CuKernelArg_* params[7];
params[0] = (CuKernelArg_*)&pw1; params[1] = (CuKernelArg_*)&pg1; params[2] = (CuKernelArg_*)&nn1;
params[3] = (CuKernelArg_*)&pw2; params[4] = (CuKernelArg_*)&pg2; params[5] = (CuKernelArg_*)&nn2;
params[6] = (CuKernelArg_*)&rate;
return cuda_launch_(K_SGD2, n1 + n2, ¶ms[0]);
}
}
int cuda_sum(struct CudaBuf a, struct CudaBuf out1) {
unsigned long n = a.len;
if (!cuda_ensure_() || !cuda_fits_(a, n) || !cuda_fits_(out1, 1UL) || n > 4294967295UL) return 0;
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)a.dev;
CuDevPtr_* po = (CuDevPtr_*)out1.dev;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[3];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&po; params[2] = (CuKernelArg_*)&nn;
return cuda_reduce_launch_(K_SUM, n, out1.dev, ¶ms[0]);
}
}
int cuda_mse_grad(struct CudaBuf y, struct CudaBuf t, struct CudaBuf dy, struct CudaBuf loss1) {
unsigned long n = y.len;
if (!cuda_ensure_() || !cuda_fits_(y, n) || !cuda_fits_(t, n) || !cuda_fits_(dy, n) ||
!cuda_fits_(loss1, 1UL) || n > 4294967295UL) return 0;
unsafe {
CuDevPtr_* py = (CuDevPtr_*)y.dev;
CuDevPtr_* pt = (CuDevPtr_*)t.dev;
CuDevPtr_* pd = (CuDevPtr_*)dy.dev;
CuDevPtr_* pl = (CuDevPtr_*)loss1.dev;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[5];
params[0] = (CuKernelArg_*)&py; params[1] = (CuKernelArg_*)&pt; params[2] = (CuKernelArg_*)&pd;
params[3] = (CuKernelArg_*)&pl; params[4] = (CuKernelArg_*)&nn;
return cuda_reduce_launch_(K_MSE, n, loss1.dev, ¶ms[0]);
}
}
int cuda_copy(struct CudaBuf src, unsigned long srcOff, struct CudaBuf dst, unsigned long dstOff, unsigned long n) {
if (!cuda_ensure_() || n == 0UL || !cuda_fits_(src, srcOff + n) || !cuda_fits_(dst, dstOff + n)) return 0;
unsafe {
return cuMemcpyDtoDAsync_v2((unsigned long long)(dst.dev + dstOff * sizeof(float)),
(unsigned long long)(src.dev + srcOff * sizeof(float)),
n * sizeof(float), g_cuda_stream) == CUDA_OK;
}
}
int cuda_upsample2x_nearest(struct CudaBuf in, struct CudaBuf out,
unsigned long planes, unsigned long h, unsigned long w) {
unsigned long n = planes * h * w * 4UL;
if (!cuda_ensure_() || n == 0UL || n > 4294967295UL ||
!cuda_fits_(in, planes * h * w) || !cuda_fits_(out, n)) return 0;
unsafe {
CuDevPtr_* pi = (CuDevPtr_*)in.dev;
CuDevPtr_* po = (CuDevPtr_*)out.dev;
unsigned int np = (unsigned int)planes;
unsigned int nh = (unsigned int)h;
unsigned int nw = (unsigned int)w;
CuKernelArg_* params[5];
params[0] = (CuKernelArg_*)π params[1] = (CuKernelArg_*)&po; params[2] = (CuKernelArg_*)&np;
params[3] = (CuKernelArg_*)&nh; params[4] = (CuKernelArg_*)&nw;
return cuda_launch_(K_UPSAMPLE2X, planes * h * w, ¶ms[0]); // a thread per input
}
}
int cuda_copy_rows(struct CudaBuf src, unsigned long srcOff, unsigned long srcStride,
struct CudaBuf dst, unsigned long dstOff, unsigned long dstStride,
unsigned long rowLen, unsigned long rows) {
if (!cuda_ensure_() || rowLen == 0UL || rows == 0UL || rowLen > 4294967295UL || rows > 4294967295UL ||
(rows > 1UL && (rowLen > srcStride || rowLen > dstStride)) ||
!cuda_fits_(src, srcOff + (rows - 1UL) * srcStride + rowLen) ||
!cuda_fits_(dst, dstOff + (rows - 1UL) * dstStride + rowLen)) return 0;
// float4 when every offset, stride and the row length are multiples of 4
// (CudaBuf bases are 256-byte aligned)
int v4 = ((srcOff | srcStride | dstOff | dstStride | rowLen) & 3UL) == 0UL;
unsigned long q = v4 ? 4UL : 1UL;
unsafe {
CuDevPtr_* ps = (CuDevPtr_*)src.dev;
CuDevPtr_* pd = (CuDevPtr_*)dst.dev;
unsigned long long so = (unsigned long long)(srcOff / q);
unsigned long long ss = (unsigned long long)(srcStride / q);
unsigned long long dof = (unsigned long long)(dstOff / q);
unsigned long long ds = (unsigned long long)(dstStride / q);
unsigned int len = (unsigned int)(rowLen / q);
unsigned int nr = (unsigned int)rows;
CuKernelArg_* params[8];
params[0] = (CuKernelArg_*)&ps; params[1] = (CuKernelArg_*)&pd; params[2] = (CuKernelArg_*)&so; params[3] = (CuKernelArg_*)&ss;
params[4] = (CuKernelArg_*)&dof; params[5] = (CuKernelArg_*)&ds; params[6] = (CuKernelArg_*)&len; params[7] = (CuKernelArg_*)&nr;
return cuda_launch_(v4 ? K_COPY_ROWS4 : K_COPY_ROWS, (rowLen / q) * rows, ¶ms[0]);
}
}
// ── Interop ──────────────────────────────────────────────────────────────────
?&CudaNativeStream cuda_stream_handle() {
if (!cuda_ensure_()) return null;
unsafe { return (?&CudaNativeStream)g_cuda_stream; }
}
int cuda_tf32_enabled() { return g_cuda_tf32; }
int cuda_compute_capability() {
if (!cuda_ensure_()) return 0;
return g_cuda_cc;
}
int cuda_sm_count() {
if (!cuda_ensure_()) return 0;
return g_cuda_sms;
}
?&CudaNativeModule cuda_load_module(const &char source, const &char name) {
if (!cuda_ensure_()) return null;
unsafe { return (?&CudaNativeModule)cuda_compile_module_((const char*)source, (const char*)name); }
}
// ── Graphs ───────────────────────────────────────────────────────────────────
int cuda_capture_begin() {
if (!cuda_ensure_()) return 0;
unsafe { return cuStreamBeginCapture_v2(g_cuda_stream, CUDA_CAPTURE_GLOBAL) == CUDA_OK; }
}
struct CudaGraph cuda_capture_end() {
struct CudaGraph g;
g.exec = 0UL;
if (!cuda_ensure_()) return g;
unsafe {
CuGraph_* graph = (CuGraph_*)0;
if (cuStreamEndCapture(g_cuda_stream, &graph) != CUDA_OK) return g;
CuGraphExec_* exec = (CuGraphExec_*)0;
int ok = cuGraphInstantiateWithFlags(&exec, graph, 0ULL) == CUDA_OK;
cuGraphDestroy(graph);
if (ok) g.exec = (unsigned long)exec;
}
return g;
}
int cuda_graph_launch(struct CudaGraph g) {
if (g.exec == 0UL || !cuda_ensure_()) return 0;
unsafe { return cuGraphLaunch((CuGraphExec_*)g.exec, g_cuda_stream) == CUDA_OK; }
}
void CudaGraph::free() {
if (self.exec != 0UL) { unsafe { cuGraphExecDestroy((CuGraphExec_*)self.exec); } }
self.exec = 0UL;
}
// ── Host-array convenience ops ───────────────────────────────────────────────
static int cuda_host_binary_(int kernel, const float* a, const float* b, float* out, unsigned long n) {
if (n == 0UL || n > 4294967295UL || !cuda_ensure_()) return 0;
unsigned long bytes = checked_mul_size(n, sizeof(float));
unsigned long long da = cuda_alloc_bytes_(bytes);
unsigned long long db = cuda_alloc_bytes_(bytes);
unsigned long long dout = cuda_alloc_bytes_(bytes);
int ok = 0;
unsafe {
ok = da != 0ULL && db != 0ULL && dout != 0ULL &&
cuda_write_bytes_(da, (const unsigned char*)a, bytes) && cuda_write_bytes_(db, (const unsigned char*)b, bytes) &&
cuda_binary_raw_(kernel, (unsigned long)da, (unsigned long)db, (unsigned long)dout, n) &&
cuda_read_bytes_((unsigned char*)out, dout, bytes);
if (da != 0ULL) cuMemFree_v2(da);
if (db != 0ULL) cuMemFree_v2(db);
if (dout != 0ULL) cuMemFree_v2(dout);
}
return ok;
}
static int cuda_host_unary_(int kernel, const float* a, float k, float* out, unsigned long n) {
if (n == 0UL || n > 4294967295UL || !cuda_ensure_()) return 0;
unsigned long bytes = checked_mul_size(n, sizeof(float));
unsigned long long da = cuda_alloc_bytes_(bytes);
unsigned long long dout = cuda_alloc_bytes_(bytes);
int ok = 0;
unsafe {
ok = da != 0ULL && dout != 0ULL && cuda_write_bytes_(da, (const unsigned char*)a, bytes) &&
cuda_unary_raw_(kernel, (unsigned long)da, k, (unsigned long)dout, n) &&
cuda_read_bytes_((unsigned char*)out, dout, bytes);
if (da != 0ULL) cuMemFree_v2(da);
if (dout != 0ULL) cuMemFree_v2(dout);
}
return ok;
}
int cuda_add_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return cuda_host_binary_(K_ADD, a, b, out, n); }
int cuda_add_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_add_f32_n_(a, b, out, n);
}
int cuda_sub_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return cuda_host_binary_(K_SUB, a, b, out, n); }
int cuda_sub_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_sub_f32_n_(a, b, out, n);
}
int cuda_mul_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return cuda_host_binary_(K_MUL, a, b, out, n); }
int cuda_mul_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_mul_f32_n_(a, b, out, n);
}
int cuda_div_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return cuda_host_binary_(K_DIV, a, b, out, n); }
int cuda_div_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_div_f32_n_(a, b, out, n);
}
int cuda_pow_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return cuda_host_binary_(K_POW, a, b, out, n); }
int cuda_pow_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_pow_f32_n_(a, b, out, n);
}
int cuda_relu_f32_n_(const &float a, &float out, unsigned long n) { return cuda_host_unary_(K_RELU, a, 0.0f, out, n); }
int cuda_relu_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_relu_f32_n_(a, out, n);
}
int cuda_log_f32_n_(const &float a, &float out, unsigned long n) { return cuda_host_unary_(K_LOG, a, 0.0f, out, n); }
int cuda_log_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_log_f32_n_(a, out, n);
}
int cuda_exp_f32_n_(const &float a, &float out, unsigned long n) { return cuda_host_unary_(K_EXP, a, 0.0f, out, n); }
int cuda_exp_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_exp_f32_n_(a, out, n);
}
int cuda_sqrt_f32_n_(const &float a, &float out, unsigned long n) { return cuda_host_unary_(K_SQRT, a, 0.0f, out, n); }
int cuda_sqrt_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_sqrt_f32_n_(a, out, n);
}
int cuda_scale_f32_n_(const &float a, float k, &float out, unsigned long n) { return cuda_host_unary_(K_SCALE, a, k, out, n); }
int cuda_scale_f32([]const float a_s, float k, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_scale_f32_n_(a, k, out, n);
}
int cuda_sum_f32_n_(const &float a, &float out, unsigned long n) {
if (n == 0UL || n > 4294967295UL || !cuda_ensure_()) return 0;
unsigned long bytes = checked_mul_size(n, sizeof(float));
unsigned long long da = cuda_alloc_bytes_(bytes);
unsigned long long dout = cuda_alloc_bytes_(sizeof(float));
int ok = 0;
unsafe {
if (da != 0ULL && dout != 0ULL && cuda_write_bytes_(da, (const unsigned char*)a, bytes)) {
CuDevPtr_* pa = (CuDevPtr_*)(unsigned long)da;
CuDevPtr_* po = (CuDevPtr_*)(unsigned long)dout;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[3];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&po; params[2] = (CuKernelArg_*)&nn;
ok = cuda_reduce_launch_(K_SUM, n, (unsigned long)dout, ¶ms[0]) &&
cuda_read_bytes_((unsigned char*)out, dout, sizeof(float));
}
if (da != 0ULL) cuMemFree_v2(da);
if (dout != 0ULL) cuMemFree_v2(dout);
}
return ok;
}
int cuda_sum_f32([]const float a_s, []float out_s) {
unsigned long n = a_s.len;
if (out_s.len < 1UL) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_sum_f32_n_(a, out, n);
}
// Launches sgemm_nn on raw device addresses.
static int cuda_sgemm_kernel_(unsigned long a, unsigned long b, unsigned long c,
unsigned long M, unsigned long K, unsigned long N) {
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)a;
CuDevPtr_* pb = (CuDevPtr_*)b;
CuDevPtr_* pc = (CuDevPtr_*)c;
int m = (int)M;
int n = (int)N;
int k = (int)K;
CuKernelArg_* params[6];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&pb; params[2] = (CuKernelArg_*)&pc;
params[3] = (CuKernelArg_*)&m; params[4] = (CuKernelArg_*)&n; params[5] = (CuKernelArg_*)&k;
return cuLaunchKernel(g_cuda_fn[K_SGEMM],
(unsigned int)((N + 63UL) / 64UL), (unsigned int)((M + 63UL) / 64UL), 1U,
CUDA_BLOCK, 1U, 1U, 0U, g_cuda_stream, ¶ms[0], (CuKernelArg_**)0) == CUDA_OK;
}
}
// Uploads a and b, runs the custom kernel (useBlas = 0) or cuBLAS, reads out.
static int cuda_host_matmul_(const float* a, const float* b, float* out,
unsigned long M, unsigned long K, unsigned long N, int useBlas) {
if (!cuda_gemm_dims_ok_(M, N, K) || !cuda_ensure_()) return 0;
unsigned long bytesA = checked_mul_size(checked_mul_size(M, K), sizeof(float));
unsigned long bytesB = checked_mul_size(checked_mul_size(K, N), sizeof(float));
unsigned long bytesC = checked_mul_size(checked_mul_size(M, N), sizeof(float));
unsigned long long da = cuda_alloc_bytes_(bytesA);
unsigned long long db = cuda_alloc_bytes_(bytesB);
unsigned long long dc = cuda_alloc_bytes_(bytesC);
int ok = 0;
unsafe {
if (da != 0ULL && db != 0ULL && dc != 0ULL &&
cuda_write_bytes_(da, (const unsigned char*)a, bytesA) && cuda_write_bytes_(db, (const unsigned char*)b, bytesB)) {
if (useBlas) ok = cuda_gemm_raw_((unsigned long)da, 0, (unsigned long)db, 0, (unsigned long)dc, M, N, K);
else ok = cuda_sgemm_kernel_((unsigned long)da, (unsigned long)db, (unsigned long)dc, M, K, N);
ok = ok && cuda_read_bytes_((unsigned char*)out, dc, bytesC);
}
if (da != 0ULL) cuMemFree_v2(da);
if (db != 0ULL) cuMemFree_v2(db);
if (dc != 0ULL) cuMemFree_v2(dc);
}
return ok;
}
int cuda_matmul_f32_n_(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
return cuda_host_matmul_(a, b, out, M, K, N, 0);
}
int cuda_matmul_f32([]const float a_s, []const float b_s, []float out_s, unsigned long M, unsigned long K, unsigned long N) {
if (a_s.len < M * K || b_s.len < K * N || out_s.len < M * N) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_matmul_f32_n_(a, b, out, M, K, N);
}
int cuda_matmul_f32_blas_n_(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
return cuda_host_matmul_(a, b, out, M, K, N, 1);
}
int cuda_matmul_f32_blas([]const float a_s, []const float b_s, []float out_s, unsigned long M, unsigned long K, unsigned long N) {
if (a_s.len < M * K || b_s.len < K * N || out_s.len < M * N) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return cuda_matmul_f32_blas_n_(a, b, out, M, K, N);
}
// ── fp16 / bf16 buffers (CudaHalfBuf) ────────────────────────────────────────
static int cuda_half_ok_(struct CudaHalfBuf b, unsigned long n) {
return b.dev != 0UL && b.len >= n && (b.dtype == CUDA_FP16 || b.dtype == CUDA_BF16);
}
// Kernel index for a CudaHalfBuf op: the fp16 kernel or its bf16 twin.
static int cuda_k16_(int fp16Kernel, int dtype) { return dtype == CUDA_BF16 ? fp16Kernel + 1 : fp16Kernel; }
static int cuda_dtype_type_(int dtype) { return dtype == CUDA_BF16 ? CUDA_R_16BF_ : CUDA_R_16F_; }
struct CudaHalfBuf CudaHalfBuf::alloc(unsigned long n, int dtype) {
struct CudaHalfBuf b;
b.dev = 0UL;
b.len = 0UL;
b.dtype = dtype;
if (n == 0UL || (dtype != CUDA_FP16 && dtype != CUDA_BF16)) return b;
unsigned long long p = cuda_alloc_bytes_(checked_mul_size(n, 2UL));
if (p == 0ULL) return b;
b.dev = (unsigned long)p;
b.len = n;
return b;
}
struct CudaHalfBuf CudaHalfBuf::upload_n_(const &float host, unsigned long n, int dtype) {
struct CudaHalfBuf b = CudaHalfBuf::alloc(n, dtype);
if (b.dev == 0UL) return b;
struct CudaBuf tmp = CudaBuf::upload_n_(host, n);
int ok = tmp.dev != 0UL && cuda_to_half(tmp, b);
if (ok) ok = cuda_sync(); // tmp is freed next
tmp.free();
if (!ok) b.free();
return b;
}
struct CudaHalfBuf CudaHalfBuf::upload([]const float host_s, int dtype) {
unsigned long n = host_s.len;
float host_d_[1];
const &float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (const &float)host_s.ptr; } }
return CudaHalfBuf::upload_n_(host, n, dtype);
}
int cuda_half_write_n_(struct CudaHalfBuf dst, const &unsigned short bits, unsigned long n) {
if (!cuda_half_ok_(dst, n) || n == 0UL || !cuda_ensure_()) return 0;
unsafe { return cuda_write_bytes_((unsigned long long)dst.dev, (const unsigned char*)bits, n * 2UL); }
}
int cuda_half_write(struct CudaHalfBuf dst, []const unsigned short bits_s) {
unsigned long n = bits_s.len;
unsigned short bits_d_[1];
const &unsigned short bits = &bits_d_[0];
if (bits_s.len > 0UL) { unsafe { bits = (const &unsigned short)bits_s.ptr; } }
return cuda_half_write_n_(dst, bits, n);
}
int cuda_half_read_n_(struct CudaHalfBuf src, &unsigned short bits, unsigned long n) {
if (!cuda_half_ok_(src, n) || n == 0UL || !cuda_ensure_()) return 0;
unsafe { return cuda_read_bytes_((unsigned char*)bits, (unsigned long long)src.dev, n * 2UL); }
}
int cuda_half_read(struct CudaHalfBuf src, []unsigned short bits_s) {
unsigned long n = bits_s.len;
unsigned short bits_d_[1];
&unsigned short bits = &bits_d_[0];
if (bits_s.len > 0UL) { unsafe { bits = (&unsigned short)bits_s.ptr; } }
return cuda_half_read_n_(src, bits, n);
}
void CudaHalfBuf::free() {
if (self.dev != 0UL) { unsafe { cuMemFree_v2((unsigned long long)self.dev); } }
self.dev = 0UL;
self.len = 0UL;
}
// Launches a two-pointer CudaHalfBuf kernel over n elements: (a, out, n).
static int cuda_launch2_(int kernel, unsigned long a, unsigned long out, unsigned long n) {
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)a;
CuDevPtr_* po = (CuDevPtr_*)out;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[3];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&po; params[2] = (CuKernelArg_*)&nn;
return cuda_launch_(kernel, n, ¶ms[0]);
}
}
int cuda_to_half(struct CudaBuf src, struct CudaHalfBuf dst) {
unsigned long n = dst.len;
if (!cuda_ensure_() || !cuda_half_ok_(dst, n) || !cuda_fits_(src, n) || n > 4294967295UL) return 0;
return cuda_launch2_(cuda_k16_(K_TO_F16, dst.dtype), src.dev, dst.dev, n);
}
int cuda_from_half(struct CudaHalfBuf src, struct CudaBuf dst) {
unsigned long n = src.len;
if (!cuda_ensure_() || !cuda_half_ok_(src, n) || !cuda_fits_(dst, n) || n > 4294967295UL) return 0;
return cuda_launch2_(cuda_k16_(K_FROM_F16, src.dtype), src.dev, dst.dev, n);
}
// c = op(a) op(b) through cublasGemmEx: a, b of 'dtype', c of cType, fp32
// accumulation (row-major via the column-major swap of cuda_gemm_raw_).
static int cuda_gemm_ex_(unsigned long a, int ta, unsigned long b, int tb, unsigned long c, int abType, int cType,
unsigned long m, unsigned long n, unsigned long k) {
float one = 1.0f;
float zero = 0.0f;
unsafe {
return cublasGemmEx(g_cuda_blas, tb ? CUDA_BLAS_T : CUDA_BLAS_N, ta ? CUDA_BLAS_T : CUDA_BLAS_N,
(int)n, (int)m, (int)k, (const float*)&one,
(const CuDevPtr_*)b, abType, tb ? (int)k : (int)n,
(const CuDevPtr_*)a, abType, ta ? (int)m : (int)k,
(const float*)&zero, (CuDevPtr_*)c, cType, (int)n,
CUDA_COMPUTE_32F, CUDA_GEMM_DEFAULT_) == 0;
}
}
static int cuda_half_gemm_ok_(struct CudaHalfBuf a, struct CudaHalfBuf b, unsigned long m, unsigned long n,
unsigned long k) {
return cuda_ensure_() && cuda_gemm_dims_ok_(m, n, k) && a.dtype == b.dtype &&
cuda_half_ok_(a, m * k) && cuda_half_ok_(b, k * n);
}
int cuda_half_gemm_strided(struct CudaHalfBuf a, int ta, unsigned long strideA, struct CudaHalfBuf b, int tb,
unsigned long strideB, struct CudaHalfBuf c, unsigned long strideC,
unsigned long m, unsigned long n, unsigned long k, unsigned long batch) {
if (!cuda_ensure_() || !cuda_gemm_dims_ok_(m, n, k) || batch == 0UL || batch > 2147483647UL ||
a.dtype != b.dtype || c.dtype != a.dtype || strideC < m * n) return 0;
unsigned long last = batch - 1UL;
if (!cuda_half_ok_(a, last * strideA + m * k) || !cuda_half_ok_(b, last * strideB + k * n) ||
!cuda_half_ok_(c, last * strideC + m * n)) return 0;
int t = cuda_dtype_type_(a.dtype);
float one = 1.0f;
float zero = 0.0f;
// Row-major, so each product goes to cuBLAS as C^T = op(B)^T op(A)^T (see cuda_gemm_raw_).
unsafe {
return cublasGemmStridedBatchedEx(g_cuda_blas, tb ? CUDA_BLAS_T : CUDA_BLAS_N, ta ? CUDA_BLAS_T : CUDA_BLAS_N,
(int)n, (int)m, (int)k, (const float*)&one,
(const CuDevPtr_*)b.dev, t, tb ? (int)k : (int)n, (long long)strideB,
(const CuDevPtr_*)a.dev, t, ta ? (int)m : (int)k, (long long)strideA,
(const float*)&zero, (CuDevPtr_*)c.dev, t, (int)n, (long long)strideC,
(int)batch, CUDA_COMPUTE_32F, CUDA_GEMM_DEFAULT_) == 0;
}
}
int cuda_half_gemm(struct CudaHalfBuf a, int ta, struct CudaHalfBuf b, int tb, struct CudaHalfBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!cuda_half_gemm_ok_(a, b, m, n, k) || c.dtype != a.dtype || !cuda_half_ok_(c, m * n)) return 0;
int t = cuda_dtype_type_(a.dtype);
return cuda_gemm_ex_(a.dev, ta, b.dev, tb, c.dev, t, t, m, n, k);
}
int cuda_half_gemm_f32(struct CudaHalfBuf a, int ta, struct CudaHalfBuf b, int tb, struct CudaBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!cuda_half_gemm_ok_(a, b, m, n, k) || !cuda_fits_(c, m * n)) return 0;
return cuda_gemm_ex_(a.dev, ta, b.dev, tb, c.dev, cuda_dtype_type_(a.dtype), CUDA_R_32F_, m, n, k);
}
int cuda_half_gemm_relu(struct CudaHalfBuf a, int ta, struct CudaHalfBuf b, int tb, struct CudaHalfBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!cuda_half_gemm_ok_(a, b, m, n, k) || c.dtype != a.dtype || !cuda_half_ok_(c, m * n)) return 0;
int t = cuda_dtype_type_(a.dtype);
int plan = cuda_lt_plan_(ta, tb, m, n, k, LT_EPILOGUE_RELU, t);
if (plan >= 0) {
float one = 1.0f;
float zero = 0.0f;
unsafe {
if (cublasLtMatmul(g_cuda_lt, g_lt_desc[plan], (const float*)&one,
(const CuDevPtr_*)b.dev, g_lt_la[plan], (const CuDevPtr_*)a.dev, g_lt_lb[plan],
(const float*)&zero, (const CuDevPtr_*)c.dev, g_lt_lc[plan],
(CuDevPtr_*)c.dev, g_lt_lc[plan],
(const unsigned long*)&g_lt_algo[plan * LT_RESULT_WORDS],
(CuDevPtr_*)(unsigned long)g_cuda_ws, CUDA_WORKSPACE_BYTES, g_cuda_stream) == 0)
return 1;
}
}
return cuda_gemm_ex_(a.dev, ta, b.dev, tb, c.dev, t, t, m, n, k) &&
cuda_launch2_(cuda_k16_(K_RELU_F16, a.dtype), c.dev, c.dev, m * n);
}
// One to four warps per row (more on short outputs, so every SM has rows
// to stream), 128-thread blocks of 4 / wpr rows each.
int cuda_half_gemv(struct CudaHalfBuf w, struct CudaBuf x, struct CudaBuf y, unsigned long n, unsigned long k) {
if (!cuda_ensure_() || n == 0UL || k == 0UL || n > 2147483647UL || k > 4294967295UL) return 0;
if (!cuda_half_ok_(w, n * k) || !cuda_fits_(x, k) || !cuda_fits_(y, n)) return 0;
int vec = k % 8UL == 0UL && w.dev % 16UL == 0UL && x.dev % 16UL == 0UL;
unsigned long sms = g_cuda_sms > 0 ? (unsigned long)g_cuda_sms : 1UL;
int wpr = 1;
if (n < sms * 16UL) wpr = n < sms * 8UL || k >= 4096UL ? 4 : 2;
if (!vec || k < 512UL) wpr = 1; // too short to split
unsigned long rpb = 4UL / (unsigned long)wpr;
// One resident wave (16 blocks per SM) grid-strided when the full grid
// would leave a partial second wave; past two waves the full grid streams better.
unsigned long blocks = (n + rpb - 1UL) / rpb;
if (blocks < sms * 32UL && blocks > sms * 16UL) blocks = sms * 16UL;
unsafe {
CuDevPtr_* pw = (CuDevPtr_*)w.dev;
CuDevPtr_* px = (CuDevPtr_*)x.dev;
CuDevPtr_* py = (CuDevPtr_*)y.dev;
unsigned int nn = (unsigned int)n;
unsigned int kk = (unsigned int)k;
int split = wpr;
CuKernelArg_* params[6];
params[0] = (CuKernelArg_*)&pw; params[1] = (CuKernelArg_*)&px; params[2] = (CuKernelArg_*)&py; params[3] = (CuKernelArg_*)&nn;
params[4] = (CuKernelArg_*)&kk; params[5] = (CuKernelArg_*)&split;
return cuLaunchKernel(g_cuda_fn[cuda_k16_(vec ? K_GEMV_F16 : K_GEMV1_F16, w.dtype)],
(unsigned int)blocks, 1U, 1U, 128U, 1U, 1U, 0U,
g_cuda_stream, ¶ms[0], (CuKernelArg_**)0) == CUDA_OK;
}
}
int cuda_half_relu(struct CudaHalfBuf a, struct CudaHalfBuf out) {
unsigned long n = a.len;
if (!cuda_ensure_() || !cuda_half_ok_(a, n) || !cuda_half_ok_(out, n) || a.dtype != out.dtype ||
n > 4294967295UL) return 0;
return cuda_launch2_(cuda_k16_(K_RELU_F16, a.dtype), a.dev, out.dev, n);
}
int cuda_half_relu_grad(struct CudaHalfBuf z, struct CudaHalfBuf g, struct CudaHalfBuf out) {
unsigned long n = z.len;
if (!cuda_ensure_() || !cuda_half_ok_(z, n) || !cuda_half_ok_(g, n) || !cuda_half_ok_(out, n) ||
z.dtype != g.dtype || z.dtype != out.dtype || n > 4294967295UL) return 0;
unsafe {
CuDevPtr_* pz = (CuDevPtr_*)z.dev;
CuDevPtr_* pg = (CuDevPtr_*)g.dev;
CuDevPtr_* po = (CuDevPtr_*)out.dev;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[4];
params[0] = (CuKernelArg_*)&pz; params[1] = (CuKernelArg_*)&pg; params[2] = (CuKernelArg_*)&po; params[3] = (CuKernelArg_*)&nn;
return cuda_launch_(cuda_k16_(K_RELUG_F16, z.dtype), n, ¶ms[0]);
}
}
int cuda_half_mse_grad(struct CudaHalfBuf y, struct CudaHalfBuf t, struct CudaHalfBuf dy, struct CudaBuf loss1) {
unsigned long n = y.len;
if (!cuda_ensure_() || !cuda_half_ok_(y, n) || !cuda_half_ok_(t, n) || !cuda_half_ok_(dy, n) ||
y.dtype != t.dtype || y.dtype != dy.dtype || !cuda_fits_(loss1, 1UL) || n > 4294967295UL) return 0;
unsafe {
CuDevPtr_* py = (CuDevPtr_*)y.dev;
CuDevPtr_* pt = (CuDevPtr_*)t.dev;
CuDevPtr_* pd = (CuDevPtr_*)dy.dev;
CuDevPtr_* pl = (CuDevPtr_*)loss1.dev;
unsigned int nn = (unsigned int)n;
CuKernelArg_* params[5];
params[0] = (CuKernelArg_*)&py; params[1] = (CuKernelArg_*)&pt; params[2] = (CuKernelArg_*)&pd;
params[3] = (CuKernelArg_*)&pl; params[4] = (CuKernelArg_*)&nn;
return cuda_reduce_launch_(cuda_k16_(K_MSE_F16, y.dtype), n, loss1.dev, ¶ms[0]);
}
}
int cuda_half_sgd(struct CudaHalfBuf w, struct CudaHalfBuf g, float lr) {
unsigned long n = w.len;
if (!cuda_ensure_() || !cuda_half_ok_(w, n) || !cuda_half_ok_(g, n) || w.dtype != g.dtype ||
n > 4294967295UL) return 0;
return cuda_sgd_raw_(w.dtype == CUDA_BF16 ? K_SGD_BF16 : K_SGD_F16, w.dev, g.dev, lr, n);
}
int cuda_half_sgd2(struct CudaHalfBuf w1, struct CudaHalfBuf g1, struct CudaHalfBuf w2, struct CudaHalfBuf g2,
float lr) {
unsigned long n1 = w1.len;
unsigned long n2 = w2.len;
if (!cuda_ensure_() || !cuda_half_ok_(w1, n1) || !cuda_half_ok_(g1, n1) || !cuda_half_ok_(w2, n2) ||
!cuda_half_ok_(g2, n2) || w1.dtype != g1.dtype || w1.dtype != w2.dtype || w1.dtype != g2.dtype ||
n1 + n2 > 4294967295UL) return 0;
unsafe {
CuDevPtr_* pw1 = (CuDevPtr_*)w1.dev;
CuDevPtr_* pg1 = (CuDevPtr_*)g1.dev;
CuDevPtr_* pw2 = (CuDevPtr_*)w2.dev;
CuDevPtr_* pg2 = (CuDevPtr_*)g2.dev;
unsigned int nn1 = (unsigned int)n1;
unsigned int nn2 = (unsigned int)n2;
float rate = lr;
CuKernelArg_* params[7];
params[0] = (CuKernelArg_*)&pw1; params[1] = (CuKernelArg_*)&pg1; params[2] = (CuKernelArg_*)&nn1;
params[3] = (CuKernelArg_*)&pw2; params[4] = (CuKernelArg_*)&pg2; params[5] = (CuKernelArg_*)&nn2;
params[6] = (CuKernelArg_*)&rate;
return cuda_launch_(cuda_k16_(K_SGD2_F16, w1.dtype), n1 + n2, ¶ms[0]);
}
}
// ── Mixed precision ──────────────────────────────────────────────────────────
extern unsigned char* malloc(unsigned long size);
extern void free(unsigned char* ptr);
int cuda_zero(struct CudaBuf b) {
if (!cuda_ensure_() || b.dev == 0UL) return 0;
unsafe { return cuMemsetD32Async((unsigned long long)b.dev, 0U, (size_t)b.len, g_cuda_stream) == CUDA_OK; }
}
int cuda_half_convert(struct CudaHalfBuf src, struct CudaHalfBuf dst) {
unsigned long n = src.len;
if (!cuda_ensure_() || !cuda_half_ok_(src, n) || !cuda_half_ok_(dst, n) || n > 4294967295UL) return 0;
if (src.dtype == dst.dtype) {
unsafe {
return cuMemcpyDtoDAsync_v2((unsigned long long)dst.dev, (unsigned long long)src.dev, (size_t)(n * 2UL),
g_cuda_stream) == CUDA_OK;
}
}
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)src.dev;
CuDevPtr_* po = (CuDevPtr_*)dst.dev;
unsigned int nn = (unsigned int)n;
int toBf = dst.dtype == CUDA_BF16 ? 1 : 0;
CuKernelArg_* params[4];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&po; params[2] = (CuKernelArg_*)&nn; params[3] = (CuKernelArg_*)&toBf;
return cuda_launch_(K_HALF_CVT, n, ¶ms[0]);
}
}
// The kernel's CvtTab: 32 source / destination pointers and lengths.
#define CUDA_CVT_TAB 32UL
struct CudaCvtTab_ {
unsigned long src[32];
unsigned long dst[32];
unsigned int n[32];
};
// One launch per 32 entries; mode as in the cvt_many kernel.
static int cuda_cvt_many_(const unsigned long* src, const unsigned long* dst, const unsigned long* len,
unsigned long count, int mode, int bf, float scale, unsigned long flag) {
unsigned long base = 0UL;
while (base < count) {
struct CudaCvtTab_ t;
unsigned long c = count - base < CUDA_CVT_TAB ? count - base : CUDA_CVT_TAB;
unsigned long most = 1UL;
unsafe {
unsigned long i = 0UL;
while (i < CUDA_CVT_TAB) {
t.src[i] = i < c ? src[base + i] : 0UL;
t.dst[i] = i < c ? dst[base + i] : 0UL;
t.n[i] = i < c ? (unsigned int)len[base + i] : 0U;
if (i < c && len[base + i] > most) most = len[base + i];
i = i + 1UL;
}
unsigned long blocks = (most + (unsigned long)CUDA_BLOCK - 1UL) / (unsigned long)CUDA_BLOCK;
if (blocks > 1024UL) blocks = 1024UL;
int m = mode;
int b = bf;
float sc = scale;
CuDevPtr_* pf = (CuDevPtr_*)flag;
int vec = 1;
i = 0UL;
while (i < c) { if ((t.src[i] | t.dst[i]) % 16UL != 0UL) vec = 0; i = i + 1UL; }
CuKernelArg_* params[6];
params[0] = (CuKernelArg_*)&t; params[1] = (CuKernelArg_*)&m; params[2] = (CuKernelArg_*)&b; params[3] = (CuKernelArg_*)≻
params[4] = (CuKernelArg_*)&pf; params[5] = (CuKernelArg_*)&vec;
if (cuLaunchKernel(g_cuda_fn[K_CVT_MANY], (unsigned int)blocks, (unsigned int)c, 1U, CUDA_BLOCK, 1U, 1U,
0U, g_cuda_stream, ¶ms[0], (CuKernelArg_**)0) != CUDA_OK) return 0;
}
base = base + c;
}
return 1;
}
// Pointer / length tables for cuda_cvt_many_ (heap, freed by the caller).
struct CudaCvtLists_ {
unsigned long* src;
unsigned long* dst;
unsigned long* len;
int ok;
int bf;
};
static struct CudaCvtLists_ cuda_cvt_lists_(unsigned long count) {
struct CudaCvtLists_ l;
unsafe {
unsigned long bytes = sizeof(unsigned long) * (count == 0UL ? 1UL : count);
l.src = (unsigned long*)malloc(bytes);
l.dst = (unsigned long*)malloc(bytes);
l.len = (unsigned long*)malloc(bytes);
}
l.ok = 1;
l.bf = -1;
return l;
}
static void cuda_cvt_lists_free_(&CudaCvtLists_ l) {
unsafe { free((unsigned char*)l.src); free((unsigned char*)l.dst); free((unsigned char*)l.len); }
}
// One 16-bit dtype per call (the kernel's bf flag); mixed lists are refused.
static void cuda_cvt_dtype_(&CudaCvtLists_ l, int dtype) {
int bf = dtype == CUDA_BF16 ? 1 : 0;
if (l.bf < 0) l.bf = bf;
else if (l.bf != bf) l.ok = 0;
}
int cuda_to_half_many(const &Vec<struct CudaBuf> src, const &Vec<struct CudaHalfBuf> dst) {
unsigned long n = src.length();
if (!cuda_ensure_() || n != dst.length()) return 0;
if (n == 0UL) return 1;
struct CudaCvtLists_ l = cuda_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct CudaBuf s = *src.at(i);
struct CudaHalfBuf d = *dst.at(i);
if (!cuda_fits_(s, d.len) || !cuda_half_ok_(d, d.len) || d.len > 4294967295UL) l.ok = 0;
cuda_cvt_dtype_(&l, d.dtype);
l.src[i] = s.dev; l.dst[i] = d.dev; l.len[i] = d.len;
i = i + 1UL;
}
}
int ok = l.ok && cuda_cvt_many_((const unsigned long*)l.src, (const unsigned long*)l.dst, (const unsigned long*)l.len,
n, 0, l.bf, (float)1.0, 0UL);
cuda_cvt_lists_free_(&l);
return ok;
}
int cuda_from_half_many(const &Vec<struct CudaHalfBuf> src, const &Vec<struct CudaBuf> dst) {
unsigned long n = src.length();
if (!cuda_ensure_() || n != dst.length()) return 0;
if (n == 0UL) return 1;
struct CudaCvtLists_ l = cuda_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct CudaHalfBuf s = *src.at(i);
struct CudaBuf d = *dst.at(i);
if (!cuda_half_ok_(s, d.len) || !cuda_fits_(d, d.len) || d.len > 4294967295UL) l.ok = 0;
cuda_cvt_dtype_(&l, s.dtype);
l.src[i] = s.dev; l.dst[i] = d.dev; l.len[i] = d.len;
i = i + 1UL;
}
}
int ok = l.ok && cuda_cvt_many_((const unsigned long*)l.src, (const unsigned long*)l.dst, (const unsigned long*)l.len,
n, 1, l.bf, (float)1.0, 0UL);
cuda_cvt_lists_free_(&l);
return ok;
}
int cuda_half_unscale_many(const &Vec<struct CudaHalfBuf> src, const &Vec<struct CudaBuf> dst, float scale,
struct CudaBuf flag) {
unsigned long n = src.length();
if (!cuda_ensure_() || n != dst.length() || !cuda_fits_(flag, 1UL)) return 0;
if (n == 0UL) return 1;
struct CudaCvtLists_ l = cuda_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct CudaHalfBuf s = *src.at(i);
struct CudaBuf d = *dst.at(i);
if (!cuda_half_ok_(s, s.len) || !cuda_fits_(d, s.len) || s.len > 4294967295UL) l.ok = 0;
cuda_cvt_dtype_(&l, s.dtype);
l.src[i] = s.dev; l.dst[i] = d.dev; l.len[i] = s.len;
i = i + 1UL;
}
}
int ok = l.ok && cuda_cvt_many_((const unsigned long*)l.src, (const unsigned long*)l.dst, (const unsigned long*)l.len,
n, 1, l.bf, scale, flag.dev);
cuda_cvt_lists_free_(&l);
return ok;
}
int cuda_unscale_many(const &Vec<struct CudaBuf> src, const &Vec<struct CudaBuf> dst, float scale,
struct CudaBuf flag) {
unsigned long n = src.length();
if (!cuda_ensure_() || n != dst.length() || !cuda_fits_(flag, 1UL)) return 0;
if (n == 0UL) return 1;
struct CudaCvtLists_ l = cuda_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct CudaBuf s = *src.at(i);
struct CudaBuf d = *dst.at(i);
if (!cuda_fits_(s, s.len) || !cuda_fits_(d, s.len) || s.len > 4294967295UL) l.ok = 0;
l.src[i] = s.dev; l.dst[i] = d.dev; l.len[i] = s.len;
i = i + 1UL;
}
}
int ok = l.ok && cuda_cvt_many_((const unsigned long*)l.src, (const unsigned long*)l.dst, (const unsigned long*)l.len,
n, 2, 0, scale, flag.dev);
cuda_cvt_lists_free_(&l);
return ok;
}
int cuda_mp_sgd(struct CudaBuf w, struct CudaBuf g, struct CudaBuf vel, struct CudaHalfBuf wHalf,
float lr, float momentum, float weightDecay, struct CudaBuf skip) {
unsigned long n = w.len;
if (!cuda_ensure_() || !cuda_fits_(w, n) || !cuda_fits_(g, n) || n > 4294967295UL) return 0;
if (vel.dev != 0UL && !cuda_fits_(vel, n)) return 0;
if (wHalf.dev != 0UL && !cuda_half_ok_(wHalf, n)) return 0;
unsafe {
CuDevPtr_* pw = (CuDevPtr_*)w.dev;
CuDevPtr_* pg = (CuDevPtr_*)g.dev;
CuDevPtr_* pv = (CuDevPtr_*)vel.dev;
CuDevPtr_* ph = (CuDevPtr_*)wHalf.dev;
CuDevPtr_* ps = (CuDevPtr_*)skip.dev;
float l = lr;
float mo = momentum;
float wd = weightDecay;
unsigned int nn = (unsigned int)n;
int bf = wHalf.dtype == CUDA_BF16 ? 1 : 0;
CuKernelArg_* params[10];
params[0] = (CuKernelArg_*)&pw; params[1] = (CuKernelArg_*)&pg; params[2] = (CuKernelArg_*)&pv; params[3] = (CuKernelArg_*)&ph;
params[4] = (CuKernelArg_*)&l; params[5] = (CuKernelArg_*)&mo; params[6] = (CuKernelArg_*)&wd; params[7] = (CuKernelArg_*)&nn;
params[8] = (CuKernelArg_*)&bf; params[9] = (CuKernelArg_*)&ps;
return cuda_launch_(K_MP_SGD, n, ¶ms[0]);
}
}
// 1 if an Adam step on these buffers can move 4 values per thread (w, m, v, g
// 16-byte aligned, the 16-bit copy 8-byte aligned or absent).
static int cuda_adam_vec_(unsigned long w, unsigned long m, unsigned long v, unsigned long g, unsigned long wh) {
return (w | m | v | g) % 16UL == 0UL && wh % 8UL == 0UL;
}
int cuda_mp_adam(struct CudaBuf w, struct CudaBuf m, struct CudaBuf v, struct CudaBuf g, struct CudaHalfBuf wHalf,
float lr, float beta1, float beta2, float eps, float weightDecay, int decoupled,
unsigned long step, struct CudaBuf skip) {
unsigned long n = w.len;
if (!cuda_ensure_() || !cuda_fits_(w, n) || !cuda_fits_(m, n) || !cuda_fits_(v, n) || !cuda_fits_(g, n) ||
n > 4294967295UL || step == 0UL) return 0;
if (wHalf.dev != 0UL && !cuda_half_ok_(wHalf, n)) return 0;
float bc1 = (float)1.0 - pow_f(beta1, (float)step);
float bc2 = (float)1.0 - pow_f(beta2, (float)step);
unsafe {
CuDevPtr_* pw = (CuDevPtr_*)w.dev;
CuDevPtr_* pm = (CuDevPtr_*)m.dev;
CuDevPtr_* pv = (CuDevPtr_*)v.dev;
CuDevPtr_* pg = (CuDevPtr_*)g.dev;
CuDevPtr_* ph = (CuDevPtr_*)wHalf.dev;
CuDevPtr_* ps = (CuDevPtr_*)skip.dev;
float l = lr; float b1 = beta1; float b2 = beta2; float e = eps; float wd = weightDecay;
float c1 = bc1; float c2 = bc2;
int dec = decoupled ? 1 : 0;
unsigned int nn = (unsigned int)n;
int bf = wHalf.dtype == CUDA_BF16 ? 1 : 0;
int vec = cuda_adam_vec_(w.dev, m.dev, v.dev, g.dev, wHalf.dev);
CuKernelArg_* params[17];
params[0] = (CuKernelArg_*)&pw; params[1] = (CuKernelArg_*)± params[2] = (CuKernelArg_*)&pv; params[3] = (CuKernelArg_*)&pg;
params[4] = (CuKernelArg_*)&ph; params[5] = (CuKernelArg_*)&l; params[6] = (CuKernelArg_*)&b1; params[7] = (CuKernelArg_*)&b2;
params[8] = (CuKernelArg_*)&e; params[9] = (CuKernelArg_*)&wd; params[10] = (CuKernelArg_*)&c1; params[11] = (CuKernelArg_*)&c2;
params[12] = (CuKernelArg_*)&dec; params[13] = (CuKernelArg_*)&nn; params[14] = (CuKernelArg_*)&bf; params[15] = (CuKernelArg_*)&ps;
params[16] = (CuKernelArg_*)&vec;
return cuda_launch_(K_MP_ADAM, vec ? (n + 3UL) / 4UL : n, ¶ms[0]);
}
}
// The kernels' AdamArgs and AdamTab (24 tensors per launch).
#define CUDA_ADAM_TAB 24UL
struct CudaAdamArgs_ {
float lr; float b1; float b2; float eps; float wd; float bc1; float bc2;
int decoupled;
};
struct CudaAdamTab_ {
unsigned long w[24];
unsigned long m[24];
unsigned long v[24];
unsigned long g[24];
unsigned long wh[24];
unsigned int n[24];
};
int cuda_mp_adam_many(const &Vec<struct CudaBuf> w, const &Vec<struct CudaBuf> m, const &Vec<struct CudaBuf> v,
const &Vec<struct CudaBuf> g, const &Vec<struct CudaHalfBuf> wHalf, float lr, float beta1,
float beta2, float eps, float weightDecay, int decoupled, unsigned long step,
struct CudaBuf skip) {
unsigned long count = w.length();
unsigned long nh = wHalf.length();
if (!cuda_ensure_() || step == 0UL || m.length() != count || v.length() != count || g.length() != count ||
(nh != 0UL && nh != count)) return 0;
int bf = -1;
unsigned long i = 0UL;
while (i < count) {
unsigned long n = (*w.at(i)).len;
if (!cuda_fits_(*w.at(i), n) || !cuda_fits_(*m.at(i), n) || !cuda_fits_(*v.at(i), n) ||
!cuda_fits_(*g.at(i), n) || n == 0UL || n > 4294967295UL) return 0;
if (nh != 0UL) {
struct CudaHalfBuf h = *wHalf.at(i);
if (!cuda_half_ok_(h, n)) return 0;
int b = h.dtype == CUDA_BF16 ? 1 : 0;
if (bf >= 0 && b != bf) return 0; // one 16-bit dtype per call
bf = b;
}
i = i + 1UL;
}
struct CudaAdamArgs_ a;
a.lr = lr; a.b1 = beta1; a.b2 = beta2; a.eps = eps; a.wd = weightDecay;
a.bc1 = (float)1.0 - pow_f(beta1, (float)step);
a.bc2 = (float)1.0 - pow_f(beta2, (float)step);
a.decoupled = decoupled ? 1 : 0;
if (bf < 0) bf = 0;
unsigned long base = 0UL;
while (base < count) {
struct CudaAdamTab_ t;
unsigned long c = count - base < CUDA_ADAM_TAB ? count - base : CUDA_ADAM_TAB;
unsigned long most = 1UL;
int vec = 1;
unsafe {
unsigned long j = 0UL;
while (j < CUDA_ADAM_TAB) {
int in = j < c;
t.w[j] = in ? (*w.at(base + j)).dev : 0UL;
t.m[j] = in ? (*m.at(base + j)).dev : 0UL;
t.v[j] = in ? (*v.at(base + j)).dev : 0UL;
t.g[j] = in ? (*g.at(base + j)).dev : 0UL;
t.wh[j] = in && nh != 0UL ? (*wHalf.at(base + j)).dev : 0UL;
t.n[j] = in ? (unsigned int)(*w.at(base + j)).len : 0U;
if (in) {
if (!cuda_adam_vec_(t.w[j], t.m[j], t.v[j], t.g[j], t.wh[j])) vec = 0;
if ((*w.at(base + j)).len > most) most = (*w.at(base + j)).len;
}
j = j + 1UL;
}
unsigned long per = vec ? (unsigned long)CUDA_BLOCK * 4UL : (unsigned long)CUDA_BLOCK;
unsigned long blocks = (most + per - 1UL) / per;
if (blocks > 2048UL) blocks = 2048UL;
int b = bf;
CuDevPtr_* ps = (CuDevPtr_*)skip.dev;
CuKernelArg_* params[5];
params[0] = (CuKernelArg_*)&t; params[1] = (CuKernelArg_*)&a; params[2] = (CuKernelArg_*)&b; params[3] = (CuKernelArg_*)&ps;
params[4] = (CuKernelArg_*)&vec;
if (cuLaunchKernel(g_cuda_fn[K_MP_ADAM_MANY], (unsigned int)blocks, (unsigned int)c, 1U, CUDA_BLOCK, 1U, 1U,
0U, g_cuda_stream, ¶ms[0], (CuKernelArg_**)0) != CUDA_OK) return 0;
}
base = base + c;
}
return 1;
}
static int cuda_half_ew_(struct CudaHalfBuf a, struct CudaHalfBuf b, struct CudaHalfBuf out, float k, int op) {
unsigned long n = a.len;
if (!cuda_ensure_() || !cuda_half_ok_(a, n) || !cuda_half_ok_(out, n) || a.dtype != out.dtype ||
n > 4294967295UL) return 0;
if (op != 3 && (!cuda_half_ok_(b, n) || b.dtype != a.dtype)) return 0;
unsafe {
CuDevPtr_* pa = (CuDevPtr_*)a.dev;
CuDevPtr_* pb = (CuDevPtr_*)b.dev;
CuDevPtr_* po = (CuDevPtr_*)out.dev;
float kk = k;
unsigned int nn = (unsigned int)n;
int o = op;
int bf = a.dtype == CUDA_BF16 ? 1 : 0;
int vec = (a.dev | b.dev | out.dev) % 16UL == 0UL;
CuKernelArg_* params[8];
params[0] = (CuKernelArg_*)&pa; params[1] = (CuKernelArg_*)&pb; params[2] = (CuKernelArg_*)&po; params[3] = (CuKernelArg_*)&kk;
params[4] = (CuKernelArg_*)&nn; params[5] = (CuKernelArg_*)&o; params[6] = (CuKernelArg_*)&bf; params[7] = (CuKernelArg_*)&vec;
return cuda_launch_(K_HALF_EW, vec ? (n + 7UL) / 8UL : n, ¶ms[0]);
}
}
int cuda_half_add(struct CudaHalfBuf a, struct CudaHalfBuf b, struct CudaHalfBuf out) { return cuda_half_ew_(a, b, out, (float)0.0, 0); }
int cuda_half_sub(struct CudaHalfBuf a, struct CudaHalfBuf b, struct CudaHalfBuf out) { return cuda_half_ew_(a, b, out, (float)0.0, 1); }
int cuda_half_mul(struct CudaHalfBuf a, struct CudaHalfBuf b, struct CudaHalfBuf out) { return cuda_half_ew_(a, b, out, (float)0.0, 2); }
int cuda_half_scale(struct CudaHalfBuf a, float k, struct CudaHalfBuf out) { return cuda_half_ew_(a, a, out, k, 3); }
int cuda_half_axpy(struct CudaHalfBuf a, float k, struct CudaHalfBuf b, struct CudaHalfBuf out) {
return cuda_half_ew_(a, b, out, k, 4);
}
int cuda_half_bias(struct CudaHalfBuf y, struct CudaHalfBuf bias, unsigned long rows, unsigned long channels,
unsigned long plane, int relu) {
unsigned long total = rows * channels * plane;
if (!cuda_ensure_() || total == 0UL || total > 4294967295UL || !cuda_half_ok_(y, total) ||
!cuda_half_ok_(bias, channels) || y.dtype != bias.dtype) return 0;
unsafe {
CuDevPtr_* py = (CuDevPtr_*)y.dev;
CuDevPtr_* pb = (CuDevPtr_*)bias.dev;
unsigned int t = (unsigned int)total;
unsigned int c = (unsigned int)channels;
unsigned int p = (unsigned int)plane;
int r = relu ? 1 : 0;
int bf = y.dtype == CUDA_BF16 ? 1 : 0;
CuKernelArg_* params[7];
params[0] = (CuKernelArg_*)&py; params[1] = (CuKernelArg_*)&pb; params[2] = (CuKernelArg_*)&t; params[3] = (CuKernelArg_*)&c;
params[4] = (CuKernelArg_*)&p; params[5] = (CuKernelArg_*)&r; params[6] = (CuKernelArg_*)&bf;
return cuda_launch_(K_HALF_BIAS, total, ¶ms[0]);
}
}
static int cuda_half_pool_(struct CudaHalfBuf x, struct CudaHalfBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride, int avg) {
if (!cuda_ensure_() || kernel == 0UL || stride == 0UL || height < kernel || width < kernel || x.dtype != y.dtype)
return 0;
unsigned long oh = (height - kernel) / stride + 1UL;
unsigned long ow = (width - kernel) / stride + 1UL;
unsigned long rows = batch * channels;
unsigned long total = rows * oh * ow;
if (total == 0UL || total > 4294967295UL || height * width > 4294967295UL ||
!cuda_half_ok_(x, rows * height * width) || !cuda_half_ok_(y, total)) return 0;
unsafe {
CuDevPtr_* px = (CuDevPtr_*)x.dev;
CuDevPtr_* py = (CuDevPtr_*)y.dev;
unsigned int r = (unsigned int)rows;
unsigned int hh = (unsigned int)height; unsigned int ww = (unsigned int)width;
unsigned int ohh = (unsigned int)oh; unsigned int oww = (unsigned int)ow;
unsigned int kk = (unsigned int)kernel; unsigned int ss = (unsigned int)stride;
int a = avg;
int bf = x.dtype == CUDA_BF16 ? 1 : 0;
CuKernelArg_* params[11];
params[0] = (CuKernelArg_*)&px; params[1] = (CuKernelArg_*)&py; params[2] = (CuKernelArg_*)&r; params[3] = (CuKernelArg_*)&hh;
params[4] = (CuKernelArg_*)&ww; params[5] = (CuKernelArg_*)&ohh; params[6] = (CuKernelArg_*)&oww; params[7] = (CuKernelArg_*)&kk;
params[8] = (CuKernelArg_*)&ss; params[9] = (CuKernelArg_*)&a; params[10] = (CuKernelArg_*)&bf;
return cuda_launch_(K_HALF_POOL, total, ¶ms[0]);
}
}
int cuda_half_maxpool2d(struct CudaHalfBuf x, struct CudaHalfBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return cuda_half_pool_(x, y, batch, channels, height, width, kernel, stride, 0);
}
int cuda_half_avgpool2d(struct CudaHalfBuf x, struct CudaHalfBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return cuda_half_pool_(x, y, batch, channels, height, width, kernel, stride, 1);
}
} // namespace std