// SafeC Standard Library — AMD GPU backend (ROCm / HIP). See gpu_rocm.h.
//
// Kernels are HIP C++ source embedded below and compiled on first use with
// hiprtc for the selected GPU's own architecture (gfx908, gfx90a, gfx942,
// gfx1100, ...), then loaded with hipModuleLoadData — so one build of this
// library runs on any ROCm GPU. Compilation takes ~0.2 s once per process.
// safeguard: link linux -L${ROCM_PATH:-/opt/rocm}/lib -lamdhip64 -lhiprtc -lrocblas
#pragma once
#include <std/ml/gpu_rocm.h>
#include <std/ml/rocm_ffi.h>
#include <std/math.h>
#include <std/mem.h>
#include <std/time.h>
namespace std {
// ── HIP runtime / hiprtc / rocBLAS C ABI (hip_runtime_api.h, hiprtc.h, rocblas.h) ──
// The HIP runtime, hiprtc and rocBLAS are declared in rocm_ffi.h.
#define ROCM_HIP_OK 0
#define ROCM_ATTR_CU_COUNT 63 // hipDeviceAttributeMultiprocessorCount
#define ROCM_ATTR_INTEGRATED 16 // hipDeviceAttributeIntegrated
#define ROCM_H2D 1 // hipMemcpyHostToDevice
#define ROCM_D2H 2 // hipMemcpyDeviceToHost
#define ROCM_D2D 3 // hipMemcpyDeviceToDevice
#define ROCM_CAPTURE_GLOBAL 0 // hipStreamCaptureModeGlobal
#define ROCM_BLAS_N 111 // rocblas_operation_none
#define ROCM_BLAS_T 112 // rocblas_operation_transpose
#define ROCM_BLAS_F16 150 // rocblas_datatype_f16_r
#define ROCM_BLAS_F32 151 // rocblas_datatype_f32_r
#define ROCM_BLAS_BF16 168 // rocblas_datatype_bf16_r
#define ROCM_BLAS_ALGO_STANDARD 0 // rocblas_gemm_algo_standard
#define ROCM_BLAS_ALGO_INDEX 1 // rocblas_gemm_algo_solution_index
#define ROCM_BLOCK 256U
#define ROCM_MAX_BLOCKS 4096UL
// ── Kernel source ────────────────────────────────────────────────────────────
// Elementwise kernels use grid-stride loops; reduce_sum and mse_grad reduce
// per wavefront with shuffles, then per block through LDS, then one atomic
// per block. sgemm_nn is a 64x64-tile LDS kernel with 4x4 register tiles.
static const char* rocm_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_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_gelu, 0.5f * x * (1.0f + tanhf(0.7978845608028654f * (x + 0.044715f * x * x * x))))\n"
"EW_UNARY(ew_scale, x * k)\n"
"\n"
"// relu backward: g_in = g_out where x > 0\n"
"extern \"C\" __global__ void ew_relu_grad(const float* __restrict__ x, const float* __restrict__ g,\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) o[i] = x[i] > 0.0f ? g[i] : 0.0f;\n"
"}\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"
"// Sum: each block reduces a grid-stride slice with wave shuffles, then one\n"
"// atomicAdd per block. *out must be zeroed first.\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"
" for (int off = warpSize / 2; off > 0; off >>= 1) s += __shfl_down(s, off);\n"
" __shared__ float part[1024 / 32];\n"
" int lane = threadIdx.x % warpSize, wid = threadIdx.x / warpSize;\n"
" if (lane == 0) part[wid] = s;\n"
" __syncthreads();\n"
" if (wid == 0) {\n"
" int nw = (blockDim.x + warpSize - 1) / warpSize;\n"
" s = lane < nw ? part[lane] : 0.0f;\n"
" for (int off = warpSize / 2; off > 0; off >>= 1) s += __shfl_down(s, off);\n"
" if (lane == 0) atomicAdd(out, s);\n"
" }\n"
"}\n"
"\n"
"// ── SGEMM: C[M,N] = op(A) * op(B) (+ optional ReLU epilogue), row-major ─────\n"
"// 64x64 output tile per 256-thread block, BK=16, each thread a 4x4 micro-tile.\n"
"// TA/TB: operand stored transposed (A is KxM / B is NxK).\n"
"template <int TA, int TB, int RELU>\n"
"__device__ void sgemm_tile(const float* __restrict__ A, const float* __restrict__ B,\n"
" float* __restrict__ C, int M, int N, int K) {\n"
" constexpr 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; // 16x16 threads, 4x4 each\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"
" // load A tile (BM x BK) -> As[k][m]\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"
" float v = 0.0f;\n"
" if (gm < M && gk < K) v = TA ? A[(size_t)gk * M + gm] : A[(size_t)gm * K + gk];\n"
" As[k][m] = v;\n"
" }\n"
" // load B tile (BK x BN) -> Bs[k][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"
" float v = 0.0f;\n"
" if (gk < K && gn < N) v = TB ? B[(size_t)gn * K + gk] : B[(size_t)gk * N + gn];\n"
" Bs[k][n] = v;\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) { float v = acc[i][j]; C[(size_t)gm * N + gn] = RELU ? (v > 0.0f ? v : 0.0f) : v; }\n"
" }\n"
" }\n"
"}\n"
"\n"
"extern \"C\" __global__ __launch_bounds__(256) void sgemm_nn(const float* A, const float* B, float* C, int M, int N, int K) { sgemm_tile<0,0,0>(A,B,C,M,N,K); }\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* y, const float* t, float* dy, float* loss, unsigned n) {\n"
" float s = 0.0f;\n"
" unsigned i = blockIdx.x * blockDim.x + threadIdx.x;\n"
" for (; i < n; i += gridDim.x * blockDim.x) { float d = y[i] - t[i]; dy[i] = 2.0f * d; s += d * d; }\n"
" for (int off = warpSize / 2; off > 0; off >>= 1) s += __shfl_down(s, off);\n"
" __shared__ float part[16];\n"
" int lane = threadIdx.x % warpSize, wid = threadIdx.x / warpSize;\n"
" if (lane == 0) part[wid] = s;\n"
" __syncthreads();\n"
" if (wid == 0) {\n"
" s = lane < (int)((blockDim.x + warpSize - 1) / warpSize) ? part[lane] : 0.0f;\n"
" for (int off = warpSize / 2; off > 0; off >>= 1) s += __shfl_down(s, off);\n"
" if (lane == 0) atomicAdd(loss, s);\n"
" }\n"
"}\n"
// planes x h x w -> planes x 2h x 2w, out[p][y][x] = in[p][y/2][x/2]
"extern \"C\" __global__ void upsample2x(const float* __restrict__ x, float* __restrict__ y,\n"
" unsigned planes, unsigned h, unsigned w) {\n"
" unsigned w2 = 2u * w, h2 = 2u * h, total = planes * h2 * w2;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {\n"
" unsigned ox = i % w2, t = i / w2, oy = t % h2, p = t / h2;\n"
" y[i] = x[((size_t)p * h + oy / 2u) * w + ox / 2u];\n"
" }\n"
"}\n"
// NCHW y[b][c][:] += bias[c], then relu if asked: the conv epilogue.
"extern \"C\" __global__ void bias_act(float* __restrict__ y, const float* __restrict__ b,\n"
" unsigned total, unsigned plane, unsigned channels, int relu) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {\n"
" float v = y[i] + b[(i / plane) % channels];\n"
" y[i] = (relu && v < 0.0f) ? 0.0f : v;\n"
" }\n"
"}\n"
// out[b] = (a[b], c[b]) for 'batch' rows of na and nc elements: channel
// concat of NCHW tensors in one launch.
"extern \"C\" __global__ void concat_rows(const float* __restrict__ a, const float* __restrict__ c,\n"
" float* __restrict__ o, unsigned batch, unsigned na, unsigned nc) {\n"
" unsigned row = na + nc, total = batch * row;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {\n"
" unsigned b = i / row, r = i % row;\n"
" o[i] = r < na ? a[(size_t)b * na + r] : c[(size_t)b * nc + (r - na)];\n"
" }\n"
"}\n"
// out[b] = (upsample2x(x[b]), skip[b]) along channels: the U-Net decoder's
// upsample and skip concat in one pass, without the upsampled tensor.
// x: [batch, cx, h, w]; skip: [batch, cs, 2h, 2w]; out: [batch, cx+cs, 2h, 2w].
"extern \"C\" __global__ void upcat(const float* __restrict__ x, const float* __restrict__ s,\n"
" float* __restrict__ o, unsigned batch, unsigned cx, unsigned cs,\n"
" unsigned h, unsigned w) {\n"
" unsigned w2 = 2u * w, h2 = 2u * h, plane = h2 * w2, row = (cx + cs) * plane, total = batch * row;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {\n"
" unsigned b = i / row, r = i % row, c = r / plane, p = r % plane;\n"
" if (c < cx) {\n"
" unsigned oy = p / w2, ox = p % w2;\n"
" o[i] = x[(((size_t)b * cx + c) * h + oy / 2u) * w + ox / 2u];\n"
" } else {\n"
" o[i] = s[((size_t)b * cs + (c - cx)) * plane + p];\n"
" }\n"
" }\n"
"}\n"
// fp16 / bf16 storage. Both conversions round to nearest even and give
// the same bits as float16.h: _Float16 for fp16, the integer
// round-half-even trick for bf16 (NaN keeps its top bits plus 0x40).
"__device__ unsigned short f2h(float x, int bf) {\n"
" if (!bf) return __builtin_bit_cast(unsigned short, (_Float16)x);\n"
" unsigned u = __builtin_bit_cast(unsigned, x);\n"
" if ((u & 0x7fffffffu) > 0x7f800000u) return (unsigned short)((u >> 16) | 0x40u);\n"
" return (unsigned short)((u + 0x7fffu + ((u >> 16) & 1u)) >> 16);\n"
"}\n"
"__device__ float h2f(unsigned short h, int bf) {\n"
" return bf ? __builtin_bit_cast(float, (unsigned)h << 16) : (float)__builtin_bit_cast(_Float16, h);\n"
"}\n"
"extern \"C\" __global__ void to_half(const float* __restrict__ a, unsigned short* __restrict__ o,\n"
" unsigned n, int bf) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
" o[i] = f2h(a[i], bf);\n"
"}\n"
"extern \"C\" __global__ void from_half(const unsigned short* __restrict__ a, float* __restrict__ o,\n"
" unsigned n, int bf) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
" o[i] = h2f(a[i], bf);\n"
"}\n"
// w = round(w - lr * g), computed in fp32. w == g is allowed.
"extern \"C\" __global__ void half_sgd(unsigned short* w, const unsigned short* g, float lr,\n"
" unsigned n, int bf) {\n"
// no FMA contraction: same rounding as the CPU's separate * and -
"#pragma clang fp contract(off)\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
" w[i] = f2h(h2f(w[i], bf) - lr * h2f(g[i], bf), bf);\n"
"}\n"
// Training ops on 16-bit buffers: fp32 arithmetic, one rounding per
// stored value (relu and relu_grad only copy bits).
"extern \"C\" __global__ void half_relu(const unsigned short* a, unsigned short* o, unsigned n, int bf) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
" o[i] = h2f(a[i], bf) > 0.0f ? a[i] : (unsigned short)0;\n"
"}\n"
"extern \"C\" __global__ void half_relu_grad(const unsigned short* z, const unsigned short* g,\n"
" unsigned short* o, unsigned n, int bf) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
" o[i] = h2f(z[i], bf) > 0.0f ? g[i] : (unsigned short)0;\n"
"}\n"
// dy = round(2 (y - t)), loss += sum (y - t)^2 in fp32
"extern \"C\" __global__ void half_mse(const unsigned short* y, const unsigned short* t, unsigned short* dy,\n"
" float* loss, unsigned n, int bf) {\n"
" float s = 0.0f;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n"
" float d = h2f(y[i], bf) - h2f(t[i], bf); dy[i] = f2h(2.0f * d, bf); s += d * d;\n"
" }\n"
" for (int off = warpSize / 2; off > 0; off >>= 1) s += __shfl_down(s, off);\n"
" __shared__ float part[16];\n"
" int lane = threadIdx.x % warpSize, wid = threadIdx.x / warpSize;\n"
" if (lane == 0) part[wid] = s;\n"
" __syncthreads();\n"
" if (wid == 0) {\n"
" s = lane < (int)((blockDim.x + warpSize - 1) / warpSize) ? part[lane] : 0.0f;\n"
" for (int off = warpSize / 2; off > 0; off >>= 1) s += __shfl_down(s, off);\n"
" if (lane == 0) atomicAdd(loss, s);\n"
" }\n"
"}\n"
// Two SGD updates in one launch: indices [0, n1) update w1, the rest w2.
"extern \"C\" __global__ void half_sgd2(unsigned short* w1, const unsigned short* g1, unsigned n1,\n"
" unsigned short* w2, const unsigned short* g2, unsigned n2,\n"
" float lr, int bf) {\n"
"#pragma clang fp contract(off)\n"
" unsigned n = n1 + n2;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n"
" if (i < n1) w1[i] = f2h(h2f(w1[i], bf) - lr * h2f(g1[i], bf), bf);\n"
" else w2[i - n1] = f2h(h2f(w2[i - n1], bf) - lr * h2f(g2[i - n1], bf), bf);\n"
" }\n"
"}\n"
// fp16 <-> bf16, through fp32 (one rounding).
"extern \"C\" __global__ void half_to_half(const unsigned short* a, unsigned short* o, unsigned n, int toBf) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)\n"
" o[i] = f2h(h2f(a[i], !toBf), toBf);\n"
"}\n"
// Up to 32 buffers per launch, blockIdx.y picking one. mode 0: fp32 ->
// 16-bit; 1: 16-bit -> fp32 times scale; 2: fp32 -> fp32 times scale (in
// place allowed). Modes 1 and 2 set *flag = 1 when a result is not finite.
"struct CvtTab { const unsigned char* src[32]; unsigned char* dst[32]; unsigned n[32]; };\n"
"extern \"C\" __global__ void cvt_many(CvtTab t, int mode, int bf, float scale, float* flag) {\n"
" unsigned b = blockIdx.y, n = t.n[b];\n"
" int bad = 0;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n"
" if (mode == 0) { ((unsigned short*)t.dst[b])[i] = f2h(((const float*)t.src[b])[i], bf); continue; }\n"
" float v = mode == 1 ? h2f(((const unsigned short*)t.src[b])[i], bf) : ((const float*)t.src[b])[i];\n"
" v = v * scale;\n"
" if (!__builtin_isfinite(v)) bad = 1;\n"
" ((float*)t.dst[b])[i] = v;\n"
" }\n"
" if (bad && flag) *flag = 1.0f;\n"
"}\n"
// Mixed-precision optimizer steps on fp32 master weights, also writing
// the 16-bit copy wh (when not null); nothing happens when *skip != 0
// (a gradient overflowed). Same arithmetic as nn.h's Optimizer::step.
"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"
"#pragma clang fp contract(off)\n"
" if (skip && *skip != 0.0f) return;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n"
" float gi = g[i] + wd * w[i];\n"
" if (vel) { float v = mom * vel[i] + gi; vel[i] = v; gi = v; }\n"
" float nw = w[i] - lr * gi;\n"
" w[i] = nw;\n"
" if (wh) wh[i] = f2h(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) {\n"
"#pragma clang fp contract(off)\n"
" if (skip && *skip != 0.0f) return;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n"
" float gi = g[i];\n"
" if (!decoupled) gi = gi + wd * w[i];\n"
" float mi = b1 * m[i] + (1.0f - b1) * gi;\n"
" float vi = b2 * v[i] + (1.0f - b2) * gi * gi;\n"
" m[i] = mi; v[i] = vi;\n"
" float upd = (mi / bc1) / (sqrtf(vi / bc2) + eps);\n"
" if (decoupled) upd = upd + wd * w[i];\n"
" float nw = w[i] - lr * upd;\n"
" w[i] = nw;\n"
" if (wh) wh[i] = f2h(nw, bf);\n"
" }\n"
"}\n"
// 16-bit elementwise: 0 a+b, 1 a-b, 2 a*b, 3 a*k, 4 a + k*b; fp32 arithmetic.
"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) {\n"
"#pragma clang fp contract(off)\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {\n"
" float x = h2f(a[i], bf), y = op == 3 ? 0.0f : h2f(b[i], bf);\n"
" float r = op == 0 ? x + y : op == 1 ? x - y : op == 2 ? x * y : op == 3 ? x * k : x + k * y;\n"
" o[i] = f2h(r, bf);\n"
" }\n"
"}\n"
// y[i] += bias[(i / plane) % channels], then max(0, .) if relu: row bias
// (plane 1) or NCHW channel bias.
"extern \"C\" __global__ void half_bias(unsigned short* y, const unsigned short* bias, unsigned total,\n"
" unsigned channels, unsigned plane, int relu, int bf) {\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {\n"
" float r = h2f(y[i], bf) + h2f(bias[(i / plane) % channels], bf);\n"
" if (relu && r < 0.0f) r = 0.0f;\n"
" y[i] = f2h(r, bf);\n"
" }\n"
"}\n"
// 16-bit max / avg pooling, kernel x kernel window, no padding, over
// planes [rows][h][w]; fp32 accumulation, one rounding (max copies bits).
"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, total = rows * plane;\n"
" for (unsigned i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += gridDim.x * blockDim.x) {\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 = h2f(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 = h2f(b, bf);\n"
" acc += v;\n"
" if (v > bv) { bv = v; best = b; }\n"
" }\n"
" y[i] = avg ? f2h(acc / (float)(k * k), bf) : best;\n"
" }\n"
"}\n"
"// Flash attention (forward): one block = 64 query rows of one head, 4 lanes\n"
"// per row, each lane holding a quarter of the row's q and output (dims\n"
"// part, part + 4, ...). K/V tiles of 32 rows are staged in LDS; per tile the\n"
"// lanes compute the scores, then rescale the running max / sum / output\n"
"// once (Dao et al.'s online softmax). The [seq, kv] score matrix is never\n"
"// formed. Q/K/V/O are [heads][rows][d] row-major; q is pre-scaled.\n"
"#define FA_BR 64\n"
"#define FA_BC 32\n"
"#define FA_TPR 4\n"
"#define FA_DMAX 128\n"
"template <int DM>\n"
"__device__ __forceinline__ void flash_impl(const float* __restrict__ Q, const float* __restrict__ K,\n"
" const float* __restrict__ V, float* __restrict__ O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) {\n"
" constexpr int NT = DM / FA_TPR; // dims per lane\n"
" __shared__ float Ks[FA_BC][DM + 1];\n"
" __shared__ float Vs[FA_BC][DM + 1];\n"
" const unsigned h = blockIdx.y;\n"
" const unsigned row = blockIdx.x * FA_BR + threadIdx.x / FA_TPR;\n"
" const unsigned part = threadIdx.x % FA_TPR;\n"
" const bool active = row < seq;\n"
" const float* Kh = K + (unsigned long long)h * kv * d;\n"
" const float* Vh = V + (unsigned long long)h * kv * d;\n"
" float qr[NT], acc[NT];\n"
"#pragma unroll\n"
" for (int t = 0; t < NT; ++t) {\n"
" unsigned c = part + FA_TPR * t;\n"
" qr[t] = (active && c < d) ? Q[((unsigned long long)h * seq + row) * d + c] * scale : 0.0f;\n"
" acc[t] = 0.0f;\n"
" }\n"
" float m = -__builtin_inff(), l = 0.0f;\n"
" const int off = (int)kv - (int)seq; // causal: key j visible to row i if j <= i + off\n"
" unsigned kvEnd = kv;\n"
" if (causal) {\n"
" int lim = (int)(blockIdx.x * FA_BR + FA_BR - 1) + off + 1;\n"
" kvEnd = lim <= 0 ? 0u : (lim < (int)kv ? (unsigned)lim : kv);\n"
" }\n"
" for (unsigned j0 = 0; j0 < kvEnd; j0 += FA_BC) {\n"
" for (unsigned e = threadIdx.x; e < FA_BC * d; e += blockDim.x) {\n"
" unsigned r = e / d, c = e % d, j = j0 + r;\n"
" Ks[r][c] = j < kv ? Kh[(unsigned long long)j * d + c] : 0.0f;\n"
" Vs[r][c] = j < kv ? Vh[(unsigned long long)j * d + c] : 0.0f;\n"
" }\n"
" __syncthreads();\n"
" const unsigned jn = kv - j0 < FA_BC ? kv - j0 : FA_BC;\n"
" float s[FA_BC];\n"
" float mt = -__builtin_inff();\n"
"#pragma unroll\n"
" for (int jj = 0; jj < FA_BC; ++jj) {\n"
" float x = 0.0f;\n"
"#pragma unroll\n"
" for (int t = 0; t < NT; ++t) x = fmaf(qr[t], Ks[jj][part + FA_TPR * t], x);\n"
" x += __shfl_xor(x, 1);\n"
" x += __shfl_xor(x, 2);\n"
" bool vis = (unsigned)jj < jn && (!causal || (int)(j0 + jj) <= (int)row + off);\n"
" s[jj] = vis ? x : -__builtin_inff();\n"
" mt = fmaxf(mt, s[jj]);\n"
" }\n"
" if (mt > -__builtin_inff()) {\n"
" float mn = fmaxf(m, mt);\n"
" float corr = __expf(m - mn);\n"
" l *= corr;\n"
"#pragma unroll\n"
" for (int t = 0; t < NT; ++t) acc[t] *= corr;\n"
"#pragma unroll\n"
" for (int jj = 0; jj < FA_BC; ++jj) {\n"
" float p = __expf(s[jj] - mn);\n"
" l += p;\n"
"#pragma unroll\n"
" for (int t = 0; t < NT; ++t) acc[t] = fmaf(p, Vs[jj][part + FA_TPR * t], acc[t]);\n"
" }\n"
" m = mn;\n"
" }\n"
" __syncthreads();\n"
" }\n"
" if (active) {\n"
" float inv = l > 0.0f ? 1.0f / l : 0.0f;\n"
"#pragma unroll\n"
" for (int t = 0; t < NT; ++t) {\n"
" unsigned c = part + FA_TPR * t;\n"
" if (c < d) O[((unsigned long long)h * seq + row) * d + c] = acc[t] * inv;\n"
" }\n"
" }\n"
"}\n"
"extern \"C\" __global__ __launch_bounds__(256) void flash_attn_d64(const float* Q, const float* K, const float* V, float* O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) { flash_impl<64>(Q, K, V, O, seq, kv, d, scale, causal); }\n"
"extern \"C\" __global__ __launch_bounds__(256) void flash_attn_d128(const float* Q, const float* K, const float* V, float* O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) { flash_impl<128>(Q, K, V, O, seq, kv, d, scale, causal); }\n"
"\n"
"// Linear attention, chunked (32 rows): phi(x) = elu(x) + 1.\n"
"// state: S_c = sum_{j in chunk c} phi(k_j)^T [v_j | 1] ([d][d+1] per chunk)\n"
"// prefix: causal -> exclusive prefix over chunks; else the total, in chunk 0\n"
"// out: row i of chunk c: (phi(q_i) P + sum_{j<=i in c} (phi(q_i).phi(k_j)) [v_j|1]) split as\n"
"// numerator columns 0..d-1 over the denominator column d\n"
"#define LA_C 32\n"
"__device__ __forceinline__ float la_phi(float x) { return x > 0.0f ? x + 1.0f : __expf(x); }\n"
"extern \"C\" __global__ __launch_bounds__(256) void lin_attn_state(const float* __restrict__ K, const float* __restrict__ V,\n"
" float* __restrict__ S, unsigned kv, unsigned d, unsigned nchunks) {\n"
" __shared__ float pk[LA_C][FA_DMAX];\n"
" __shared__ float vv[LA_C][FA_DMAX + 1];\n"
" const unsigned c = blockIdx.x, h = blockIdx.y, w = d + 1;\n"
" const unsigned j0 = c * LA_C;\n"
" for (unsigned e = threadIdx.x; e < LA_C * w; e += blockDim.x) {\n"
" unsigned r = e / w, col = e % w, j = j0 + r;\n"
" bool in = j < kv;\n"
" if (col < d) {\n"
" pk[r][col] = in ? la_phi(K[((unsigned long long)h * kv + j) * d + col]) : 0.0f;\n"
" vv[r][col] = in ? V[((unsigned long long)h * kv + j) * d + col] : 0.0f;\n"
" } else vv[r][d] = in ? 1.0f : 0.0f;\n"
" }\n"
" __syncthreads();\n"
" float* Sc = S + ((unsigned long long)h * nchunks + c) * d * w;\n"
" for (unsigned e = threadIdx.x; e < d * w; e += blockDim.x) {\n"
" unsigned r = e / w, col = e % w;\n"
" float a = 0.0f;\n"
"#pragma unroll 8\n"
" for (int jj = 0; jj < LA_C; ++jj) a = fmaf(pk[jj][r], vv[jj][col], a);\n"
" Sc[e] = a;\n"
" }\n"
"}\n"
"extern \"C\" __global__ __launch_bounds__(256) void lin_attn_prefix(float* __restrict__ S, unsigned d, unsigned nchunks, int causal) {\n"
" const unsigned w = d + 1, n = d * w, h = blockIdx.y;\n"
" const unsigned e = blockIdx.x * blockDim.x + threadIdx.x;\n"
" if (e >= n) return;\n"
" float* base = S + (unsigned long long)h * nchunks * n + e;\n"
" float run = 0.0f;\n"
" for (unsigned c = 0; c < nchunks; ++c) {\n"
" float t = base[(unsigned long long)c * n];\n"
" if (causal) base[(unsigned long long)c * n] = run;\n"
" run += t;\n"
" }\n"
" if (!causal) base[0] = run;\n"
"}\n"
"extern \"C\" __global__ __launch_bounds__(256) void lin_attn_out(const float* __restrict__ Q, const float* __restrict__ K,\n"
" const float* __restrict__ V, const float* __restrict__ S, float* __restrict__ O,\n"
" unsigned seq, unsigned kv, unsigned d, unsigned nchunks, int causal) {\n"
" __shared__ float pq[LA_C][FA_DMAX];\n"
" __shared__ float pk[LA_C][FA_DMAX + 1]; // later reused for the [LA_C][d+1] results\n"
" __shared__ float vv[LA_C][FA_DMAX + 1];\n"
" __shared__ float A[LA_C][LA_C + 1];\n"
" const unsigned c = blockIdx.x, h = blockIdx.y, w = d + 1;\n"
" const unsigned i0 = c * LA_C;\n"
" const float* P = S + ((unsigned long long)h * nchunks + (causal ? c : 0u)) * d * w;\n"
" for (unsigned e = threadIdx.x; e < LA_C * w; e += blockDim.x) {\n"
" unsigned r = e / w, col = e % w, i = i0 + r;\n"
" if (col < d) {\n"
" pq[r][col] = i < seq ? la_phi(Q[((unsigned long long)h * seq + i) * d + col]) : 0.0f;\n"
" if (causal) {\n"
" pk[r][col] = i < kv ? la_phi(K[((unsigned long long)h * kv + i) * d + col]) : 0.0f;\n"
" vv[r][col] = i < kv ? V[((unsigned long long)h * kv + i) * d + col] : 0.0f;\n"
" }\n"
" } else if (causal) vv[r][d] = i < kv ? 1.0f : 0.0f;\n"
" }\n"
" __syncthreads();\n"
" if (causal) {\n"
" for (unsigned e = threadIdx.x; e < LA_C * LA_C; e += blockDim.x) {\n"
" unsigned i = e / LA_C, j = e % LA_C;\n"
" float a = 0.0f;\n"
" if (j <= i) for (unsigned r = 0; r < d; ++r) a = fmaf(pq[i][r], pk[j][r], a);\n"
" A[i][j] = a;\n"
" }\n"
" __syncthreads();\n"
" }\n"
" float res[17];\n"
" int nres = 0;\n"
" for (unsigned e = threadIdx.x; e < LA_C * w; e += blockDim.x) {\n"
" unsigned i = e / w, col = e % w;\n"
" float a = 0.0f;\n"
" for (unsigned r = 0; r < d; ++r) a = fmaf(pq[i][r], P[(unsigned long long)r * w + col], a);\n"
" if (causal)\n"
" for (unsigned j = 0; j <= i; ++j) a = fmaf(A[i][j], vv[j][col], a);\n"
" res[nres++] = a;\n"
" }\n"
" __syncthreads(); // pk no longer needed: results go there\n"
" nres = 0;\n"
" for (unsigned e = threadIdx.x; e < LA_C * w; e += blockDim.x) pk[e / w][e % w] = res[nres++];\n"
" __syncthreads();\n"
" for (unsigned e = threadIdx.x; e < LA_C * d; e += blockDim.x) {\n"
" unsigned i = e / d, col = e % d;\n"
" if (i0 + i < seq) {\n"
" float den = pk[i][d];\n"
" O[((unsigned long long)h * seq + i0 + i) * d + col] = den != 0.0f ? pk[i][col] / den : 0.0f;\n"
" }\n"
" }\n"
"}\n"
"// Decode-step ops (one token at a time, M = 1).\n"
"// gemv_half: y = W x with W [n][k] fp16/bf16 (row-major) and fp32 x, y; one\n"
"// wave per output row, 16-byte weight loads when k % 8 == 0. Memory-bound:\n"
"// ~1 TB/s on an MI100, where a GEMM with one row (rocBLAS) is far slower.\n"
"__device__ __forceinline__ float gv_wsum(float a) {\n"
" a += __shfl_xor(a, 32); a += __shfl_xor(a, 16); a += __shfl_xor(a, 8);\n"
" a += __shfl_xor(a, 4); a += __shfl_xor(a, 2); a += __shfl_xor(a, 1);\n"
" return a;\n"
"}\n"
"template <int BF>\n"
"__device__ __forceinline__ void gemv_half_impl(const unsigned short* __restrict__ W, const float* __restrict__ x,\n"
" float* __restrict__ y, unsigned n, unsigned k) {\n"
" const unsigned lane = threadIdx.x & 63u;\n"
" const unsigned row = blockIdx.x * (blockDim.x >> 6) + (threadIdx.x >> 6);\n"
" if (row >= n) return;\n"
" const unsigned short* w = W + (unsigned long long)row * k;\n"
" float a = 0.0f;\n"
" if ((k & 7u) == 0u) {\n"
"#pragma unroll 2\n"
" for (unsigned kk = lane * 8u; kk < k; kk += 512u) {\n"
" const uint4 v = *(const uint4*)(w + kk);\n"
" const float4 x0 = *(const float4*)(x + kk), x1 = *(const float4*)(x + kk + 4u);\n"
" a = fmaf(h2f((unsigned short)(v.x & 0xffffu), BF), x0.x, a); a = fmaf(h2f((unsigned short)(v.x >> 16), BF), x0.y, a);\n"
" a = fmaf(h2f((unsigned short)(v.y & 0xffffu), BF), x0.z, a); a = fmaf(h2f((unsigned short)(v.y >> 16), BF), x0.w, a);\n"
" a = fmaf(h2f((unsigned short)(v.z & 0xffffu), BF), x1.x, a); a = fmaf(h2f((unsigned short)(v.z >> 16), BF), x1.y, a);\n"
" a = fmaf(h2f((unsigned short)(v.w & 0xffffu), BF), x1.z, a); a = fmaf(h2f((unsigned short)(v.w >> 16), BF), x1.w, a);\n"
" }\n"
" } else {\n"
" for (unsigned kk = lane; kk < k; kk += 64u) a = fmaf(h2f(w[kk], BF), x[kk], a);\n"
" }\n"
" a = gv_wsum(a);\n"
" if (lane == 0u) y[row] = a;\n"
"}\n"
"extern \"C\" __global__ __launch_bounds__(256) void gemv_half(const unsigned short* W, const float* x, float* y,\n"
" unsigned n, unsigned k, int bf) {\n"
" if (bf) gemv_half_impl<1>(W, x, y, n, k);\n"
" else gemv_half_impl<0>(W, x, y, n, k);\n"
"}\n"
"// rmsnorm_rows: y[r] = x[r] * rsqrt(mean(x[r]^2) + eps) * w (w == 0: no\n"
"// weight); one block per row; y may alias x.\n"
"extern \"C\" __global__ __launch_bounds__(256) void rmsnorm_rows(const float* x, const float* __restrict__ w,\n"
" float* y, unsigned n, float eps) {\n"
" __shared__ float red[4];\n"
" const float* xr = x + (unsigned long long)blockIdx.x * n;\n"
" float* yr = y + (unsigned long long)blockIdx.x * n;\n"
" float s = 0.0f;\n"
" for (unsigned i = threadIdx.x; i < n; i += 256u) { float v = xr[i]; s = fmaf(v, v, s); }\n"
" s = gv_wsum(s);\n"
" if ((threadIdx.x & 63u) == 0u) red[threadIdx.x >> 6] = s;\n"
" __syncthreads();\n"
" const float inv = rsqrtf((red[0] + red[1] + red[2] + red[3]) / (float)n + eps);\n"
" for (unsigned i = threadIdx.x; i < n; i += 256u) yr[i] = xr[i] * inv * (w ? w[i] : 1.0f);\n"
"}\n"
"// ── Attention on matrix cores (CDNA: gfx908 / gfx90a / gfx94x / gfx950) ──────\n"
"// fp32 MFMA versions of the flash and linear attention kernels above, used\n"
"// when the GPU has matrix cores (mfma_probe). On other GPUs the kernels\n"
"// compile to empty bodies and the CUDA-core kernels above run instead.\n"
"#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) || defined(__gfx950__)\n"
"#define SAFEC_MFMA 1\n"
"#else\n"
"#define SAFEC_MFMA 0\n"
"#endif\n"
"extern \"C\" __global__ void mfma_probe(float* o) { if (threadIdx.x == 0 && blockIdx.x == 0) o[0] = (float)SAFEC_MFMA; }\n"
"#if SAFEC_MFMA\n"
"typedef float fm_v16 __attribute__((ext_vector_type(16)));\n"
"typedef float fm_v4 __attribute__((ext_vector_type(4)));\n"
"// Flash attention. Grid (heads, query blocks); query blocks run last-first so\n"
"// the longest causal blocks start first. Scores are computed transposed\n"
"// (S^T = K Q^T), so each lane owns one query row: its q (pre-scaled by\n"
"// scale * log2 e) stays in registers as the MFMA B operand, the row max and\n"
"// sum need one cross-half shuffle, and P feeds the P V product straight from\n"
"// registers. The output is rescaled only when some row's max grew by\n"
"// more than 2^8.\n"
"// - head dim <= 64: 32x32x2 MFMA, 4 waves x 32 rows, K/V tiles of 64 keys\n"
"// - head dim <= 128: 16x16x4 MFMA, 4 waves x 16 rows (half the registers)\n"
"#define FM_BC 32\n"
"#define FM_T 64\n"
"#define FM_BR 128\n"
"#define FM16_BR 64\n"
"template <int DM>\n"
"__device__ __forceinline__ void flash_m32_impl(const float* __restrict__ Q, const float* __restrict__ K,\n"
" const float* __restrict__ V, float* __restrict__ O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) {\n"
" constexpr int HD = DM / 2;\n"
" constexpr int NB = DM / 32;\n"
" constexpr int D4 = DM / 4;\n"
" __shared__ float4 Ks[FM_T * D4];\n"
" __shared__ float Vs[FM_T * DM];\n"
" const unsigned h = blockIdx.x;\n"
" const unsigned qb = gridDim.y - 1u - blockIdx.y;\n"
" const unsigned lane = threadIdx.x & 63u, wave = threadIdx.x >> 6;\n"
" const unsigned col = lane & 31u, half = lane >> 5;\n"
" const unsigned q0 = qb * FM_BR + wave * 32u;\n"
" const unsigned qrow = q0 + col;\n"
" const float* Kh = K + (unsigned long long)h * kv * d;\n"
" const float* Vh = V + (unsigned long long)h * kv * d;\n"
" const float qs = scale * 1.4426950408889634f;\n"
" float qr[HD];\n"
"#pragma unroll\n"
" for (int s = 0; s < HD; ++s) {\n"
" unsigned c = half * HD + s;\n"
" qr[s] = (qrow < seq && c < d) ? Q[((unsigned long long)h * seq + qrow) * d + c] * qs : 0.0f;\n"
" }\n"
" fm_v16 acc[NB];\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) acc[b][r] = 0.0f;\n"
" float m = -__builtin_inff(), l = 0.0f;\n"
" const int off = (int)kv - (int)seq;\n"
" unsigned kvEnd = kv, wEnd = kv;\n"
" if (causal) {\n"
" int lim = (int)(qb * FM_BR + FM_BR - 1) + off + 1;\n"
" kvEnd = lim <= 0 ? 0u : (lim < (int)kv ? (unsigned)lim : kv);\n"
" int wl = (int)(q0 + 31u) + off + 1;\n"
" wEnd = wl <= 0 ? 0u : (wl < (int)kv ? (unsigned)wl : kv);\n"
" }\n"
" const bool vec = (d & 3u) == 0u;\n"
" for (unsigned jt = 0; jt < kvEnd; jt += FM_T) {\n"
" __syncthreads();\n"
" for (unsigned e = threadIdx.x; e < FM_T * D4; e += blockDim.x) {\n"
" unsigned r = e / D4, c4 = e % D4, j = jt + r, c = c4 * 4u;\n"
" float4 kq = make_float4(0.f, 0.f, 0.f, 0.f), vq = kq;\n"
" if (j < kv && c < d) {\n"
" const float* kp = Kh + (unsigned long long)j * d + c;\n"
" const float* vp = Vh + (unsigned long long)j * d + c;\n"
" if (vec) { kq = *(const float4*)kp; vq = *(const float4*)vp; }\n"
" else {\n"
" kq.x = kp[0]; vq.x = vp[0];\n"
" if (c + 1 < d) { kq.y = kp[1]; vq.y = vp[1]; }\n"
" if (c + 2 < d) { kq.z = kp[2]; vq.z = vp[2]; }\n"
" if (c + 3 < d) { kq.w = kp[3]; vq.w = vp[3]; }\n"
" }\n"
" }\n"
" Ks[r * D4 + (c4 ^ (r & 7u))] = kq;\n"
" *(float4*)&Vs[r * DM + c] = vq;\n"
" }\n"
" __syncthreads();\n"
" for (unsigned sub = 0; sub < FM_T / FM_BC; ++sub) {\n"
" const unsigned j0 = jt + sub * FM_BC;\n"
" if (j0 >= wEnd || j0 >= kvEnd) break;\n"
" fm_v16 sc;\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) sc[r] = 0.0f;\n"
" const float4* kr = Ks + (sub * FM_BC + col) * D4;\n"
"#pragma unroll\n"
" for (int s4 = 0; s4 < HD / 4; ++s4) {\n"
" float4 kk = kr[(half * (HD / 4) + s4) ^ (col & 7u)];\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.x, qr[4 * s4 + 0], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.y, qr[4 * s4 + 1], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.z, qr[4 * s4 + 2], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.w, qr[4 * s4 + 3], sc, 0, 0, 0);\n"
" }\n"
" float p[16];\n"
" float mt = -__builtin_inff();\n"
" const bool edge = j0 + FM_BC > kv || (causal && (int)(j0 + FM_BC - 1) > (int)q0 + off);\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" float x = sc[r];\n"
" if (edge) {\n"
" unsigned kj = j0 + 8u * (r / 4) + 4u * half + (r % 4);\n"
" bool vis = kj < kv && (!causal || (int)kj <= (int)qrow + off);\n"
" x = vis ? x : -__builtin_inff();\n"
" }\n"
" p[r] = x;\n"
" mt = fmaxf(mt, x);\n"
" }\n"
" mt = fmaxf(mt, __shfl_xor(mt, 32));\n"
" // keep the old max until a row's max grows by more than 2^8 (as\n"
" // FlashAttention-4): p stays <= 256, and most tiles skip the rescale\n"
" const bool grow = mt > m + 8.0f;\n"
" float mn = grow ? mt : m;\n"
" float ms = mn == -__builtin_inff() ? 0.0f : mn;\n"
" float ls = 0.0f;\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) { p[r] = exp2f(p[r] - ms); ls += p[r]; }\n"
" if (__any(grow)) {\n"
" float corr = exp2f(m - ms);\n"
" l *= corr;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) acc[b][r] *= corr;\n"
" }\n"
" l += ls;\n"
" m = mn;\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" const float* vr = Vs + sub * FM_BC * DM + (8u * (r / 4) + 4u * half + (r % 4)) * DM + col;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
" acc[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(vr[32 * b], p[r], acc[b], 0, 0, 0);\n"
" }\n"
" }\n"
" }\n"
" l += __shfl_xor(l, 32);\n"
" if (qrow < seq) {\n"
" float inv = l > 0.0f ? 1.0f / l : 0.0f;\n"
" float* orow = O + ((unsigned long long)h * seq + qrow) * d;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" unsigned c = 32u * b + 8u * (r / 4) + 4u * half + (r % 4);\n"
" if (c < d) orow[c] = acc[b][r] * inv;\n"
" }\n"
" }\n"
"}\n"
"template <int DM>\n"
"__device__ __forceinline__ void flash_m16_impl(const float* __restrict__ Q, const float* __restrict__ K,\n"
" const float* __restrict__ V, float* __restrict__ O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) {\n"
" constexpr int QD = DM / 4; // dims per lane group in QK^T\n"
" constexpr int NB = DM / 16; // 16-dim output blocks\n"
" constexpr int D4 = DM / 4;\n"
" __shared__ float4 Ks[FM_BC * D4];\n"
" __shared__ float4 Vs[FM_BC * D4];\n"
" const unsigned h = blockIdx.x;\n"
" const unsigned qb = gridDim.y - 1u - blockIdx.y;\n"
" const unsigned lane = threadIdx.x & 63u, wave = threadIdx.x >> 6;\n"
" const unsigned col = lane & 15u, grp = lane >> 4;\n"
" const unsigned q0 = qb * FM16_BR + wave * 16u;\n"
" const unsigned qrow = q0 + col;\n"
" const float* Kh = K + (unsigned long long)h * kv * d;\n"
" const float* Vh = V + (unsigned long long)h * kv * d;\n"
" const float qs = scale * 1.4426950408889634f;\n"
" float qr[QD];\n"
"#pragma unroll\n"
" for (int s = 0; s < QD; ++s) {\n"
" unsigned c = grp * QD + s;\n"
" qr[s] = (qrow < seq && c < d) ? Q[((unsigned long long)h * seq + qrow) * d + c] * qs : 0.0f;\n"
" }\n"
" fm_v4 acc[NB];\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b) acc[b] = (fm_v4){0.f, 0.f, 0.f, 0.f};\n"
" float m = -__builtin_inff(), l = 0.0f;\n"
" const int off = (int)kv - (int)seq;\n"
" unsigned kvEnd = kv, wEnd = kv;\n"
" if (causal) {\n"
" int lim = (int)(qb * FM16_BR + FM16_BR - 1) + off + 1;\n"
" kvEnd = lim <= 0 ? 0u : (lim < (int)kv ? (unsigned)lim : kv);\n"
" int wl = (int)(q0 + 15u) + off + 1;\n"
" wEnd = wl <= 0 ? 0u : (wl < (int)kv ? (unsigned)wl : kv);\n"
" }\n"
" const bool vec = (d & 3u) == 0u;\n"
" for (unsigned j0 = 0; j0 < kvEnd; j0 += FM_BC) {\n"
" __syncthreads();\n"
" for (unsigned e = threadIdx.x; e < FM_BC * D4; e += blockDim.x) {\n"
" unsigned r = e / D4, c4 = e % D4, j = j0 + r, c = c4 * 4u;\n"
" float4 kq = make_float4(0.f, 0.f, 0.f, 0.f), vq = kq;\n"
" if (j < kv && c < d) {\n"
" const float* kp = Kh + (unsigned long long)j * d + c;\n"
" const float* vp = Vh + (unsigned long long)j * d + c;\n"
" if (vec) { kq = *(const float4*)kp; vq = *(const float4*)vp; }\n"
" else {\n"
" kq.x = kp[0]; vq.x = vp[0];\n"
" if (c + 1 < d) { kq.y = kp[1]; vq.y = vp[1]; }\n"
" if (c + 2 < d) { kq.z = kp[2]; vq.z = vp[2]; }\n"
" if (c + 3 < d) { kq.w = kp[3]; vq.w = vp[3]; }\n"
" }\n"
" }\n"
" Ks[r * D4 + (c4 ^ (r & 7u))] = kq;\n"
" Vs[r * D4 + (c4 ^ (((r >> 2) & 1u) << 2))] = vq;\n"
" }\n"
" __syncthreads();\n"
" if (j0 >= wEnd) continue;\n"
" fm_v4 sc[2];\n"
"#pragma unroll\n"
" for (int t = 0; t < 2; ++t) {\n"
" sc[t] = (fm_v4){0.f, 0.f, 0.f, 0.f};\n"
" const unsigned kr = 16u * t + col;\n"
" const float4* krow = Ks + kr * D4;\n"
"#pragma unroll\n"
" for (int s4 = 0; s4 < QD / 4; ++s4) {\n"
" float4 kk = krow[(grp * (QD / 4) + s4) ^ (kr & 7u)];\n"
" sc[t] = __builtin_amdgcn_mfma_f32_16x16x4f32(kk.x, qr[4 * s4 + 0], sc[t], 0, 0, 0);\n"
" sc[t] = __builtin_amdgcn_mfma_f32_16x16x4f32(kk.y, qr[4 * s4 + 1], sc[t], 0, 0, 0);\n"
" sc[t] = __builtin_amdgcn_mfma_f32_16x16x4f32(kk.z, qr[4 * s4 + 2], sc[t], 0, 0, 0);\n"
" sc[t] = __builtin_amdgcn_mfma_f32_16x16x4f32(kk.w, qr[4 * s4 + 3], sc[t], 0, 0, 0);\n"
" }\n"
" }\n"
" float p[8];\n"
" float mt = -__builtin_inff();\n"
" const bool edge = j0 + FM_BC > kv || (causal && (int)(j0 + FM_BC - 1) > (int)q0 + off);\n"
"#pragma unroll\n"
" for (int t = 0; t < 2; ++t)\n"
"#pragma unroll\n"
" for (int r = 0; r < 4; ++r) {\n"
" float x = sc[t][r];\n"
" if (edge) {\n"
" unsigned kj = j0 + 16u * t + 4u * grp + r;\n"
" bool vis = kj < kv && (!causal || (int)kj <= (int)qrow + off);\n"
" x = vis ? x : -__builtin_inff();\n"
" }\n"
" p[4 * t + r] = x;\n"
" mt = fmaxf(mt, x);\n"
" }\n"
" mt = fmaxf(mt, __shfl_xor(mt, 16));\n"
" mt = fmaxf(mt, __shfl_xor(mt, 32));\n"
" // keep the old max until a row's max grows by more than 2^8 (as\n"
" // FlashAttention-4): p stays <= 256, and most tiles skip the rescale\n"
" const bool grow = mt > m + 8.0f;\n"
" float mn = grow ? mt : m;\n"
" float ms = mn == -__builtin_inff() ? 0.0f : mn;\n"
" float ls = 0.0f;\n"
"#pragma unroll\n"
" for (int r = 0; r < 8; ++r) { p[r] = exp2f(p[r] - ms); ls += p[r]; }\n"
" if (__any(grow)) {\n"
" float corr = exp2f(m - ms);\n"
" l *= corr;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b) acc[b] *= corr;\n"
" }\n"
" l += ls;\n"
" m = mn;\n"
" const float* vf = (const float*)Vs;\n"
"#pragma unroll\n"
" for (int t = 0; t < 2; ++t)\n"
"#pragma unroll\n"
" for (int r = 0; r < 4; ++r) {\n"
" const unsigned key = 16u * t + 4u * grp + r;\n"
" const float* vr = vf + key * DM;\n"
" const unsigned sw = ((key >> 2) & 1u) << 4;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
" acc[b] = __builtin_amdgcn_mfma_f32_16x16x4f32(vr[(16u * b + col) ^ sw], p[4 * t + r], acc[b], 0, 0, 0);\n"
" }\n"
" }\n"
" l += __shfl_xor(l, 16);\n"
" l += __shfl_xor(l, 32);\n"
" if (qrow < seq) {\n"
" float inv = l > 0.0f ? 1.0f / l : 0.0f;\n"
" float* orow = O + ((unsigned long long)h * seq + qrow) * d;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 4; ++r) {\n"
" unsigned c = 16u * b + 4u * grp + r;\n"
" if (c < d) orow[c] = acc[b][r] * inv;\n"
" }\n"
" }\n"
"}\n"
"// Linear attention, chunk-parallel, three passes (rows in chunks of LA_R):\n"
"// la_state: W[c] = phi(K_c)^T V_c and Z[c] = sum phi(K_c) per chunk (MFMA)\n"
"// la_scan: causal: exclusive prefix over chunks; else the total, in chunk 0\n"
"// la_out: out = (phi(Q_c) W + A V_c) / (phi(Q_c) . Z + rowsum A), A the\n"
"// causal intra-chunk scores phi(Q_c) phi(K_c)^T (j <= i); phi(Q)\n"
"// rows sit in registers as the MFMA operand, like flash's q.\n"
"#define LA_T 32\n"
"#define LA_R 128\n"
"// K/V tile staging, split so the next tile's global loads can be issued\n"
"// before the current tile is used: la_fetch reads rows [j0, j0 + LA_T) of\n"
"// phi(K) (all DM dims) and of V (cols c0 .. c0 + VC) into registers, la_put\n"
"// writes them to Ks (swizzled float4, as the flash kernel) and Vs ([LA_T][VC]).\n"
"template <int DM, int VC> struct LaTile {\n"
" static constexpr int KE = LA_T * DM / 4 / 256, VE = LA_T * VC / 4 / 256;\n"
" float4 k[KE], v[VE];\n"
"};\n"
"__device__ __forceinline__ float4 la_ld4(const float* p, unsigned c, unsigned d, bool vec, bool phi) {\n"
" float4 x = make_float4(0.f, 0.f, 0.f, 0.f);\n"
" if (vec) x = *(const float4*)p;\n"
" else {\n"
" x.x = p[0];\n"
" if (c + 1 < d) x.y = p[1];\n"
" if (c + 2 < d) x.z = p[2];\n"
" if (c + 3 < d) x.w = p[3];\n"
" }\n"
" if (phi) {\n"
" x.x = la_phi(x.x); x.y = c + 1 < d ? la_phi(x.y) : 0.f;\n"
" x.z = c + 2 < d ? la_phi(x.z) : 0.f; x.w = c + 3 < d ? la_phi(x.w) : 0.f;\n"
" }\n"
" return x;\n"
"}\n"
"template <int DM, int VC>\n"
"__device__ __forceinline__ void la_fetch(LaTile<DM, VC>& t, const float* __restrict__ Kh, const float* __restrict__ Vh,\n"
" unsigned j0, unsigned jend, unsigned d, unsigned c0) {\n"
" constexpr int D4 = DM / 4, V4 = VC / 4;\n"
" const bool vec = (d & 3u) == 0u;\n"
"#pragma unroll\n"
" for (int i = 0; i < t.KE; ++i) {\n"
" unsigned e = threadIdx.x + 256u * i, r = e / D4, c = (e % D4) * 4u, j = j0 + r;\n"
" t.k[i] = (j < jend && c < d) ? la_ld4(Kh + (unsigned long long)j * d + c, c, d, vec, true) : make_float4(0.f, 0.f, 0.f, 0.f);\n"
" }\n"
"#pragma unroll\n"
" for (int i = 0; i < t.VE; ++i) {\n"
" unsigned e = threadIdx.x + 256u * i, r = e / V4, c = c0 + (e % V4) * 4u, j = j0 + r;\n"
" t.v[i] = (j < jend && c < d) ? la_ld4(Vh + (unsigned long long)j * d + c, c, d, vec, false) : make_float4(0.f, 0.f, 0.f, 0.f);\n"
" }\n"
"}\n"
"template <int DM, int VC>\n"
"__device__ __forceinline__ void la_put(const LaTile<DM, VC>& t, float4* Ks, float* Vs) {\n"
" constexpr int D4 = DM / 4, V4 = VC / 4;\n"
"#pragma unroll\n"
" for (int i = 0; i < t.KE; ++i) {\n"
" unsigned e = threadIdx.x + 256u * i, r = e / D4, c4 = e % D4;\n"
" Ks[r * D4 + (c4 ^ (r & 7u))] = t.k[i];\n"
" }\n"
"#pragma unroll\n"
" for (int i = 0; i < t.VE; ++i) {\n"
" unsigned e = threadIdx.x + 256u * i, r = e / V4;\n"
" *(float4*)&Vs[r * VC + (e % V4) * 4u] = t.v[i];\n"
" }\n"
"}\n"
"\n"
"// Pass 1: W[h][c] = phi(K_c)^T V_c ([DM][DM]) and Z[h][c] = sum phi(K_c) ([DM])\n"
"// over rows [c*rows, (c+1)*rows). Wave w owns the 32x32 blocks w, w+4, ...\n"
"template <int DM>\n"
"__device__ __forceinline__ void la_state_impl(const float* __restrict__ K, const float* __restrict__ V,\n"
" float* __restrict__ W, float* __restrict__ Z,\n"
" unsigned kv, unsigned d, unsigned rows, unsigned nc) {\n"
" constexpr int D4 = DM / 4, NBL = (DM / 32) * (DM / 32), PW = (NBL + 3) / 4;\n"
" __shared__ float4 Ks[LA_T * D4];\n"
" __shared__ float Vs[LA_T * DM];\n"
" const unsigned c = blockIdx.x, h = blockIdx.y;\n"
" const unsigned lane = threadIdx.x & 63u, wave = threadIdx.x >> 6, col = lane & 31u, half = lane >> 5;\n"
" const float* Kh = K + (unsigned long long)h * kv * d;\n"
" const float* Vh = V + (unsigned long long)h * kv * d;\n"
" const unsigned r0 = c * rows, r1 = min(kv, r0 + rows);\n"
" fm_v16 acc[PW];\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) acc[b][r] = 0.0f;\n"
" float z = 0.0f;\n"
" const float* kf = (const float*)Ks;\n"
" LaTile<DM, DM> tile;\n"
" if (r0 < r1) la_fetch<DM, DM>(tile, Kh, Vh, r0, r1, d, 0u);\n"
" for (unsigned j0 = r0; j0 < r1; j0 += LA_T) {\n"
" __syncthreads();\n"
" la_put<DM, DM>(tile, Ks, Vs);\n"
" __syncthreads();\n"
" if (j0 + LA_T < r1) la_fetch<DM, DM>(tile, Kh, Vh, j0 + LA_T, r1, d, 0u);\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b) {\n"
" const int blk = wave + 4 * b;\n"
" if (blk < NBL) {\n"
" const unsigned bi = blk / (DM / 32), bj = blk % (DM / 32);\n"
" const unsigned dk = 32u * bi + col;\n"
"#pragma unroll\n"
" for (int s = 0; s < LA_T / 2; ++s) {\n"
" const unsigned key = 2u * s + half;\n"
" float a = kf[key * DM + ((((dk >> 2) ^ (key & 7u)) << 2) | (dk & 3u))];\n"
" float bb = Vs[key * DM + 32u * bj + col];\n"
" acc[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(a, bb, acc[b], 0, 0, 0);\n"
" }\n"
" }\n"
" }\n"
" if (threadIdx.x < DM) {\n"
" const unsigned dk = threadIdx.x;\n"
"#pragma unroll 8\n"
" for (int key = 0; key < LA_T; ++key) z += kf[key * DM + ((((dk >> 2) ^ (key & 7u)) << 2) | (dk & 3u))];\n"
" }\n"
" }\n"
" float* w = W + ((unsigned long long)h * nc + c) * DM * DM;\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b) {\n"
" const int blk = wave + 4 * b;\n"
" if (blk < NBL) {\n"
" const unsigned bi = blk / (DM / 32), bj = blk % (DM / 32);\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r)\n"
" w[(32u * bi + 8u * (r / 4) + 4u * half + (r % 4)) * DM + 32u * bj + col] = acc[b][r];\n"
" }\n"
" }\n"
" if (threadIdx.x < DM) Z[((unsigned long long)h * nc + c) * DM + threadIdx.x] = z;\n"
"}\n"
"\n"
"// Pass 2: causal -> exclusive prefix over chunks, in place; else the total, in chunk 0.\n"
"template <int DM>\n"
"__device__ __forceinline__ void la_scan_impl(float* __restrict__ W, float* __restrict__ Z, unsigned nc, int causal) {\n"
" const unsigned h = blockIdx.y, e = blockIdx.x * blockDim.x + threadIdx.x;\n"
" if (e >= DM * DM + DM) return;\n"
" float* p;\n"
" unsigned long long stride;\n"
" if (e < DM * DM) { p = W + (unsigned long long)h * nc * DM * DM + e; stride = DM * DM; }\n"
" else { p = Z + (unsigned long long)h * nc * DM + (e - DM * DM); stride = DM; }\n"
" float run = 0.0f;\n"
" for (unsigned c0 = 0; c0 < nc; c0 += 8) {\n"
" float v[8];\n"
"#pragma unroll\n"
" for (int u = 0; u < 8; ++u) v[u] = c0 + u < nc ? p[(c0 + u) * stride] : 0.0f;\n"
" if (causal) {\n"
"#pragma unroll\n"
" for (int u = 0; u < 8; ++u)\n"
" if (c0 + u < nc) { p[(c0 + u) * stride] = run; run += v[u]; }\n"
" } else {\n"
"#pragma unroll\n"
" for (int u = 0; u < 8; ++u) run += v[u];\n"
" }\n"
" }\n"
" if (!causal) p[0] = run;\n"
"}\n"
"\n"
"// Pass 3: rows [blk*LA_R, +LA_R) of the output, columns c0 .. c0 + DO\n"
"// (c0 = blockIdx.z * DO). Wave w: rows +32w. Lane: row col, phi(q) dims\n"
"// half*DM/2 .. in registers.\n"
"template <int DM, int DO>\n"
"__device__ __forceinline__ void la_out_impl(const float* __restrict__ Q, const float* __restrict__ K,\n"
" const float* __restrict__ V, const float* __restrict__ W,\n"
" const float* __restrict__ Z, float* __restrict__ O,\n"
" unsigned seq, unsigned kv, unsigned d, unsigned nc, int causal) {\n"
" constexpr int HD = DM / 2, NB = DO / 32, D4 = DM / 4;\n"
" constexpr int WR = DM < 64 ? DM : 64; // state rows per pass (both half-waves)\n"
" constexpr int NP = DM / WR, SP = HD / NP; // passes, dims per half-wave per pass\n"
" constexpr int P1 = LA_T * DM + LA_T * DO, P2 = WR * DO, P3 = 4 * 32 * DO;\n"
" constexpr int POOL = P1 > P2 ? (P1 > P3 ? P1 : P3) : (P2 > P3 ? P2 : P3); // floats, phases reuse it\n"
" __shared__ float4 pool[POOL / 4];\n"
" float4* Ks = pool;\n"
" float* Vs = (float*)(pool + LA_T * D4);\n"
" float* Wl = (float*)pool;\n"
" const unsigned blk = blockIdx.x, h = blockIdx.y, c0 = blockIdx.z * DO;\n"
" const unsigned lane = threadIdx.x & 63u, wave = threadIdx.x >> 6, col = lane & 31u, half = lane >> 5;\n"
" const unsigned row = blk * LA_R + wave * 32u + col;\n"
" float qr[HD];\n"
" {\n"
" // phi(Q) for the block's LA_R rows, loaded coalesced through LDS in\n"
" // pieces of 32 x DM per wave-round (rows of all waves together)\n"
" float4* qs = pool; // [4 waves][32 rows][D4], col4 ^ (row & 7)\n"
" constexpr int QCH = POOL / (4 * 32 * DM) >= 1 ? 32 : 32 * POOL / (4 * 32 * DM); // rows per wave per round\n"
" static_assert(4 * QCH * DM <= POOL, \"pool\");\n"
" for (unsigned r0 = 0; r0 < 32u; r0 += QCH) {\n"
" __syncthreads();\n"
" for (unsigned e = threadIdx.x; e < 4u * QCH * D4; e += blockDim.x) {\n"
" unsigned wr = e / (QCH * D4), rr = (e / D4) % QCH, c4 = e % D4;\n"
" unsigned grow = blk * LA_R + wr * 32u + r0 + rr, c = 4u * c4;\n"
" float4 x = make_float4(0.f, 0.f, 0.f, 0.f);\n"
" if (grow < seq && c < d)\n"
" x = la_ld4(Q + ((unsigned long long)h * seq + grow) * d + c, c, d, (d & 3u) == 0u, true);\n"
" qs[(wr * QCH + rr) * D4 + (c4 ^ (rr & 7u))] = x;\n"
" }\n"
" __syncthreads();\n"
" if (col >= r0 && col < r0 + QCH) {\n"
" const unsigned rr = col - r0;\n"
"#pragma unroll\n"
" for (int s4 = 0; s4 < HD / 4; ++s4) {\n"
" float4 x = qs[(wave * QCH + rr) * D4 + ((half * (HD / 4) + s4) ^ (rr & 7u))];\n"
" qr[4 * s4] = x.x; qr[4 * s4 + 1] = x.y; qr[4 * s4 + 2] = x.z; qr[4 * s4 + 3] = x.w;\n"
" }\n"
" }\n"
" }\n"
" }\n"
" const float* Ws = W + ((unsigned long long)h * nc + (causal ? blk : 0u)) * DM * DM;\n"
" const float* Zs = Z + ((unsigned long long)h * nc + (causal ? blk : 0u)) * DM;\n"
" fm_v16 acc[NB];\n"
" float den = 0.0f;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) acc[b][r] = 0.0f;\n"
" if (causal) {\n"
" const float* Kh = K + (unsigned long long)h * kv * d;\n"
" const float* Vh = V + (unsigned long long)h * kv * d;\n"
" const unsigned cs = blk * LA_R;\n"
" LaTile<DM, DO> tile;\n"
" if (cs < kv) la_fetch<DM, DO>(tile, Kh, Vh, cs, kv, d, c0);\n"
" for (unsigned t = 0; t < LA_R / LA_T; ++t) {\n"
" const unsigned j0 = cs + t * LA_T;\n"
" if (j0 >= kv) break;\n"
" __syncthreads();\n"
" la_put<DM, DO>(tile, Ks, Vs);\n"
" __syncthreads();\n"
" if (t + 1 < LA_R / LA_T && j0 + LA_T < kv) la_fetch<DM, DO>(tile, Kh, Vh, j0 + LA_T, kv, d, c0);\n"
" if (t > wave) continue;\n"
" fm_v16 sc;\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) sc[r] = 0.0f;\n"
" const float4* kr = Ks + col * D4;\n"
"#pragma unroll\n"
" for (int s4 = 0; s4 < HD / 4; ++s4) {\n"
" float4 kk = kr[(half * (HD / 4) + s4) ^ (col & 7u)];\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.x, qr[4 * s4 + 0], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.y, qr[4 * s4 + 1], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.z, qr[4 * s4 + 2], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.w, qr[4 * s4 + 3], sc, 0, 0, 0);\n"
" }\n"
" float a[16];\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" unsigned kl = 8u * (r / 4) + 4u * half + (r % 4);\n"
" a[r] = (t < wave || kl <= col) ? sc[r] : 0.0f;\n"
" den += a[r];\n"
" }\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" const float* vr = Vs + (8u * (r / 4) + 4u * half + (r % 4)) * DO + col;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
" acc[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(vr[32 * b], a[r], acc[b], 0, 0, 0);\n"
" }\n"
" }\n"
" }\n"
" // inter-chunk: acc^T += S^T phi(Q)^T, den += phi(q) . z\n"
"#pragma unroll\n"
" for (int s = 0; s < HD; ++s) den = fmaf(qr[s], Zs[half * HD + s], den);\n"
"#pragma unroll\n"
" for (int p = 0; p < NP; ++p) {\n"
" __syncthreads();\n"
" for (unsigned e = threadIdx.x; e < WR * DO / 4; e += blockDim.x) {\n"
" unsigned rr = e / (DO / 4), c4 = e % (DO / 4);\n"
" unsigned grow = (rr / SP) * HD + p * SP + (rr % SP);\n"
" *(float4*)&Wl[rr * DO + 4u * c4] = *(const float4*)(Ws + (unsigned long long)grow * DM + c0 + 4u * c4);\n"
" }\n"
" __syncthreads();\n"
"#pragma unroll\n"
" for (int s = 0; s < SP; ++s) {\n"
" const float* wr = Wl + (half * SP + s) * DO + col;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
" acc[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(wr[32 * b], qr[p * SP + s], acc[b], 0, 0, 0);\n"
" }\n"
" }\n"
" den += __shfl_xor(den, 32);\n"
" const float inv = den != 0.0f ? 1.0f / den : 0.0f;\n"
" // stage the wave's 32 x DO tile in LDS (all of Ks/Vs/Wl is free now), then\n"
" // write whole row segments\n"
" __syncthreads();\n"
" float4* st = pool + wave * 32u * (DO / 4); // [32][DO/4] float4, column index ^ (row & 15)\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r4 = 0; r4 < 4; ++r4) {\n"
" unsigned c4 = 8u * b + 2u * r4 + half;\n"
" st[col * (DO / 4) + (c4 ^ (col & 15u))] = make_float4(acc[b][4 * r4] * inv, acc[b][4 * r4 + 1] * inv,\n"
" acc[b][4 * r4 + 2] * inv, acc[b][4 * r4 + 3] * inv);\n"
" }\n"
" __syncthreads();\n"
" const unsigned rbase = blk * LA_R + wave * 32u;\n"
" if ((d & 3u) == 0u) {\n"
" for (unsigned e = lane; e < 32u * (DO / 4); e += 64u) {\n"
" unsigned rr = e / (DO / 4), c = 4u * (e % (DO / 4));\n"
" if (rbase + rr < seq && c0 + c < d)\n"
" *(float4*)(O + ((unsigned long long)h * seq + rbase + rr) * d + c0 + c) = st[rr * (DO / 4) + ((c / 4u) ^ (rr & 15u))];\n"
" }\n"
" } else {\n"
" for (unsigned e = lane; e < 32u * DO; e += 64u) {\n"
" unsigned rr = e / DO, c = e % DO;\n"
" if (rbase + rr < seq && c0 + c < d)\n"
" O[((unsigned long long)h * seq + rbase + rr) * d + c0 + c] = ((const float*)&st[rr * (DO / 4) + ((c / 4u) ^ (rr & 15u))])[c & 3u];\n"
" }\n"
" }\n"
"}\n"
"\n"
"// Causal linear attention in one pass. Block = (chunk c of LA_R rows, head h,\n"
"// column slice z), taken in chunk order from an atomic ticket. It loads the\n"
"// chunk's K/V tiles once and uses them twice: for the masked intra-chunk term\n"
"// and for the chunk's own state S_c = phi(K_c)^T V_c (slice z). The prefix of\n"
"// earlier chunks comes from a decoupled look-back (Merrill & Garland): each\n"
"// block publishes S_c (flag 1) right away, then sums predecessors' states\n"
"// back to the first one that has published an inclusive prefix (flag 2), and\n"
"// publishes its own inclusive prefix. Blocks only ever wait on lower tickets,\n"
"// which are already running, so the waits always end.\n"
"__device__ __forceinline__ float la_ldg(const float* p) {\n"
" return __hip_atomic_load((float*)p, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);\n"
"}\n"
"template <int DM, int DO>\n"
"__device__ __forceinline__ void la_causal_impl(const float* __restrict__ Q, const float* __restrict__ K,\n"
" const float* __restrict__ V, float* __restrict__ R,\n"
" int* __restrict__ F, float* __restrict__ O,\n"
" unsigned seq, unsigned heads, unsigned d, unsigned nc) {\n"
" constexpr int HD = DM / 2, NB = DO / 32, D4 = DM / 4, NZ = DM / DO;\n"
" constexpr int RS = DM * DO + DM; // one record: state slice, then z\n"
" constexpr int BJ = DO / 32, NBL = (DM / 32) * BJ, PW = (NBL + 3) / 4;\n"
" constexpr int P1 = LA_T * DM + LA_T * DO, P2 = DM * DO, P3 = 4 * 32 * DO;\n"
" constexpr int PM = P1 > P2 ? (P1 > P3 ? P1 : P3) : (P2 > P3 ? P2 : P3);\n"
" constexpr int POOL = PM >= P2 + DM + 4 ? PM : P2 + DM + 4; // Pl, then Zp and two flags after it\n"
" __shared__ float4 pool[POOL / 4];\n"
" float* Zp = (float*)pool + P2;\n"
" unsigned& s_tk = *(unsigned*)((float*)pool + P2 + DM);\n"
" int& s_flag = *(int*)((float*)pool + P2 + DM + 1);\n"
" float4* Ks = pool;\n"
" float* Vs = (float*)(pool + LA_T * D4);\n"
" float* Pl = (float*)pool;\n"
" const float* kf = (const float*)Ks;\n"
" if (threadIdx.x == 0) s_tk = atomicAdd((unsigned*)(F + heads * NZ * nc), 1u);\n"
" __syncthreads();\n"
" const unsigned tk = s_tk % (nc * heads * NZ); // wraps if the kernel is replayed without the reset\n"
" __syncthreads(); // the pool is reused below\n"
" const unsigned c = tk / (heads * NZ), hz = tk % (heads * NZ), h = hz / NZ, z = hz % NZ, c0 = z * DO;\n"
" const unsigned kv = seq;\n"
" const unsigned lane = threadIdx.x & 63u, wave = threadIdx.x >> 6, col = lane & 31u, half = lane >> 5;\n"
" const unsigned row = c * LA_R + wave * 32u + col;\n"
" float qr[HD];\n"
" {\n"
" float4* qs = pool;\n"
" constexpr int QCH = POOL / (4 * 32 * DM) >= 1 ? 32 : 32 * POOL / (4 * 32 * DM);\n"
" for (unsigned r0 = 0; r0 < 32u; r0 += QCH) {\n"
" __syncthreads();\n"
" for (unsigned e = threadIdx.x; e < 4u * QCH * D4; e += blockDim.x) {\n"
" unsigned wr = e / (QCH * D4), rr = (e / D4) % QCH, c4 = e % D4;\n"
" unsigned grow = c * LA_R + wr * 32u + r0 + rr, cc = 4u * c4;\n"
" float4 x = make_float4(0.f, 0.f, 0.f, 0.f);\n"
" if (grow < seq && cc < d)\n"
" x = la_ld4(Q + ((unsigned long long)h * seq + grow) * d + cc, cc, d, (d & 3u) == 0u, true);\n"
" qs[(wr * QCH + rr) * D4 + (c4 ^ (rr & 7u))] = x;\n"
" }\n"
" __syncthreads();\n"
" if (col >= r0 && col < r0 + QCH) {\n"
" const unsigned rr = col - r0;\n"
"#pragma unroll\n"
" for (int s4 = 0; s4 < HD / 4; ++s4) {\n"
" float4 x = qs[(wave * QCH + rr) * D4 + ((half * (HD / 4) + s4) ^ (rr & 7u))];\n"
" qr[4 * s4] = x.x; qr[4 * s4 + 1] = x.y; qr[4 * s4 + 2] = x.z; qr[4 * s4 + 3] = x.w;\n"
" }\n"
" }\n"
" }\n"
" }\n"
" fm_v16 acc[NB], st[PW];\n"
" float den = 0.0f, zl = 0.0f;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) acc[b][r] = 0.0f;\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) st[b][r] = 0.0f;\n"
" {\n"
" const float* Kh = K + (unsigned long long)h * kv * d;\n"
" const float* Vh = V + (unsigned long long)h * kv * d;\n"
" const unsigned cs = c * LA_R;\n"
" LaTile<DM, DO> tile;\n"
" la_fetch<DM, DO>(tile, Kh, Vh, cs, kv, d, c0);\n"
" for (unsigned t = 0; t < LA_R / LA_T; ++t) {\n"
" const unsigned j0 = cs + t * LA_T;\n"
" if (j0 >= kv) break;\n"
" __syncthreads();\n"
" la_put<DM, DO>(tile, Ks, Vs);\n"
" __syncthreads();\n"
" if (t + 1 < LA_R / LA_T && j0 + LA_T < kv) la_fetch<DM, DO>(tile, Kh, Vh, j0 + LA_T, kv, d, c0);\n"
" // the chunk's state slice: S[dk][dv] += sum_key phi(k)[dk] v[dv]\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b) {\n"
" const int bl = wave + 4 * b;\n"
" if (bl < NBL) {\n"
" const unsigned dk = 32u * (bl / BJ) + col, dv = 32u * (bl % BJ) + col;\n"
"#pragma unroll\n"
" for (int s = 0; s < LA_T / 2; ++s) {\n"
" const unsigned key = 2u * s + half;\n"
" float a = kf[key * DM + ((((dk >> 2) ^ (key & 7u)) << 2) | (dk & 3u))];\n"
" st[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(a, Vs[key * DO + dv], st[b], 0, 0, 0);\n"
" }\n"
" }\n"
" }\n"
" if (threadIdx.x < DM) {\n"
" const unsigned dk = threadIdx.x;\n"
"#pragma unroll 8\n"
" for (int key = 0; key < LA_T; ++key) zl += kf[key * DM + ((((dk >> 2) ^ (key & 7u)) << 2) | (dk & 3u))];\n"
" }\n"
" if (t <= wave) {\n"
" fm_v16 sc;\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) sc[r] = 0.0f;\n"
" const float4* kr = Ks + col * D4;\n"
"#pragma unroll\n"
" for (int s4 = 0; s4 < HD / 4; ++s4) {\n"
" float4 kk = kr[(half * (HD / 4) + s4) ^ (col & 7u)];\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.x, qr[4 * s4 + 0], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.y, qr[4 * s4 + 1], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.z, qr[4 * s4 + 2], sc, 0, 0, 0);\n"
" sc = __builtin_amdgcn_mfma_f32_32x32x2f32(kk.w, qr[4 * s4 + 3], sc, 0, 0, 0);\n"
" }\n"
" float a[16];\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" unsigned kl = 8u * (r / 4) + 4u * half + (r % 4);\n"
" a[r] = (t < wave || kl <= col) ? sc[r] : 0.0f;\n"
" den += a[r];\n"
" }\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" const float* vr = Vs + (8u * (r / 4) + 4u * half + (r % 4)) * DO + col;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
" acc[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(vr[32 * b], a[r], acc[b], 0, 0, 0);\n"
" }\n"
" }\n"
" }\n"
" }\n"
" // publish S_c: inclusive already for chunk 0\n"
" float* Rh = R + (unsigned long long)(h * NZ + z) * nc * 2u * RS;\n"
" int* Fh = F + (h * NZ + z) * nc;\n"
" {\n"
" float* rp = Rh + ((unsigned long long)c * 2u + (c == 0 ? 1u : 0u)) * RS;\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b) {\n"
" const int bl = wave + 4 * b;\n"
" if (bl < NBL)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r)\n"
" rp[(32u * (bl / BJ) + 8u * (r / 4) + 4u * half + (r % 4)) * DO + 32u * (bl % BJ) + col] = st[b][r];\n"
" }\n"
" if (threadIdx.x < DM) rp[DM * DO + threadIdx.x] = zl;\n"
" __threadfence();\n"
" __syncthreads();\n"
" if (threadIdx.x == 0) __hip_atomic_store(&Fh[c], c == 0 ? 2 : 1, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT);\n"
" }\n"
" // look back: Pl, Zp = sums over chunks < c\n"
" for (unsigned e = threadIdx.x; e < DM * DO; e += blockDim.x) Pl[e] = 0.0f;\n"
" if (threadIdx.x < DM) Zp[threadIdx.x] = 0.0f;\n"
" __syncthreads();\n"
" for (int u = (int)c - 1; u >= 0; --u) {\n"
" if (threadIdx.x == 0) {\n"
" int f;\n"
" while ((f = __hip_atomic_load(&Fh[u], __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_AGENT)) == 0) __builtin_amdgcn_s_sleep(1);\n"
" s_flag = f;\n"
" }\n"
" __syncthreads();\n"
" const int f = s_flag;\n"
" const float* rp = Rh + ((unsigned long long)u * 2u + (f == 2 ? 1u : 0u)) * RS;\n"
" for (unsigned e = threadIdx.x; e < DM * DO; e += blockDim.x) Pl[e] += la_ldg(rp + e);\n"
" if (threadIdx.x < DM) Zp[threadIdx.x] += la_ldg(rp + DM * DO + threadIdx.x);\n"
" __syncthreads();\n"
" if (f == 2) break;\n"
" }\n"
" if (c > 0) {\n"
" float* rp = Rh + ((unsigned long long)c * 2u + 1u) * RS;\n"
"#pragma unroll\n"
" for (int b = 0; b < PW; ++b) {\n"
" const int bl = wave + 4 * b;\n"
" if (bl < NBL)\n"
"#pragma unroll\n"
" for (int r = 0; r < 16; ++r) {\n"
" const unsigned i = (32u * (bl / BJ) + 8u * (r / 4) + 4u * half + (r % 4)) * DO + 32u * (bl % BJ) + col;\n"
" rp[i] = Pl[i] + st[b][r];\n"
" }\n"
" }\n"
" if (threadIdx.x < DM) rp[DM * DO + threadIdx.x] = Zp[threadIdx.x] + zl;\n"
" __threadfence();\n"
" __syncthreads();\n"
" if (threadIdx.x == 0) __hip_atomic_store(&Fh[c], 2, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT);\n"
" }\n"
" // inter-chunk: acc^T += P^T phi(Q)^T, den += phi(q) . z\n"
"#pragma unroll\n"
" for (int s = 0; s < HD; ++s) {\n"
" den = fmaf(qr[s], Zp[half * HD + s], den);\n"
" const float* wr = Pl + (half * HD + s) * DO + col;\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
" acc[b] = __builtin_amdgcn_mfma_f32_32x32x2f32(wr[32 * b], qr[s], acc[b], 0, 0, 0);\n"
" }\n"
" den += __shfl_xor(den, 32);\n"
" const float inv = den != 0.0f ? 1.0f / den : 0.0f;\n"
" __syncthreads();\n"
" constexpr unsigned SWM = (DO / 4 < 16 ? DO / 4 : 16) - 1; // swizzle within a staged row\n"
" float4* sto = pool + wave * 32u * (DO / 4);\n"
"#pragma unroll\n"
" for (int b = 0; b < NB; ++b)\n"
"#pragma unroll\n"
" for (int r4 = 0; r4 < 4; ++r4) {\n"
" unsigned c4 = 8u * b + 2u * r4 + half;\n"
" sto[col * (DO / 4) + (c4 ^ (col & SWM))] = make_float4(acc[b][4 * r4] * inv, acc[b][4 * r4 + 1] * inv,\n"
" acc[b][4 * r4 + 2] * inv, acc[b][4 * r4 + 3] * inv);\n"
" }\n"
" __syncthreads();\n"
" const unsigned rbase = c * LA_R + wave * 32u;\n"
" if ((d & 3u) == 0u) {\n"
" for (unsigned e = lane; e < 32u * (DO / 4); e += 64u) {\n"
" unsigned rr = e / (DO / 4), cc = 4u * (e % (DO / 4));\n"
" if (rbase + rr < seq && c0 + cc < d)\n"
" *(float4*)(O + ((unsigned long long)h * seq + rbase + rr) * d + c0 + cc) = sto[rr * (DO / 4) + ((cc / 4u) ^ (rr & SWM))];\n"
" }\n"
" } else {\n"
" for (unsigned e = lane; e < 32u * DO; e += 64u) {\n"
" unsigned rr = e / DO, cc = e % DO;\n"
" if (rbase + rr < seq && c0 + cc < d)\n"
" O[((unsigned long long)h * seq + rbase + rr) * d + c0 + cc] = ((const float*)&sto[rr * (DO / 4) + ((cc / 4u) ^ (rr & SWM))])[cc & 3u];\n"
" }\n"
" }\n"
"}\n"
"#endif\n"
"#if SAFEC_MFMA\n"
"#define FM_ATTR(minb) __launch_bounds__(256, minb)\n"
"#else\n"
"#define FM_ATTR(minb) __launch_bounds__(256)\n"
"#endif\n"
"extern \"C\" __global__ FM_ATTR(2) void flash_mfma_d64(const float* Q, const float* K, const float* V, float* O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) {\n"
"#if SAFEC_MFMA\n"
" flash_m32_impl<64>(Q, K, V, O, seq, kv, d, scale, causal);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ FM_ATTR(2) void flash_mfma_d128(const float* Q, const float* K, const float* V, float* O,\n"
" unsigned seq, unsigned kv, unsigned d, float scale, int causal) {\n"
"#if SAFEC_MFMA\n"
" flash_m16_impl<128>(Q, K, V, O, seq, kv, d, scale, causal);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ FM_ATTR(3) void la_state_d64(const float* K, const float* V, float* W, float* Z,\n"
" unsigned kv, unsigned d, unsigned rows, unsigned nc) {\n"
"#if SAFEC_MFMA\n"
" la_state_impl<64>(K, V, W, Z, kv, d, rows, nc);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ FM_ATTR(2) void la_state_d128(const float* K, const float* V, float* W, float* Z,\n"
" unsigned kv, unsigned d, unsigned rows, unsigned nc) {\n"
"#if SAFEC_MFMA\n"
" la_state_impl<128>(K, V, W, Z, kv, d, rows, nc);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ __launch_bounds__(256) void la_scan_d64(float* W, float* Z, unsigned nc, int causal) {\n"
"#if SAFEC_MFMA\n"
" la_scan_impl<64>(W, Z, nc, causal);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ __launch_bounds__(256) void la_scan_d128(float* W, float* Z, unsigned nc, int causal) {\n"
"#if SAFEC_MFMA\n"
" la_scan_impl<128>(W, Z, nc, causal);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ FM_ATTR(2) void la_out_d64(const float* Q, const float* K, const float* V, const float* W,\n"
" const float* Z, float* O, unsigned seq, unsigned kv, unsigned d, unsigned nc, int causal) {\n"
"#if SAFEC_MFMA\n"
" la_out_impl<64, 64>(Q, K, V, W, Z, O, seq, kv, d, nc, causal);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ FM_ATTR(2) void la_out_d128(const float* Q, const float* K, const float* V, const float* W,\n"
" const float* Z, float* O, unsigned seq, unsigned kv, unsigned d, unsigned nc, int causal) {\n"
"#if SAFEC_MFMA\n"
" la_out_impl<128, 64>(Q, K, V, W, Z, O, seq, kv, d, nc, causal);\n"
"#endif\n"
"}\n"
"extern \"C\" __global__ FM_ATTR(2) void la_causal_d64(const float* Q, const float* K, const float* V, float* R,\n"
" int* F, float* O, unsigned seq, unsigned heads, unsigned d, unsigned nc) {\n"
"#if SAFEC_MFMA\n"
" la_causal_impl<64, 64>(Q, K, V, R, F, O, seq, heads, d, nc);\n"
"#endif\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_UPSAMPLE2X 15
#define K_BIAS_ACT 16
#define K_TO_HALF 17
#define K_FROM_HALF 18
#define K_HALF_SGD 19
#define K_CONCAT 20
#define K_UPCAT 21
#define K_HALF_RELU 22
#define K_HALF_RELUG 23
#define K_HALF_MSE 24
#define K_HALF_SGD2 25
#define K_FLASH64 26
#define K_FLASH128 27
#define K_LIN_STATE 28
#define K_LIN_PREFIX 29
#define K_LIN_OUT 30
#define K_MFMA_PROBE 31
#define K_FLASH_M64 32
#define K_FLASH_M128 33
#define K_LA_STATE64 34
#define K_LA_STATE128 35
#define K_LA_SCAN64 36
#define K_LA_SCAN128 37
#define K_LA_OUT64 38
#define K_LA_OUT128 39
#define K_LA_CAUSAL64 40
#define K_GEMV_HALF 41
#define K_RMSNORM 42
#define K_GELU 43
#define K_HALF_CVT 44
#define K_CVT_MANY 45
#define K_MP_SGD 46
#define K_MP_ADAM 47
#define K_HALF_EW 48
#define K_HALF_BIAS 49
#define K_HALF_POOL 50
#define K_COUNT 51
static const char* rocm_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_UPSAMPLE2X) return "upsample2x";
if (k == K_BIAS_ACT) return "bias_act";
if (k == K_TO_HALF) return "to_half";
if (k == K_FROM_HALF) return "from_half";
if (k == K_HALF_SGD) return "half_sgd";
if (k == K_CONCAT) return "concat_rows";
if (k == K_UPCAT) return "upcat";
if (k == K_HALF_RELU) return "half_relu";
if (k == K_HALF_RELUG) return "half_relu_grad";
if (k == K_HALF_MSE) return "half_mse";
if (k == K_HALF_SGD2) return "half_sgd2";
if (k == K_FLASH64) return "flash_attn_d64";
if (k == K_FLASH128) return "flash_attn_d128";
if (k == K_LIN_STATE) return "lin_attn_state";
if (k == K_LIN_PREFIX) return "lin_attn_prefix";
if (k == K_LIN_OUT) return "lin_attn_out";
if (k == K_MFMA_PROBE) return "mfma_probe";
if (k == K_FLASH_M64) return "flash_mfma_d64";
if (k == K_FLASH_M128) return "flash_mfma_d128";
if (k == K_LA_STATE64) return "la_state_d64";
if (k == K_LA_STATE128) return "la_state_d128";
if (k == K_LA_SCAN64) return "la_scan_d64";
if (k == K_LA_SCAN128) return "la_scan_d128";
if (k == K_LA_OUT64) return "la_out_d64";
if (k == K_LA_OUT128) return "la_out_d128";
if (k == K_LA_CAUSAL64) return "la_causal_d64";
if (k == K_GEMV_HALF) return "gemv_half";
if (k == K_RMSNORM) return "rmsnorm_rows";
if (k == K_GELU) return "ew_gelu";
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_HALF_EW) return "half_ew";
if (k == K_HALF_BIAS) return "half_bias";
if (k == K_HALF_POOL) return "half_pool";
return "mse_grad";
}
// ── Backend state (one per process) ──────────────────────────────────────────
static int g_rocm_state = 0; // 0 = not tried, 1 = ready, -1 = unavailable
static RocmNativeStream* g_rocm_stream = null;
static RocblasHandle_* g_rocm_blas = null;
static HipFunction_* g_rocm_fn[K_COUNT];
static char g_rocm_name[256];
static int g_rocm_mfma = 0; // 1 = the MFMA attention kernels run (matrix cores, not switched off)
static int g_rocm_mfma_has = 0; // 1 = the GPU has matrix cores (CDNA)
// Initializes everything once: picks the GPU with the most compute units
// (discrete before integrated), creates the stream and rocBLAS handle, and
// compiles + loads the kernels with hiprtc.
static int rocm_ensure_() {
if (g_rocm_state != 0) return g_rocm_state == 1;
g_rocm_state = -1;
unsafe {
g_rocm_name[0] = (char)0;
if (hipInit(0U) != ROCM_HIP_OK) return 0;
int count = 0;
if (hipGetDeviceCount(&count) != ROCM_HIP_OK || count <= 0) return 0;
int best = -1;
long long bestScore = -1LL;
int d = 0;
while (d < count) {
int cu = 0;
int integrated = 0;
hipDeviceGetAttribute(&cu, ROCM_ATTR_CU_COUNT, d);
hipDeviceGetAttribute(&integrated, ROCM_ATTR_INTEGRATED, d);
long long score = (long long)cu + (integrated != 0 ? 0LL : 1000000LL);
if (score > bestScore) { bestScore = score; best = d; }
d = d + 1;
}
if (best < 0 || hipSetDevice(best) != ROCM_HIP_OK) return 0;
hipDeviceGetName(&g_rocm_name[0], 256, best);
if (hipStreamCreate(&g_rocm_stream) != ROCM_HIP_OK) return 0;
if (rocblas_create_handle(&g_rocm_blas) != 0) return 0;
rocblas_set_stream(g_rocm_blas, g_rocm_stream);
HiprtcProgram_* prog = null;
if (hiprtcCreateProgram(&prog, rocm_kernel_source_(), "safec_rocm.hip", 0,
(const char**)0, (const char**)0) != 0) return 0;
const char* opts[1];
opts[0] = "-O3";
if (hiprtcCompileProgram(prog, 1, &opts[0]) != 0) { hiprtcDestroyProgram(&prog); return 0; }
unsigned long codeSize = 0UL;
hiprtcGetCodeSize(prog, &codeSize);
char* code = (char*)alloc_raw_(codeSize);
if (code == (char*)0) { hiprtcDestroyProgram(&prog); return 0; }
hiprtcGetCode(prog, code);
hiprtcDestroyProgram(&prog);
HipModule_* module = null;
int loaded = hipModuleLoadData(&module, (const char*)code);
dealloc(code);
if (loaded != ROCM_HIP_OK) return 0;
int k = 0;
while (k < K_COUNT) {
if (hipModuleGetFunction(&g_rocm_fn[k], module, rocm_kernel_name_(k)) != ROCM_HIP_OK) return 0;
k = k + 1;
}
// Ask the compiled code whether this GPU has matrix cores (the kernel
// source knows the target; the host API doesn't say).
HipDevPtr_* flag = null;
if (hipMalloc(&flag, 4UL) == ROCM_HIP_OK) {
HipKernelArg_* pp[1];
pp[0] = (HipKernelArg_*)&flag;
float has = 0.0f;
if (hipModuleLaunchKernel(g_rocm_fn[K_MFMA_PROBE], 1U, 1U, 1U, 64U, 1U, 1U, 0U, g_rocm_stream,
&pp[0], (HipKernelArg_**)0) == ROCM_HIP_OK &&
hipMemcpyAsync((unsigned char*)&has, (const unsigned char*)flag, 4UL, ROCM_D2H, g_rocm_stream) == ROCM_HIP_OK &&
hipStreamSynchronize(g_rocm_stream) == ROCM_HIP_OK)
g_rocm_mfma_has = has == 1.0f ? 1 : 0;
g_rocm_mfma = g_rocm_mfma_has;
hipFree(flag);
}
}
g_rocm_state = 1;
return 1;
}
int rocm_available() { return rocm_ensure_(); }
int rocm_attention_use_mfma(int on) {
if (!rocm_ensure_()) return 0;
g_rocm_mfma = on && g_rocm_mfma_has ? 1 : 0;
return g_rocm_mfma;
}
const &char rocm_device_name() {
rocm_ensure_();
unsafe { return (const &char)&g_rocm_name[0]; }
}
// ── Buffers ──────────────────────────────────────────────────────────────────
struct RocmBuf RocmBuf::alloc(unsigned long n) {
struct RocmBuf b;
b.dev = 0UL;
b.len = 0UL;
if (n == 0UL || !rocm_ensure_()) return b;
unsafe {
HipDevPtr_* p = (HipDevPtr_*)0;
if (hipMalloc(&p, checked_mul_size(n, sizeof(float))) != ROCM_HIP_OK) return b;
b.dev = (unsigned long)p;
}
b.len = n;
return b;
}
int rocm_write_n_(struct RocmBuf dst, const &float host, unsigned long n) {
if (dst.dev == 0UL || n > dst.len || !rocm_ensure_()) return 0;
unsafe {
return hipMemcpyAsync((unsigned char*)dst.dev, (const unsigned char*)host, n * sizeof(float),
ROCM_H2D, g_rocm_stream) == ROCM_HIP_OK;
}
}
int rocm_write(struct RocmBuf 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 rocm_write_n_(dst, host, n);
}
struct RocmBuf RocmBuf::upload_n_(const &float host, unsigned long n) {
struct RocmBuf b = RocmBuf::alloc(n);
if (b.dev != 0UL && !rocm_write_n_(b, host, n)) b.free();
return b;
}
struct RocmBuf RocmBuf::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 RocmBuf::upload_n_(host, n);
}
int rocm_read_n_(struct RocmBuf src, &float host, unsigned long n) {
if (src.dev == 0UL || n > src.len || !rocm_ensure_()) return 0;
unsafe {
if (hipMemcpyAsync((unsigned char*)host, (const unsigned char*)src.dev, n * sizeof(float),
ROCM_D2H, g_rocm_stream) != ROCM_HIP_OK) return 0;
return hipStreamSynchronize(g_rocm_stream) == ROCM_HIP_OK;
}
}
int rocm_read(struct RocmBuf 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 rocm_read_n_(src, host, n);
}
void RocmBuf::free() {
if (self.dev != 0UL) { unsafe { hipFree((HipDevPtr_*)self.dev); } }
self.dev = 0UL;
self.len = 0UL;
}
int rocm_sync() {
if (!rocm_ensure_()) return 0;
unsafe { return hipStreamSynchronize(g_rocm_stream) == ROCM_HIP_OK; }
}
?&RocmNativeStream rocm_stream_handle() {
if (!rocm_ensure_()) return null;
unsafe { return (?&RocmNativeStream)g_rocm_stream; }
}
int rocm_copy(struct RocmBuf src, unsigned long srcOff, struct RocmBuf dst, unsigned long dstOff, unsigned long n) {
if (!rocm_ensure_() || src.dev == 0UL || dst.dev == 0UL || n == 0UL ||
srcOff + n > src.len || dstOff + n > dst.len) return 0;
unsafe {
return hipMemcpyAsync((unsigned char*)(dst.dev + dstOff * sizeof(float)),
(const unsigned char*)(src.dev + srcOff * sizeof(float)),
n * sizeof(float), ROCM_D2D, g_rocm_stream) == ROCM_HIP_OK;
}
}
// ── Launch helpers ───────────────────────────────────────────────────────────
static int rocm_launch_(int kernel, unsigned long n, HipKernelArg_** params) {
unsigned long blocks = (n + (unsigned long)ROCM_BLOCK - 1UL) / (unsigned long)ROCM_BLOCK;
if (blocks > ROCM_MAX_BLOCKS) blocks = ROCM_MAX_BLOCKS;
if (blocks == 0UL) blocks = 1UL;
unsafe {
return hipModuleLaunchKernel(g_rocm_fn[kernel], (unsigned int)blocks, 1U, 1U,
ROCM_BLOCK, 1U, 1U, 0U, g_rocm_stream,
params, (HipKernelArg_**)0) == ROCM_HIP_OK;
}
}
static int rocm_fits_(struct RocmBuf b, unsigned long n) { return b.dev != 0UL && b.len >= n; }
static int rocm_binary_(int kernel, struct RocmBuf a, struct RocmBuf b, struct RocmBuf out) {
unsigned long n = out.len;
if (!rocm_ensure_() || !rocm_fits_(a, n) || !rocm_fits_(b, n) || !rocm_fits_(out, n) ||
n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)a.dev;
HipDevPtr_* pb = (HipDevPtr_*)b.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int nn = (unsigned int)n;
HipKernelArg_* params[4];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&pb; params[2] = (HipKernelArg_*)&po; params[3] = (HipKernelArg_*)&nn;
return rocm_launch_(kernel, n, ¶ms[0]);
}
}
static int rocm_unary_(int kernel, struct RocmBuf a, float k, struct RocmBuf out) {
unsigned long n = out.len;
if (!rocm_ensure_() || !rocm_fits_(a, n) || !rocm_fits_(out, n) || n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)a.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int nn = (unsigned int)n;
float kk = k;
HipKernelArg_* params[4];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&po; params[2] = (HipKernelArg_*)&nn; params[3] = (HipKernelArg_*)&kk;
return rocm_launch_(kernel, n, ¶ms[0]);
}
}
// ── Queued ops ───────────────────────────────────────────────────────────────
// ── gemm_ex solution choice ──────────────────────────────────────────────────
// rocBLAS's default solution for a gemm_ex shape can be several times slower
// than its best on some GPUs (on an MI100: 3.5x for a 128x256x1024 fp16
// product, 2-3x for small bf16 ones). So the first call for each (types,
// transposes, size) outside graph capture times every solution rocBLAS
// offers on that call's own buffers, like MIOpen's Find, and keeps the
// fastest (0 = the default, if nothing beats it). During capture an untuned
// shape uses the default. The cache holds ROCM_GEMM_PLANS shapes (a deep
// convolutional net uses a few hundred: each layer's forward and both gradients,
// per batch size).
#define ROCM_GEMM_PLANS 1024
#define ROCM_GEMM_MAX_SOLS 512
#define ROCM_GEMM_REPS 8
static int g_rocm_capturing = 0;
static int g_gplan_used[ROCM_GEMM_PLANS];
static unsigned long g_gplan_key[ROCM_GEMM_PLANS]; // packed types, transposes
static int g_gplan_m[ROCM_GEMM_PLANS];
static int g_gplan_n[ROCM_GEMM_PLANS];
static int g_gplan_k[ROCM_GEMM_PLANS];
static int g_gplan_sol[ROCM_GEMM_PLANS];
static int g_gplan_next = 0;
static int g_gplan_list[ROCM_GEMM_MAX_SOLS];
// One column-major rocBLAS gemm_ex call, D = C = op(A) op(B); 'sol' 0 uses
// the default solution.
static int rocm_gemm_ex_call_(int opA, int opB, int m, int n, int k,
unsigned long a, int aType, int lda, unsigned long b, int ldb,
unsigned long c, int cType, int ldc, int sol) {
float one = 1.0f;
float zero = 0.0f;
unsafe {
return rocblas_gemm_ex(g_rocm_blas, opA, opB, m, n, k, (const float*)&one,
(const HipDevPtr_*)a, aType, lda, (const HipDevPtr_*)b, aType, ldb,
(const float*)&zero, (const HipDevPtr_*)c, cType, ldc, (HipDevPtr_*)c, cType, ldc,
ROCM_BLAS_F32, sol == 0 ? ROCM_BLAS_ALGO_STANDARD : ROCM_BLAS_ALGO_INDEX,
sol, 0U) == 0;
}
}
static double rocm_gemm_time_(int opA, int opB, int m, int n, int k,
unsigned long a, int aType, int lda, unsigned long b, int ldb,
unsigned long c, int cType, int ldc, int sol) {
if (!rocm_gemm_ex_call_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, sol)) return -1.0;
unsafe { hipStreamSynchronize(g_rocm_stream); }
long long t0 = time_mono_ns();
int r = 0;
while (r < ROCM_GEMM_REPS) {
rocm_gemm_ex_call_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, sol);
r = r + 1;
}
unsafe { hipStreamSynchronize(g_rocm_stream); }
return (double)(time_mono_ns() - t0);
}
static int rocm_gemm_ex_(int opA, int opB, int m, int n, int k,
unsigned long a, int aType, int lda, unsigned long b, int ldb,
unsigned long c, int cType, int ldc) {
unsigned long key = (unsigned long)aType | ((unsigned long)cType << 16) |
((unsigned long)opA << 32) | ((unsigned long)opB << 48);
int i = 0;
while (i < ROCM_GEMM_PLANS) {
if (g_gplan_used[i] && g_gplan_key[i] == key && g_gplan_m[i] == m && g_gplan_n[i] == n &&
g_gplan_k[i] == k)
return rocm_gemm_ex_call_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, g_gplan_sol[i]);
i = i + 1;
}
if (g_rocm_capturing) return rocm_gemm_ex_call_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, 0);
int best = 0;
double bestT = rocm_gemm_time_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, 0);
if (bestT < 0.0) return 0;
int count = ROCM_GEMM_MAX_SOLS;
float one = 1.0f;
float zero = 0.0f;
unsafe {
if (rocblas_gemm_ex_get_solutions(g_rocm_blas, opA, opB, m, n, k, (const float*)&one,
(const HipDevPtr_*)a, aType, lda, (const HipDevPtr_*)b, aType, ldb,
(const float*)&zero, (const HipDevPtr_*)c, cType, ldc, (HipDevPtr_*)c, cType, ldc,
ROCM_BLAS_F32, ROCM_BLAS_ALGO_INDEX, 0U,
&g_gplan_list[0], &count) != 0) count = 0;
}
if (count > ROCM_GEMM_MAX_SOLS) count = ROCM_GEMM_MAX_SOLS;
int s = 0;
while (s < count) {
int sol = g_gplan_list[s];
if (sol != 0) {
double t = rocm_gemm_time_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, sol);
// Prefer the default unless a solution is clearly faster.
if (t >= 0.0 && t < bestT * 0.95) { bestT = t; best = sol; }
}
s = s + 1;
}
int slot = -1;
i = 0;
while (i < ROCM_GEMM_PLANS && slot < 0) { if (!g_gplan_used[i]) slot = i; i = i + 1; }
if (slot < 0) { slot = g_gplan_next; g_gplan_next = (g_gplan_next + 1) % ROCM_GEMM_PLANS; }
g_gplan_used[slot] = 1;
g_gplan_key[slot] = key;
g_gplan_m[slot] = m;
g_gplan_n[slot] = n;
g_gplan_k[slot] = k;
g_gplan_sol[slot] = best;
// The timing runs already computed C; run the chosen solution once more
// so the result comes from the solution later calls will use.
return rocm_gemm_ex_call_(opA, opB, m, n, k, a, aType, lda, b, ldb, c, cType, ldc, best);
}
int rocm_gemm(struct RocmBuf a, int ta, struct RocmBuf b, int tb, struct RocmBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!rocm_ensure_() || !rocm_fits_(a, m * k) || !rocm_fits_(b, k * n) || !rocm_fits_(c, m * n))
return 0;
// Row-major C = op(A) op(B) is column-major C^T = op(B)^T op(A)^T, so
// hand rocBLAS the operands swapped; each operand's leading dimension
// is its stored row length.
return rocm_gemm_ex_(tb ? ROCM_BLAS_T : ROCM_BLAS_N, ta ? ROCM_BLAS_T : ROCM_BLAS_N,
(int)n, (int)m, (int)k,
b.dev, ROCM_BLAS_F32, tb ? (int)k : (int)n,
a.dev, ta ? (int)m : (int)k,
c.dev, ROCM_BLAS_F32, (int)n);
}
int rocm_add(struct RocmBuf a, struct RocmBuf b, struct RocmBuf out) { return rocm_binary_(K_ADD, a, b, out); }
int rocm_sub(struct RocmBuf a, struct RocmBuf b, struct RocmBuf out) { return rocm_binary_(K_SUB, a, b, out); }
int rocm_mul(struct RocmBuf a, struct RocmBuf b, struct RocmBuf out) { return rocm_binary_(K_MUL, a, b, out); }
int rocm_div(struct RocmBuf a, struct RocmBuf b, struct RocmBuf out) { return rocm_binary_(K_DIV, a, b, out); }
int rocm_pow(struct RocmBuf a, struct RocmBuf b, struct RocmBuf out) { return rocm_binary_(K_POW, a, b, out); }
int rocm_relu_grad(struct RocmBuf z, struct RocmBuf g, struct RocmBuf out) { return rocm_binary_(K_RELU_GRAD, z, g, out); }
int rocm_scale(struct RocmBuf a, float k, struct RocmBuf out) { return rocm_unary_(K_SCALE, a, k, out); }
int rocm_relu(struct RocmBuf a, struct RocmBuf out) { return rocm_unary_(K_RELU, a, 0.0f, out); }
int rocm_exp(struct RocmBuf a, struct RocmBuf out) { return rocm_unary_(K_EXP, a, 0.0f, out); }
int rocm_log(struct RocmBuf a, struct RocmBuf out) { return rocm_unary_(K_LOG, a, 0.0f, out); }
int rocm_sqrt(struct RocmBuf a, struct RocmBuf out) { return rocm_unary_(K_SQRT, a, 0.0f, out); }
int rocm_gelu(struct RocmBuf a, struct RocmBuf out) { return rocm_unary_(K_GELU, a, 0.0f, out); }
int rocm_rmsnorm(struct RocmBuf x, struct RocmBuf w, struct RocmBuf y, unsigned long rows, unsigned long n, float eps) {
if (!rocm_ensure_() || rows == 0UL || n == 0UL || rows > 2147483647UL || n > 4294967295UL) return 0;
if (!rocm_fits_(x, rows * n) || !rocm_fits_(y, rows * n)) return 0;
if (w.dev != 0UL && w.len < n) return 0;
unsafe {
HipDevPtr_* px = (HipDevPtr_*)x.dev;
HipDevPtr_* pw = (HipDevPtr_*)w.dev;
HipDevPtr_* py = (HipDevPtr_*)y.dev;
unsigned int nn = (unsigned int)n;
float e = eps;
HipKernelArg_* params[5];
params[0] = (HipKernelArg_*)&px; params[1] = (HipKernelArg_*)&pw; params[2] = (HipKernelArg_*)&py; params[3] = (HipKernelArg_*)&nn; params[4] = (HipKernelArg_*)&e;
return hipModuleLaunchKernel(g_rocm_fn[K_RMSNORM], (unsigned int)rows, 1U, 1U, ROCM_BLOCK, 1U, 1U, 0U,
g_rocm_stream, ¶ms[0], (HipKernelArg_**)0) == ROCM_HIP_OK;
}
}
// ── fp16 / bf16 ─────────────────────────────────────────────────────────────
static int rocm_half_dtype_ok_(int dtype) { return dtype == ROCM_FP16 || dtype == ROCM_BF16; }
static int rocm_half_fits_(struct RocmHalfBuf b, unsigned long n) {
return b.dev != 0UL && b.len >= n && rocm_half_dtype_ok_(b.dtype);
}
struct RocmHalfBuf RocmHalfBuf::alloc(unsigned long n, int dtype) {
struct RocmHalfBuf b;
b.dev = 0UL;
b.len = 0UL;
b.dtype = dtype;
if (n == 0UL || !rocm_half_dtype_ok_(dtype) || !rocm_ensure_()) return b;
unsafe {
HipDevPtr_* p = (HipDevPtr_*)0;
if (hipMalloc(&p, checked_mul_size(n, 2UL)) != ROCM_HIP_OK) return b;
b.dev = (unsigned long)p;
}
b.len = n;
return b;
}
void RocmHalfBuf::free() {
if (self.dev != 0UL) { unsafe { hipFree((HipDevPtr_*)self.dev); } }
self.dev = 0UL;
self.len = 0UL;
}
int rocm_half_write_n_(struct RocmHalfBuf dst, const &unsigned short bits, unsigned long n) {
if (dst.dev == 0UL || n > dst.len || !rocm_ensure_()) return 0;
unsafe {
return hipMemcpyAsync((unsigned char*)dst.dev, (const unsigned char*)bits, n * 2UL, ROCM_H2D, g_rocm_stream) == ROCM_HIP_OK;
}
}
int rocm_half_write(struct RocmHalfBuf 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 rocm_half_write_n_(dst, bits, n);
}
int rocm_half_read_n_(struct RocmHalfBuf src, &unsigned short bits, unsigned long n) {
if (src.dev == 0UL || n > src.len || !rocm_ensure_()) return 0;
unsafe {
if (hipMemcpyAsync((unsigned char*)bits, (const unsigned char*)src.dev, n * 2UL, ROCM_D2H, g_rocm_stream) != ROCM_HIP_OK)
return 0;
return hipStreamSynchronize(g_rocm_stream) == ROCM_HIP_OK;
}
}
int rocm_half_read(struct RocmHalfBuf 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 rocm_half_read_n_(src, bits, n);
}
static int rocm_convert_(int kernel, unsigned long a, unsigned long o, unsigned long n, int bf) {
if (n == 0UL || n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)a;
HipDevPtr_* po = (HipDevPtr_*)o;
unsigned int nn = (unsigned int)n;
int b = bf;
HipKernelArg_* params[4];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&po; params[2] = (HipKernelArg_*)&nn; params[3] = (HipKernelArg_*)&b;
return rocm_launch_(kernel, n, ¶ms[0]);
}
}
int rocm_to_half(struct RocmBuf src, struct RocmHalfBuf dst) {
unsigned long n = dst.len;
if (!rocm_ensure_() || !rocm_fits_(src, n) || !rocm_half_fits_(dst, n)) return 0;
return rocm_convert_(K_TO_HALF, src.dev, dst.dev, n, dst.dtype == ROCM_BF16);
}
int rocm_from_half(struct RocmHalfBuf src, struct RocmBuf dst) {
unsigned long n = dst.len;
if (!rocm_ensure_() || !rocm_half_fits_(src, n) || !rocm_fits_(dst, n)) return 0;
return rocm_convert_(K_FROM_HALF, src.dev, dst.dev, n, src.dtype == ROCM_BF16);
}
struct RocmHalfBuf RocmHalfBuf::upload_n_(const &float host, unsigned long n, int dtype) {
struct RocmHalfBuf h = RocmHalfBuf::alloc(n, dtype);
if (h.dev == 0UL) return h;
struct RocmBuf tmp = RocmBuf::upload_n_(host, n);
int ok = tmp.dev != 0UL && rocm_to_half(tmp, h) && rocm_sync();
tmp.free();
if (!ok) h.free();
return h;
}
struct RocmHalfBuf RocmHalfBuf::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 RocmHalfBuf::upload_n_(host, n, dtype);
}
// rocBLAS gemm_ex: 16-bit A and B, fp32 accumulation; C/D are 16-bit
// (outType = the inputs' type) or fp32. Row-major via the same operand swap
// as rocm_gemm.
static int rocm_half_gemm_(struct RocmHalfBuf a, int ta, struct RocmHalfBuf b, int tb,
unsigned long cdev, int outType, unsigned long m, unsigned long n, unsigned long k) {
int t = a.dtype == ROCM_BF16 ? ROCM_BLAS_BF16 : ROCM_BLAS_F16;
return rocm_gemm_ex_(tb ? ROCM_BLAS_T : ROCM_BLAS_N, ta ? ROCM_BLAS_T : ROCM_BLAS_N,
(int)n, (int)m, (int)k,
b.dev, t, tb ? (int)k : (int)n,
a.dev, ta ? (int)m : (int)k,
cdev, outType == 0 ? t : outType, (int)n);
}
int rocm_half_gemm(struct RocmHalfBuf a, int ta, struct RocmHalfBuf b, int tb, struct RocmHalfBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!rocm_ensure_() || !rocm_half_fits_(a, m * k) || !rocm_half_fits_(b, k * n) ||
!rocm_half_fits_(c, m * n) || a.dtype != b.dtype || a.dtype != c.dtype) return 0;
return rocm_half_gemm_(a, ta, b, tb, c.dev, 0, m, n, k);
}
int rocm_half_gemm_f32(struct RocmHalfBuf a, int ta, struct RocmHalfBuf b, int tb, struct RocmBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!rocm_ensure_() || !rocm_half_fits_(a, m * k) || !rocm_half_fits_(b, k * n) ||
!rocm_fits_(c, m * n) || a.dtype != b.dtype) return 0;
return rocm_half_gemm_(a, ta, b, tb, c.dev, ROCM_BLAS_F32, m, n, k);
}
int rocm_half_sgd(struct RocmHalfBuf w, struct RocmHalfBuf g, float lr) {
unsigned long n = w.len;
if (!rocm_ensure_() || !rocm_half_fits_(w, n) || !rocm_half_fits_(g, n) || w.dtype != g.dtype ||
n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pw = (HipDevPtr_*)w.dev;
HipDevPtr_* pg = (HipDevPtr_*)g.dev;
float l = lr;
unsigned int nn = (unsigned int)n;
int bf = w.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[5];
params[0] = (HipKernelArg_*)&pw; params[1] = (HipKernelArg_*)&pg; params[2] = (HipKernelArg_*)&l; params[3] = (HipKernelArg_*)&nn;
params[4] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_SGD, n, ¶ms[0]);
}
}
int rocm_half_relu(struct RocmHalfBuf a, struct RocmHalfBuf out) {
unsigned long n = a.len;
if (!rocm_ensure_() || !rocm_half_fits_(a, n) || !rocm_half_fits_(out, n) || a.dtype != out.dtype ||
n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)a.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int nn = (unsigned int)n;
int bf = a.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[4];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&po; params[2] = (HipKernelArg_*)&nn; params[3] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_RELU, n, ¶ms[0]);
}
}
int rocm_half_gemv(struct RocmHalfBuf w, struct RocmBuf x, struct RocmBuf y, unsigned long n, unsigned long k) {
if (!rocm_ensure_() || n == 0UL || k == 0UL || n > 2147483647UL || k > 4294967295UL) return 0;
if (!rocm_half_fits_(w, n * k) || !rocm_fits_(x, k) || !rocm_fits_(y, n)) return 0;
unsafe {
HipDevPtr_* pw = (HipDevPtr_*)w.dev;
HipDevPtr_* px = (HipDevPtr_*)x.dev;
HipDevPtr_* py = (HipDevPtr_*)y.dev;
unsigned int nn = (unsigned int)n;
unsigned int kk = (unsigned int)k;
int bf = w.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[6];
params[0] = (HipKernelArg_*)&pw; params[1] = (HipKernelArg_*)&px; params[2] = (HipKernelArg_*)&py; params[3] = (HipKernelArg_*)&nn;
params[4] = (HipKernelArg_*)&kk; params[5] = (HipKernelArg_*)&bf;
return hipModuleLaunchKernel(g_rocm_fn[K_GEMV_HALF], (unsigned int)((n + 3UL) / 4UL), 1U, 1U, ROCM_BLOCK, 1U, 1U, 0U,
g_rocm_stream, ¶ms[0], (HipKernelArg_**)0) == ROCM_HIP_OK;
}
}
// rocBLAS has no ReLU epilogue: the GEMM, then ReLU in place.
int rocm_half_gemm_relu(struct RocmHalfBuf a, int ta, struct RocmHalfBuf b, int tb, struct RocmHalfBuf c,
unsigned long m, unsigned long n, unsigned long k) {
if (!rocm_half_gemm(a, ta, b, tb, c, m, n, k)) return 0;
struct RocmHalfBuf view = c;
view.len = m * n;
return rocm_half_relu(view, view);
}
int rocm_half_relu_grad(struct RocmHalfBuf z, struct RocmHalfBuf g, struct RocmHalfBuf out) {
unsigned long n = z.len;
if (!rocm_ensure_() || !rocm_half_fits_(z, n) || !rocm_half_fits_(g, n) || !rocm_half_fits_(out, n) ||
z.dtype != g.dtype || z.dtype != out.dtype || n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pz = (HipDevPtr_*)z.dev;
HipDevPtr_* pg = (HipDevPtr_*)g.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int nn = (unsigned int)n;
int bf = z.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[5];
params[0] = (HipKernelArg_*)&pz; params[1] = (HipKernelArg_*)&pg; params[2] = (HipKernelArg_*)&po; params[3] = (HipKernelArg_*)&nn;
params[4] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_RELUG, n, ¶ms[0]);
}
}
int rocm_half_mse_grad(struct RocmHalfBuf y, struct RocmHalfBuf t, struct RocmHalfBuf dy, struct RocmBuf loss1) {
unsigned long n = y.len;
if (!rocm_ensure_() || !rocm_half_fits_(y, n) || !rocm_half_fits_(t, n) || !rocm_half_fits_(dy, n) ||
y.dtype != t.dtype || y.dtype != dy.dtype || !rocm_fits_(loss1, 1UL) || n > 4294967295UL) return 0;
unsafe {
if (hipMemsetAsync((HipDevPtr_*)loss1.dev, 0, sizeof(float), g_rocm_stream) != ROCM_HIP_OK) return 0;
HipDevPtr_* py = (HipDevPtr_*)y.dev;
HipDevPtr_* pt = (HipDevPtr_*)t.dev;
HipDevPtr_* pd = (HipDevPtr_*)dy.dev;
HipDevPtr_* pl = (HipDevPtr_*)loss1.dev;
unsigned int nn = (unsigned int)n;
int bf = y.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[6];
params[0] = (HipKernelArg_*)&py; params[1] = (HipKernelArg_*)&pt; params[2] = (HipKernelArg_*)&pd;
params[3] = (HipKernelArg_*)&pl; params[4] = (HipKernelArg_*)&nn; params[5] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_MSE, n, ¶ms[0]);
}
}
int rocm_half_sgd2(struct RocmHalfBuf w1, struct RocmHalfBuf g1, struct RocmHalfBuf w2, struct RocmHalfBuf g2,
float lr) {
unsigned long n1 = w1.len;
unsigned long n2 = w2.len;
if (!rocm_ensure_() || !rocm_half_fits_(w1, n1) || !rocm_half_fits_(g1, n1) || !rocm_half_fits_(w2, n2) ||
!rocm_half_fits_(g2, n2) || w1.dtype != g1.dtype || w1.dtype != w2.dtype || w1.dtype != g2.dtype ||
n1 + n2 > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pw1 = (HipDevPtr_*)w1.dev;
HipDevPtr_* pg1 = (HipDevPtr_*)g1.dev;
HipDevPtr_* pw2 = (HipDevPtr_*)w2.dev;
HipDevPtr_* pg2 = (HipDevPtr_*)g2.dev;
unsigned int nn1 = (unsigned int)n1;
unsigned int nn2 = (unsigned int)n2;
float rate = lr;
int bf = w1.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[8];
params[0] = (HipKernelArg_*)&pw1; params[1] = (HipKernelArg_*)&pg1; params[2] = (HipKernelArg_*)&nn1;
params[3] = (HipKernelArg_*)&pw2; params[4] = (HipKernelArg_*)&pg2; params[5] = (HipKernelArg_*)&nn2;
params[6] = (HipKernelArg_*)&rate; params[7] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_SGD2, n1 + n2, ¶ms[0]);
}
}
int rocm_concat(struct RocmBuf a, unsigned long na, struct RocmBuf c, unsigned long nc,
struct RocmBuf out, unsigned long batch) {
unsigned long total = batch * (na + nc);
if (!rocm_ensure_() || batch == 0UL || na == 0UL || nc == 0UL || total > 4294967295UL ||
!rocm_fits_(a, batch * na) || !rocm_fits_(c, batch * nc) || !rocm_fits_(out, total)) return 0;
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)a.dev;
HipDevPtr_* pc = (HipDevPtr_*)c.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int nb = (unsigned int)batch;
unsigned int n1 = (unsigned int)na;
unsigned int n2 = (unsigned int)nc;
HipKernelArg_* params[6];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&pc; params[2] = (HipKernelArg_*)&po;
params[3] = (HipKernelArg_*)&nb; params[4] = (HipKernelArg_*)&n1; params[5] = (HipKernelArg_*)&n2;
return rocm_launch_(K_CONCAT, total, ¶ms[0]);
}
}
int rocm_upsample2x_concat(struct RocmBuf x, unsigned long xChannels, struct RocmBuf skip,
unsigned long skipChannels, struct RocmBuf out,
unsigned long batch, unsigned long h, unsigned long w) {
unsigned long plane = 4UL * h * w;
unsigned long total = batch * (xChannels + skipChannels) * plane;
if (!rocm_ensure_() || total == 0UL || total > 4294967295UL ||
!rocm_fits_(x, batch * xChannels * h * w) || !rocm_fits_(skip, batch * skipChannels * plane) ||
!rocm_fits_(out, total)) return 0;
unsafe {
HipDevPtr_* px = (HipDevPtr_*)x.dev;
HipDevPtr_* ps = (HipDevPtr_*)skip.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int nb = (unsigned int)batch;
unsigned int ncx = (unsigned int)xChannels;
unsigned int ncs = (unsigned int)skipChannels;
unsigned int nh = (unsigned int)h;
unsigned int nw = (unsigned int)w;
HipKernelArg_* params[8];
params[0] = (HipKernelArg_*)&px; params[1] = (HipKernelArg_*)&ps; params[2] = (HipKernelArg_*)&po; params[3] = (HipKernelArg_*)&nb;
params[4] = (HipKernelArg_*)&ncx; params[5] = (HipKernelArg_*)&ncs; params[6] = (HipKernelArg_*)&nh; params[7] = (HipKernelArg_*)&nw;
return rocm_launch_(K_UPCAT, total, ¶ms[0]);
}
}
int rocm_upsample2x_nearest(struct RocmBuf in, struct RocmBuf out,
unsigned long planes, unsigned long h, unsigned long w) {
unsigned long total = planes * 4UL * h * w;
if (!rocm_ensure_() || total == 0UL || total > 4294967295UL ||
!rocm_fits_(in, planes * h * w) || !rocm_fits_(out, total)) return 0;
unsafe {
HipDevPtr_* px = (HipDevPtr_*)in.dev;
HipDevPtr_* py = (HipDevPtr_*)out.dev;
unsigned int np = (unsigned int)planes;
unsigned int nh = (unsigned int)h;
unsigned int nw = (unsigned int)w;
HipKernelArg_* params[5];
params[0] = (HipKernelArg_*)&px; params[1] = (HipKernelArg_*)&py; params[2] = (HipKernelArg_*)&np;
params[3] = (HipKernelArg_*)&nh; params[4] = (HipKernelArg_*)&nw;
return rocm_launch_(K_UPSAMPLE2X, total, ¶ms[0]);
}
}
int rocm_bias_channels(struct RocmBuf y, struct RocmBuf bias, unsigned long batch, unsigned long channels,
unsigned long plane, int relu) {
unsigned long total = batch * channels * plane;
if (!rocm_ensure_() || total == 0UL || total > 4294967295UL ||
!rocm_fits_(y, total) || !rocm_fits_(bias, channels)) return 0;
unsafe {
HipDevPtr_* py = (HipDevPtr_*)y.dev;
HipDevPtr_* pb = (HipDevPtr_*)bias.dev;
unsigned int nt = (unsigned int)total;
unsigned int npl = (unsigned int)plane;
unsigned int nc = (unsigned int)channels;
int r = relu ? 1 : 0;
HipKernelArg_* params[6];
params[0] = (HipKernelArg_*)&py; params[1] = (HipKernelArg_*)&pb; params[2] = (HipKernelArg_*)&nt;
params[3] = (HipKernelArg_*)&npl; params[4] = (HipKernelArg_*)&nc; params[5] = (HipKernelArg_*)&r;
return rocm_launch_(K_BIAS_ACT, total, ¶ms[0]);
}
}
int rocm_sgd(struct RocmBuf w, struct RocmBuf g, float lr) {
unsigned long n = w.len;
if (!rocm_ensure_() || !rocm_fits_(w, n) || !rocm_fits_(g, n) || n > 4294967295UL) return 0;
unsafe {
HipDevPtr_* pw = (HipDevPtr_*)w.dev;
HipDevPtr_* pg = (HipDevPtr_*)g.dev;
float rate = lr;
unsigned int nn = (unsigned int)n;
HipKernelArg_* params[4];
params[0] = (HipKernelArg_*)&pw; params[1] = (HipKernelArg_*)&pg; params[2] = (HipKernelArg_*)&rate; params[3] = (HipKernelArg_*)&nn;
return rocm_launch_(K_SGD, n, ¶ms[0]);
}
}
int rocm_sum(struct RocmBuf a, struct RocmBuf out1) {
unsigned long n = a.len;
if (!rocm_ensure_() || !rocm_fits_(a, n) || !rocm_fits_(out1, 1UL) || n > 4294967295UL) return 0;
unsafe {
if (hipMemsetAsync((HipDevPtr_*)out1.dev, 0, sizeof(float), g_rocm_stream) != ROCM_HIP_OK) return 0;
HipDevPtr_* pa = (HipDevPtr_*)a.dev;
HipDevPtr_* po = (HipDevPtr_*)out1.dev;
unsigned int nn = (unsigned int)n;
HipKernelArg_* params[3];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&po; params[2] = (HipKernelArg_*)&nn;
return rocm_launch_(K_SUM, n, ¶ms[0]);
}
}
int rocm_mse_grad(struct RocmBuf y, struct RocmBuf t, struct RocmBuf dy, struct RocmBuf loss1) {
unsigned long n = y.len;
if (!rocm_ensure_() || !rocm_fits_(t, n) || !rocm_fits_(dy, n) || !rocm_fits_(loss1, 1UL) ||
n > 4294967295UL) return 0;
unsafe {
if (hipMemsetAsync((HipDevPtr_*)loss1.dev, 0, sizeof(float), g_rocm_stream) != ROCM_HIP_OK) return 0;
HipDevPtr_* py = (HipDevPtr_*)y.dev;
HipDevPtr_* pt = (HipDevPtr_*)t.dev;
HipDevPtr_* pd = (HipDevPtr_*)dy.dev;
HipDevPtr_* pl = (HipDevPtr_*)loss1.dev;
unsigned int nn = (unsigned int)n;
HipKernelArg_* params[5];
params[0] = (HipKernelArg_*)&py; params[1] = (HipKernelArg_*)&pt; params[2] = (HipKernelArg_*)&pd;
params[3] = (HipKernelArg_*)&pl; params[4] = (HipKernelArg_*)&nn;
return rocm_launch_(K_MSE, n, ¶ms[0]);
}
}
// ── Graphs ───────────────────────────────────────────────────────────────────
int rocm_capture_begin() {
if (!rocm_ensure_()) return 0;
unsafe {
if (hipStreamBeginCapture(g_rocm_stream, ROCM_CAPTURE_GLOBAL) != ROCM_HIP_OK) return 0;
}
g_rocm_capturing = 1;
return 1;
}
struct RocmGraph rocm_capture_end() {
struct RocmGraph g;
g.exec = 0UL;
if (!rocm_ensure_()) return g;
unsafe {
HipGraph_* graph = null;
g_rocm_capturing = 0;
if (hipStreamEndCapture(g_rocm_stream, &graph) != ROCM_HIP_OK) return g;
HipGraphExec_* exec = null;
int ok = hipGraphInstantiate(&exec, graph, null, (char*)0, 0UL) == ROCM_HIP_OK;
hipGraphDestroy(graph);
if (ok) g.exec = (unsigned long)exec;
}
return g;
}
int rocm_graph_launch(struct RocmGraph g) {
if (g.exec == 0UL || !rocm_ensure_()) return 0;
unsafe { return hipGraphLaunch((HipGraphExec_*)g.exec, g_rocm_stream) == ROCM_HIP_OK; }
}
void RocmGraph::free() {
if (self.exec != 0UL) { unsafe { hipGraphExecDestroy((HipGraphExec_*)self.exec); } }
self.exec = 0UL;
}
// ── Host-array convenience ops ───────────────────────────────────────────────
static int rocm_host_binary_(int kernel, const &float a, const &float b, &float out, unsigned long n) {
struct RocmBuf da = RocmBuf::upload_n_(a, n);
struct RocmBuf db = RocmBuf::upload_n_(b, n);
struct RocmBuf dout = RocmBuf::alloc(n);
int ok = da.dev != 0UL && db.dev != 0UL && dout.dev != 0UL &&
rocm_binary_(kernel, da, db, dout) && rocm_read_n_(dout, out, n);
da.free();
db.free();
dout.free();
return ok;
}
static int rocm_host_unary_(int kernel, const &float a, float k, &float out, unsigned long n) {
struct RocmBuf da = RocmBuf::upload_n_(a, n);
struct RocmBuf dout = RocmBuf::alloc(n);
int ok = da.dev != 0UL && dout.dev != 0UL &&
rocm_unary_(kernel, da, k, dout) && rocm_read_n_(dout, out, n);
da.free();
dout.free();
return ok;
}
int rocm_add_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return rocm_host_binary_(K_ADD, a, b, out, n); }
int rocm_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 rocm_add_f32_n_(a, b, out, n);
}
int rocm_sub_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return rocm_host_binary_(K_SUB, a, b, out, n); }
int rocm_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 rocm_sub_f32_n_(a, b, out, n);
}
int rocm_mul_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return rocm_host_binary_(K_MUL, a, b, out, n); }
int rocm_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 rocm_mul_f32_n_(a, b, out, n);
}
int rocm_div_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return rocm_host_binary_(K_DIV, a, b, out, n); }
int rocm_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 rocm_div_f32_n_(a, b, out, n);
}
int rocm_pow_f32_n_(const &float a, const &float b, &float out, unsigned long n) { return rocm_host_binary_(K_POW, a, b, out, n); }
int rocm_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 rocm_pow_f32_n_(a, b, out, n);
}
int rocm_relu_f32_n_(const &float a, &float out, unsigned long n) { return rocm_host_unary_(K_RELU, a, 0.0f, out, n); }
int rocm_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 rocm_relu_f32_n_(a, out, n);
}
int rocm_log_f32_n_(const &float a, &float out, unsigned long n) { return rocm_host_unary_(K_LOG, a, 0.0f, out, n); }
int rocm_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 rocm_log_f32_n_(a, out, n);
}
int rocm_exp_f32_n_(const &float a, &float out, unsigned long n) { return rocm_host_unary_(K_EXP, a, 0.0f, out, n); }
int rocm_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 rocm_exp_f32_n_(a, out, n);
}
int rocm_sqrt_f32_n_(const &float a, &float out, unsigned long n) { return rocm_host_unary_(K_SQRT, a, 0.0f, out, n); }
int rocm_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 rocm_sqrt_f32_n_(a, out, n);
}
int rocm_scale_f32_n_(const &float a, float k, &float out, unsigned long n) { return rocm_host_unary_(K_SCALE, a, k, out, n); }
int rocm_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 rocm_scale_f32_n_(a, k, out, n);
}
int rocm_sum_f32_n_(const &float a, &float out1, unsigned long n) {
struct RocmBuf da = RocmBuf::upload_n_(a, n);
struct RocmBuf dout = RocmBuf::alloc(1UL);
int ok = da.dev != 0UL && dout.dev != 0UL && rocm_sum(da, dout) && rocm_read_n_(dout, out1, 1UL);
da.free();
dout.free();
return ok;
}
int rocm_sum_f32([]const float a_s, &float out1) {
unsigned long n = a_s.len;
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
return rocm_sum_f32_n_(a, out1, n);
}
// ── Attention ────────────────────────────────────────────────────────────────
static int rocm_launch_grid3_(int kernel, unsigned long gx, unsigned long gy, unsigned long gz, HipKernelArg_** params) {
if (gx == 0UL || gy == 0UL || gz == 0UL || gx > 2147483647UL || gy > 65535UL || gz > 65535UL) return 0;
unsafe {
return hipModuleLaunchKernel(g_rocm_fn[kernel], (unsigned int)gx, (unsigned int)gy, (unsigned int)gz,
ROCM_BLOCK, 1U, 1U, 0U, g_rocm_stream,
params, (HipKernelArg_**)0) == ROCM_HIP_OK;
}
}
static int rocm_launch_grid_(int kernel, unsigned long gx, unsigned long gy, HipKernelArg_** params) {
return rocm_launch_grid3_(kernel, gx, gy, 1UL, params);
}
static int rocm_attn_shape_ok_(struct RocmBuf q, struct RocmBuf k, struct RocmBuf v, struct RocmBuf out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen, unsigned long headDim) {
if (heads == 0UL || seqLen == 0UL || kvLen == 0UL || headDim == 0UL || headDim > 128UL) return 0;
if (heads > 65535UL || seqLen > 2147483647UL || kvLen > 2147483647UL) return 0;
return rocm_fits_(q, heads * seqLen * headDim) && rocm_fits_(k, heads * kvLen * headDim) &&
rocm_fits_(v, heads * kvLen * headDim) && rocm_fits_(out, heads * seqLen * headDim);
}
int rocm_flash_attention(struct RocmBuf q, struct RocmBuf k, struct RocmBuf v, struct RocmBuf out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen, unsigned long headDim,
float scale, int causal) {
if (!rocm_ensure_() || !rocm_attn_shape_ok_(q, k, v, out, heads, seqLen, kvLen, headDim)) return 0;
if (causal && kvLen < seqLen) return 0;
unsafe {
HipDevPtr_* pq = (HipDevPtr_*)q.dev;
HipDevPtr_* pk = (HipDevPtr_*)k.dev;
HipDevPtr_* pv = (HipDevPtr_*)v.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
unsigned int s = (unsigned int)seqLen;
unsigned int kn = (unsigned int)kvLen;
unsigned int d = (unsigned int)headDim;
float sc = scale > 0.0f ? scale : (float)(1.0 / sqrt_d((double)headDim));
int ca = causal ? 1 : 0;
HipKernelArg_* params[9];
params[0] = (HipKernelArg_*)&pq; params[1] = (HipKernelArg_*)&pk; params[2] = (HipKernelArg_*)&pv; params[3] = (HipKernelArg_*)&po;
params[4] = (HipKernelArg_*)&s; params[5] = (HipKernelArg_*)&kn; params[6] = (HipKernelArg_*)&d; params[7] = (HipKernelArg_*)≻
params[8] = (HipKernelArg_*)&ca;
if (g_rocm_mfma) {
// matrix cores: grid (heads, query blocks of 128 / 64 rows)
if (headDim <= 64UL) return rocm_launch_grid_(K_FLASH_M64, heads, (seqLen + 127UL) / 128UL, ¶ms[0]);
return rocm_launch_grid_(K_FLASH_M128, heads, (seqLen + 63UL) / 64UL, ¶ms[0]);
}
return rocm_launch_grid_(headDim <= 64UL ? K_FLASH64 : K_FLASH128, (seqLen + 63UL) / 64UL, heads,
¶ms[0]);
}
}
// Per-chunk states for rocm_linear_attention, kept and grown as needed
// (not during graph capture: growing would allocate).
static struct RocmBuf g_rocm_lin_scratch;
static int g_rocm_lin_init = 0;
// Grows the linear-attention scratch to 'need' floats (not during capture).
static int rocm_lin_scratch_(unsigned long need) {
if (!g_rocm_lin_init) { g_rocm_lin_scratch.dev = 0UL; g_rocm_lin_scratch.len = 0UL; g_rocm_lin_init = 1; }
if (g_rocm_lin_scratch.len >= need) return 1;
if (g_rocm_capturing) return 0;
g_rocm_lin_scratch.free();
g_rocm_lin_scratch = RocmBuf::alloc(need);
if (g_rocm_lin_scratch.dev == 0UL) { g_rocm_lin_scratch.len = 0UL; return 0; }
return 1;
}
// Matrix-core linear attention. Causal with head dim <= 64: one fused pass
// (la_causal_d64, decoupled look-back over 128-row chunks). Otherwise three
// passes: per-chunk states (128-row chunks; non-causal uses longer row ranges
// once there would be more than 64 of them), a scan over chunks, then the
// output pass. Head dims are padded to 64 or 128.
static int rocm_linear_attention_mfma_(struct RocmBuf q, struct RocmBuf k, struct RocmBuf v, struct RocmBuf out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim, int causal) {
unsigned long dm = headDim <= 64UL ? 64UL : 128UL;
if (causal && dm == 64UL) {
// one pass: per-chunk records (state, then inclusive prefix) and a
// flag per chunk plus a ticket counter after them, zeroed per call
unsigned long ncc = (seqLen + 127UL) / 128UL;
unsigned long rs = 64UL * 64UL + 64UL;
unsigned long recs = heads * ncc * 2UL * rs;
if (!rocm_lin_scratch_(recs + heads * ncc + 1UL)) return 0;
unsafe {
HipDevPtr_* pq = (HipDevPtr_*)q.dev;
HipDevPtr_* pk = (HipDevPtr_*)k.dev;
HipDevPtr_* pv = (HipDevPtr_*)v.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
HipDevPtr_* pr = (HipDevPtr_*)g_rocm_lin_scratch.dev;
HipDevPtr_* pf = (HipDevPtr_*)(g_rocm_lin_scratch.dev + recs * 4UL);
if (hipMemsetAsync(pf, 0, (heads * ncc + 1UL) * 4UL, g_rocm_stream) != ROCM_HIP_OK) return 0;
unsigned int s = (unsigned int)seqLen;
unsigned int hn = (unsigned int)heads;
unsigned int d = (unsigned int)headDim;
unsigned int nc = (unsigned int)ncc;
HipKernelArg_* pc[10];
pc[0] = (HipKernelArg_*)&pq; pc[1] = (HipKernelArg_*)&pk; pc[2] = (HipKernelArg_*)&pv; pc[3] = (HipKernelArg_*)≺ pc[4] = (HipKernelArg_*)&pf;
pc[5] = (HipKernelArg_*)&po; pc[6] = (HipKernelArg_*)&s; pc[7] = (HipKernelArg_*)&hn; pc[8] = (HipKernelArg_*)&d; pc[9] = (HipKernelArg_*)&nc;
return rocm_launch_grid_(K_LA_CAUSAL64, ncc * heads, 1UL, &pc[0]);
}
}
unsigned long rows = 128UL;
if (!causal) rows = 128UL * (((kvLen + 127UL) / 128UL + 63UL) / 64UL);
unsigned long nchunks = (kvLen + rows - 1UL) / rows;
if (!rocm_lin_scratch_(heads * nchunks * dm * (dm + 1UL))) return 0;
int big = dm == 128UL;
unsafe {
HipDevPtr_* pq = (HipDevPtr_*)q.dev;
HipDevPtr_* pk = (HipDevPtr_*)k.dev;
HipDevPtr_* pv = (HipDevPtr_*)v.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
HipDevPtr_* pw = (HipDevPtr_*)g_rocm_lin_scratch.dev;
HipDevPtr_* pz = (HipDevPtr_*)(g_rocm_lin_scratch.dev + heads * nchunks * dm * dm * 4UL);
unsigned int s = (unsigned int)seqLen;
unsigned int kn = (unsigned int)kvLen;
unsigned int d = (unsigned int)headDim;
unsigned int rw = (unsigned int)rows;
unsigned int nc = (unsigned int)nchunks;
int ca = causal ? 1 : 0;
HipKernelArg_* p1[8];
p1[0] = (HipKernelArg_*)&pk; p1[1] = (HipKernelArg_*)&pv; p1[2] = (HipKernelArg_*)&pw; p1[3] = (HipKernelArg_*)&pz; p1[4] = (HipKernelArg_*)&kn;
p1[5] = (HipKernelArg_*)&d; p1[6] = (HipKernelArg_*)&rw; p1[7] = (HipKernelArg_*)&nc;
if (!rocm_launch_grid_(big ? K_LA_STATE128 : K_LA_STATE64, nchunks, heads, &p1[0])) return 0;
if (causal || nchunks > 1UL) {
HipKernelArg_* p2[4];
p2[0] = (HipKernelArg_*)&pw; p2[1] = (HipKernelArg_*)&pz; p2[2] = (HipKernelArg_*)&nc; p2[3] = (HipKernelArg_*)&ca;
if (!rocm_launch_grid_(big ? K_LA_SCAN128 : K_LA_SCAN64, (dm * dm + dm + 255UL) / 256UL, heads,
&p2[0])) return 0;
}
HipKernelArg_* p3[11];
p3[0] = (HipKernelArg_*)&pq; p3[1] = (HipKernelArg_*)&pk; p3[2] = (HipKernelArg_*)&pv; p3[3] = (HipKernelArg_*)&pw; p3[4] = (HipKernelArg_*)&pz;
p3[5] = (HipKernelArg_*)&po; p3[6] = (HipKernelArg_*)&s; p3[7] = (HipKernelArg_*)&kn; p3[8] = (HipKernelArg_*)&d; p3[9] = (HipKernelArg_*)&nc;
p3[10] = (HipKernelArg_*)&ca;
return rocm_launch_grid3_(big ? K_LA_OUT128 : K_LA_OUT64, (seqLen + 127UL) / 128UL, heads, dm / 64UL,
&p3[0]);
}
}
int rocm_linear_attention(struct RocmBuf q, struct RocmBuf k, struct RocmBuf v, struct RocmBuf out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen, unsigned long headDim,
int causal) {
if (!rocm_ensure_() || !rocm_attn_shape_ok_(q, k, v, out, heads, seqLen, kvLen, headDim)) return 0;
if (causal && kvLen != seqLen) return 0;
if (g_rocm_mfma) return rocm_linear_attention_mfma_(q, k, v, out, heads, seqLen, kvLen, headDim, causal);
unsigned long nchunks = (kvLen + 31UL) / 32UL;
if (!rocm_lin_scratch_(heads * nchunks * headDim * (headDim + 1UL))) return 0;
unsafe {
HipDevPtr_* pq = (HipDevPtr_*)q.dev;
HipDevPtr_* pk = (HipDevPtr_*)k.dev;
HipDevPtr_* pv = (HipDevPtr_*)v.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
HipDevPtr_* ps = (HipDevPtr_*)g_rocm_lin_scratch.dev;
unsigned int s = (unsigned int)seqLen;
unsigned int kn = (unsigned int)kvLen;
unsigned int d = (unsigned int)headDim;
unsigned int nc = (unsigned int)nchunks;
int ca = causal ? 1 : 0;
HipKernelArg_* p1[6];
p1[0] = (HipKernelArg_*)&pk; p1[1] = (HipKernelArg_*)&pv; p1[2] = (HipKernelArg_*)&ps; p1[3] = (HipKernelArg_*)&kn; p1[4] = (HipKernelArg_*)&d;
p1[5] = (HipKernelArg_*)&nc;
if (!rocm_launch_grid_(K_LIN_STATE, nchunks, heads, &p1[0])) return 0;
HipKernelArg_* p2[4];
p2[0] = (HipKernelArg_*)&ps; p2[1] = (HipKernelArg_*)&d; p2[2] = (HipKernelArg_*)&nc; p2[3] = (HipKernelArg_*)&ca;
unsigned long elems = headDim * (headDim + 1UL);
if (!rocm_launch_grid_(K_LIN_PREFIX, (elems + 255UL) / 256UL, heads, &p2[0])) return 0;
HipKernelArg_* p3[10];
p3[0] = (HipKernelArg_*)&pq; p3[1] = (HipKernelArg_*)&pk; p3[2] = (HipKernelArg_*)&pv; p3[3] = (HipKernelArg_*)&ps; p3[4] = (HipKernelArg_*)&po;
p3[5] = (HipKernelArg_*)&s; p3[6] = (HipKernelArg_*)&kn; p3[7] = (HipKernelArg_*)&d; p3[8] = (HipKernelArg_*)&nc; p3[9] = (HipKernelArg_*)&ca;
return rocm_launch_grid_(K_LIN_OUT, (seqLen + 31UL) / 32UL, heads, &p3[0]);
}
}
static int rocm_attention_f32_(const &float q, const &float k, const &float v, &float out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim, float scale, int causal, int linear) {
unsigned long nq = heads * seqLen * headDim;
unsigned long nk = heads * kvLen * headDim;
struct RocmBuf dq = RocmBuf::upload_n_(q, nq);
struct RocmBuf dk = RocmBuf::upload_n_(k, nk);
struct RocmBuf dv = RocmBuf::upload_n_(v, nk);
struct RocmBuf dout = RocmBuf::alloc(nq);
int ok = dq.dev != 0UL && dk.dev != 0UL && dv.dev != 0UL && dout.dev != 0UL;
if (ok) {
ok = linear ? rocm_linear_attention(dq, dk, dv, dout, heads, seqLen, kvLen, headDim, causal)
: rocm_flash_attention(dq, dk, dv, dout, heads, seqLen, kvLen, headDim, scale, causal);
ok = ok && rocm_read_n_(dout, out, nq);
}
dq.free();
dk.free();
dv.free();
dout.free();
return ok;
}
int rocm_flash_attention_f32(const &float q, const &float k, const &float v, &float out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim, float scale, int causal) {
return rocm_attention_f32_(q, k, v, out, heads, seqLen, kvLen, headDim, scale, causal, 0);
}
int rocm_linear_attention_f32(const &float q, const &float k, const &float v, &float out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim, int causal) {
return rocm_attention_f32_(q, k, v, out, heads, seqLen, kvLen, headDim, 0.0f, causal, 1);
}
int rocm_matmul_f32_n_(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
if (M > 2147483647UL || N > 2147483647UL || K > 2147483647UL) return 0;
struct RocmBuf da = RocmBuf::upload_n_(a, M * K);
struct RocmBuf db = RocmBuf::upload_n_(b, K * N);
struct RocmBuf dc = RocmBuf::alloc(M * N);
int ok = da.dev != 0UL && db.dev != 0UL && dc.dev != 0UL;
if (ok) {
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)da.dev;
HipDevPtr_* pb = (HipDevPtr_*)db.dev;
HipDevPtr_* pc = (HipDevPtr_*)dc.dev;
int m = (int)M;
int n = (int)N;
int k = (int)K;
HipKernelArg_* params[6];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&pb; params[2] = (HipKernelArg_*)&pc;
params[3] = (HipKernelArg_*)&m; params[4] = (HipKernelArg_*)&n; params[5] = (HipKernelArg_*)&k;
ok = hipModuleLaunchKernel(g_rocm_fn[K_SGEMM],
(unsigned int)((N + 63UL) / 64UL), (unsigned int)((M + 63UL) / 64UL), 1U,
ROCM_BLOCK, 1U, 1U, 0U, g_rocm_stream,
¶ms[0], (HipKernelArg_**)0) == ROCM_HIP_OK;
}
ok = ok && rocm_read_n_(dc, out, M * N);
}
da.free();
db.free();
dc.free();
return ok;
}
int rocm_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 rocm_matmul_f32_n_(a, b, out, M, K, N);
}
int rocm_matmul_f32_blas_n_(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
struct RocmBuf da = RocmBuf::upload_n_(a, M * K);
struct RocmBuf db = RocmBuf::upload_n_(b, K * N);
struct RocmBuf dc = RocmBuf::alloc(M * N);
int ok = da.dev != 0UL && db.dev != 0UL && dc.dev != 0UL &&
rocm_gemm(da, 0, db, 0, dc, M, N, K) && rocm_read_n_(dc, out, M * N);
da.free();
db.free();
dc.free();
return ok;
}
int rocm_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 rocm_matmul_f32_blas_n_(a, b, out, M, K, N);
}
// ── Mixed precision ──────────────────────────────────────────────────────────
extern unsigned char* malloc(unsigned long size);
extern void free(unsigned char* ptr);
int rocm_zero(struct RocmBuf b) {
if (!rocm_ensure_() || b.dev == 0UL) return 0;
unsafe { return hipMemsetAsync((HipDevPtr_*)b.dev, 0, b.len * sizeof(float), g_rocm_stream) == ROCM_HIP_OK; }
}
int rocm_half_convert(struct RocmHalfBuf src, struct RocmHalfBuf dst) {
unsigned long n = src.len;
if (!rocm_ensure_() || !rocm_half_fits_(src, n) || !rocm_half_fits_(dst, n) || n > 4294967295UL) return 0;
if (src.dtype == dst.dtype) {
unsafe {
return hipMemcpyAsync((unsigned char*)dst.dev, (const unsigned char*)src.dev, n * 2UL, ROCM_D2D, g_rocm_stream) == ROCM_HIP_OK;
}
}
return rocm_convert_(K_HALF_CVT, src.dev, dst.dev, n, dst.dtype == ROCM_BF16);
}
// The kernel's CvtTab: 32 source / destination pointers and lengths.
#define ROCM_CVT_TAB 32UL
struct RocmCvtTab_ {
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 rocm_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 RocmCvtTab_ t;
unsigned long c = count - base < ROCM_CVT_TAB ? count - base : ROCM_CVT_TAB;
unsigned long most = 1UL;
unsafe {
unsigned long i = 0UL;
while (i < ROCM_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)ROCM_BLOCK - 1UL) / (unsigned long)ROCM_BLOCK;
if (blocks > 1024UL) blocks = 1024UL;
int m = mode;
int b = bf;
float sc = scale;
HipDevPtr_* pf = (HipDevPtr_*)flag;
HipKernelArg_* params[5];
params[0] = (HipKernelArg_*)&t; params[1] = (HipKernelArg_*)&m; params[2] = (HipKernelArg_*)&b; params[3] = (HipKernelArg_*)≻
params[4] = (HipKernelArg_*)&pf;
if (hipModuleLaunchKernel(g_rocm_fn[K_CVT_MANY], (unsigned int)blocks, (unsigned int)c, 1U, ROCM_BLOCK, 1U, 1U,
0U, g_rocm_stream, ¶ms[0], (HipKernelArg_**)0) != ROCM_HIP_OK) return 0;
}
base = base + c;
}
return 1;
}
// Pointer / length tables for rocm_cvt_many_ (heap, freed by the caller).
struct RocmCvtLists_ {
unsigned long* src;
unsigned long* dst;
unsigned long* len;
int ok;
int bf;
};
static struct RocmCvtLists_ rocm_cvt_lists_(unsigned long count) {
struct RocmCvtLists_ 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 rocm_cvt_lists_free_(&RocmCvtLists_ 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 rocm_cvt_dtype_(&RocmCvtLists_ l, int dtype) {
int bf = dtype == ROCM_BF16 ? 1 : 0;
if (l.bf < 0) l.bf = bf;
else if (l.bf != bf) l.ok = 0;
}
int rocm_to_half_many(const &Vec<struct RocmBuf> src, const &Vec<struct RocmHalfBuf> dst) {
unsigned long n = src.length();
if (!rocm_ensure_() || n != dst.length()) return 0;
if (n == 0UL) return 1;
struct RocmCvtLists_ l = rocm_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct RocmBuf s = *src.at(i);
struct RocmHalfBuf d = *dst.at(i);
if (!rocm_fits_(s, d.len) || !rocm_half_fits_(d, d.len) || d.len > 4294967295UL) l.ok = 0;
rocm_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 && rocm_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);
rocm_cvt_lists_free_(&l);
return ok;
}
int rocm_from_half_many(const &Vec<struct RocmHalfBuf> src, const &Vec<struct RocmBuf> dst) {
unsigned long n = src.length();
if (!rocm_ensure_() || n != dst.length()) return 0;
if (n == 0UL) return 1;
struct RocmCvtLists_ l = rocm_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct RocmHalfBuf s = *src.at(i);
struct RocmBuf d = *dst.at(i);
if (!rocm_half_fits_(s, d.len) || !rocm_fits_(d, d.len) || d.len > 4294967295UL) l.ok = 0;
rocm_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 && rocm_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);
rocm_cvt_lists_free_(&l);
return ok;
}
int rocm_half_unscale_many(const &Vec<struct RocmHalfBuf> src, const &Vec<struct RocmBuf> dst, float scale,
struct RocmBuf flag) {
unsigned long n = src.length();
if (!rocm_ensure_() || n != dst.length() || !rocm_fits_(flag, 1UL)) return 0;
if (n == 0UL) return 1;
struct RocmCvtLists_ l = rocm_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct RocmHalfBuf s = *src.at(i);
struct RocmBuf d = *dst.at(i);
if (!rocm_half_fits_(s, s.len) || !rocm_fits_(d, s.len) || s.len > 4294967295UL) l.ok = 0;
rocm_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 && rocm_cvt_many_((const unsigned long*)l.src, (const unsigned long*)l.dst, (const unsigned long*)l.len,
n, 1, l.bf, scale, flag.dev);
rocm_cvt_lists_free_(&l);
return ok;
}
int rocm_unscale_many(const &Vec<struct RocmBuf> src, const &Vec<struct RocmBuf> dst, float scale,
struct RocmBuf flag) {
unsigned long n = src.length();
if (!rocm_ensure_() || n != dst.length() || !rocm_fits_(flag, 1UL)) return 0;
if (n == 0UL) return 1;
struct RocmCvtLists_ l = rocm_cvt_lists_(n);
unsafe {
unsigned long i = 0UL;
while (i < n) {
struct RocmBuf s = *src.at(i);
struct RocmBuf d = *dst.at(i);
if (!rocm_fits_(s, s.len) || !rocm_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 && rocm_cvt_many_((const unsigned long*)l.src, (const unsigned long*)l.dst, (const unsigned long*)l.len,
n, 2, 0, scale, flag.dev);
rocm_cvt_lists_free_(&l);
return ok;
}
int rocm_mp_sgd(struct RocmBuf w, struct RocmBuf g, struct RocmBuf vel, struct RocmHalfBuf wHalf,
float lr, float momentum, float weightDecay, struct RocmBuf skip) {
unsigned long n = w.len;
if (!rocm_ensure_() || !rocm_fits_(w, n) || !rocm_fits_(g, n) || n > 4294967295UL) return 0;
if (vel.dev != 0UL && !rocm_fits_(vel, n)) return 0;
if (wHalf.dev != 0UL && !rocm_half_fits_(wHalf, n)) return 0;
unsafe {
HipDevPtr_* pw = (HipDevPtr_*)w.dev;
HipDevPtr_* pg = (HipDevPtr_*)g.dev;
HipDevPtr_* pv = (HipDevPtr_*)vel.dev;
HipDevPtr_* ph = (HipDevPtr_*)wHalf.dev;
HipDevPtr_* ps = (HipDevPtr_*)skip.dev;
float l = lr;
float mo = momentum;
float wd = weightDecay;
unsigned int nn = (unsigned int)n;
int bf = wHalf.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[10];
params[0] = (HipKernelArg_*)&pw; params[1] = (HipKernelArg_*)&pg; params[2] = (HipKernelArg_*)&pv; params[3] = (HipKernelArg_*)&ph;
params[4] = (HipKernelArg_*)&l; params[5] = (HipKernelArg_*)&mo; params[6] = (HipKernelArg_*)&wd; params[7] = (HipKernelArg_*)&nn;
params[8] = (HipKernelArg_*)&bf; params[9] = (HipKernelArg_*)&ps;
return rocm_launch_(K_MP_SGD, n, ¶ms[0]);
}
}
int rocm_mp_adam(struct RocmBuf w, struct RocmBuf m, struct RocmBuf v, struct RocmBuf g, struct RocmHalfBuf wHalf,
float lr, float beta1, float beta2, float eps, float weightDecay, int decoupled,
unsigned long step, struct RocmBuf skip) {
unsigned long n = w.len;
if (!rocm_ensure_() || !rocm_fits_(w, n) || !rocm_fits_(m, n) || !rocm_fits_(v, n) || !rocm_fits_(g, n) ||
n > 4294967295UL || step == 0UL) return 0;
if (wHalf.dev != 0UL && !rocm_half_fits_(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 {
HipDevPtr_* pw = (HipDevPtr_*)w.dev;
HipDevPtr_* pm = (HipDevPtr_*)m.dev;
HipDevPtr_* pv = (HipDevPtr_*)v.dev;
HipDevPtr_* pg = (HipDevPtr_*)g.dev;
HipDevPtr_* ph = (HipDevPtr_*)wHalf.dev;
HipDevPtr_* ps = (HipDevPtr_*)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 == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[16];
params[0] = (HipKernelArg_*)&pw; params[1] = (HipKernelArg_*)± params[2] = (HipKernelArg_*)&pv; params[3] = (HipKernelArg_*)&pg;
params[4] = (HipKernelArg_*)&ph; params[5] = (HipKernelArg_*)&l; params[6] = (HipKernelArg_*)&b1; params[7] = (HipKernelArg_*)&b2;
params[8] = (HipKernelArg_*)&e; params[9] = (HipKernelArg_*)&wd; params[10] = (HipKernelArg_*)&c1; params[11] = (HipKernelArg_*)&c2;
params[12] = (HipKernelArg_*)&dec; params[13] = (HipKernelArg_*)&nn; params[14] = (HipKernelArg_*)&bf; params[15] = (HipKernelArg_*)&ps;
return rocm_launch_(K_MP_ADAM, n, ¶ms[0]);
}
}
static int rocm_half_ew_(struct RocmHalfBuf a, struct RocmHalfBuf b, struct RocmHalfBuf out, float k, int op) {
unsigned long n = a.len;
if (!rocm_ensure_() || !rocm_half_fits_(a, n) || !rocm_half_fits_(out, n) || a.dtype != out.dtype ||
n > 4294967295UL) return 0;
if (op != 3 && (!rocm_half_fits_(b, n) || b.dtype != a.dtype)) return 0;
unsafe {
HipDevPtr_* pa = (HipDevPtr_*)a.dev;
HipDevPtr_* pb = (HipDevPtr_*)b.dev;
HipDevPtr_* po = (HipDevPtr_*)out.dev;
float kk = k;
unsigned int nn = (unsigned int)n;
int o = op;
int bf = a.dtype == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[7];
params[0] = (HipKernelArg_*)&pa; params[1] = (HipKernelArg_*)&pb; params[2] = (HipKernelArg_*)&po; params[3] = (HipKernelArg_*)&kk;
params[4] = (HipKernelArg_*)&nn; params[5] = (HipKernelArg_*)&o; params[6] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_EW, n, ¶ms[0]);
}
}
int rocm_half_add(struct RocmHalfBuf a, struct RocmHalfBuf b, struct RocmHalfBuf out) { return rocm_half_ew_(a, b, out, (float)0.0, 0); }
int rocm_half_sub(struct RocmHalfBuf a, struct RocmHalfBuf b, struct RocmHalfBuf out) { return rocm_half_ew_(a, b, out, (float)0.0, 1); }
int rocm_half_mul(struct RocmHalfBuf a, struct RocmHalfBuf b, struct RocmHalfBuf out) { return rocm_half_ew_(a, b, out, (float)0.0, 2); }
int rocm_half_scale(struct RocmHalfBuf a, float k, struct RocmHalfBuf out) { return rocm_half_ew_(a, a, out, k, 3); }
int rocm_half_axpy(struct RocmHalfBuf a, float k, struct RocmHalfBuf b, struct RocmHalfBuf out) {
return rocm_half_ew_(a, b, out, k, 4);
}
int rocm_half_bias(struct RocmHalfBuf y, struct RocmHalfBuf bias, unsigned long rows, unsigned long channels,
unsigned long plane, int relu) {
unsigned long total = rows * channels * plane;
if (!rocm_ensure_() || total == 0UL || total > 4294967295UL || !rocm_half_fits_(y, total) ||
!rocm_half_fits_(bias, channels) || y.dtype != bias.dtype) return 0;
unsafe {
HipDevPtr_* py = (HipDevPtr_*)y.dev;
HipDevPtr_* pb = (HipDevPtr_*)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 == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[7];
params[0] = (HipKernelArg_*)&py; params[1] = (HipKernelArg_*)&pb; params[2] = (HipKernelArg_*)&t; params[3] = (HipKernelArg_*)&c;
params[4] = (HipKernelArg_*)&p; params[5] = (HipKernelArg_*)&r; params[6] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_BIAS, total, ¶ms[0]);
}
}
static int rocm_half_pool_(struct RocmHalfBuf x, struct RocmHalfBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride, int avg) {
if (!rocm_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 ||
!rocm_half_fits_(x, rows * height * width) || !rocm_half_fits_(y, total)) return 0;
unsafe {
HipDevPtr_* px = (HipDevPtr_*)x.dev;
HipDevPtr_* py = (HipDevPtr_*)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 == ROCM_BF16 ? 1 : 0;
HipKernelArg_* params[11];
params[0] = (HipKernelArg_*)&px; params[1] = (HipKernelArg_*)&py; params[2] = (HipKernelArg_*)&r; params[3] = (HipKernelArg_*)&hh;
params[4] = (HipKernelArg_*)&ww; params[5] = (HipKernelArg_*)&ohh; params[6] = (HipKernelArg_*)&oww; params[7] = (HipKernelArg_*)&kk;
params[8] = (HipKernelArg_*)&ss; params[9] = (HipKernelArg_*)&a; params[10] = (HipKernelArg_*)&bf;
return rocm_launch_(K_HALF_POOL, total, ¶ms[0]);
}
}
int rocm_half_maxpool2d(struct RocmHalfBuf x, struct RocmHalfBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return rocm_half_pool_(x, y, batch, channels, height, width, kernel, stride, 0);
}
int rocm_half_avgpool2d(struct RocmHalfBuf x, struct RocmHalfBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return rocm_half_pool_(x, y, batch, channels, height, width, kernel, stride, 1);
}
} // namespace std