// Generated by gen_device_kernels.sh from device_kernels.cu -- do not edit.
#pragma once
namespace std {
#define DK_EW 0
#define DK_EWOP_0 1
#define DK_EWOP_1 2
#define DK_EWOP_2 3
#define DK_EWOP_3 4
#define DK_EWOP_4 5
#define DK_EWOP_5 6
#define DK_EWOP_6 7
#define DK_EWOP_7 8
#define DK_EWOP_8 9
#define DK_EWOP_9 10
#define DK_EWOP_10 11
#define DK_EWOP_11 12
#define DK_EWOP_12 13
#define DK_EWOP_13 14
#define DK_EWOP_14 15
#define DK_EWOP_20 16
#define DK_EWOP_21 17
#define DK_EWOP_22 18
#define DK_EWOP_23 19
#define DK_EWOP_24 20
#define DK_EWOP_25 21
#define DK_EWOP_26 22
#define DK_EWOP_27 23
#define DK_EWOP_28 24
#define DK_EWOP_29 25
#define DK_EWOP_30 26
#define DK_EWOP_31 27
#define DK_EWOP_32 28
#define DK_EWOP_33 29
#define DK_EWOP_40 30
#define DK_EWOP_41 31
#define DK_EWOP_42 32
#define DK_EWOP_43 33
#define DK_EWOP_44 34
#define DK_EWOP_45 35
#define DK_EWOP_46 36
#define DK_EWOP_47 37
#define DK_EWOP_48 38
#define DK_EWOP_49 39
#define DK_EWOP_50 40
#define DK_EWOP_51 41
#define DK_SUM 42
#define DK_BIAS_ROWS 43
#define DK_MUL_ROW 44
#define DK_COLSUM_ACC 45
#define DK_MUL_ROW_ACC 46
#define DK_ROWSUM 47
#define DK_ROWBCAST_ACC 48
#define DK_CONCAT 49
#define DK_SLICE 50
#define DK_GATHER_ROWS 51
#define DK_SCATTER_ROWS 52
#define DK_SOFTMAX_ROWS 53
#define DK_SOFTMAX_BWD 54
#define DK_TRANSPOSE 55
#define DK_ATTN_FWD 56
#define DK_ATTN_BWD_Q 57
#define DK_ATTN_BWD_KV 58
#define DK_CE_FWD 59
#define DK_CE_BWD 60
#define DK_EWLOSS_FWD 61
#define DK_EWLOSS_BWD 62
#define DK_LAYERNORM 63
#define DK_LAYERNORM_BWD 64
#define DK_OPT 65
#define DK_OPT_MULTI 66
#define DK_CLIP_SCALE 67
#define DK_IM2COL 68
#define DK_COL2IM 69
#define DK_SWAP01 70
#define DK_BIAS_CHANNELS 71
#define DK_BIAS_CHANNELS_BWD 72
#define DK_BN_PARTIAL 73
#define DK_BN_FINISH 74
#define DK_GN_STATS 75
#define DK_NORM_APPLY 76
#define DK_NORM_SUMS 77
#define DK_BN_BWD_DX 78
#define DK_GN_BWD 79
#define DK_POOL 80
#define DK_POOL_BWD 81
#define DK_GPOOL 82
#define DK_UPSAMPLE 83
#define DK_RESIZE 84
#define DK_RESIZE_BWD 85
#define DK_AMP_UNSCALE 86
#define DK_AMP_UNSCALE_MULTI 87
#define DK_AMP_UPDATE 88
#define DK_SPMM 89
#define DK_SPMM_TGRAD 90
#define DK_SPMM_VGRAD 91
#define DK_COUNT 92
#define DK_SOURCE_CRC 0xbc6a9ac7U
// The kernel running elementwise op code op alone (dk_ewop_<op>), or DK_EW.
static int device_ew_kernel_(int op) {
if (op == 0) return DK_EWOP_0;
if (op == 1) return DK_EWOP_1;
if (op == 2) return DK_EWOP_2;
if (op == 3) return DK_EWOP_3;
if (op == 4) return DK_EWOP_4;
if (op == 5) return DK_EWOP_5;
if (op == 6) return DK_EWOP_6;
if (op == 7) return DK_EWOP_7;
if (op == 8) return DK_EWOP_8;
if (op == 9) return DK_EWOP_9;
if (op == 10) return DK_EWOP_10;
if (op == 11) return DK_EWOP_11;
if (op == 12) return DK_EWOP_12;
if (op == 13) return DK_EWOP_13;
if (op == 14) return DK_EWOP_14;
if (op == 20) return DK_EWOP_20;
if (op == 21) return DK_EWOP_21;
if (op == 22) return DK_EWOP_22;
if (op == 23) return DK_EWOP_23;
if (op == 24) return DK_EWOP_24;
if (op == 25) return DK_EWOP_25;
if (op == 26) return DK_EWOP_26;
if (op == 27) return DK_EWOP_27;
if (op == 28) return DK_EWOP_28;
if (op == 29) return DK_EWOP_29;
if (op == 30) return DK_EWOP_30;
if (op == 31) return DK_EWOP_31;
if (op == 32) return DK_EWOP_32;
if (op == 33) return DK_EWOP_33;
if (op == 40) return DK_EWOP_40;
if (op == 41) return DK_EWOP_41;
if (op == 42) return DK_EWOP_42;
if (op == 43) return DK_EWOP_43;
if (op == 44) return DK_EWOP_44;
if (op == 45) return DK_EWOP_45;
if (op == 46) return DK_EWOP_46;
if (op == 47) return DK_EWOP_47;
if (op == 48) return DK_EWOP_48;
if (op == 49) return DK_EWOP_49;
if (op == 50) return DK_EWOP_50;
if (op == 51) return DK_EWOP_51;
return DK_EW;
}
static const char* device_kernel_name_(int k) {
if (k == 0) return "dk_ew";
if (k == 1) return "dk_ewop_0";
if (k == 2) return "dk_ewop_1";
if (k == 3) return "dk_ewop_2";
if (k == 4) return "dk_ewop_3";
if (k == 5) return "dk_ewop_4";
if (k == 6) return "dk_ewop_5";
if (k == 7) return "dk_ewop_6";
if (k == 8) return "dk_ewop_7";
if (k == 9) return "dk_ewop_8";
if (k == 10) return "dk_ewop_9";
if (k == 11) return "dk_ewop_10";
if (k == 12) return "dk_ewop_11";
if (k == 13) return "dk_ewop_12";
if (k == 14) return "dk_ewop_13";
if (k == 15) return "dk_ewop_14";
if (k == 16) return "dk_ewop_20";
if (k == 17) return "dk_ewop_21";
if (k == 18) return "dk_ewop_22";
if (k == 19) return "dk_ewop_23";
if (k == 20) return "dk_ewop_24";
if (k == 21) return "dk_ewop_25";
if (k == 22) return "dk_ewop_26";
if (k == 23) return "dk_ewop_27";
if (k == 24) return "dk_ewop_28";
if (k == 25) return "dk_ewop_29";
if (k == 26) return "dk_ewop_30";
if (k == 27) return "dk_ewop_31";
if (k == 28) return "dk_ewop_32";
if (k == 29) return "dk_ewop_33";
if (k == 30) return "dk_ewop_40";
if (k == 31) return "dk_ewop_41";
if (k == 32) return "dk_ewop_42";
if (k == 33) return "dk_ewop_43";
if (k == 34) return "dk_ewop_44";
if (k == 35) return "dk_ewop_45";
if (k == 36) return "dk_ewop_46";
if (k == 37) return "dk_ewop_47";
if (k == 38) return "dk_ewop_48";
if (k == 39) return "dk_ewop_49";
if (k == 40) return "dk_ewop_50";
if (k == 41) return "dk_ewop_51";
if (k == 42) return "dk_sum";
if (k == 43) return "dk_bias_rows";
if (k == 44) return "dk_mul_row";
if (k == 45) return "dk_colsum_acc";
if (k == 46) return "dk_mul_row_acc";
if (k == 47) return "dk_rowsum";
if (k == 48) return "dk_rowbcast_acc";
if (k == 49) return "dk_concat";
if (k == 50) return "dk_slice";
if (k == 51) return "dk_gather_rows";
if (k == 52) return "dk_scatter_rows";
if (k == 53) return "dk_softmax_rows";
if (k == 54) return "dk_softmax_bwd";
if (k == 55) return "dk_transpose";
if (k == 56) return "dk_attn_fwd";
if (k == 57) return "dk_attn_bwd_q";
if (k == 58) return "dk_attn_bwd_kv";
if (k == 59) return "dk_ce_fwd";
if (k == 60) return "dk_ce_bwd";
if (k == 61) return "dk_ewloss_fwd";
if (k == 62) return "dk_ewloss_bwd";
if (k == 63) return "dk_layernorm";
if (k == 64) return "dk_layernorm_bwd";
if (k == 65) return "dk_opt";
if (k == 66) return "dk_opt_multi";
if (k == 67) return "dk_clip_scale";
if (k == 68) return "dk_im2col";
if (k == 69) return "dk_col2im";
if (k == 70) return "dk_swap01";
if (k == 71) return "dk_bias_channels";
if (k == 72) return "dk_bias_channels_bwd";
if (k == 73) return "dk_bn_partial";
if (k == 74) return "dk_bn_finish";
if (k == 75) return "dk_gn_stats";
if (k == 76) return "dk_norm_apply";
if (k == 77) return "dk_norm_sums";
if (k == 78) return "dk_bn_bwd_dx";
if (k == 79) return "dk_gn_bwd";
if (k == 80) return "dk_pool";
if (k == 81) return "dk_pool_bwd";
if (k == 82) return "dk_gpool";
if (k == 83) return "dk_upsample";
if (k == 84) return "dk_resize";
if (k == 85) return "dk_resize_bwd";
if (k == 86) return "dk_amp_unscale";
if (k == 87) return "dk_amp_unscale_multi";
if (k == 88) return "dk_amp_update";
if (k == 89) return "dk_spmm";
if (k == 90) return "dk_spmm_tgrad";
if (k == 91) return "dk_spmm_vgrad";
return "";
}
static const char* device_kernel_source_() {
return
"// SafeC Standard Library — the device kernels behind std::ml's tensor ops on\n"
"// ROCm (compiled at run time by hiprtc), CUDA (NVRTC), Metal\n"
"// (newLibraryWithSource) and Vulkan (SPIR-V compiled ahead of time by\n"
"// slangc): one source for all four. gen_device_kernels.sh turns this file\n"
"// into device_kernels.h (the source as a string and the kernel ids) and\n"
"// device_kernels_spirv.h (the SPIR-V); edit this file, then regenerate.\n"
"//\n"
"// Every kernel takes one DevArgs (device.h): device addresses in p[],\n"
"// integers in u[], floats in f[]. It runs over n threads (n passed as u[7])\n"
"// in blocks of 256; per-row kernels use one thread per row.\n"
"\n"
"// ── Portability: CUDA / HIP and Metal ────────────────────────────────────────\n"
"// The source is CUDA C++; as Metal Shading Language (C++14 with address\n"
"// spaces) a few things differ, and the macros below absorb them, expanding\n"
"// to nothing (or to the CUDA spelling) for hiprtc and NVRTC:\n"
"// - pointers and references name the memory they point into: dk_global\n"
"// (device memory), dk_shared (a block's shared memory), dk_thread (a\n"
"// thread's own variables);\n"
"// - Metal has no 64-bit 'long long' (DK_U64) and no double (dk_double, a\n"
"// wide accumulator: float there; the reductions it serves stay\n"
"// two-pass, so fp32 holds the tolerance);\n"
"// - a kernel receives its arguments in constant memory and its thread\n"
"// indices as parameters (DK_KERNEL declares both, under the CUDA names\n"
"// threadIdx, blockIdx, blockDim and gridDim), so a helper that reads\n"
"// them takes DK_IDS after its parameters and is called with DK_IDA;\n"
"// - block reductions (block_add, block_sum_d) use simdgroup sums and a\n"
"// small threadgroup scratch the kernel declares (dk_red);\n"
"// - attention tiles are 8 x 8 instead of 16 x 16 (ATTN_T): Metal's 32 KB\n"
"// of threadgroup memory doesn't hold four 16-row tiles of 128 floats;\n"
"// a second build with DK_ATTN64 keeps 16 x 16 for head dims up to 64;\n"
"// - integer division by a run-time value costs dozens of cycles on Apple\n"
"// GPUs, so the hot index arithmetic divides with DK_UDIV / DK_UMOD: on\n"
"// Metal through a float reciprocal made exact by correction (dk_udiv,\n"
"// the reciprocal declared by DK_RECIP), elsewhere plain / and %;\n"
"// - the kernels with a thread per row (softmax, layer norm, cross-entropy,\n"
"// row sums) take a simdgroup per row on Metal (GSR, MpsBackend launches\n"
"// 32 threads per row), its lanes splitting the columns: a few hundred\n"
"// rows are too few threads to keep an Apple GPU busy.\n"
"// Vulkan builds take Metal's paths: DK_SIMD (32-wide simdgroups: block\n"
"// reductions, a simdgroup per row) and DK_TG32K (32 KB of shared memory:\n"
"// the attention tiles), and the Metal-only GEMMs and convolutions below.\n"
"#if defined(DK_VULKAN)\n"
"// Slang, for SPIR-V (gen_device_kernels.sh first rewrites the C++ spellings\n"
"// Slang lacks: unsigned int -> uint, templates -> generics, reference\n"
"// parameters -> values or inout, threadgroup pointers -> TGPtr). Device\n"
"// addresses are buffer device addresses (physical storage buffer pointers),\n"
"// a kernel's DevArgs are its push constants, and the simdgroup is a subgroup\n"
"// VulkanBackend requires to be 32 wide.\n"
"#define DK_SIMD 1\n"
"#define DK_TG32K 1\n"
"typealias ulong = uint64_t;\n"
"typealias ulong2 = vector<uint64_t, 2>;\n"
"typealias ushort = uint16_t;\n"
"#define DK_U64 uint64_t\n"
"#define dk_double float\n"
"#define dk_global\n"
"#define dk_shared\n"
"#define dk_thread\n"
"#define device\n"
"#define constant\n"
"#define __device__\n"
"#define __forceinline__ [ForceInline]\n"
"#define __shared__ static groupshared\n"
"#define __syncthreads() GroupMemoryBarrierWithGroupSync()\n"
"struct DevArgs {\n"
" uint64_t p[8];\n"
" uint u[8];\n"
" float f[8];\n"
"};\n"
"// The entry point (its DevArgs in push constants) calls the kernel's body\n"
"// with the CUDA names; DK_ENTRY gives a block size other than 256.\n"
"#define DK_KERNEL(name) DK_ENTRY(name, 256)\n"
"#define DK_ENTRY(name, T) [shader(\"compute\")] [numthreads(T, 1, 1)] \\\n"
" void name(uniform DevArgs a, uint3 dk_t : SV_GroupThreadID, uint3 dk_g : SV_GroupID) { \\\n"
" name##_body(a, dk_t, dk_g, uint3(T, 1u, 1u), WorkgroupCount(), dk_t.x % 32u, dk_t.x / 32u); } \\\n"
" void name##_body(DevArgs a, uint3 threadIdx, uint3 blockIdx, uint3 blockDim, uint3 gridDim, uint dk_lane, uint dk_simd)\n"
"#define DK_IDS , uint3 threadIdx, uint3 blockIdx, uint3 blockDim, uint3 gridDim, uint dk_lane, uint dk_simd\n"
"#define DK_IDA , threadIdx, blockIdx, blockDim, gridDim, dk_lane, dk_simd\n"
"static groupshared float dk_red[32];\n"
"float simd_sum(float v) { return WaveActiveSum(v); }\n"
"float simd_max(float v) { return WaveActiveMax(v); }\n"
"// float atomics (SPV_EXT_shader_atomic_float_add) and the C math library's float functions\n"
"void atomicAdd(float* p, float v) { InterlockedAdd(*p, v); }\n"
"float expf(float x) { return exp(x); }\n"
"float logf(float x) { return log(x); }\n"
"float tanhf(float x) { return tanh(x); }\n"
"float sqrtf(float x) { return sqrt(x); }\n"
"float rsqrtf(float x) { return rsqrt(x); }\n"
"float fabsf(float x) { return abs(x); }\n"
"float floorf(float x) { return floor(x); }\n"
"float ceilf(float x) { return ceil(x); }\n"
"float fminf(float x, float y) { return min(x, y); }\n"
"float fmaxf(float x, float y) { return max(x, y); }\n"
"float fmaf(float x, float y, float z) { return fma(x, y, z); }\n"
"uint __float_as_uint(float f) { return asuint(f); }\n"
"float __uint_as_float(uint u) { return asfloat(u); }\n"
"uint dk_udiv(uint q, uint d, float inv) {\n"
" uint r = (uint)((float)q * inv);\n"
" while (r * d > q) r--;\n"
" while ((r + 1u) * d <= q) r++;\n"
" return r;\n"
"}\n"
"#define DK_RECIP(inv, d) float inv = 1.0f / (float)(d);\n"
"#define DK_UDIV(q, d, inv) dk_udiv((q), (d), (inv))\n"
"#define DK_UMOD(q, d, inv) ((q) - dk_udiv((q), (d), (inv)) * (d))\n"
"#define GSR(r, n) for (uint r = (blockIdx.x * blockDim.x + threadIdx.x) / 32u; r < (n); r += blockDim.x * gridDim.x / 32u)\n"
"#elif defined(__METAL_VERSION__)\n"
"#define DK_SIMD 1\n"
"#define DK_TG32K 1\n"
"#include <metal_stdlib>\n"
"using namespace metal;\n"
"#define DK_U64 ulong\n"
"#define dk_double float\n"
"#define dk_global device\n"
"#define dk_shared threadgroup\n"
"#define dk_thread thread\n"
"#define __device__\n"
"#define __forceinline__ inline __attribute__((always_inline))\n"
"#define __shared__ threadgroup\n"
"#define __syncthreads() threadgroup_barrier(mem_flags::mem_threadgroup)\n"
"struct DevArgs_ {\n"
" ulong p[8];\n"
" unsigned int u[8];\n"
" float f[8];\n"
"};\n"
"// const DevArgs& names the kernel's argument struct in constant memory.\n"
"#define DevArgs constant DevArgs_\n"
"#define DK_KERNEL(name) kernel void name(constant DevArgs_& a [[buffer(0)]], \\\n"
" uint3 threadIdx [[thread_position_in_threadgroup]], uint3 blockIdx [[threadgroup_position_in_grid]], \\\n"
" uint3 blockDim [[threads_per_threadgroup]], uint3 gridDim [[threadgroups_per_grid]], \\\n"
" uint dk_lane [[thread_index_in_simdgroup]], uint dk_simd [[simdgroup_index_in_threadgroup]], \\\n"
" threadgroup float* dk_red [[threadgroup(0)]])\n"
"#define DK_IDS , uint3 threadIdx, uint3 blockIdx, uint3 blockDim, uint3 gridDim, uint dk_lane, uint dk_simd, \\\n"
" threadgroup float* dk_red\n"
"#define DK_IDA , threadIdx, blockIdx, blockDim, gridDim, dk_lane, dk_simd, dk_red\n"
"// float atomics and the C math library's float functions\n"
"inline void atomicAdd(device float* p, float v) { atomic_fetch_add_explicit((device atomic_float*)p, v, memory_order_relaxed); }\n"
"inline float expf(float x) { return metal::exp(x); }\n"
"inline float logf(float x) { return metal::log(x); }\n"
"inline float tanhf(float x) { return metal::tanh(x); }\n"
"inline float sqrtf(float x) { return metal::sqrt(x); }\n"
"inline float rsqrtf(float x) { return metal::rsqrt(x); }\n"
"inline float fabsf(float x) { return metal::fabs(x); }\n"
"inline float floorf(float x) { return metal::floor(x); }\n"
"inline float ceilf(float x) { return metal::ceil(x); }\n"
"inline float fminf(float x, float y) { return metal::fmin(x, y); }\n"
"inline float fmaxf(float x, float y) { return metal::fmax(x, y); }\n"
"inline float fmaf(float x, float y, float z) { return metal::fma(x, y, z); }\n"
"inline unsigned int __float_as_uint(float f) { return as_type<uint>(f); }\n"
"inline float __uint_as_float(unsigned int u) { return as_type<float>(u); }\n"
"inline unsigned int dk_udiv(unsigned int q, unsigned int d, float inv) {\n"
" unsigned int r = (unsigned int)((float)q * inv);\n"
" while (r * d > q) r--;\n"
" while ((r + 1u) * d <= q) r++;\n"
" return r;\n"
"}\n"
"#define DK_RECIP(inv, d) float inv = 1.0f / (float)(d);\n"
"#define DK_UDIV(q, d, inv) dk_udiv((q), (d), (inv))\n"
"#define DK_UMOD(q, d, inv) ((q) - dk_udiv((q), (d), (inv)) * (d))\n"
"// rows r for this simdgroup (a grid of 32 threads per row); dk_lane splits the columns\n"
"#define GSR(r, n) for (unsigned int r = (blockIdx.x * blockDim.x + threadIdx.x) / 32u; r < (n); r += blockDim.x * gridDim.x / 32u)\n"
"#else\n"
"#define DK_U64 unsigned long long\n"
"#define dk_double double\n"
"#define dk_global\n"
"#define dk_shared\n"
"#define dk_thread\n"
"#define DK_KERNEL(name) extern \"C\" __global__ void name(DevArgs a)\n"
"#define DK_IDS\n"
"#define DK_IDA\n"
"#define DK_RECIP(inv, d)\n"
"#define DK_UDIV(q, d, inv) q / d\n"
"#define DK_UMOD(q, d, inv) q % d\n"
"struct DevArgs {\n"
" DK_U64 p[8];\n"
" unsigned int u[8];\n"
" float f[8];\n"
"};\n"
"#endif\n"
"\n"
"#define GS(i, n) for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x)\n"
"\n"
"// ── Element types ────────────────────────────────────────────────────────────\n"
"// A buffer's element type travels in its handle's top two bits (device.h's\n"
"// DEV_TAG_*): 0 fp32, 1 fp16, 2 bf16 (activations stored in 16 bits). F(k)\n"
"// views argument k's buffer: loads widen to float, stores round to its type,\n"
"// so every kernel reads and writes any of them. Atomics target fp32 buffers\n"
"// (gradients, sums) through FA(k). The source compiles twice: with\n"
"// DK_TAGGED for launches with a 16-bit argument, and plain (views are bare\n"
"// float pointers) for the rest.\n"
"#define DEV_ADDR_MASK 0x3fffffffffffffffULL\n"
"__device__ unsigned short f2h_bits(float f);\n"
"__device__ float h2f_bits(unsigned short h);\n"
"__device__ unsigned short f2b_bits(float f);\n"
"__device__ float b2f_bits(unsigned short b);\n"
"#if defined(DK_VULKAN)\n"
"// Slang has no references: the view's subscript loads and stores (in the\n"
"// DK_TAGGED build, converting by the tag), and a view tests true when set.\n"
"struct TP {\n"
" DK_U64 p;\n"
" __init(DK_U64 h) { p = h; }\n"
"#ifdef DK_TAGGED\n"
" uint tag() { return (uint)(p >> 62); }\n"
" DK_U64 base() { return p & DEV_ADDR_MASK; }\n"
" __subscript(DK_U64 i) -> float {\n"
" get {\n"
" uint t = tag();\n"
" if (t == 0u) return ((float*)base())[i];\n"
" uint16_t h = ((uint16_t*)base())[i];\n"
" return t == 1u ? h2f_bits(h) : b2f_bits(h);\n"
" }\n"
" [nonmutating] set {\n"
" uint t = tag();\n"
" if (t == 0u) ((float*)base())[i] = newValue;\n"
" else ((uint16_t*)base())[i] = t == 1u ? f2h_bits(newValue) : f2b_bits(newValue);\n"
" }\n"
" }\n"
"#else\n"
" __subscript(DK_U64 i) -> float {\n"
" get { return ((float*)p)[i]; }\n"
" [nonmutating] set { ((float*)p)[i] = newValue; }\n"
" }\n"
"#endif\n"
"};\n"
"#ifdef DK_TAGGED\n"
"TP operator+(TP t, DK_U64 k) { return TP(t.p + k * (t.tag() == 0u ? 4ull : 2ull)); }\n"
"#else\n"
"TP operator+(TP t, DK_U64 k) { return TP(t.p + k * 4ull); }\n"
"#endif\n"
"extension bool {\n"
" __implicit_conversion(900) __init(TP t) { this = (t.p & DEV_ADDR_MASK) != 0ull; }\n"
"}\n"
"#elif defined(DK_TAGGED)\n"
"struct Ref {\n"
" DK_U64 p; // the element's address\n"
" unsigned int tag;\n"
" __device__ float get() const {\n"
" if (tag == 0u) return *(const dk_global float*)p;\n"
" unsigned short h = *(const dk_global unsigned short*)p;\n"
" return tag == 1u ? h2f_bits(h) : b2f_bits(h);\n"
" }\n"
" __device__ void set(float v) const {\n"
" if (tag == 0u) *(dk_global float*)p = v;\n"
" else *(dk_global unsigned short*)p = tag == 1u ? f2h_bits(v) : f2b_bits(v);\n"
" }\n"
" __device__ operator float() const { return get(); }\n"
" __device__ const dk_thread Ref& operator=(float v) const { set(v); return *this; }\n"
" __device__ const dk_thread Ref& operator=(const dk_thread Ref& o) const { set(o.get()); return *this; }\n"
" __device__ const dk_thread Ref& operator+=(float v) const { set(get() + v); return *this; }\n"
" __device__ const dk_thread Ref& operator-=(float v) const { set(get() - v); return *this; }\n"
"};\n"
"struct TP {\n"
" DK_U64 p; // tagged handle\n"
" __device__ TP(DK_U64 h) : p(h) {}\n"
" __device__ unsigned int tag() const { return (unsigned int)(p >> 62); }\n"
" __device__ DK_U64 base() const { return p & DEV_ADDR_MASK; }\n"
" __device__ Ref operator[](DK_U64 i) const {\n"
" Ref r; r.tag = tag(); r.p = base() + i * (r.tag == 0u ? 4ull : 2ull); return r;\n"
" }\n"
" __device__ TP operator+(DK_U64 k) const { return TP(p + k * (tag() == 0u ? 4ull : 2ull)); }\n"
" __device__ explicit operator bool() const { return base() != 0ull; }\n"
"};\n"
"#else\n"
"// The plain build (no argument is 16-bit): views are float pointers.\n"
"struct TP {\n"
" DK_U64 p;\n"
" __device__ TP(DK_U64 h) : p(h) {}\n"
" __device__ dk_global float& operator[](DK_U64 i) const { return ((dk_global float*)p)[i]; }\n"
" __device__ TP operator+(DK_U64 k) const { return TP(p + k * 4ull); }\n"
" __device__ explicit operator bool() const { return p != 0ull; }\n"
"};\n"
"#endif\n"
"#define F(k) TP(a.p[k])\n"
"#define CF(k) TP(a.p[k])\n"
"#define FA(k) ((dk_global float*)(a.p[k] & DEV_ADDR_MASK))\n"
"#define N (a.u[7])\n"
"\n"
"// Gradient kernels that add to their output take u6 = 1 to write it instead\n"
"// (it was an unwritten zero gradient): ACCUM(out, v) is out (+)= v.\n"
"#define ACCUM(o, v) do { if (a.u[6]) (o) = (v); else (o) += (v); } while (0)\n"
"\n"
"// ── Elementwise: op codes match tensor.h's EW_* ─────────────────────────────\n"
"// p0 = in a, p1 = in b, p2 = out, p3 = aux (a third operand); f0 = k, f1 = k2.\n"
"// The accumulating forms (EW_ACC .. and 40 ..) add to out, or with u1 = 1\n"
"// write it (out was all zeros: an unwritten gradient).\n"
"#define EW_ADD 0\n"
"#define EW_SUB 1\n"
"#define EW_MUL 2\n"
"#define EW_SCALE 3\n"
"#define EW_RELU 4\n"
"#define EW_FILL 5\n"
"#define EW_COPY 6\n"
"#define EW_ACC 7\n"
"#define EW_DEC 8\n"
"#define EW_ACC_MUL 9\n"
"#define EW_ACC_SCALE 10\n"
"#define EW_ACC_RELU 11\n"
"#define EW_ACC_K 12\n"
"#define EW_ACC_SIGB 13\n"
"#define EW_ACC_TANHB 14\n"
"#define EW_EXP 20\n"
"#define EW_LOG 21\n"
"#define EW_ABS 22\n"
"#define EW_LEAKY 23\n"
"#define EW_CLAMP 24\n"
"#define EW_SIGMOID 25\n"
"#define EW_TANH 26\n"
"#define EW_GELU 27\n"
"#define EW_SILU 28\n"
"#define EW_ADDK 29\n"
"#define EW_MIN 30\n"
"#define EW_ELU 31\n"
"#define EW_LERP 32 // out += k * (a - out), in place (a soft update)\n"
"#define EW_MULS 33 // out = a * b[0] (b a device scalar)\n"
"#define EW_ACC_PK 40 // out += *b (a device scalar: a seed gradient)\n"
"#define EW_ACC_EXPB 41 // out += g * y (a = g, b = y)\n"
"#define EW_ACC_LOGB 42 // out += g / x (a = g, b = x)\n"
"#define EW_ACC_ABSB 43 // out += g * sign(x)\n"
"#define EW_ACC_LEAKYB 44 // out += g * (x > 0 ? 1 : k)\n"
"#define EW_ACC_CLAMPB 45 // out += x == y ? g : 0 (aux = y)\n"
"#define EW_ACC_GELUB 46 // out += g * gelu'(x)\n"
"#define EW_ACC_SILUB 47 // out += g * silu'(x)\n"
"#define EW_ACC_MINA 48 // out += x <= other ? g : 0 (b = x, aux = other)\n"
"#define EW_ACC_MINB 49 // out += x < other ? 0 : g ... (b = x of 'a' side, aux = this side)\n"
"#define EW_ACC_ELUB 50 // out += g * (x > 0 ? 1 : y + k) (b = x, aux = y)\n"
"#define EW_ACC_PSCALE 51 // out += *b * k * ... (b device scalar seed): out += k * b[0]\n"
"\n"
"__device__ float gelu_f(float x) {\n"
" float c = 0.7978845608028654f;\n"
" return 0.5f * x * (1.0f + tanhf(c * (x + 0.044715f * x * x * x)));\n"
"}\n"
"__device__ float gelu_d(float x) {\n"
" float c = 0.7978845608028654f;\n"
" float t = tanhf(c * (x + 0.044715f * x * x * x));\n"
" float du = c * (1.0f + 0.134145f * x * x);\n"
" return 0.5f * (1.0f + t) + 0.5f * x * (1.0f - t * t) * du;\n"
"}\n"
"__device__ float sig_f(float x) { return 1.0f / (1.0f + expf(-x)); }\n"
"\n"
"// One elementwise result from the operands' values (y0: b[0] for the ops\n"
"// reading a device scalar); ow: the accumulating forms write.\n"
"__device__ __forceinline__ float ew_val(int op, float x, float y, float z, float o, float k, float k2, bool ow) {\n"
" float ob = ow ? 0.0f : o;\n"
" switch (op) {\n"
" case EW_ADD: return x + y;\n"
" case EW_SUB: return x - y;\n"
" case EW_MUL: return x * y;\n"
" case EW_SCALE: return x * k;\n"
" case EW_RELU: return x > 0.0f ? x : 0.0f;\n"
" case EW_FILL: return k;\n"
" case EW_COPY: return x;\n"
" case EW_ACC: return ob + x;\n"
" case EW_DEC: return ob - x;\n"
" case EW_ACC_MUL: return fmaf(x, y, ob);\n"
" case EW_ACC_SCALE: return fmaf(x, k, ob);\n"
" case EW_ACC_RELU: return ob + (y > 0.0f ? x : 0.0f);\n"
" case EW_ACC_K: return ob + k;\n"
" case EW_ACC_SIGB: return fmaf(x * y, 1.0f - y, ob);\n"
" case EW_ACC_TANHB: return fmaf(x, 1.0f - y * y, ob);\n"
" case EW_EXP: return expf(x);\n"
" case EW_LOG: return logf(x);\n"
" case EW_ABS: return fabsf(x);\n"
" case EW_LEAKY: return x > 0.0f ? x : k * x;\n"
" case EW_CLAMP: return fminf(fmaxf(x, k), k2);\n"
" case EW_SIGMOID: return sig_f(x);\n"
" case EW_TANH: return tanhf(x);\n"
" case EW_GELU: return gelu_f(x);\n"
" case EW_SILU: return x * sig_f(x);\n"
" case EW_ADDK: return x + k;\n"
" case EW_MIN: return x <= y ? x : y;\n"
" case EW_ELU: return x > 0.0f ? x : k * (expf(x) - 1.0f);\n"
" case EW_LERP: return fmaf(k, x - o, o);\n"
" case EW_MULS: return x * y;\n"
" case EW_ACC_PK: return ob + y;\n"
" case EW_ACC_EXPB: return fmaf(x, y, ob);\n"
" case EW_ACC_LOGB: return ob + x / y;\n"
" case EW_ACC_ABSB: return ob + (y > 0.0f ? x : (y < 0.0f ? -x : 0.0f));\n"
" case EW_ACC_LEAKYB: return ob + x * (y > 0.0f ? 1.0f : k);\n"
" case EW_ACC_CLAMPB: return ob + (y == z ? x : 0.0f);\n"
" case EW_ACC_GELUB: return fmaf(x, gelu_d(y), ob);\n"
" case EW_ACC_SILUB: { float s = sig_f(y); return fmaf(x, s + y * s * (1.0f - s), ob); }\n"
" case EW_ACC_MINA: return ob + (y <= z ? x : 0.0f);\n"
" case EW_ACC_MINB: return ob + (!(z <= y) ? x : 0.0f);\n"
" case EW_ACC_ELUB: return ob + x * (y > 0.0f ? 1.0f : z + k);\n"
" case EW_ACC_PSCALE: return ob + k * y;\n"
" }\n"
" return 0.0f;\n"
"}\n"
"// Which operands an op reads: x, y per element, y as a device scalar, z, o.\n"
"__device__ __forceinline__ bool ew_uses_x(int op) { return op != EW_FILL && op != EW_ACC_K && op != EW_ACC_PK && op != EW_ACC_PSCALE; }\n"
"__device__ __forceinline__ bool ew_y_scalar(int op) { return op == EW_MULS || op == EW_ACC_PK || op == EW_ACC_PSCALE; }\n"
"__device__ __forceinline__ bool ew_uses_y(int op) {\n"
" return op == EW_ADD || op == EW_SUB || op == EW_MUL || op == EW_ACC_MUL || op == EW_ACC_RELU || op == EW_ACC_SIGB ||\n"
" op == EW_ACC_TANHB || op == EW_MIN || (op >= EW_ACC_EXPB && op <= EW_ACC_ELUB);\n"
"}\n"
"__device__ __forceinline__ bool ew_uses_z(int op) { return op == EW_ACC_CLAMPB || op == EW_ACC_MINA || op == EW_ACC_MINB || op == EW_ACC_ELUB; }\n"
"__device__ __forceinline__ bool ew_reads_o(int op, bool ow) {\n"
" return op == EW_LERP || (!ow && ((op >= EW_ACC && op <= EW_ACC_TANHB) || (op >= EW_ACC_PK && op <= EW_ACC_PSCALE)));\n"
"}\n"
"\n"
"#ifdef DK_TAGGED\n"
"// Two 16-bit values of type t from one 32-bit word, and back.\n"
"__device__ __forceinline__ float h16_lo(unsigned int w, unsigned int t) { unsigned short h = (unsigned short)(w & 0xffffu); return t == 1u ? h2f_bits(h) : b2f_bits(h); }\n"
"__device__ __forceinline__ float h16_hi(unsigned int w, unsigned int t) { unsigned short h = (unsigned short)(w >> 16); return t == 1u ? h2f_bits(h) : b2f_bits(h); }\n"
"__device__ __forceinline__ unsigned int h16_pack(float lo, float hi, unsigned int t) {\n"
" unsigned int a = t == 1u ? f2h_bits(lo) : f2b_bits(lo), b = t == 1u ? f2h_bits(hi) : f2b_bits(hi);\n"
" return a | (b << 16);\n"
"}\n"
"// Every element operand of type t (the output's) and 4-byte aligned: two\n"
"// values per thread through 32-bit words.\n"
"__device__ bool ew_pairable(const DevArgs& a, int op, bool ow, unsigned int t) {\n"
" if (t == 0u) return false;\n"
" DK_U64 ps[4];\n"
" bool use[4];\n"
" ps[0] = a.p[0]; use[0] = ew_uses_x(op);\n"
" ps[1] = a.p[1]; use[1] = ew_uses_y(op);\n"
" ps[2] = a.p[2]; use[2] = true;\n"
" ps[3] = a.p[3]; use[3] = ew_uses_z(op);\n"
" for (int q = 0; q < 4; q++)\n"
" if (use[q] && ((unsigned int)(ps[q] >> 62) != t || (ps[q] & 3ull) != 0ull)) return false;\n"
" return true;\n"
"}\n"
"#endif\n"
"\n"
"// One instantiation per op: a tight loop, no per-element dispatch.\n"
"template <int OPC>\n"
"__device__ __forceinline__ void ew_run(const DevArgs& a DK_IDS) {\n"
" const int OP = OPC >= 0 ? OPC : (int)a.u[0];\n"
" TP x = CF(0); TP y = CF(1); TP o = F(2); TP z = CF(3);\n"
" float k = a.f[0], k2 = a.f[1];\n"
" bool ow = a.u[1] != 0;\n"
" const bool ux = ew_uses_x(OP), uy = ew_uses_y(OP), ys = ew_y_scalar(OP), uz = ew_uses_z(OP);\n"
" bool uo = ew_reads_o(OP, ow);\n"
" float y0 = ys ? (float)y[0] : 0.0f;\n"
"#ifdef DK_TAGGED\n"
" unsigned int t = (unsigned int)(a.p[2] >> 62);\n"
" if (ew_pairable(a, OP, ow, t)) {\n"
" const dk_global unsigned int* xw = (const dk_global unsigned int*)(a.p[0] & DEV_ADDR_MASK);\n"
" const dk_global unsigned int* yw = (const dk_global unsigned int*)(a.p[1] & DEV_ADDR_MASK);\n"
" const dk_global unsigned int* zw = (const dk_global unsigned int*)(a.p[3] & DEV_ADDR_MASK);\n"
" dk_global unsigned int* ow_ = (dk_global unsigned int*)(a.p[2] & DEV_ADDR_MASK);\n"
" unsigned int n2 = N / 2u;\n"
" GS(j, n2) {\n"
" unsigned int xv = ux ? xw[j] : 0u, yv = uy ? yw[j] : 0u, zv = uz ? zw[j] : 0u, ov = uo ? ow_[j] : 0u;\n"
" float lo, hi;\n"
" if (t == 1u) {\n"
" lo = ew_val(OP, h16_lo(xv, 1u), uy ? h16_lo(yv, 1u) : y0, h16_lo(zv, 1u), h16_lo(ov, 1u), k, k2, ow);\n"
" hi = ew_val(OP, h16_hi(xv, 1u), uy ? h16_hi(yv, 1u) : y0, h16_hi(zv, 1u), h16_hi(ov, 1u), k, k2, ow);\n"
" ow_[j] = h16_pack(lo, hi, 1u);\n"
" } else {\n"
" lo = ew_val(OP, h16_lo(xv, 2u), uy ? h16_lo(yv, 2u) : y0, h16_lo(zv, 2u), h16_lo(ov, 2u), k, k2, ow);\n"
" hi = ew_val(OP, h16_hi(xv, 2u), uy ? h16_hi(yv, 2u) : y0, h16_hi(zv, 2u), h16_hi(ov, 2u), k, k2, ow);\n"
" ow_[j] = h16_pack(lo, hi, 2u);\n"
" }\n"
" }\n"
" if ((N & 1u) && blockIdx.x == 0 && threadIdx.x == 0) {\n"
" unsigned int i = N - 1u;\n"
" o[i] = ew_val(OP, ux ? (float)x[i] : 0.0f, uy ? (float)y[i] : y0, uz ? (float)z[i] : 0.0f,\n"
" uo ? (float)o[i] : 0.0f, k, k2, ow);\n"
" }\n"
" return;\n"
" }\n"
"#endif\n"
" GS(i, N) {\n"
" o[i] = ew_val(OP, ux ? (float)x[i] : 0.0f, uy ? (float)y[i] : y0, uz ? (float)z[i] : 0.0f,\n"
" uo ? (float)o[i] : 0.0f, k, k2, ow);\n"
" }\n"
"}\n"
"\n"
"// Any op (u0); dk_ewop_<code> below run one op each (device.sc picks them).\n"
"DK_KERNEL(dk_ew) { ew_run<-1>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_0) { ew_run<0>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_1) { ew_run<1>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_2) { ew_run<2>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_3) { ew_run<3>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_4) { ew_run<4>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_5) { ew_run<5>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_6) { ew_run<6>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_7) { ew_run<7>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_8) { ew_run<8>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_9) { ew_run<9>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_10) { ew_run<10>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_11) { ew_run<11>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_12) { ew_run<12>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_13) { ew_run<13>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_14) { ew_run<14>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_20) { ew_run<20>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_21) { ew_run<21>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_22) { ew_run<22>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_23) { ew_run<23>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_24) { ew_run<24>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_25) { ew_run<25>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_26) { ew_run<26>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_27) { ew_run<27>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_28) { ew_run<28>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_29) { ew_run<29>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_30) { ew_run<30>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_31) { ew_run<31>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_32) { ew_run<32>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_33) { ew_run<33>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_40) { ew_run<40>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_41) { ew_run<41>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_42) { ew_run<42>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_43) { ew_run<43>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_44) { ew_run<44>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_45) { ew_run<45>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_46) { ew_run<46>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_47) { ew_run<47>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_48) { ew_run<48>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_49) { ew_run<49>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_50) { ew_run<50>(a DK_IDA); }\n"
"DK_KERNEL(dk_ewop_51) { ew_run<51>(a DK_IDA); }\n"
"\n"
"// ── Reductions (256 threads per block, one atomic per block) ─────────────────\n"
"__device__ void block_add(float v, dk_global float* out DK_IDS) {\n"
"#if defined(DK_SIMD)\n"
" // each simdgroup's sum, then the first simdgroup adds those\n"
" v = simd_sum(v);\n"
" if (dk_lane == 0u) dk_red[dk_simd] = v;\n"
" __syncthreads();\n"
" if (dk_simd == 0u) {\n"
" float s = simd_sum(threadIdx.x < (blockDim.x + 31u) / 32u ? dk_red[threadIdx.x] : 0.0f);\n"
" if (threadIdx.x == 0u) atomicAdd(out, s);\n"
" }\n"
"#else\n"
" __shared__ float s[256];\n"
" s[threadIdx.x] = v;\n"
" __syncthreads();\n"
" for (unsigned int w = 128; w > 0; w >>= 1) {\n"
" if (threadIdx.x < w) s[threadIdx.x] += s[threadIdx.x + w];\n"
" __syncthreads();\n"
" }\n"
" if (threadIdx.x == 0) atomicAdd(out, s[0]);\n"
"#endif\n"
"}\n"
"\n"
"// p0 = in, p1 = out[0] (zeroed by the caller); u0: 0 sum, 1 sum of squares\n"
"DK_KERNEL(dk_sum) {\n"
" TP x = CF(0);\n"
" float acc = 0.0f;\n"
" GS(i, N) { float v = x[i]; acc += a.u[0] ? v * v : v; }\n"
" block_add(acc, FA(1) DK_IDA);\n"
"}\n"
"\n"
"// ── Rows and columns: u0 = cols (rows = n / cols for row-shaped n) ──────────\n"
"// out[r,c] = x[r,c] + b[c] n = rows*cols\n"
"DK_KERNEL(dk_bias_rows) {\n"
" unsigned int C = a.u[0];\n"
" GS(i, N) F(2)[i] = CF(0)[i] + CF(1)[i % C];\n"
"}\n"
"// out[r,c] = x[r,c] * g[c] n = rows*cols\n"
"DK_KERNEL(dk_mul_row) {\n"
" unsigned int C = a.u[0];\n"
" GS(i, N) F(2)[i] = CF(0)[i] * CF(1)[i % C];\n"
"}\n"
"// gb[c] += sum_r d[r,c] (* x[r,c] when p2 != 0); u0 = cols, u1 = rows, u2 = S\n"
"// row slices (0 or 1: one). Units are (256 columns, row slice); a thread per\n"
"// column. n = ceil(cols/256)*S*256\n"
"DK_KERNEL(dk_colsum_acc) {\n"
" unsigned int C = a.u[0], R = a.u[1], S = a.u[2] > 1 ? a.u[2] : 1;\n"
" TP d = CF(0); dk_global float* gb = FA(1); TP x = CF(2);\n"
" for (unsigned int u = blockIdx.x; u < ((C + 255) / 256) * S; u += gridDim.x) {\n"
" unsigned int c = (u / S) * 256 + threadIdx.x, sl = u % S;\n"
" if (c < C) {\n"
" float s = 0.0f;\n"
" if (x) { for (unsigned int r = sl; r < R; r += S) s += d[r * C + c] * x[r * C + c]; }\n"
" else { for (unsigned int r = sl; r < R; r += S) s += d[r * C + c]; }\n"
" if (S > 1) atomicAdd(&gb[c], s); else gb[c] += s;\n"
" }\n"
" }\n"
"}\n"
"// dx[r,c] += d[r,c] * g[c] n = rows*cols\n"
"DK_KERNEL(dk_mul_row_acc) {\n"
" unsigned int C = a.u[0];\n"
" GS(i, N) ACCUM(F(2)[i], CF(0)[i] * CF(1)[i % C]);\n"
"}\n"
"// out[r] = sum_c x[r,c] n = rows\n"
"DK_KERNEL(dk_rowsum) {\n"
" unsigned int C = a.u[0];\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" float s = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) s += CF(0)[r * C + c];\n"
" s = simd_sum(s);\n"
" if (dk_lane == 0u) F(1)[r] = s;\n"
" }\n"
"#else\n"
" GS(r, N) { float s = 0.0f; for (unsigned int c = 0; c < C; c++) s += CF(0)[r * C + c]; F(1)[r] = s; }\n"
"#endif\n"
"}\n"
"// dx[r,c] += g[r] n = rows*cols\n"
"DK_KERNEL(dk_rowbcast_acc) {\n"
" unsigned int C = a.u[0];\n"
" GS(i, N) ACCUM(F(1)[i], CF(0)[i / C]);\n"
"}\n"
"// out [rows, p+q] = x [rows, p] ++ y [rows, q]; u0 = p, u1 = q n = rows*(p+q)\n"
"DK_KERNEL(dk_concat) {\n"
" unsigned int P = a.u[0], Q = a.u[1], W = P + Q;\n"
" GS(i, N) { unsigned int r = i / W, c = i % W; F(2)[i] = c < P ? CF(0)[r * P + c] : CF(1)[r * Q + c - P]; }\n"
"}\n"
"// Columns [start, start+len) of a wide [rows, cols] tensor and a narrow\n"
"// [rows, len] one; u0 = cols, u1 = start, u2 = len, u3 = mode:\n"
"// 0 narrow (p1) = wide (p0) columns (slice)\n"
"// 1 wide (p1) += narrow (p0) (slice backward)\n"
"// 2 narrow (p1) += wide (p0) columns (concat backward) n = rows*len\n"
"DK_KERNEL(dk_slice) {\n"
" unsigned int C = a.u[0], S = a.u[1], L = a.u[2];\n"
" GS(i, N) {\n"
" unsigned int r = i / L, c = i % L;\n"
" if (a.u[3] == 1) F(1)[r * C + S + c] += CF(0)[i];\n"
" else if (a.u[3] == 2) ACCUM(F(1)[i], CF(0)[r * C + S + c]);\n"
" else F(1)[i] = CF(0)[r * C + S + c];\n"
" }\n"
"}\n"
"// out[i, :] = table[idx[i], :]; u0 = d n = m*d\n"
"DK_KERNEL(dk_gather_rows) {\n"
" unsigned int D = a.u[0];\n"
" GS(i, N) { unsigned int r = i / D, c = i % D; F(2)[i] = CF(0)[(unsigned int)CF(1)[r] * D + c]; }\n"
"}\n"
"// table.grad[idx[i], :] += g[i, :] n = m*d\n"
"DK_KERNEL(dk_scatter_rows) {\n"
" unsigned int D = a.u[0];\n"
" GS(i, N) { unsigned int r = i / D, c = i % D; atomicAdd(&FA(2)[(unsigned int)CF(1)[r] * D + c], CF(0)[i]); }\n"
"}\n"
"\n"
"// ── Attention masks (attention.h) ────────────────────────────────────────────\n"
"// Query row i sits at position pos = i + (keys - queries); flags bit 0\n"
"// (causal) hides keys j > pos, bit 1 (a sliding window of w) keys j <= pos - w.\n"
"__device__ __forceinline__ bool attn_vis(unsigned int flags, unsigned int w, int pos, int j) {\n"
" if ((flags & 1u) && j > pos) return false;\n"
" if ((flags & 2u) && j <= pos - (int)w) return false;\n"
" return true;\n"
"}\n"
"\n"
"// ── Softmax family: one thread per row, u0 = cols ────────────────────────────\n"
"// u1: 0 softmax, 1 log-softmax. A masked softmax (u1 = 0): u2 = mask flags,\n"
"// u3 = window, u4 = (int) keys - queries; hidden entries and rows with no\n"
"// visible key come out 0.\n"
"DK_KERNEL(dk_softmax_rows) {\n"
" unsigned int C = a.u[0];\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" TP x = CF(0) + r * C; TP o = F(1) + r * C;\n"
" unsigned int fl = a.u[2];\n"
" int pos = (int)r + (int)a.u[4];\n"
" float mx = -3.0e38f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) if (!fl || attn_vis(fl, a.u[3], pos, (int)c)) mx = fmaxf(mx, x[c]);\n"
" mx = simd_max(mx);\n"
" float s = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) if (!fl || attn_vis(fl, a.u[3], pos, (int)c)) s += expf(x[c] - mx);\n"
" s = simd_sum(s);\n"
" if (fl) {\n"
" float inv = s > 0.0f ? 1.0f / s : 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) o[c] = attn_vis(fl, a.u[3], pos, (int)c) ? expf(x[c] - mx) * inv : 0.0f;\n"
" } else if (a.u[1]) {\n"
" float l = mx + logf(s);\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) o[c] = x[c] - l;\n"
" } else {\n"
" float inv = 1.0f / s;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) o[c] = expf(x[c] - mx) * inv;\n"
" }\n"
" }\n"
"#else\n"
" GS(r, N) {\n"
" TP x = CF(0) + r * C; TP o = F(1) + r * C;\n"
" if (a.u[2]) {\n"
" int pos = (int)r + (int)a.u[4];\n"
" float mx = -3.0e38f;\n"
" for (unsigned int c = 0; c < C; c++) if (attn_vis(a.u[2], a.u[3], pos, (int)c)) mx = fmaxf(mx, x[c]);\n"
" float s = 0.0f;\n"
" for (unsigned int c = 0; c < C; c++) if (attn_vis(a.u[2], a.u[3], pos, (int)c)) s += expf(x[c] - mx);\n"
" float inv = s > 0.0f ? 1.0f / s : 0.0f;\n"
" for (unsigned int c = 0; c < C; c++) o[c] = attn_vis(a.u[2], a.u[3], pos, (int)c) ? expf(x[c] - mx) * inv : 0.0f;\n"
" continue;\n"
" }\n"
" float mx = x[0]; for (unsigned int c = 1; c < C; c++) mx = fmaxf(mx, x[c]);\n"
" float s = 0.0f; for (unsigned int c = 0; c < C; c++) s += expf(x[c] - mx);\n"
" if (a.u[1]) { float l = mx + logf(s); for (unsigned int c = 0; c < C; c++) o[c] = x[c] - l; }\n"
" else { float inv = 1.0f / s; for (unsigned int c = 0; c < C; c++) o[c] = expf(x[c] - mx) * inv; }\n"
" }\n"
"#endif\n"
"}\n"
"// dx += y * (g - sum(g*y)) (softmax), dx += g - exp(y) * sum(g) (log-softmax)\n"
"// p0 = g, p1 = y (the forward output), p2 = dx\n"
"DK_KERNEL(dk_softmax_bwd) {\n"
" unsigned int C = a.u[0];\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" TP g = CF(0) + r * C; TP y = CF(1) + r * C; TP dx = F(2) + r * C;\n"
" float s = 0.0f;\n"
" if (a.u[1]) {\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) s += g[c];\n"
" s = simd_sum(s);\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) dx[c] += g[c] - expf(y[c]) * s;\n"
" } else {\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) s += g[c] * y[c];\n"
" s = simd_sum(s);\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) dx[c] += y[c] * (g[c] - s);\n"
" }\n"
" }\n"
"#else\n"
" GS(r, N) {\n"
" TP g = CF(0) + r * C; TP y = CF(1) + r * C; TP dx = F(2) + r * C;\n"
" if (a.u[1]) {\n"
" float s = 0.0f; for (unsigned int c = 0; c < C; c++) s += g[c];\n"
" for (unsigned int c = 0; c < C; c++) dx[c] += g[c] - expf(y[c]) * s;\n"
" } else {\n"
" float s = 0.0f; for (unsigned int c = 0; c < C; c++) s += g[c] * y[c];\n"
" for (unsigned int c = 0; c < C; c++) dx[c] += y[c] * (g[c] - s);\n"
" }\n"
" }\n"
"#endif\n"
"}\n"
"\n"
"// ── Transpose ────────────────────────────────────────────────────────────────\n"
"// out [C, R] = in [R, C]^T, or with u2 = 1 out += it (u6 = 1: write); u0 = R,\n"
"// u1 = C n = R*C\n"
"DK_KERNEL(dk_transpose) {\n"
" unsigned int R = a.u[0], C = a.u[1];\n"
" GS(i, N) {\n"
" unsigned int r = i / C, c = i % C;\n"
" float v = CF(0)[i];\n"
" if (a.u[2]) ACCUM(F(1)[c * R + r], v); else F(1)[c * R + r] = v;\n"
" }\n"
"}\n"
"\n"
"// ── Fused multi-head attention (attention.h) ─────────────────────────────────\n"
"// q [S, H*D], k [L, H*D], v [L, H*E], out [S, H*E] (head h is columns\n"
"// [h*D, (h+1)*D) of each row): out = softmax(scale q k^T, masked) v per head\n"
"// without forming the score matrix -- FlashAttention-2's tiling (Dao 2023):\n"
"// a block of 256 threads takes one head's tile of 16 rows and walks the\n"
"// other side in tiles of 16 through shared memory, 16 x 16 scores at a time,\n"
"// with an online softmax; the row's log-sum-exp is kept for the backward\n"
"// pass, which recomputes the scores from it. Thread t works on row t / 16\n"
"// of its tile and, across the head's columns, on every 16th one from t % 16\n"
"// (D, E <= ATTN_MAXD). Tiles a causal or windowed mask hides entirely are\n"
"// skipped. Arguments shared by the three kernels: u0 = S, u1 = L, u2 = H,\n"
"// u3 = D, u4 = E, u5 = flags (bits 0-1 the mask, as attn_vis; bits 2-4: dq,\n"
"// dk, dv are unwritten zero gradients: write them), u6 = window, f0 = scale;\n"
"// p4 = lse [H*S] (fp32), p5 = dout, p7 = D [H*S] (fp32, rowsum(dout * out),\n"
"// written by dk_attn_bwd_q for dk_attn_bwd_kv). Launched over\n"
"// 256 * H * ceil(rows / 16) threads (rows: S, or L for dk_attn_bwd_kv).\n"
"#if defined(DK_TG32K) && defined(DK_ATTN64)\n"
"#define ATTN_MAXD 64 // a second Metal build for head dims up to 64, where 16 x 16 tiles fit\n"
"#define ATTN_T 16\n"
"#elif defined(DK_TG32K)\n"
"#define ATTN_MAXD 128\n"
"#define ATTN_T 8 // 64 threads per block (see Portability above)\n"
"#else\n"
"#define ATTN_MAXD 128\n"
"#define ATTN_T 16\n"
"#endif\n"
"#define ATTN_LD (ATTN_MAXD + 1)\n"
"#define ATTN_PER (ATTN_MAXD / ATTN_T)\n"
"#define ATTN_HID (-3.0e38f)\n"
"// The key range a tile of rows [i0, i0 + 16) can see: [lo, hi).\n"
"__device__ __forceinline__ void attn_range(const DevArgs& a, unsigned int i0, dk_thread int& lo, dk_thread int& hi) {\n"
" int off = (int)a.u[1] - (int)a.u[0];\n"
" lo = 0; hi = (int)a.u[1];\n"
" if (a.u[5] & 1u) { int m = (int)i0 + ATTN_T - 1 + off + 1; if (m < hi) hi = m; }\n"
" if (a.u[5] & 2u) { int m = (int)i0 + off - (int)a.u[6] + 1; if (m > lo) lo = m; }\n"
" if (lo < 0) lo = 0;\n"
"}\n"
"// rows [r0, r0 + 16) of a [rows, H*W] tensor's head h -> sh [16][ATTN_LD] (0 past the end)\n"
"__device__ __forceinline__ void attn_load(TP x, dk_shared float* sh, unsigned int r0, unsigned int rows, unsigned int H,\n"
" unsigned int h, unsigned int W DK_IDS) {\n"
" for (unsigned int k = threadIdx.x; k < ATTN_T * W; k += blockDim.x) {\n"
" unsigned int r = k / W, c = k % W;\n"
" sh[r * ATTN_LD + c] = r0 + r < rows ? (float)x[((r0 + r) * H + h) * W + c] : 0.0f;\n"
" }\n"
"}\n"
"// p0 q, p1 k, p2 v, p3 out, p4 lse (written)\n"
"DK_KERNEL(dk_attn_fwd) {\n"
" __shared__ float Qs[ATTN_T * ATTN_LD], Ks[ATTN_T * ATTN_LD], Vs[ATTN_T * ATTN_LD], Ps[ATTN_T * (ATTN_T + 1)];\n"
" unsigned int S = a.u[0], L = a.u[1], H = a.u[2], D = a.u[3], E = a.u[4], fl = a.u[5] & 3u;\n"
" int off = (int)L - (int)S;\n"
" unsigned int tiles = (S + ATTN_T - 1) / ATTN_T;\n"
" unsigned int r = threadIdx.x / ATTN_T, ln = threadIdx.x % ATTN_T;\n"
" for (unsigned int u = blockIdx.x; u < H * tiles; u += gridDim.x) {\n"
" unsigned int h = u / tiles, i0 = (u % tiles) * ATTN_T, i = i0 + r;\n"
" attn_load(CF(0), Qs, i0, S, H, h, D DK_IDA);\n"
" float acc[ATTN_PER];\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) acc[t] = 0.0f;\n"
" float m = ATTN_HID, l = 0.0f;\n"
" int lo, hi;\n"
" attn_range(a, i0, lo, hi);\n"
" for (int j0 = (lo / ATTN_T) * ATTN_T; j0 < hi; j0 += ATTN_T) {\n"
" __syncthreads();\n"
" attn_load(CF(1), Ks, (unsigned int)j0, L, H, h, D DK_IDA);\n"
" attn_load(CF(2), Vs, (unsigned int)j0, L, H, h, E DK_IDA);\n"
" __syncthreads();\n"
" unsigned int j = (unsigned int)j0 + ln;\n"
" float sc = ATTN_HID;\n"
" if (i < S && j < L && attn_vis(fl, a.u[6], (int)i + off, (int)j)) {\n"
" float s = 0.0f;\n"
" for (unsigned int d = 0; d < D; d++) s += Qs[r * ATTN_LD + d] * Ks[ln * ATTN_LD + d];\n"
" sc = s * a.f[0];\n"
" }\n"
" Ps[r * (ATTN_T + 1) + ln] = sc;\n"
" __syncthreads();\n"
" float mt = ATTN_HID;\n"
" for (unsigned int c = 0; c < ATTN_T; c++) mt = fmaxf(mt, Ps[r * (ATTN_T + 1) + c]);\n"
" if (mt > ATTN_HID) {\n"
" float mn = fmaxf(m, mt), cr = expf(m - mn);\n"
" float pv[ATTN_T], ps = 0.0f;\n"
" for (unsigned int c = 0; c < ATTN_T; c++) {\n"
" float sv = Ps[r * (ATTN_T + 1) + c];\n"
" pv[c] = sv > ATTN_HID ? expf(sv - mn) : 0.0f;\n"
" ps += pv[c];\n"
" }\n"
" l = l * cr + ps;\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) {\n"
" unsigned int e = ln + t * ATTN_T;\n"
" if (e < E) {\n"
" float x = acc[t] * cr;\n"
" for (unsigned int c = 0; c < ATTN_T; c++) x += pv[c] * Vs[c * ATTN_LD + e];\n"
" acc[t] = x;\n"
" }\n"
" }\n"
" m = mn;\n"
" }\n"
" }\n"
" if (i < S) {\n"
" float inv = l > 0.0f ? 1.0f / l : 0.0f;\n"
" TP o = F(3) + (i * H + h) * E;\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) { unsigned int e = ln + t * ATTN_T; if (e < E) o[e] = acc[t] * inv; }\n"
" if (ln == 0) FA(4)[h * S + i] = l > 0.0f ? m + logf(l) : 3.0e38f; // no visible key: every p is 0\n"
" }\n"
" __syncthreads();\n"
" }\n"
"}\n"
"// p0 q, p1 k, p2 v, p3 out, p6 dq (+)= scale * sum_j ds_ij k_j (0: none),\n"
"// ds_ij = p_ij (dout_i . v_j - D_i)\n"
"DK_KERNEL(dk_attn_bwd_q) {\n"
" __shared__ float Qs[ATTN_T * ATTN_LD], Gs[ATTN_T * ATTN_LD], Ks[ATTN_T * ATTN_LD], Vs[ATTN_T * ATTN_LD];\n"
" __shared__ float Ds[ATTN_T * (ATTN_T + 1)], Rd[ATTN_T], Rl[ATTN_T];\n"
" unsigned int S = a.u[0], L = a.u[1], H = a.u[2], D = a.u[3], E = a.u[4], fl = a.u[5] & 3u;\n"
" int off = (int)L - (int)S;\n"
" unsigned int tiles = (S + ATTN_T - 1) / ATTN_T;\n"
" unsigned int r = threadIdx.x / ATTN_T, ln = threadIdx.x % ATTN_T;\n"
" for (unsigned int u = blockIdx.x; u < H * tiles; u += gridDim.x) {\n"
" unsigned int h = u / tiles, i0 = (u % tiles) * ATTN_T, i = i0 + r;\n"
" attn_load(CF(5), Gs, i0, S, H, h, E DK_IDA);\n"
" attn_load(CF(3), Ks, i0, S, H, h, E DK_IDA); // the output rows, for D (Ks is free until the key tiles)\n"
" __syncthreads();\n"
" float part = 0.0f;\n"
" for (unsigned int e = ln; e < E; e += ATTN_T) part += Gs[r * ATTN_LD + e] * Ks[r * ATTN_LD + e];\n"
" Ds[r * (ATTN_T + 1) + ln] = part;\n"
" __syncthreads();\n"
" if (ln == 0) {\n"
" float Di = 0.0f;\n"
" for (unsigned int c = 0; c < ATTN_T; c++) Di += Ds[r * (ATTN_T + 1) + c];\n"
" Rd[r] = Di;\n"
" Rl[r] = i < S ? FA(4)[h * S + i] : 0.0f;\n"
" if (i < S) FA(7)[h * S + i] = Di;\n"
" }\n"
" if (!a.p[6]) { __syncthreads(); continue; }\n"
" attn_load(CF(0), Qs, i0, S, H, h, D DK_IDA);\n"
" float dq[ATTN_PER];\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) dq[t] = 0.0f;\n"
" int lo, hi;\n"
" attn_range(a, i0, lo, hi);\n"
" for (int j0 = (lo / ATTN_T) * ATTN_T; j0 < hi; j0 += ATTN_T) {\n"
" __syncthreads();\n"
" attn_load(CF(1), Ks, (unsigned int)j0, L, H, h, D DK_IDA);\n"
" attn_load(CF(2), Vs, (unsigned int)j0, L, H, h, E DK_IDA);\n"
" __syncthreads();\n"
" unsigned int j = (unsigned int)j0 + ln;\n"
" float ds = 0.0f;\n"
" if (i < S && j < L && attn_vis(fl, a.u[6], (int)i + off, (int)j)) {\n"
" float s = 0.0f, dp = 0.0f;\n"
" for (unsigned int d = 0; d < D; d++) s += Qs[r * ATTN_LD + d] * Ks[ln * ATTN_LD + d];\n"
" for (unsigned int e = 0; e < E; e++) dp += Gs[r * ATTN_LD + e] * Vs[ln * ATTN_LD + e];\n"
" ds = expf(s * a.f[0] - Rl[r]) * (dp - Rd[r]);\n"
" }\n"
" Ds[r * (ATTN_T + 1) + ln] = ds;\n"
" __syncthreads();\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) {\n"
" unsigned int d = ln + t * ATTN_T;\n"
" if (d < D) {\n"
" float x = dq[t];\n"
" for (unsigned int c = 0; c < ATTN_T; c++) x += Ds[r * (ATTN_T + 1) + c] * Ks[c * ATTN_LD + d];\n"
" dq[t] = x;\n"
" }\n"
" }\n"
" }\n"
" if (i < S) {\n"
" TP o = F(6) + (i * H + h) * D;\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) {\n"
" unsigned int d = ln + t * ATTN_T;\n"
" if (d < D) { if (a.u[5] & 4u) o[d] = a.f[0] * dq[t]; else o[d] += a.f[0] * dq[t]; }\n"
" }\n"
" }\n"
" __syncthreads();\n"
" }\n"
"}\n"
"// p0 q, p1 k, p2 v, p3 dk (+)= scale * sum_i ds_ij q_i, p6 dv (+)= sum_i p_ij\n"
"// dout_i (either 0: none); after dk_attn_bwd_q. A block takes 16 keys.\n"
"DK_KERNEL(dk_attn_bwd_kv) {\n"
" __shared__ float Ks[ATTN_T * ATTN_LD], Vs[ATTN_T * ATTN_LD], Qs[ATTN_T * ATTN_LD], Gs[ATTN_T * ATTN_LD];\n"
" __shared__ float Ps[ATTN_T * (ATTN_T + 1)], Ds[ATTN_T * (ATTN_T + 1)], Rd[ATTN_T], Rl[ATTN_T];\n"
" unsigned int S = a.u[0], L = a.u[1], H = a.u[2], D = a.u[3], E = a.u[4], fl = a.u[5] & 3u;\n"
" int off = (int)L - (int)S;\n"
" unsigned int tiles = (L + ATTN_T - 1) / ATTN_T;\n"
" unsigned int r = threadIdx.x / ATTN_T, ln = threadIdx.x % ATTN_T;\n"
" for (unsigned int u = blockIdx.x; u < H * tiles; u += gridDim.x) {\n"
" unsigned int h = u / tiles, j0 = (u % tiles) * ATTN_T, j = j0 + r;\n"
" attn_load(CF(1), Ks, j0, L, H, h, D DK_IDA);\n"
" attn_load(CF(2), Vs, j0, L, H, h, E DK_IDA);\n"
" float dk[ATTN_PER], dv[ATTN_PER];\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) { dk[t] = 0.0f; dv[t] = 0.0f; }\n"
" // the queries that can see keys [j0, j0 + 16): [qlo, qhi)\n"
" int qlo = 0, qhi = (int)S;\n"
" if (fl & 1u) { int m2 = (int)j0 - off; if (m2 > qlo) qlo = m2; }\n"
" if (fl & 2u) { int m2 = (int)j0 + ATTN_T - 1 - off + (int)a.u[6]; if (m2 < qhi) qhi = m2; }\n"
" if (qlo < 0) qlo = 0;\n"
" for (int i0 = (qlo / ATTN_T) * ATTN_T; i0 < qhi; i0 += ATTN_T) {\n"
" __syncthreads();\n"
" attn_load(CF(0), Qs, (unsigned int)i0, S, H, h, D DK_IDA);\n"
" attn_load(CF(5), Gs, (unsigned int)i0, S, H, h, E DK_IDA);\n"
" if (threadIdx.x < ATTN_T) {\n"
" unsigned int i = (unsigned int)i0 + threadIdx.x;\n"
" Rl[threadIdx.x] = i < S ? FA(4)[h * S + i] : 0.0f;\n"
" Rd[threadIdx.x] = i < S ? FA(7)[h * S + i] : 0.0f;\n"
" }\n"
" __syncthreads();\n"
" // thread (key r, query ln)\n"
" unsigned int i = (unsigned int)i0 + ln;\n"
" float p = 0.0f, ds = 0.0f;\n"
" if (i < S && j < L && attn_vis(fl, a.u[6], (int)i + off, (int)j)) {\n"
" float s = 0.0f, dp = 0.0f;\n"
" for (unsigned int d = 0; d < D; d++) s += Qs[ln * ATTN_LD + d] * Ks[r * ATTN_LD + d];\n"
" for (unsigned int e = 0; e < E; e++) dp += Gs[ln * ATTN_LD + e] * Vs[r * ATTN_LD + e];\n"
" p = expf(s * a.f[0] - Rl[ln]);\n"
" ds = p * (dp - Rd[ln]);\n"
" }\n"
" Ps[r * (ATTN_T + 1) + ln] = p;\n"
" Ds[r * (ATTN_T + 1) + ln] = ds;\n"
" __syncthreads();\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) {\n"
" unsigned int d = ln + t * ATTN_T;\n"
" if (d < E) {\n"
" float x = dv[t];\n"
" for (unsigned int c = 0; c < ATTN_T; c++) x += Ps[r * (ATTN_T + 1) + c] * Gs[c * ATTN_LD + d];\n"
" dv[t] = x;\n"
" }\n"
" if (d < D) {\n"
" float x = dk[t];\n"
" for (unsigned int c = 0; c < ATTN_T; c++) x += Ds[r * (ATTN_T + 1) + c] * Qs[c * ATTN_LD + d];\n"
" dk[t] = x;\n"
" }\n"
" }\n"
" }\n"
" if (j < L) {\n"
" if (a.p[3]) {\n"
" TP o = F(3) + (j * H + h) * D;\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) {\n"
" unsigned int d = ln + t * ATTN_T;\n"
" if (d < D) { if (a.u[5] & 8u) o[d] = a.f[0] * dk[t]; else o[d] += a.f[0] * dk[t]; }\n"
" }\n"
" }\n"
" if (a.p[6]) {\n"
" TP o = F(6) + (j * H + h) * E;\n"
" for (unsigned int t = 0; t < ATTN_PER; t++) {\n"
" unsigned int e = ln + t * ATTN_T;\n"
" if (e < E) { if (a.u[5] & 16u) o[e] = dv[t]; else o[e] += dv[t]; }\n"
" }\n"
" }\n"
" }\n"
" __syncthreads();\n"
" }\n"
"}\n"
"\n"
"// ── Fused losses -> loss[0] (zeroed by the caller), averaged ─────────────────\n"
"// Cross-entropy on class indices: p0 = logits [rows, C], p1 = labels [rows],\n"
"// p2 = loss[0]; u0 = C; f0 = 1/rows n = rows\n"
"DK_KERNEL(dk_ce_fwd) {\n"
" unsigned int C = a.u[0];\n"
" float acc = 0.0f;\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" TP x = CF(0) + r * C;\n"
" float mx = -3.0e38f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) mx = fmaxf(mx, x[c]);\n"
" mx = simd_max(mx);\n"
" float s = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) s += expf(x[c] - mx);\n"
" s = simd_sum(s);\n"
" unsigned int lab = (unsigned int)CF(1)[r];\n"
" if (dk_lane == 0u) acc += (mx + logf(s) - x[lab]) * a.f[0];\n"
" }\n"
"#else\n"
" GS(r, N) {\n"
" TP x = CF(0) + r * C;\n"
" float mx = x[0]; for (unsigned int c = 1; c < C; c++) mx = fmaxf(mx, x[c]);\n"
" float s = 0.0f; for (unsigned int c = 0; c < C; c++) s += expf(x[c] - mx);\n"
" unsigned int lab = (unsigned int)CF(1)[r];\n"
" acc += (mx + logf(s) - x[lab]) * a.f[0];\n"
" }\n"
"#endif\n"
" block_add(acc, FA(2) DK_IDA);\n"
"}\n"
"// dlogits += seed[0] * (softmax - onehot) / rows; p3 = seed (the loss's grad)\n"
"DK_KERNEL(dk_ce_bwd) {\n"
" unsigned int C = a.u[0];\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" TP x = CF(0) + r * C; TP d = F(2) + r * C;\n"
" float sc = CF(3)[0] * a.f[0];\n"
" float mx = -3.0e38f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) mx = fmaxf(mx, x[c]);\n"
" mx = simd_max(mx);\n"
" float s = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) s += expf(x[c] - mx);\n"
" s = simd_sum(s);\n"
" unsigned int lab = (unsigned int)CF(1)[r];\n"
" float inv = 1.0f / s;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) ACCUM(d[c], sc * (expf(x[c] - mx) * inv - (c == lab ? 1.0f : 0.0f)));\n"
" }\n"
"#else\n"
" GS(r, N) {\n"
" TP x = CF(0) + r * C; TP d = F(2) + r * C;\n"
" float sc = CF(3)[0] * a.f[0];\n"
" float mx = x[0]; for (unsigned int c = 1; c < C; c++) mx = fmaxf(mx, x[c]);\n"
" float s = 0.0f; for (unsigned int c = 0; c < C; c++) s += expf(x[c] - mx);\n"
" unsigned int lab = (unsigned int)CF(1)[r];\n"
" float inv = 1.0f / s;\n"
" for (unsigned int c = 0; c < C; c++) ACCUM(d[c], sc * (expf(x[c] - mx) * inv - (c == lab ? 1.0f : 0.0f)));\n"
" }\n"
"#endif\n"
"}\n"
"// Elementwise losses: u0 = 0 BCE on logits, 1 Huber (f1 = delta); p0 = x,\n"
"// p1 = target, p2 = loss[0], f0 = 1/n n = elements\n"
"DK_KERNEL(dk_ewloss_fwd) {\n"
" float acc = 0.0f;\n"
" GS(i, N) {\n"
" float x = CF(0)[i], t = CF(1)[i], v;\n"
" if (a.u[0] == 0) v = fmaxf(x, 0.0f) - x * t + logf(1.0f + expf(-fabsf(x)));\n"
" else { float d = x - t, ad = fabsf(d), k = a.f[1]; v = ad <= k ? 0.5f * d * d : k * (ad - 0.5f * k); }\n"
" acc += v * a.f[0];\n"
" }\n"
" block_add(acc, FA(2) DK_IDA);\n"
"}\n"
"// dx += seed[0] / n * (sigmoid(x) - t) | clamp(x - t, -delta, delta); p2 = dx, p3 = seed\n"
"DK_KERNEL(dk_ewloss_bwd) {\n"
" float sc = CF(3)[0] * a.f[0];\n"
" GS(i, N) {\n"
" float x = CF(0)[i], t = CF(1)[i], g;\n"
" if (a.u[0] == 0) g = sig_f(x) - t;\n"
" else { float d = x - t, k = a.f[1]; g = fminf(fmaxf(d, -k), k); }\n"
" ACCUM(F(2)[i], sc * g);\n"
" }\n"
"}\n"
"\n"
"// ── Layer norm: one thread per row, u0 = cols, f0 = eps ──────────────────────\n"
"// p0 = x, p1 = xhat out, p2 = inv std [rows]\n"
"DK_KERNEL(dk_layernorm) {\n"
" unsigned int C = a.u[0];\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" TP x = CF(0) + r * C; TP o = F(1) + r * C;\n"
" float m = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) m += x[c];\n"
" m = simd_sum(m) / (float)C;\n"
" float v = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) { float d = x[c] - m; v += d * d; }\n"
" float is = rsqrtf(simd_sum(v) / (float)C + a.f[0]);\n"
" if (dk_lane == 0u) F(2)[r] = is;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) o[c] = (x[c] - m) * is;\n"
" }\n"
"#else\n"
" GS(r, N) {\n"
" TP x = CF(0) + r * C; TP o = F(1) + r * C;\n"
" float m = 0.0f; for (unsigned int c = 0; c < C; c++) m += x[c]; m /= (float)C;\n"
" float v = 0.0f; for (unsigned int c = 0; c < C; c++) { float d = x[c] - m; v += d * d; }\n"
" float is = rsqrtf(v / (float)C + a.f[0]);\n"
" F(2)[r] = is;\n"
" for (unsigned int c = 0; c < C; c++) o[c] = (x[c] - m) * is;\n"
" }\n"
"#endif\n"
"}\n"
"// dx += inv * (g - mean(g) - xhat * mean(g * xhat)); p0 = g, p1 = xhat, p2 = inv, p3 = dx\n"
"DK_KERNEL(dk_layernorm_bwd) {\n"
" unsigned int C = a.u[0];\n"
"#if defined(DK_SIMD)\n"
" GSR(r, N) {\n"
" TP g = CF(0) + r * C; TP y = CF(1) + r * C; TP dx = F(3) + r * C;\n"
" float mg = 0.0f, mgy = 0.0f;\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) { mg += g[c]; mgy += g[c] * y[c]; }\n"
" mg = simd_sum(mg) / (float)C; mgy = simd_sum(mgy) / (float)C;\n"
" float is = CF(2)[r];\n"
" for (unsigned int c = dk_lane; c < C; c += 32u) dx[c] += is * (g[c] - mg - y[c] * mgy);\n"
" }\n"
"#else\n"
" GS(r, N) {\n"
" TP g = CF(0) + r * C; TP y = CF(1) + r * C; TP dx = F(3) + r * C;\n"
" float mg = 0.0f, mgy = 0.0f;\n"
" for (unsigned int c = 0; c < C; c++) { mg += g[c]; mgy += g[c] * y[c]; }\n"
" mg /= (float)C; mgy /= (float)C;\n"
" float is = CF(2)[r];\n"
" for (unsigned int c = 0; c < C; c++) dx[c] += is * (g[c] - mg - y[c] * mgy);\n"
" }\n"
"#endif\n"
"}\n"
"\n"
"// ── Optimizer step over one parameter (nn.h's OPT_* arithmetic) ─────────────\n"
"// p0 = w, p1 = g, p2 = m, p3 = v, p4 = scale[0] (gradient clip factor, or 0 = none)\n"
"// u0 = kind; f0 lr, f1 wd, f2 momentum, f3 beta1, f4 beta2, f5 eps, f6 1/bc1, f7 1/sqrt(bc2)\n"
"#define OPT_SGD 0\n"
"#define OPT_MOMENTUM 1\n"
"#define OPT_NESTEROV 2\n"
"#define OPT_RMSPROP 3\n"
"#define OPT_ADAM 4\n"
"#define OPT_ADAMW 5\n"
"#define OPT_ADAGRAD 6\n"
"#define OPT_LION 7\n"
"__device__ void opt_one(const DevArgs& a, dk_global float* w, const dk_global float* gr, dk_global float* m, dk_global float* v,\n"
" unsigned int i, float cs) {\n"
" int kind = (int)a.u[0];\n"
" float lr = a.f[0], wd = a.f[1], mu = a.f[2], b1 = a.f[3], b2 = a.f[4], eps = a.f[5];\n"
" {\n"
" float g0 = gr[i] * cs;\n"
" if (kind == OPT_SGD) { w[i] -= lr * (g0 + wd * w[i]); }\n"
" else if (kind == OPT_MOMENTUM || kind == OPT_NESTEROV) {\n"
" float g = g0 + wd * w[i]; float mi = mu * m[i] + g; m[i] = mi;\n"
" w[i] -= lr * (kind == OPT_NESTEROV ? g + mu * mi : mi);\n"
" } else if (kind == OPT_ADAGRAD) {\n"
" float g = g0 + wd * w[i]; float vi = v[i] + g * g; v[i] = vi;\n"
" w[i] -= lr * g / (sqrtf(vi) + eps);\n"
" } else if (kind == OPT_LION) {\n"
" float c = b1 * m[i] + (1.0f - b1) * g0;\n"
" float sg = c > 0.0f ? 1.0f : (c < 0.0f ? -1.0f : 0.0f);\n"
" w[i] -= lr * (sg + wd * w[i]);\n"
" m[i] = b2 * m[i] + (1.0f - b2) * g0;\n"
" } else if (kind == OPT_RMSPROP) {\n"
" float g = g0 + wd * w[i]; float vi = b2 * v[i] + (1.0f - b2) * g * g; v[i] = vi;\n"
" w[i] -= lr * g / (sqrtf(vi) + eps);\n"
" } else {\n"
" float l2 = kind == OPT_ADAMW ? 0.0f : wd, dec = kind == OPT_ADAMW ? wd : 0.0f;\n"
" float g = g0 + l2 * w[i];\n"
" float mi = b1 * m[i] + (1.0f - b1) * g, vi = b2 * v[i] + (1.0f - b2) * g * g;\n"
" m[i] = mi; v[i] = vi;\n"
" w[i] -= lr * ((mi * a.f[6]) / (sqrtf(vi) * a.f[7] + eps) + dec * w[i]);\n"
" }\n"
" }\n"
"}\n"
"// p5 (when set): the loss scale's state; an overflowed step changes nothing.\n"
"DK_KERNEL(dk_opt) {\n"
" TP clip = CF(4);\n"
" if (a.p[5] && CF(5)[2] != 0.0f) return;\n"
" float cs = clip ? clip[0] : 1.0f;\n"
" GS(i, N) opt_one(a, FA(0), FA(1), FA(2), FA(3), i, cs);\n"
"}\n"
"// Every parameter's step in one launch: p0 = a table of u1 entries\n"
"// {w, g, m, v, n, start} (u64 each; start = the entry's first thread),\n"
"// otherwise dk_opt's arguments n = total values\n"
"DK_KERNEL(dk_opt_multi) {\n"
" const dk_global DK_U64* T = (const dk_global DK_U64*)a.p[0];\n"
" unsigned int cnt = a.u[1];\n"
" if (a.p[5] && CF(5)[2] != 0.0f) return;\n"
" GS(i, N) {\n"
" unsigned int lo = 0, hi = cnt - 1;\n"
" while (lo < hi) { unsigned int mid = (lo + hi + 1) / 2; if (T[mid * 6 + 5] <= i) lo = mid; else hi = mid - 1; }\n"
" const dk_global DK_U64* e = T + lo * 6;\n"
" opt_one(a, (dk_global float*)e[0], (const dk_global float*)e[1], (dk_global float*)e[2], (dk_global float*)e[3],\n"
" i - (unsigned int)e[5], 1.0f);\n"
" }\n"
"}\n"
"// scale[0] = sumsq[0] > maxNorm^2 ? maxNorm / sqrt(sumsq[0]) : 1; p0 = sumsq, p1 = scale, f0 = maxNorm\n"
"DK_KERNEL(dk_clip_scale) {\n"
" if (blockIdx.x == 0 && threadIdx.x == 0) {\n"
" float n = sqrtf(CF(0)[0]);\n"
" F(1)[0] = (n > a.f[0] && n > 0.0f) ? a.f[0] / n : 1.0f;\n"
" }\n"
"}\n"
"\n"
"// ── Convolution helpers (convnet.h's layout: [batch, C*H*W] rows, NCHW) ─────\n"
"// u0 C, u1 H, u2 W, u3 kH, u4 kW, u5 OH, u6 OW; f0 strideH, f1 strideW,\n"
"// f2 padH, f3 padW, f4 dilH, f5 dilW (integers carried as floats)\n"
"// Batched: col [C*kH*kW, B*OH*OW] for B images (f6 = B); p0 = x [B, C*H*W], p1 = col\n"
"// n = C*kH*kW*B*OH*OW\n"
"#define CONV_DECODE \\\n"
" unsigned int C = a.u[0], H = a.u[1], W = a.u[2], KH = a.u[3], KW = a.u[4], OH = a.u[5], OW = a.u[6]; \\\n"
" unsigned int B = (unsigned int)a.f[6]; \\\n"
" int sh = (int)a.f[0], sw = (int)a.f[1], ph = (int)a.f[2], pw = (int)a.f[3], dh = (int)a.f[4], dw = (int)a.f[5];\n"
"DK_KERNEL(dk_im2col) {\n"
" CONV_DECODE\n"
" unsigned int OP = OH * OW, BP = B * OP;\n"
" DK_RECIP(iBP, BP) DK_RECIP(iOP, OP) DK_RECIP(iOW, OW) DK_RECIP(iKW, KW) DK_RECIP(iKH, KH) DK_RECIP(iKK, KW * KH)\n"
" GS(i, N) {\n"
" unsigned int q = DK_UMOD(i, BP, iBP), r = DK_UDIV(i, BP, iBP); // column, row of col\n"
" unsigned int b = DK_UDIV(q, OP, iOP), p = DK_UMOD(q, OP, iOP), oh = DK_UDIV(p, OW, iOW), ow = DK_UMOD(p, OW, iOW);\n"
" unsigned int kw = DK_UMOD(r, KW, iKW), kh = DK_UMOD((DK_UDIV(r, KW, iKW)), KH, iKH), c = DK_UDIV(r, (KW * KH), iKK);\n"
" int y = (int)oh * sh - ph + (int)kh * dh, x = (int)ow * sw - pw + (int)kw * dw;\n"
" F(1)[i] = (y >= 0 && y < (int)H && x >= 0 && x < (int)W)\n"
" ? CF(0)[((b * C + c) * H + (unsigned int)y) * W + (unsigned int)x] : 0.0f;\n"
" }\n"
"}\n"
"// x [B, C*H*W] += col2im(col): the adjoint of dk_im2col (atomics: windows overlap)\n"
"DK_KERNEL(dk_col2im) {\n"
" CONV_DECODE\n"
" unsigned int OP = OH * OW, BP = B * OP;\n"
" DK_RECIP(iBP, BP) DK_RECIP(iOP, OP) DK_RECIP(iOW, OW) DK_RECIP(iKW, KW) DK_RECIP(iKH, KH) DK_RECIP(iKK, KW * KH)\n"
" GS(i, N) {\n"
" unsigned int q = DK_UMOD(i, BP, iBP), r = DK_UDIV(i, BP, iBP);\n"
" unsigned int b = DK_UDIV(q, OP, iOP), p = DK_UMOD(q, OP, iOP), oh = DK_UDIV(p, OW, iOW), ow = DK_UMOD(p, OW, iOW);\n"
" unsigned int kw = DK_UMOD(r, KW, iKW), kh = DK_UMOD((DK_UDIV(r, KW, iKW)), KH, iKH), c = DK_UDIV(r, (KW * KH), iKK);\n"
" int y = (int)oh * sh - ph + (int)kh * dh, x = (int)ow * sw - pw + (int)kw * dw;\n"
" if (y >= 0 && y < (int)H && x >= 0 && x < (int)W)\n"
" atomicAdd(&FA(1)[((b * C + c) * H + (unsigned int)y) * W + (unsigned int)x], CF(0)[i]);\n"
" }\n"
"}\n"
"// out [B2, A, P] (+)= in [A, B2, P] with the first two axes swapped; u0 A, u1 B2, u2 P,\n"
"// u3 = 1 to accumulate n = A*B2*P\n"
"DK_KERNEL(dk_swap01) {\n"
" unsigned int A = a.u[0], B2 = a.u[1], P = a.u[2];\n"
" DK_RECIP(iP, P) DK_RECIP(iB2, B2)\n"
" GS(i, N) {\n"
" unsigned int p = DK_UMOD(i, P, iP), t = DK_UDIV(i, P, iP), bi = DK_UMOD(t, B2, iB2), ai = DK_UDIV(t, B2, iB2); // i indexes 'in'\n"
" unsigned int o = (bi * A + ai) * P + p;\n"
" if (a.u[3]) F(1)[o] += CF(0)[i]; else F(1)[o] = CF(0)[i];\n"
" }\n"
"}\n"
"// q / d through a float reciprocal (inv = 1/d), made exact by correction:\n"
"// integer division costs dozens of instructions per element.\n"
"__device__ unsigned int udiv(unsigned int q, unsigned int d, float inv) {\n"
" unsigned int r = (unsigned int)((float)q * inv);\n"
" while (r * d > q) r--;\n"
" while ((r + 1u) * d <= q) r++;\n"
" return r;\n"
"}\n"
"// The channel of value i in [batch, C, P] (iP = 1/P, iC = 1/C).\n"
"__device__ unsigned int chan_of(unsigned int i, unsigned int C, unsigned int P, float iP, float iC) {\n"
" unsigned int t = udiv(i, P, iP);\n"
" return t - udiv(t, C, iC) * C;\n"
"}\n"
"#ifdef DK_TAGGED\n"
"// Pairs of 16-bit values (type T) through 32-bit words: word j holds\n"
"// elements 2j and 2j + 1.\n"
"template <int T> __device__ __forceinline__ float pk_lo(unsigned int w) { return T == 1 ? h2f_bits((unsigned short)(w & 0xffffu)) : b2f_bits((unsigned short)(w & 0xffffu)); }\n"
"template <int T> __device__ __forceinline__ float pk_hi(unsigned int w) { return T == 1 ? h2f_bits((unsigned short)(w >> 16)) : b2f_bits((unsigned short)(w >> 16)); }\n"
"template <int T> __device__ __forceinline__ unsigned int pk(float lo, float hi) {\n"
" return T == 1 ? ((unsigned int)f2h_bits(lo) | ((unsigned int)f2h_bits(hi) << 16))\n"
" : ((unsigned int)f2b_bits(lo) | ((unsigned int)f2b_bits(hi) << 16));\n"
"}\n"
"__device__ __forceinline__ const dk_global unsigned int* pw(DK_U64 p) { return (const dk_global unsigned int*)(p & DEV_ADDR_MASK); }\n"
"__device__ __forceinline__ dk_global unsigned int* pwm(DK_U64 p) { return (dk_global unsigned int*)(p & DEV_ADDR_MASK); }\n"
"// The 16-bit type every listed handle shares (4-byte aligned, nonzero), or 0.\n"
"__device__ unsigned int pair_tag(DK_U64 p0, DK_U64 p1, DK_U64 p2, DK_U64 p3) {\n"
" unsigned int t = (unsigned int)(p0 >> 62);\n"
" DK_U64 ps[4] = { p0, p1, p2, p3 };\n"
" if (t == 0u) return 0u;\n"
" for (int q = 0; q < 4; q++) {\n"
" if (ps[q] == 0ull) continue;\n"
" if ((unsigned int)(ps[q] >> 62) != t || (ps[q] & 3ull) != 0ull) return 0u;\n"
" }\n"
" return t;\n"
"}\n"
"#endif\n"
"\n"
"// out[b, c, p] += bias[c]; u0 = channels, u1 = plane n = batch*C*plane\n"
"#ifdef DK_TAGGED\n"
"template <int T> __device__ void bias_ch_pairs(const DevArgs& a DK_IDS) {\n"
" unsigned int C = a.u[0], P = a.u[1];\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" dk_global unsigned int* o = pwm(a.p[0]);\n"
" GS(j, N / 2u) {\n"
" float b = FA(1)[chan_of(2u * j, C, P, iP, iC)];\n"
" unsigned int w = o[j];\n"
" o[j] = pk<T>(pk_lo<T>(w) + b, pk_hi<T>(w) + b);\n"
" }\n"
"}\n"
"#endif\n"
"DK_KERNEL(dk_bias_channels) {\n"
" unsigned int C = a.u[0], P = a.u[1];\n"
"#ifdef DK_TAGGED\n"
" if ((P & 1u) == 0u && (N & 1u) == 0u) {\n"
" unsigned int t = pair_tag(a.p[0], 0ull, 0ull, 0ull);\n"
" if (t == 1u) { bias_ch_pairs<1>(a DK_IDA); return; }\n"
" if (t == 2u) { bias_ch_pairs<2>(a DK_IDA); return; }\n"
" }\n"
"#endif\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" GS(i, N) F(0)[i] += CF(1)[chan_of(i, C, P, iP, iC)];\n"
"}\n"
"// Block-per-unit kernels run n = units * 256 threads: block b takes units\n"
"// b, b + gridDim.x, ...; its threads split the unit's values.\n"
"#define UNITS(u, U) for (unsigned int u = blockIdx.x; u < (U); u += gridDim.x)\n"
"#define TSPLIT(q, M) for (unsigned int q = threadIdx.x; q < (M); q += blockDim.x)\n"
"// The block's total of v, returned to every thread.\n"
"__device__ dk_double block_sum_d(dk_double v DK_IDS) {\n"
"#if defined(DK_SIMD)\n"
" __syncthreads(); // the previous call's readers are done with dk_red\n"
" v = simd_sum(v);\n"
" if (dk_lane == 0u) dk_red[dk_simd] = v;\n"
" __syncthreads();\n"
" float s = 0.0f;\n"
" for (unsigned int k = 0; k < (blockDim.x + 31u) / 32u; k++) s += dk_red[k];\n"
" return s;\n"
"#else\n"
" __shared__ dk_double s[256];\n"
" __syncthreads();\n"
" s[threadIdx.x] = v;\n"
" __syncthreads();\n"
" for (unsigned int w = 128; w > 0; w >>= 1) {\n"
" if (threadIdx.x < w) s[threadIdx.x] += s[threadIdx.x + w];\n"
" __syncthreads();\n"
" }\n"
" return s[0];\n"
"#endif\n"
"}\n"
"// Per-channel reductions over [batch, C, plane] split each channel's M =\n"
"// batch*plane values into S slices (u3), a block each: slice s takes the\n"
"// 256-value runs r with r % S == s. Units are (channel, slice) = (u / S, u % S).\n"
"#define SLICE_LOOP(q, M, s, S) for (unsigned int q = (s) * blockDim.x + threadIdx.x; q < (M); q += (S) * blockDim.x)\n"
"// The index of the q-th value of channel c (needs iP = 1/P in scope).\n"
"__device__ unsigned int chan_at(unsigned int q, unsigned int c, unsigned int C, unsigned int P, float iP) {\n"
" unsigned int b = udiv(q, P, iP);\n"
" return (b * C + c) * P + (q - b * P);\n"
"}\n"
"#define CHAN_AT(q, c, C, P) chan_at((q), (c), (C), (P), iP)\n"
"// gb[c] += sum over batch and plane of d; u0 = C, u1 = plane, u2 = batch, u3 = S n = C*S*256\n"
"DK_KERNEL(dk_bias_channels_bwd) {\n"
" unsigned int C = a.u[0], P = a.u[1], M = a.u[2] * a.u[1], S = a.u[3];\n"
" float iP = 1.0f / (float)P;\n"
"#ifdef DK_TAGGED\n"
" unsigned int t = (P & 1u) == 0u ? pair_tag(a.p[0], 0ull, 0ull, 0ull) : 0u;\n"
" if (t != 0u) {\n"
" // pairs: the channel's M / 2 words\n"
" const dk_global unsigned int* d = pw(a.p[0]);\n"
" UNITS(u, C * S) {\n"
" unsigned int c = u / S, sl = u % S;\n"
" dk_double s = 0.0;\n"
" SLICE_LOOP(q, M / 2u, sl, S) {\n"
" unsigned int w = d[CHAN_AT(2u * q, c, C, P) / 2u];\n"
" s += t == 1u ? (dk_double)pk_lo<1>(w) + (dk_double)pk_hi<1>(w) : (dk_double)pk_lo<2>(w) + (dk_double)pk_hi<2>(w);\n"
" }\n"
" s = block_sum_d(s DK_IDA);\n"
" if (threadIdx.x == 0) atomicAdd(&FA(1)[c], (float)s);\n"
" }\n"
" return;\n"
" }\n"
"#endif\n"
" UNITS(u, C * S) {\n"
" unsigned int c = u / S, sl = u % S;\n"
" dk_double s = 0.0;\n"
" SLICE_LOOP(q, M, sl, S) s += (dk_double)CF(0)[CHAN_AT(q, c, C, P)];\n"
" s = block_sum_d(s DK_IDA);\n"
" if (threadIdx.x == 0) atomicAdd(&FA(1)[c], (float)s);\n"
" }\n"
"}\n"
"\n"
"// ── Batch / group normalization (convnet.h's layout [batch, C*plane]) ───────\n"
"// Batch norm statistics, pass 1: per (channel, slice) the count, mean and sum\n"
"// of squared deviations, as doubles in p1 [C*S*3]. p0 x; u0 C, u1 plane, u2 batch, u3 S\n"
"DK_KERNEL(dk_bn_partial) {\n"
" unsigned int C = a.u[0], P = a.u[1], M = a.u[2] * a.u[1], S = a.u[3];\n"
" float iP = 1.0f / (float)P;\n"
" dk_global dk_double* part = (dk_global dk_double*)a.p[1];\n"
"#ifdef DK_TAGGED\n"
" unsigned int t = (P & 1u) == 0u ? pair_tag(a.p[0], 0ull, 0ull, 0ull) : 0u;\n"
" if (t != 0u) {\n"
" const dk_global unsigned int* xw = pw(a.p[0]);\n"
" UNITS(u, C * S) {\n"
" unsigned int c = u / S, sl = u % S;\n"
" dk_double s = 0.0, cnt = 0.0;\n"
" SLICE_LOOP(q, M / 2u, sl, S) {\n"
" unsigned int w = xw[CHAN_AT(2u * q, c, C, P) / 2u];\n"
" float lo = t == 1u ? pk_lo<1>(w) : pk_lo<2>(w), hi = t == 1u ? pk_hi<1>(w) : pk_hi<2>(w);\n"
" s += (dk_double)lo + (dk_double)hi; cnt += 2.0;\n"
" }\n"
" cnt = block_sum_d(cnt DK_IDA);\n"
" dk_double tot = block_sum_d(s DK_IDA);\n"
" dk_double mu = cnt > 0.0 ? tot / cnt : 0.0;\n"
" dk_double sq = 0.0;\n"
" SLICE_LOOP(q, M / 2u, sl, S) {\n"
" unsigned int w = xw[CHAN_AT(2u * q, c, C, P) / 2u];\n"
" dk_double lo = (dk_double)(t == 1u ? pk_lo<1>(w) : pk_lo<2>(w)) - mu, hi = (dk_double)(t == 1u ? pk_hi<1>(w) : pk_hi<2>(w)) - mu;\n"
" sq += lo * lo + hi * hi;\n"
" }\n"
" sq = block_sum_d(sq DK_IDA);\n"
" if (threadIdx.x == 0) { part[u * 3] = cnt; part[u * 3 + 1] = mu; part[u * 3 + 2] = sq; }\n"
" }\n"
" return;\n"
" }\n"
"#endif\n"
" UNITS(u, C * S) {\n"
" unsigned int c = u / S, sl = u % S;\n"
" dk_double s = 0.0, cnt = 0.0;\n"
" SLICE_LOOP(q, M, sl, S) { s += (dk_double)CF(0)[CHAN_AT(q, c, C, P)]; cnt += 1.0; }\n"
" cnt = block_sum_d(cnt DK_IDA);\n"
" dk_double tot = block_sum_d(s DK_IDA);\n"
" dk_double mu = cnt > 0.0 ? tot / cnt : 0.0;\n"
" dk_double sq = 0.0;\n"
" SLICE_LOOP(q, M, sl, S) { dk_double d = (dk_double)CF(0)[CHAN_AT(q, c, C, P)] - mu; sq += d * d; }\n"
" sq = block_sum_d(sq DK_IDA);\n"
" if (threadIdx.x == 0) { part[u * 3] = cnt; part[u * 3 + 1] = mu; part[u * 3 + 2] = sq; }\n"
" }\n"
"}\n"
"// Pass 2, one thread per channel: the slices combined (Chan et al.), the\n"
"// running statistics updated (training), mean and invstd. p0 partials, p1\n"
"// mean[C], p2 invstd[C], p3 running mean, p4 running var; u0 C, u3 S,\n"
"// u4 training; f0 momentum, f1 eps n = C\n"
"DK_KERNEL(dk_bn_finish) {\n"
" unsigned int S = a.u[3];\n"
" const dk_global dk_double* part = (const dk_global dk_double*)a.p[0];\n"
" GS(c, N) {\n"
" float mean, var;\n"
" if (a.u[4]) {\n"
" dk_double n = 0.0, mu = 0.0, m2 = 0.0;\n"
" for (unsigned int s = 0; s < S; s++) {\n"
" const dk_global dk_double* p = part + (c * S + s) * 3;\n"
" if (p[0] == 0.0) continue;\n"
" dk_double nn = n + p[0], d = p[1] - mu;\n"
" mu += d * p[0] / nn;\n"
" m2 += p[2] + d * d * n * p[0] / nn;\n"
" n = nn;\n"
" }\n"
" mean = (float)mu;\n"
" var = (float)(m2 / n);\n"
" float unb = n > 1.0 ? (float)(m2 / (n - 1.0)) : var;\n"
" F(3)[c] = (1.0f - a.f[0]) * F(3)[c] + a.f[0] * mean;\n"
" F(4)[c] = (1.0f - a.f[0]) * F(4)[c] + a.f[0] * unb;\n"
" } else {\n"
" mean = CF(3)[c];\n"
" var = CF(4)[c];\n"
" }\n"
" F(1)[c] = mean;\n"
" F(2)[c] = 1.0f / sqrtf(var + a.f[1]);\n"
" }\n"
"}\n"
"// Group norm statistics per (image, group) unit of M contiguous values:\n"
"// p0 x, p1 mean[U], p2 invstd[U]; u0 M, u1 U; f0 eps n = U*256\n"
"DK_KERNEL(dk_gn_stats) {\n"
" unsigned int M = a.u[0];\n"
" UNITS(u, a.u[1]) {\n"
" TP x = CF(0) + (DK_U64)u * M;\n"
" dk_double s = 0.0;\n"
" TSPLIT(q, M) s += (dk_double)x[q];\n"
" dk_double mu = block_sum_d(s DK_IDA) / (dk_double)M;\n"
" dk_double sq = 0.0;\n"
" TSPLIT(q, M) { dk_double d = (dk_double)x[q] - mu; sq += d * d; }\n"
" sq = block_sum_d(sq DK_IDA);\n"
" if (threadIdx.x == 0) { F(1)[u] = (float)mu; F(2)[u] = 1.0f / sqrtf((float)(sq / (dk_double)M) + a.f[0]); }\n"
" }\n"
"}\n"
"// The statistic of value i: its channel (batch norm, G = 0) or its (image, group).\n"
"__device__ unsigned int norm_stat(unsigned int i, unsigned int C, unsigned int P, unsigned int G, float iP, float iC) {\n"
" unsigned int t = udiv(i, P, iP), img = udiv(t, C, iC), c = t - img * C;\n"
" return G == 0 ? c : img * G + c / (C / G);\n"
"}\n"
"// A normalization fused with the activation after it: u5 = 0 none, 1 ReLU,\n"
"// 2 leaky ReLU (slope f1). The backward kernels take the fused output in p7\n"
"// and use dy * act'(y) for dy.\n"
"__device__ float fused_act(const DevArgs& a, float v) {\n"
" return a.u[5] == 0 || v > 0.0f ? v : (a.u[5] == 1 ? 0.0f : a.f[1] * v);\n"
"}\n"
"__device__ float fused_dy(const DevArgs& a, DK_U64 i) {\n"
" float dy = TP(a.p[0])[i];\n"
" if (a.u[5] == 0) return dy;\n"
" return TP(a.p[7])[i] > 0.0f ? dy : (a.u[5] == 1 ? 0.0f : a.f[1] * dy);\n"
"}\n"
"// out = gamma[c] (x - mean) invstd + beta[c]: p0 x, p1 out, p2 mean, p3 invstd,\n"
"// p4 gamma, p5 beta; u0 C, u1 plane, u2 groups (0: batch norm) n = values\n"
"#ifdef DK_TAGGED\n"
"template <int T> __device__ void norm_apply_pairs(const DevArgs& a DK_IDS) {\n"
" unsigned int C = a.u[0], P = a.u[1], G = a.u[2];\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" const dk_global unsigned int* x = pw(a.p[0]);\n"
" dk_global unsigned int* o = pwm(a.p[1]);\n"
" GS(j, N / 2u) {\n"
" unsigned int i = 2u * j, c = chan_of(i, C, P, iP, iC), st = norm_stat(i, C, P, G, iP, iC);\n"
" float g = FA(4)[c], b = FA(5)[c], m = FA(2)[st], is = FA(3)[st];\n"
" unsigned int w = x[j];\n"
" o[j] = pk<T>(fused_act(a, g * (pk_lo<T>(w) - m) * is + b), fused_act(a, g * (pk_hi<T>(w) - m) * is + b));\n"
" }\n"
"}\n"
"#endif\n"
"DK_KERNEL(dk_norm_apply) {\n"
" unsigned int C = a.u[0], P = a.u[1], G = a.u[2];\n"
"#ifdef DK_TAGGED\n"
" if ((P & 1u) == 0u && (N & 1u) == 0u) {\n"
" unsigned int t = pair_tag(a.p[0], a.p[1], 0ull, 0ull);\n"
" if (t == 1u) { norm_apply_pairs<1>(a DK_IDA); return; }\n"
" if (t == 2u) { norm_apply_pairs<2>(a DK_IDA); return; }\n"
" }\n"
"#endif\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" GS(i, N) {\n"
" unsigned int c = chan_of(i, C, P, iP, iC), st = norm_stat(i, C, P, G, iP, iC);\n"
" F(1)[i] = fused_act(a, CF(4)[c] * (CF(0)[i] - CF(2)[st]) * CF(3)[st] + CF(5)[c]);\n"
" }\n"
"}\n"
"// Per channel: dgamma[c] += sum dy xh, dbeta[c] += sum dy, and (p6, when set,\n"
"// zeroed by the caller) sums[c] += sum dy, sums[C + c] += sum dy xh. p0 dy,\n"
"// p1 x, p2 mean, p3 invstd, p4 dgamma, p5 dbeta (either may be 0); u0 C,\n"
"// u1 plane, u2 groups, u3 batch, u4 S slices n = C*S*256\n"
"#ifdef DK_TAGGED\n"
"// dy * act'(y) for a pair\n"
"template <int T> __device__ __forceinline__ void fused_dy2(const DevArgs& a, unsigned int j, dk_thread float& lo, dk_thread float& hi) {\n"
" unsigned int w = pw(a.p[0])[j];\n"
" lo = pk_lo<T>(w); hi = pk_hi<T>(w);\n"
" if (a.u[5] != 0u) {\n"
" unsigned int y = pw(a.p[7])[j];\n"
" float s = a.u[5] == 1u ? 0.0f : a.f[1];\n"
" if (!(pk_lo<T>(y) > 0.0f)) lo *= s;\n"
" if (!(pk_hi<T>(y) > 0.0f)) hi *= s;\n"
" }\n"
"}\n"
"template <int T> __device__ void norm_sums_pairs(const DevArgs& a DK_IDS) {\n"
" unsigned int C = a.u[0], P = a.u[1], G = a.u[2], M = a.u[3] * a.u[1], S = a.u[4];\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" const dk_global unsigned int* xw = pw(a.p[1]);\n"
" UNITS(u, C * S) {\n"
" unsigned int c = u / S, sl = u % S;\n"
" dk_double sd = 0.0, sx = 0.0;\n"
" SLICE_LOOP(q, M / 2u, sl, S) {\n"
" unsigned int i = CHAN_AT(2u * q, c, C, P), st = G == 0 ? c : norm_stat(i, C, P, G, iP, iC);\n"
" float dlo, dhi;\n"
" fused_dy2<T>(a, i / 2u, dlo, dhi);\n"
" float m = FA(2)[st], is = FA(3)[st];\n"
" unsigned int w = xw[i / 2u];\n"
" sd += (dk_double)dlo + (dk_double)dhi;\n"
" sx += (dk_double)(dlo * (pk_lo<T>(w) - m) * is) + (dk_double)(dhi * (pk_hi<T>(w) - m) * is);\n"
" }\n"
" sd = block_sum_d(sd DK_IDA);\n"
" sx = block_sum_d(sx DK_IDA);\n"
" if (threadIdx.x == 0) {\n"
" if (a.p[4]) atomicAdd(&FA(4)[c], (float)sx);\n"
" if (a.p[5]) atomicAdd(&FA(5)[c], (float)sd);\n"
" if (a.p[6]) { atomicAdd(&FA(6)[c], (float)sd); atomicAdd(&FA(6)[C + c], (float)sx); }\n"
" }\n"
" }\n"
"}\n"
"#endif\n"
"DK_KERNEL(dk_norm_sums) {\n"
" unsigned int C = a.u[0], P = a.u[1], G = a.u[2], M = a.u[3] * a.u[1], S = a.u[4];\n"
"#ifdef DK_TAGGED\n"
" if ((P & 1u) == 0u) {\n"
" unsigned int t = pair_tag(a.p[0], a.p[1], a.u[5] != 0u ? a.p[7] : 0ull, 0ull);\n"
" if (t == 1u) { norm_sums_pairs<1>(a DK_IDA); return; }\n"
" if (t == 2u) { norm_sums_pairs<2>(a DK_IDA); return; }\n"
" }\n"
"#endif\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" UNITS(u, C * S) {\n"
" unsigned int c = u / S, sl = u % S;\n"
" dk_double sd = 0.0, sx = 0.0;\n"
" SLICE_LOOP(q, M, sl, S) {\n"
" unsigned int i = CHAN_AT(q, c, C, P), st = G == 0 ? c : norm_stat(i, C, P, G, iP, iC);\n"
" float dy = fused_dy(a, i);\n"
" sd += (dk_double)dy;\n"
" sx += (dk_double)(dy * (CF(1)[i] - CF(2)[st]) * CF(3)[st]);\n"
" }\n"
" sd = block_sum_d(sd DK_IDA);\n"
" sx = block_sum_d(sx DK_IDA);\n"
" if (threadIdx.x == 0) {\n"
" if (a.p[4]) atomicAdd(&FA(4)[c], (float)sx);\n"
" if (a.p[5]) atomicAdd(&FA(5)[c], (float)sd);\n"
" if (a.p[6]) { atomicAdd(&FA(6)[c], (float)sd); atomicAdd(&FA(6)[C + c], (float)sx); }\n"
" }\n"
" }\n"
"}\n"
"// Batch norm: dx += gamma invstd (dy - m1 - xh m2), m1 = sum dy / N, m2 = sum dy xh / N\n"
"// in training (0 in eval). p0 dy, p1 x, p2 mean, p3 invstd, p4 gamma, p5 sums,\n"
"// p6 dx; u0 C, u1 plane, u2 training; f0 1/N n = values\n"
"#ifdef DK_TAGGED\n"
"template <int T> __device__ void bn_bwd_dx_pairs(const DevArgs& a DK_IDS) {\n"
" unsigned int C = a.u[0], P = a.u[1];\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" const dk_global unsigned int* xw = pw(a.p[1]);\n"
" dk_global unsigned int* dx = pwm(a.p[6]);\n"
" GS(j, N / 2u) {\n"
" unsigned int c = chan_of(2u * j, C, P, iP, iC);\n"
" float inv = FA(3)[c], mean = FA(2)[c], g = FA(4)[c] * inv;\n"
" float m1 = a.u[2] ? FA(5)[c] * a.f[0] : 0.0f, m2 = a.u[2] ? FA(5)[C + c] * a.f[0] : 0.0f;\n"
" float dlo, dhi;\n"
" fused_dy2<T>(a, j, dlo, dhi);\n"
" unsigned int w = xw[j];\n"
" float xlo = (pk_lo<T>(w) - mean) * inv, xhi = (pk_hi<T>(w) - mean) * inv;\n"
" float vlo = g * (dlo - m1 - xlo * m2), vhi = g * (dhi - m1 - xhi * m2);\n"
" if (!a.u[6]) { unsigned int o = dx[j]; vlo += pk_lo<T>(o); vhi += pk_hi<T>(o); }\n"
" dx[j] = pk<T>(vlo, vhi);\n"
" }\n"
"}\n"
"#endif\n"
"DK_KERNEL(dk_bn_bwd_dx) {\n"
" unsigned int C = a.u[0], P = a.u[1];\n"
"#ifdef DK_TAGGED\n"
" if ((P & 1u) == 0u && (N & 1u) == 0u) {\n"
" unsigned int t = pair_tag(a.p[0], a.p[1], a.p[6], a.u[5] != 0u ? a.p[7] : 0ull);\n"
" if (t == 1u) { bn_bwd_dx_pairs<1>(a DK_IDA); return; }\n"
" if (t == 2u) { bn_bwd_dx_pairs<2>(a DK_IDA); return; }\n"
" }\n"
"#endif\n"
" float iP = 1.0f / (float)P, iC = 1.0f / (float)C;\n"
" GS(i, N) {\n"
" unsigned int c = chan_of(i, C, P, iP, iC);\n"
" float inv = CF(3)[c];\n"
" float m1 = a.u[2] ? CF(5)[c] * a.f[0] : 0.0f, m2 = a.u[2] ? CF(5)[C + c] * a.f[0] : 0.0f;\n"
" float xh = (CF(1)[i] - CF(2)[c]) * inv;\n"
" ACCUM(F(6)[i], CF(4)[c] * inv * (fused_dy(a, i) - m1 - xh * m2));\n"
" }\n"
"}\n"
"// Group norm: per (image, group) unit, dxh = dy gamma[c],\n"
"// dx += invstd (dxh - mean(dxh) - xh mean(dxh xh)). p0 dy, p1 x, p2 mean,\n"
"// p3 invstd, p4 gamma, p5 dx; u0 C, u1 plane, u2 groups, u3 units n = U*256\n"
"DK_KERNEL(dk_gn_bwd) {\n"
" unsigned int C = a.u[0], P = a.u[1], G = a.u[2];\n"
" unsigned int Cg = C / G, M = Cg * P;\n"
" float iP = 1.0f / (float)P;\n"
" UNITS(u, a.u[3]) {\n"
" DK_U64 base = (DK_U64)u * M;\n"
" unsigned int g = u % G;\n"
" float mean = CF(2)[u], inv = CF(3)[u];\n"
" dk_double s1 = 0.0, s2 = 0.0;\n"
" TSPLIT(q, M) {\n"
" float dxh = fused_dy(a, base + q) * CF(4)[g * Cg + udiv(q, P, iP)];\n"
" s1 += (dk_double)dxh;\n"
" s2 += (dk_double)(dxh * (CF(1)[base + q] - mean) * inv);\n"
" }\n"
" float m1 = (float)(block_sum_d(s1 DK_IDA) / (dk_double)M), m2 = (float)(block_sum_d(s2 DK_IDA) / (dk_double)M);\n"
" TSPLIT(q, M) {\n"
" float dxh = fused_dy(a, base + q) * CF(4)[g * Cg + udiv(q, P, iP)];\n"
" float xh = (CF(1)[base + q] - mean) * inv;\n"
" F(5)[base + q] += inv * (dxh - m1 - xh * m2);\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── Pooling ([batch, C*H*W] planes) ──────────────────────────────────────────\n"
"// u0 C, u1 H, u2 W, u3 oH, u4 oW, u5 mode (0 max, 1 avg over the padded\n"
"// window, 2 avg over the part inside the image); f0..f5 kH, kW, sH, sW, pH, pW.\n"
"struct PoolWin { unsigned int h0, h1, w0, w1, padded; };\n"
"__device__ PoolWin pool_win(const DevArgs& a, unsigned int oh, unsigned int ow) {\n"
" unsigned int H = a.u[1], W = a.u[2];\n"
" int kH = (int)a.f[0], kW = (int)a.f[1], sH = (int)a.f[2], sW = (int)a.f[3], pH = (int)a.f[4], pW = (int)a.f[5];\n"
" int h0 = (int)oh * sH - pH, w0 = (int)ow * sW - pW, h1 = h0 + kH, w1 = w0 + kW;\n"
" int hp = h1 < (int)H + pH ? h1 : (int)H + pH, wp = w1 < (int)W + pW ? w1 : (int)W + pW;\n"
" PoolWin r;\n"
" r.padded = (unsigned int)((hp - h0) * (wp - w0));\n"
" r.h0 = h0 < 0 ? 0 : h0; r.w0 = w0 < 0 ? 0 : w0;\n"
" r.h1 = h1 > (int)H ? H : h1; r.w1 = w1 > (int)W ? W : w1;\n"
" return r;\n"
"}\n"
"// p0 x, p1 out [batch, C*oH*oW] (forward); p0 x, p1 dout, p2 dx (backward: u6 = 1)\n"
"// n = batch*C*oH*oW\n"
"DK_KERNEL(dk_pool) {\n"
" unsigned int H = a.u[1], W = a.u[2], oH = a.u[3], oW = a.u[4], mode = a.u[5];\n"
" float iW = 1.0f / (float)oW, iH = 1.0f / (float)oH;\n"
" GS(i, N) {\n"
" unsigned int t = udiv(i, oW, iW), ow = i - t * oW, pl = udiv(t, oH, iH), oh = t - pl * oH;\n"
" PoolWin r = pool_win(a, oh, ow);\n"
" TP img = CF(0) + (DK_U64)pl * H * W;\n"
" if (mode == 0) {\n"
" unsigned int best = r.h0 * W + r.w0;\n"
" float bv = img[best];\n"
" for (unsigned int y = r.h0; y < r.h1; y++)\n"
" for (unsigned int x = r.w0; x < r.w1; x++) if (img[y * W + x] > bv) { bv = img[y * W + x]; best = y * W + x; }\n"
" if (a.u[6]) atomicAdd(&FA(2)[(DK_U64)pl * H * W + best], CF(1)[i]); else F(1)[i] = bv;\n"
" } else {\n"
" float div = mode == 1 ? (float)r.padded : (float)((r.h1 - r.h0) * (r.w1 - r.w0));\n"
" if (a.u[6]) {\n"
" float share = CF(1)[i] / div;\n"
" for (unsigned int y = r.h0; y < r.h1; y++)\n"
" for (unsigned int x = r.w0; x < r.w1; x++) atomicAdd(&FA(2)[(DK_U64)pl * H * W + y * W + x], share);\n"
" } else {\n"
" float acc = 0.0f;\n"
" for (unsigned int y = r.h0; y < r.h1; y++) for (unsigned int x = r.w0; x < r.w1; x++) acc += img[y * W + x];\n"
" F(1)[i] = acc / div;\n"
" }\n"
" }\n"
" }\n"
"}\n"
"// The pooling's backward as a gather (no atomics, any element type): input\n"
"// (y, x) sums the gradients of the windows that contain it -- for max\n"
"// pooling, of those whose (first) largest value it is. p0 x, p1 dout, p2 dx;\n"
"// geometry as dk_pool; u6 = 1: dx was an unwritten zero gradient n = planes*H*W\n"
"DK_KERNEL(dk_pool_bwd) {\n"
" unsigned int H = a.u[1], W = a.u[2], oH = a.u[3], oW = a.u[4], mode = a.u[5];\n"
" int kH = (int)a.f[0], kW = (int)a.f[1], sH = (int)a.f[2], sW = (int)a.f[3], pH = (int)a.f[4], pW = (int)a.f[5];\n"
" if (kH == sH && kW == sW && pH == 0 && pW == 0 &&\n"
" (!a.u[6] || (oH * (unsigned int)kH == H && oW * (unsigned int)kW == W))) {\n"
" // non-overlapping windows partition the inputs: a thread per window\n"
" // adds into its own inputs (n counts inputs; the windows are fewer)\n"
" float iW = 1.0f / (float)oW, iH = 1.0f / (float)oH;\n"
" unsigned int planes = N / (H * W);\n"
" GS(i, planes * oH * oW) {\n"
" unsigned int t = udiv(i, oW, iW), ow = i - t * oW, pl = udiv(t, oH, iH), oh = t - pl * oH;\n"
" PoolWin r = pool_win(a, oh, ow);\n"
" TP img = CF(0) + (DK_U64)pl * H * W;\n"
" TP dx = F(2) + (DK_U64)pl * H * W;\n"
" float d = CF(1)[i];\n"
" if (mode == 0) {\n"
" unsigned int best = r.h0 * W + r.w0;\n"
" float bv = img[best];\n"
" for (unsigned int yy = r.h0; yy < r.h1; yy++)\n"
" for (unsigned int xx = r.w0; xx < r.w1; xx++) if (img[yy * W + xx] > bv) { bv = img[yy * W + xx]; best = yy * W + xx; }\n"
" if (a.u[6]) {\n"
" // exact tiling, unwritten gradient: write the window\n"
" for (unsigned int yy = r.h0; yy < r.h1; yy++)\n"
" for (unsigned int xx = r.w0; xx < r.w1; xx++) dx[yy * W + xx] = yy * W + xx == best ? d : 0.0f;\n"
" } else dx[best] += d;\n"
" } else {\n"
" float share = d / (mode == 1 ? (float)r.padded : (float)((r.h1 - r.h0) * (r.w1 - r.w0)));\n"
" for (unsigned int yy = r.h0; yy < r.h1; yy++)\n"
" for (unsigned int xx = r.w0; xx < r.w1; xx++) ACCUM(dx[yy * W + xx], share);\n"
" }\n"
" }\n"
" return;\n"
" }\n"
" float iW = 1.0f / (float)W, iH = 1.0f / (float)H, iK = 1.0f / (float)kH, iKw = 1.0f / (float)kW;\n"
" GS(i, N) {\n"
" unsigned int t = udiv(i, W, iW), x = i - t * W, pl = udiv(t, H, iH), y = t - pl * H;\n"
" TP img = CF(0) + (DK_U64)pl * H * W;\n"
" TP dout = CF(1) + (DK_U64)pl * oH * oW;\n"
" int oh0, oh1, ow0, ow1;\n"
" if (kH == sH && kW == sW && pH == 0 && pW == 0) {\n"
" // non-overlapping windows: exactly one holds (y, x)\n"
" oh0 = oh1 = (int)udiv(y, (unsigned int)kH, iK);\n"
" ow0 = ow1 = (int)udiv(x, (unsigned int)kW, iKw);\n"
" } else {\n"
" oh0 = ((int)y + pH - kH + sH) / sH; oh1 = ((int)y + pH) / sH; // windows over row y\n"
" ow0 = ((int)x + pW - kW + sW) / sW; ow1 = ((int)x + pW) / sW;\n"
" if ((int)y + pH - kH + 1 < 0) oh0 = 0;\n"
" if ((int)x + pW - kW + 1 < 0) ow0 = 0;\n"
" }\n"
" if (oh1 > (int)oH - 1) oh1 = (int)oH - 1;\n"
" if (ow1 > (int)oW - 1) ow1 = (int)oW - 1;\n"
" float s = 0.0f;\n"
" for (int oh = oh0; oh <= oh1; oh++) {\n"
" for (int ow = ow0; ow <= ow1; ow++) {\n"
" PoolWin r = pool_win(a, (unsigned int)oh, (unsigned int)ow);\n"
" if (y < r.h0 || y >= r.h1 || x < r.w0 || x >= r.w1) continue;\n"
" float d = dout[(unsigned int)oh * oW + (unsigned int)ow];\n"
" if (mode == 0) {\n"
" unsigned int best = r.h0 * W + r.w0;\n"
" float bv = img[best];\n"
" for (unsigned int yy = r.h0; yy < r.h1; yy++)\n"
" for (unsigned int xx = r.w0; xx < r.w1; xx++) if (img[yy * W + xx] > bv) { bv = img[yy * W + xx]; best = yy * W + xx; }\n"
" if (best == y * W + x) s += d;\n"
" } else {\n"
" s += d / (mode == 1 ? (float)r.padded : (float)((r.h1 - r.h0) * (r.w1 - r.w0)));\n"
" }\n"
" }\n"
" }\n"
" ACCUM(F(2)[i], s);\n"
" }\n"
"}\n"
"// Global pooling, one plane of u0 values per unit: p0 x, p1 out [planes];\n"
"// backward (u2 = 1): p1 dout, p2 dx. u1 = 0 max, 1 avg n = planes*256\n"
"DK_KERNEL(dk_gpool) {\n"
" unsigned int P = a.u[0];\n"
" UNITS(pl, N / 256) {\n"
" TP x = CF(0) + (DK_U64)pl * P;\n"
" if (a.u[1]) {\n"
" if (a.u[2]) { float share = CF(1)[pl] / (float)P; TSPLIT(q, P) ACCUM(F(2)[(DK_U64)pl * P + q], share); }\n"
" else { dk_double s = 0.0; TSPLIT(q, P) s += (dk_double)x[q]; s = block_sum_d(s DK_IDA); if (threadIdx.x == 0) F(1)[pl] = (float)(s / (dk_double)P); }\n"
" } else if (threadIdx.x == 0) {\n"
" unsigned int best = 0;\n"
" for (unsigned int q = 1; q < P; q++) if (x[q] > x[best]) best = q;\n"
" if (a.u[2]) F(2)[(DK_U64)pl * P + best] += CF(1)[pl]; else F(1)[pl] = x[best];\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── Upsampling and resizing (planes of H x W -> oH x oW) ─────────────────────\n"
"// Nearest by an integer factor u2: forward (u3 = 0) p0 x, p1 out, one thread\n"
"// per output; backward (u3 = 1) p0 dout, p1 dx, one per input. u0 W, u1 H\n"
"DK_KERNEL(dk_upsample) {\n"
" unsigned int W = a.u[0], H = a.u[1], f = a.u[2], oW = W * f, oH = H * f;\n"
" GS(i, N) {\n"
" if (a.u[3] == 0) {\n"
" unsigned int ow = i % oW, oh = (i / oW) % oH, pl = i / (oW * oH);\n"
" F(1)[i] = CF(0)[((DK_U64)pl * H + oh / f) * W + ow / f];\n"
" } else {\n"
" unsigned int w = i % W, h = (i / W) % H, pl = i / (W * H);\n"
" float s = 0.0f;\n"
" for (unsigned int y = 0; y < f; y++)\n"
" for (unsigned int x = 0; x < f; x++) s += CF(0)[((DK_U64)pl * oH + h * f + y) * oW + w * f + x];\n"
" ACCUM(F(1)[i], s);\n"
" }\n"
" }\n"
"}\n"
"// Up to four taps per axis (unused ones weigh 0), kept in registers: every\n"
"// loop over them is unrolled.\n"
"struct Taps { unsigned int i0, i1, i2, i3; float w0, w1, w2, w3; };\n"
"__device__ float rs_src(unsigned int o, unsigned int in, unsigned int out, unsigned int align, int clampZero) {\n"
" if (align) return out <= 1 ? 0.0f : (float)o * (float)(in - 1) / (float)(out - 1);\n"
" float s = ((float)o + 0.5f) * (float)in / (float)out - 0.5f;\n"
" if (clampZero && s < 0.0f) s = 0.0f;\n"
" return s;\n"
"}\n"
"__device__ unsigned int rs_clamp(int j, unsigned int in) { return j < 0 ? 0 : (j >= (int)in ? in - 1 : (unsigned int)j); }\n"
"__device__ Taps rs_taps(unsigned int o, unsigned int in, unsigned int out, unsigned int mode, unsigned int align) {\n"
" Taps t;\n"
" t.w0 = t.w1 = t.w2 = t.w3 = 0.0f;\n"
" if (mode == 0) {\n"
" unsigned int i = (unsigned int)floorf((float)o * ((float)in / (float)out));\n"
" t.i0 = t.i1 = t.i2 = t.i3 = i < in ? i : in - 1; t.w0 = 1.0f;\n"
" } else if (mode == 1) {\n"
" float s = rs_src(o, in, out, align, 1);\n"
" unsigned int i0 = (unsigned int)s;\n"
" if (i0 > in - 1) i0 = in - 1;\n"
" float l = s - (float)i0;\n"
" t.i0 = i0; t.i1 = t.i2 = t.i3 = i0 + 1 < in ? i0 + 1 : in - 1;\n"
" t.w0 = 1.0f - l; t.w1 = l;\n"
" } else {\n"
" float s = rs_src(o, in, out, align, 0), f = floorf(s), x = s - f, A = -0.75f;\n"
" float x1 = x + 1.0f, x2 = 1.0f - x, x3 = 2.0f - x;\n"
" t.w0 = ((A * x1 - 5.0f * A) * x1 + 8.0f * A) * x1 - 4.0f * A;\n"
" t.w1 = ((A + 2.0f) * x - (A + 3.0f)) * x * x + 1.0f;\n"
" t.w2 = ((A + 2.0f) * x2 - (A + 3.0f)) * x2 * x2 + 1.0f;\n"
" t.w3 = ((A * x3 - 5.0f * A) * x3 + 8.0f * A) * x3 - 4.0f * A;\n"
" int i = (int)f;\n"
" t.i0 = rs_clamp(i - 1, in); t.i1 = rs_clamp(i, in); t.i2 = rs_clamp(i + 1, in); t.i3 = rs_clamp(i + 2, in);\n"
" }\n"
" return t;\n"
"}\n"
"// sum_k w_k row[i_k]\n"
"__device__ float rs_dot(const dk_thread Taps& t, TP row) {\n"
" float s = 0.0f;\n"
" if (t.w0 != 0.0f) s += t.w0 * row[t.i0];\n"
" if (t.w1 != 0.0f) s += t.w1 * row[t.i1];\n"
" if (t.w2 != 0.0f) s += t.w2 * row[t.i2];\n"
" if (t.w3 != 0.0f) s += t.w3 * row[t.i3];\n"
" return s;\n"
"}\n"
"// convnet.h's RESIZE_* (u4), alignCorners u5: p0 x, p1 out. u0 H, u1 W, u2 oH,\n"
"// u3 oW n = planes*oH*oW\n"
"DK_KERNEL(dk_resize) {\n"
" unsigned int H = a.u[0], W = a.u[1], oH = a.u[2], oW = a.u[3];\n"
" GS(i, N) {\n"
" unsigned int ow = i % oW, oh = (i / oW) % oH, pl = i / (oW * oH);\n"
" Taps ty = rs_taps(oh, H, oH, a.u[4], a.u[5]), tx = rs_taps(ow, W, oW, a.u[4], a.u[5]);\n"
" TP img = CF(0) + (DK_U64)pl * H * W;\n"
" float acc = 0.0f;\n"
" if (ty.w0 != 0.0f) acc += ty.w0 * rs_dot(tx, img + ty.i0 * W);\n"
" if (ty.w1 != 0.0f) acc += ty.w1 * rs_dot(tx, img + ty.i1 * W);\n"
" if (ty.w2 != 0.0f) acc += ty.w2 * rs_dot(tx, img + ty.i2 * W);\n"
" if (ty.w3 != 0.0f) acc += ty.w3 * rs_dot(tx, img + ty.i3 * W);\n"
" F(1)[i] = acc;\n"
" }\n"
"}\n"
"// The weight with which output o's taps read input i along one axis (0: none).\n"
"__device__ float rs_weight(unsigned int o, unsigned int i, unsigned int in, unsigned int out, unsigned int mode,\n"
" unsigned int align) {\n"
" Taps t = rs_taps(o, in, out, mode, align);\n"
" return (t.i0 == i ? t.w0 : 0.0f) + (t.i1 == i ? t.w1 : 0.0f) + (t.i2 == i ? t.w2 : 0.0f) + (t.i3 == i ? t.w3 : 0.0f);\n"
"}\n"
"// The outputs that can read input i: their source coordinate lies in\n"
"// [i - 1, i + 1) for bilinear taps (floor(s), floor(s) + 1) and [i - 2, i + 2)\n"
"// for bicubic ones (floor(s) - 1 .. floor(s) + 2), borders clamped; a margin\n"
"// of one either side for rounding.\n"
"__device__ void rs_range(unsigned int i, unsigned int in, unsigned int out, unsigned int mode, unsigned int align,\n"
" dk_thread int* lo, dk_thread int* hi) {\n"
" float a0, a1;\n"
" float lo_s = mode == 1 ? (float)i - 1.0f : (float)i - 2.0f, hi_s = mode == 1 ? (float)i + 1.0f : (float)i + 2.0f;\n"
" if (mode == 0) { float r = (float)out / (float)in; a0 = (float)i * r; a1 = (float)(i + 1) * r; }\n"
" else if (align) {\n"
" float r = in > 1 ? (float)(out - 1) / (float)(in - 1) : (float)out;\n"
" a0 = lo_s * r; a1 = hi_s * r;\n"
" } else {\n"
" float r = (float)out / (float)in;\n"
" a0 = (lo_s + 0.5f) * r - 0.5f; a1 = (hi_s + 0.5f) * r - 0.5f;\n"
" }\n"
" int l = (int)floorf(a0) - 1, h = (int)ceilf(a1) + 1;\n"
" *lo = l < 0 ? 0 : l;\n"
" *hi = h > (int)out - 1 ? (int)out - 1 : h;\n"
"}\n"
"// The resize's backward as two gathers (no atomics), one axis at a time:\n"
"// pass 0 (f0 = 0): t[pl, oy, x] = sum over ox reading x of dout[pl, oy, ox] wx;\n"
"// pass 1 (f0 = 1): dx[pl, y, x] (+)= sum over oy reading y of t[pl, oy, x] wy.\n"
"// p0 source, p1 destination; u0 H, u1 W, u2 oH, u3 oW, u4 mode, u5 align\n"
"// n = planes*oH*W (pass 0), planes*H*W (pass 1)\n"
"DK_KERNEL(dk_resize_bwd) {\n"
" unsigned int H = a.u[0], W = a.u[1], oH = a.u[2], oW = a.u[3], mode = a.u[4], align = a.u[5];\n"
" GS(i, N) {\n"
" unsigned int x = i % W, r = (i / W);\n"
" int q0, q1;\n"
" float s = 0.0f;\n"
" if (a.f[0] == 0.0f) {\n"
" // r = pl * oH + oy: a row of dout\n"
" rs_range(x, W, oW, mode, align, &q0, &q1);\n"
" TP d = CF(0) + (DK_U64)r * oW;\n"
" for (int ox = q0; ox <= q1; ox++) {\n"
" float w = rs_weight((unsigned int)ox, x, W, oW, mode, align);\n"
" if (w != 0.0f) s += w * d[ox];\n"
" }\n"
" F(1)[i] = s;\n"
" } else {\n"
" unsigned int y = r % H, pl = r / H;\n"
" rs_range(y, H, oH, mode, align, &q0, &q1);\n"
" TP t = CF(0) + (DK_U64)pl * oH * W + x;\n"
" for (int oy = q0; oy <= q1; oy++) {\n"
" float w = rs_weight((unsigned int)oy, y, H, oH, mode, align);\n"
" if (w != 0.0f) s += w * t[(unsigned int)oy * W];\n"
" }\n"
" ACCUM(F(1)[i], s);\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── 16-bit storage (mixed precision) ─────────────────────────────────────────\n"
"// fp16 / bf16 values held as unsigned shorts, converted with round to\n"
"// nearest even (the same bits as float16.h's f32_to_fp16 / f32_to_bf16).\n"
"// Hardware conversions: v_cvt on AMD (_Float16), cvt.rn on NVIDIA, half on\n"
"// Apple GPUs; on Vulkan GLSL.std.450's PackHalf2x16 / UnpackHalf2x16, which\n"
"// need no shaderFloat16.\n"
"#if defined(DK_VULKAN)\n"
"uint16_t f2h_bits(float f) {\n"
" float2 v = float2(f, 0.0f);\n"
" return (uint16_t)spirv_asm { result:$$uint = OpExtInst glsl450 PackHalf2x16 $v };\n"
"}\n"
"float h2f_bits(uint16_t b) {\n"
" uint h = (uint)b;\n"
" float2 v = spirv_asm { result:$$float2 = OpExtInst glsl450 UnpackHalf2x16 $h };\n"
" return v.x;\n"
"}\n"
"#elif defined(__METAL_VERSION__)\n"
"unsigned short f2h_bits(float f) { return as_type<ushort>(half(f)); }\n"
"float h2f_bits(unsigned short b) { return float(as_type<half>(b)); }\n"
"#elif defined(__HIPCC_RTC__) || defined(__HIP__)\n"
"__device__ unsigned short f2h_bits(float f) { _Float16 h = (_Float16)f; return *(unsigned short*)&h; }\n"
"__device__ float h2f_bits(unsigned short b) { return (float)*(_Float16*)&b; }\n"
"#else\n"
"__device__ unsigned short f2h_bits(float f) { unsigned short r; asm(\"cvt.rn.f16.f32 %0, %1;\" : \"=h\"(r) : \"f\"(f)); return r; }\n"
"__device__ float h2f_bits(unsigned short b) { float r; asm(\"cvt.f32.f16 %0, %1;\" : \"=f\"(r) : \"h\"(b)); return r; }\n"
"#endif\n"
"__device__ unsigned short f2b_bits(float f) {\n"
" unsigned int x = __float_as_uint(f);\n"
" if ((x & 0x7f800000u) == 0x7f800000u && (x & 0x7fffffu)) return (unsigned short)((x >> 16) | 0x40u);\n"
" return (unsigned short)((x + 0x7fffu + ((x >> 16) & 1u)) >> 16);\n"
"}\n"
"__device__ float b2f_bits(unsigned short b) { return __uint_as_float((unsigned int)b << 16); }\n"
"// The dynamic loss scale's state (one per device, 4 floats): [scale,\n"
"// 1/scale, overflow flag, good steps].\n"
"// g *= 1/scale, the flag raised on a non-finite value: p0 g, p1 state n = values\n"
"DK_KERNEL(dk_amp_unscale) {\n"
" TP st = F(1);\n"
" float inv = st[1];\n"
" GS(i, N) {\n"
" float v = F(0)[i] * inv;\n"
" if (!isfinite(v)) st[2] = 1.0f;\n"
" F(0)[i] = v;\n"
" }\n"
"}\n"
"// Every gradient of a table at once: p0 = u1 entries {g, n, start} (u64),\n"
"// p1 state n = total values\n"
"DK_KERNEL(dk_amp_unscale_multi) {\n"
" const dk_global DK_U64* T = (const dk_global DK_U64*)a.p[0];\n"
" unsigned int cnt = a.u[1];\n"
" TP st = F(1);\n"
" float inv = st[1];\n"
" GS(i, N) {\n"
" unsigned int lo = 0, hi = cnt - 1;\n"
" while (lo < hi) { unsigned int mid = (lo + hi + 1) / 2; if (T[mid * 3 + 2] <= i) lo = mid; else hi = mid - 1; }\n"
" dk_global float* g = (dk_global float*)T[lo * 3];\n"
" unsigned int j = i - (unsigned int)T[lo * 3 + 2];\n"
" float v = g[j] * inv;\n"
" if (!isfinite(v)) st[2] = 1.0f;\n"
" g[j] = v;\n"
" }\n"
"}\n"
"// After a step: an overflow halves the scale, u0 good steps in a row double\n"
"// it; the flag is cleared. p0 state n = 1\n"
"DK_KERNEL(dk_amp_update) {\n"
" if (blockIdx.x == 0 && threadIdx.x == 0) {\n"
" TP st = F(0);\n"
" if (st[2] != 0.0f) { st[0] = fmaxf(st[0] * 0.5f, 1.0f); st[3] = 0.0f; }\n"
" else {\n"
" st[3] += 1.0f;\n"
" if (st[3] >= (float)a.u[0]) { st[0] = fminf(st[0] * 2.0f, 16777216.0f); st[3] = 0.0f; }\n"
" }\n"
" st[1] = 1.0f / st[0];\n"
" st[2] = 0.0f;\n"
" }\n"
"}\n"
"\n"
"// ── Sparse (CSR) x dense: sparse.h ───────────────────────────────────────────\n"
"// out [rows, n] (+)= S [rows, cols] . X [cols, n]; p0 rowPtr, p1 colIdx (as\n"
"// float32 indices), p2 values, p3 X, p4 out; u0 = n, u1 = 1 to accumulate n = rows*ncols\n"
"DK_KERNEL(dk_spmm) {\n"
" unsigned int NC = a.u[0];\n"
" GS(i, N) {\n"
" unsigned int r = i / NC, j = i % NC;\n"
" unsigned int lo = (unsigned int)CF(0)[r], hi = (unsigned int)CF(0)[r + 1];\n"
" float s = 0.0f;\n"
" for (unsigned int e = lo; e < hi; e++) s += CF(2)[e] * CF(3)[(unsigned int)CF(1)[e] * NC + j];\n"
" if (a.u[1]) F(4)[i] += s; else F(4)[i] = s;\n"
" }\n"
"}\n"
"// dX[cols, n] += S^T . G: one thread per (nonzero, column) n = nnz*ncols\n"
"// p5 = row of each nonzero\n"
"DK_KERNEL(dk_spmm_tgrad) {\n"
" unsigned int NC = a.u[0];\n"
" GS(i, N) {\n"
" unsigned int e = i / NC, j = i % NC;\n"
" unsigned int r = (unsigned int)CF(5)[e], col = (unsigned int)CF(1)[e];\n"
" atomicAdd(&FA(4)[col * NC + j], CF(2)[e] * CF(3)[r * NC + j]);\n"
" }\n"
"}\n"
"// dvalues[e] += G[row(e), :] . X[col(e), :] n = nnz\n"
"// p1 colIdx, p3 X, p4 G, p5 row of each nonzero, p6 dvalues\n"
"DK_KERNEL(dk_spmm_vgrad) {\n"
" unsigned int NC = a.u[0];\n"
" GS(e, N) {\n"
" unsigned int r = (unsigned int)CF(5)[e], col = (unsigned int)CF(1)[e];\n"
" float s = 0.0f;\n"
" for (unsigned int j = 0; j < NC; j++) s += CF(4)[r * NC + j] * CF(3)[col * NC + j];\n"
" F(6)[e] += s;\n"
" }\n"
"}\n"
"\n"
"// ── Metal and Vulkan: GEMMs, convolutions and copies (rocBLAS / cuBLAS,\n"
"// MIOpen / cuDNN and the copy engines' role on the other backends;\n"
"// MpsBackend and VulkanBackend launch these by name) ───────────────────────\n"
"#if defined(__METAL_VERSION__) || defined(DK_VULKAN)\n"
"#if defined(DK_VULKAN)\n"
"// Vulkan has no simdgroup matrices: on a 32-wide subgroup, lane l holds\n"
"// elements 2l and 2l + 1 (row-major) of each 8 x 8 matrix, so a simdgroup's\n"
"// fragments are a register tile per lane. mk_mma (below) multiplies them\n"
"// straight from the shared tiles; the rest -- loads, stores and products\n"
"// with the diagonal alpha and beta -- exchange values through the subgroup.\n"
"// Threadgroup memory is one array (its size per kernel: gen_device_kernels.sh\n"
"// sets DK_TG_FLOATS) and TGPtr an offset into it.\n"
"#ifndef DK_TG_FLOATS\n"
"#define DK_TG_FLOATS 3072\n"
"#endif\n"
"static groupshared float dk_tg[DK_TG_FLOATS];\n"
"struct TGPtr {\n"
" uint o;\n"
" __init(uint off) { o = off; }\n"
" __subscript(uint i) -> float {\n"
" get { return dk_tg[o + i]; }\n"
" [nonmutating] set { dk_tg[o + i] = newValue; }\n"
" }\n"
"};\n"
"TGPtr operator+(TGPtr p, uint k) { return TGPtr(p.o + k); }\n"
"namespace mem_flags { static const int mem_threadgroup = 1; }\n"
"void threadgroup_barrier(int flags) { GroupMemoryBarrierWithGroupSync(); }\n"
"void simdgroup_barrier(int flags) { GroupMemoryBarrierWithWaveSync(); }\n"
"struct simdgroup_float8x8 {\n"
" float2 v;\n"
" __init(float d) { // d on the diagonal, as Metal's\n"
" uint l = WaveGetLaneIndex(), r = l / 4u, c = 2u * (l % 4u);\n"
" v = float2(r == c ? d : 0.0f, r == c + 1u ? d : 0.0f);\n"
" }\n"
"};\n"
"// d = a b + c: lane (r, c0) gathers row r of a and columns c0, c0 + 1 of b\n"
"void simdgroup_multiply_accumulate(out simdgroup_float8x8 d, simdgroup_float8x8 a, simdgroup_float8x8 b,\n"
" simdgroup_float8x8 c) {\n"
" uint l = WaveGetLaneIndex(), r = l / 4u;\n"
" float2 s = c.v;\n"
" [unroll]\n"
" for (uint k = 0; k < 8u; k++) {\n"
" float2 ar = WaveShuffle(a.v, r * 4u + k / 2u);\n"
" float ak = (k & 1u) != 0u ? ar.y : ar.x;\n"
" s += ak * WaveShuffle(b.v, k * 4u + l % 4u);\n"
" }\n"
" d.v = s;\n"
"}\n"
"void simdgroup_multiply(out simdgroup_float8x8 d, simdgroup_float8x8 a, simdgroup_float8x8 b) {\n"
" simdgroup_multiply_accumulate(d, a, b, simdgroup_float8x8(0.0f));\n"
"}\n"
"void simdgroup_load(out simdgroup_float8x8 m, float* p, ulong ld) {\n"
" uint l = WaveGetLaneIndex();\n"
" float* q = p + (ulong)(l / 4u) * ld + 2u * (l % 4u);\n"
" m.v = float2(q[0], q[1]);\n"
"}\n"
"void simdgroup_store(simdgroup_float8x8 m, float* p, ulong ld) {\n"
" uint l = WaveGetLaneIndex();\n"
" float* q = p + (ulong)(l / 4u) * ld + 2u * (l % 4u);\n"
" q[0] = m.v.x;\n"
" q[1] = m.v.y;\n"
"}\n"
"void simdgroup_store(simdgroup_float8x8 m, TGPtr p, uint ld) {\n"
" uint l = WaveGetLaneIndex(), o = (l / 4u) * ld + 2u * (l % 4u);\n"
" p[o] = m.v.x;\n"
" p[o + 1u] = m.v.y;\n"
"}\n"
"#endif\n"
"// A strided batch of c_i = alpha op(a_i) op(b_i) + beta c_i (device.h's\n"
"// DevGemm), row-major: a and b fp32, fp16 or bf16 (MK_T), c any of them\n"
"// (ct), products accumulated in fp32 on the 8 x 8 simdgroup matrix units.\n"
"// A threadgroup of four simdgroups takes a BM x BN tile of c (32 x 32 or\n"
"// 64 x 64: mk_gemm<BM>_<op(a)><op(b)>_<type>), k in steps of BK through\n"
"// threadgroup memory; the grid is (n / BN, m / BM, batch). Tiles inside the\n"
"// matrices load without bounds checks (a divergent load per element costs\n"
"// 3x) and store straight from the matrix units, alpha and beta applied as\n"
"// products with diagonal matrices; edge tiles check every element.\n"
"struct DkGemm {\n"
" ulong a, b, c; // addresses\n"
" ulong lda, ldb, ldc;\n"
" ulong sa, sb, sc; // batch strides, in elements\n"
" uint m, n, k, ta, tb, at, ct;\n"
" float alpha, beta;\n"
"};\n"
"template <uint T> inline float mk_ld(ulong base, uint i) {\n"
" if (T == 0u) return ((const device float*)base)[i];\n"
" ushort h = ((const device ushort*)base)[i];\n"
" return T == 1u ? h2f_bits(h) : b2f_bits(h);\n"
"}\n"
"inline float mk_ld_any(ulong base, ulong i, uint t) {\n"
" if (t == 0u) return ((const device float*)base)[i];\n"
" ushort h = ((const device ushort*)base)[i];\n"
" return t == 1u ? h2f_bits(h) : b2f_bits(h);\n"
"}\n"
"inline void mk_st(ulong base, ulong i, uint t, float v) {\n"
" if (t == 0u) ((device float*)base)[i] = v;\n"
" else ((device ushort*)base)[i] = t == 1u ? f2h_bits(v) : f2b_bits(v);\n"
"}\n"
"// One BK step of the tile product. The tiles are As [BM][BK] and Bs [BK][BN],\n"
"// or k-major As [BK][BM] (AK) and n-major Bs [BN][BK] (BNM) where that lets\n"
"// neighbouring threads store to neighbouring words; the matrix loads\n"
"// transpose them back.\n"
"#if defined(DK_VULKAN)\n"
"// Lane (r, c0) of the simdgroup at (sr, sc) accumulates rows sr + 8i + r and\n"
"// columns sc + 8j + c0, c0 + 1 of the tile.\n"
"__generic<let BM : uint, let BN : uint, let BK : uint, let SM : uint, let SN : uint, let AK : bool = false, let BNM : bool = false>\n"
"void mk_mma(TGPtr As, TGPtr Bs, inout simdgroup_float8x8 acc[SM][SN], uint sr, uint sc) {\n"
" uint l = WaveGetLaneIndex(), r = l / 4u, c0 = 2u * (l % 4u);\n"
" [unroll]\n"
" for (uint k = 0; k < BK; k++) {\n"
" float av[SM];\n"
" float2 bv[SN];\n"
" [unroll]\n"
" for (uint i = 0; i < SM; i++) av[i] = AK ? As[k * BM + sr + i * 8u + r] : As[(sr + i * 8u + r) * BK + k];\n"
" [unroll]\n"
" for (uint j = 0; j < SN; j++) {\n"
" uint n = sc + j * 8u + c0;\n"
" bv[j] = BNM ? float2(Bs[n * BK + k], Bs[(n + 1u) * BK + k]) : float2(Bs[k * BN + n], Bs[k * BN + n + 1u]);\n"
" }\n"
" [unroll]\n"
" for (uint i = 0; i < SM; i++)\n"
" [unroll]\n"
" for (uint j = 0; j < SN; j++) acc[i][j].v += av[i] * bv[j];\n"
" }\n"
"}\n"
"#else\n"
"template <uint BM, uint BN, uint BK, uint SM, uint SN, bool AK = false, bool BNM = false>\n"
"inline void mk_mma(threadgroup float* As, threadgroup float* Bs, thread simdgroup_float8x8 (&acc)[SM][SN], uint sr, uint sc) {\n"
" #pragma clang loop unroll(full)\n"
" for (uint kk = 0; kk < BK; kk += 8u) {\n"
" simdgroup_float8x8 fa[SM], fb[SN];\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++) {\n"
" if (AK) simdgroup_load(fa[i], &As[kk * BM + sr + i * 8u], BM, ulong2(0, 0), true);\n"
" else simdgroup_load(fa[i], &As[(sr + i * 8u) * BK + kk], BK);\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" if (BNM) simdgroup_load(fb[j], &Bs[(sc + j * 8u) * BK + kk], BK, ulong2(0, 0), true);\n"
" else simdgroup_load(fb[j], &Bs[kk * BN + sc + j * 8u], BN);\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) simdgroup_multiply_accumulate(acc[i][j], fa[i], fb[j], acc[i][j]);\n"
" }\n"
"}\n"
"#endif\n"
"\n"
"template <uint T, bool TA, bool TB, uint BM, uint BN, uint BK>\n"
"inline void mk_gemm(constant DkGemm& g, uint3 tg, uint tid, uint sg, threadgroup float* As, threadgroup float* Bs) {\n"
" constexpr uint SM = BM / 16u, SN = BN / 16u; // 8 x 8 fragments per simdgroup (2 x 2 simdgroups)\n"
" constexpr ulong es = T == 0u ? 4ul : 2ul;\n"
" uint m0 = tg.y * BM, n0 = tg.x * BN;\n"
" ulong A = g.a + (ulong)tg.z * g.sa * es, B = g.b + (ulong)tg.z * g.sb * es;\n"
" ulong C = g.c + (ulong)tg.z * g.sc * (g.ct == 0u ? 4ul : 2ul);\n"
" uint sr = (sg / 2u) * SM * 8u, sc = (sg % 2u) * SN * 8u;\n"
" uint lda = (uint)g.lda, ldb = (uint)g.ldb, Mx = g.m, Nx = g.n, Kx = g.k;\n"
" bool inside = m0 + BM <= Mx && n0 + BN <= Nx;\n"
" simdgroup_float8x8 acc[SM][SN];\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) acc[i][j] = simdgroup_float8x8(0.0f);\n"
" for (uint k0 = 0; k0 < Kx; k0 += BK) {\n"
" // neighbouring threads read neighbouring addresses either way round\n"
" if (inside && k0 + BK <= Kx) {\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = TA ? e % BM : e / BK, c = TA ? e / BM : e % BK;\n"
" As[TA ? c * BM + r : r * BK + c] = mk_ld<T>(A, TA ? (k0 + c) * lda + m0 + r : (m0 + r) * lda + k0 + c);\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint e = tid + q * 128u, r = TB ? e % BK : e / BN, c = TB ? e / BK : e % BN;\n"
" Bs[TB ? c * BK + r : r * BN + c] = mk_ld<T>(B, TB ? (n0 + c) * ldb + k0 + r : (k0 + r) * ldb + n0 + c);\n"
" }\n"
" } else {\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = TA ? e % BM : e / BK, c = TA ? e / BM : e % BK;\n"
" uint gm = m0 + r, gk = k0 + c;\n"
" float v = 0.0f;\n"
" if (gm < Mx && gk < Kx) v = mk_ld<T>(A, TA ? gk * lda + gm : gm * lda + gk);\n"
" As[TA ? c * BM + r : r * BK + c] = v;\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint e = tid + q * 128u, r = TB ? e % BK : e / BN, c = TB ? e / BK : e % BN;\n"
" uint gk = k0 + r, gn = n0 + c;\n"
" float v = 0.0f;\n"
" if (gk < Kx && gn < Nx) v = mk_ld<T>(B, TB ? gn * ldb + gk : gk * ldb + gn);\n"
" Bs[TB ? c * BK + r : r * BN + c] = v;\n"
" }\n"
" }\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" mk_mma<BM, BN, BK, SM, SN, TA, TB>(As, Bs, acc, sr, sc);\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
" if (inside && g.ct == 0u) {\n"
" // c = alpha acc + beta c on the matrix units (diagonal alpha, beta)\n"
" simdgroup_float8x8 da = simdgroup_float8x8(g.alpha), db = simdgroup_float8x8(g.beta);\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" device float* o = (device float*)C + (ulong)(m0 + sr + i * 8u) * g.ldc + n0 + sc + j * 8u;\n"
" simdgroup_float8x8 r = acc[i][j];\n"
" if (g.alpha != 1.0f) simdgroup_multiply(r, r, da);\n"
" if (g.beta != 0.0f) {\n"
" simdgroup_float8x8 old;\n"
" simdgroup_load(old, o, g.ldc);\n"
" simdgroup_multiply_accumulate(r, old, db, r);\n"
" }\n"
" simdgroup_store(r, o, g.ldc);\n"
" }\n"
" return;\n"
" }\n"
" // an edge tile or a 16-bit c: each simdgroup's fragments through its own\n"
" // 8 x 8 of threadgroup memory, then element by element\n"
" threadgroup float* Cs = As + sg * 64u;\n"
" uint lane = tid % 32u;\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" simdgroup_store(acc[i][j], Cs, 8);\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < 2u; q++) {\n"
" uint e = lane + q * 32u;\n"
" uint gm = m0 + sr + i * 8u + e / 8u, gn = n0 + sc + j * 8u + e % 8u;\n"
" if (gm < Mx && gn < Nx) {\n"
" ulong o = (ulong)gm * g.ldc + gn;\n"
" float v = g.alpha * Cs[e];\n"
" if (g.beta != 0.0f) v += g.beta * mk_ld_any(C, o, g.ct); // beta 0: c is not read\n"
" mk_st(C, o, g.ct, v);\n"
" }\n"
" }\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
"}\n"
"#if defined(DK_VULKAN)\n"
"#define MK_GEMM(name, T, TA, TB, BM, BN, BK) \\\n"
"[shader(\"compute\")] [numthreads(128, 1, 1)] \\\n"
"void name(uniform DkGemm g, uint3 tg : SV_GroupID, uint tid : SV_GroupIndex) { \\\n"
" mk_gemm<T, TA, TB, BM, BN, BK>(g, tg, tid, tid / 32u, TGPtr(0u), TGPtr(BM * BK)); \\\n"
"}\n"
"#else\n"
"#define MK_GEMM(name, T, TA, TB, BM, BN, BK) \\\n"
"kernel void name(constant DkGemm& g [[buffer(0)]], uint3 tg [[threadgroup_position_in_grid]], \\\n"
" uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { \\\n"
" threadgroup float As[BM * BK], Bs[BK * BN]; \\\n"
" mk_gemm<T, TA, TB, BM, BN, BK>(g, tg, tid, sg, As, Bs); \\\n"
"}\n"
"#endif\n"
"#define MK_GEMMS(T, sfx) \\\n"
" MK_GEMM(mk_gemm32_nn_##sfx, T, false, false, 32u, 32u, 32u) \\\n"
" MK_GEMM(mk_gemm32_tn_##sfx, T, true, false, 32u, 32u, 32u) \\\n"
" MK_GEMM(mk_gemm32_nt_##sfx, T, false, true, 32u, 32u, 32u) \\\n"
" MK_GEMM(mk_gemm32_tt_##sfx, T, true, true, 32u, 32u, 32u) \\\n"
" MK_GEMM(mk_gemm64_nn_##sfx, T, false, false, 64u, 64u, 16u) \\\n"
" MK_GEMM(mk_gemm64_tn_##sfx, T, true, false, 64u, 64u, 16u) \\\n"
" MK_GEMM(mk_gemm64_nt_##sfx, T, false, true, 64u, 64u, 16u) \\\n"
" MK_GEMM(mk_gemm64_tt_##sfx, T, true, true, 64u, 64u, 16u)\n"
"MK_GEMMS(0u, f32)\n"
"MK_GEMMS(1u, f16)\n"
"MK_GEMMS(2u, bf16)\n"
"// Convolutions as implicit GEMMs (MpsBackend::conv: MIOpen / cuDNN's role),\n"
"// the column matrix gathered tile by tile instead of written out by im2col.\n"
"// In the terms of a plain convolution -- input I [B, Ci, Hi, Wi], weights\n"
"// W [Co, Ci, kH, kW], output O [B, Co, Ho, Wo] (a transposed convolution is\n"
"// one with I and O swapped) --\n"
"// mk_conv_fwd O = conv(I, W): m = co, n = (b, oy, ox), k = (ci, ky, kx)\n"
"// mk_conv_bwd_x I = conv^T(O, W): m = ci, n = (b, y, x), k = (co, ky, kx),\n"
"// a gather over the windows holding (y, x), no atomics\n"
"// mk_conv_bwd_w W += O . I-columns: m = co, n = (ci, ky, kx), k = (b, oy, ox),\n"
"// split over k (grid z), the slices added atomically (fp32)\n"
"// Gathered operands load from a clamped address and select 0 outside the\n"
"// image, so neighbouring threads never diverge on a load.\n"
"struct DkConv {\n"
" ulong x, w, y; // I, W, O\n"
" uint B, Ci, Hi, Wi, Co, Ho, Wo, kH, kW, sH, sW, pH, pW, dH, dW;\n"
" uint kSplit; // mk_conv_bwd_w: k per grid z slice\n"
"};\n"
"// q / d and q % d through a reciprocal, for q < 2^31\n"
"struct MkDiv {\n"
" uint d; float inv;\n"
" inline uint div(uint q) const { return dk_udiv(q, d, inv); }\n"
"};\n"
"inline MkDiv mk_div(uint d) { MkDiv r; r.d = d; r.inv = 1.0f / (float)d; return r; }\n"
"\n"
"// Writes the tile (m0, n0) of an [M, B x P] product into a [B, M, P] tensor\n"
"// (O or I): straight from the matrix units when the tile is whole and P a\n"
"// multiple of 8, else element by element.\n"
"template <uint T, uint BM, uint BN, uint SM, uint SN>\n"
"inline void mk_conv_store(thread simdgroup_float8x8 (&acc)[SM][SN], ulong out, uint M, uint Ntot, uint P, MkDiv dP,\n"
" uint m0, uint n0, uint sr, uint sc, uint tid, uint sg, threadgroup float* scratch) {\n"
" if (T == 0u && m0 + BM <= M && n0 + BN <= Ntot && P % 8u == 0u) {\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" uint n = n0 + sc + j * 8u, b = dP.div(n), p = n - b * P, m = m0 + sr + i * 8u;\n"
" simdgroup_store(acc[i][j], (device float*)out + ((ulong)b * M + m) * P + p, P);\n"
" }\n"
" return;\n"
" }\n"
" threadgroup float* Cs = scratch + sg * 64u;\n"
" uint lane = tid % 32u;\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" simdgroup_store(acc[i][j], Cs, 8);\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < 2u; q++) {\n"
" uint e = lane + q * 32u, m = m0 + sr + i * 8u + e / 8u, n = n0 + sc + j * 8u + e % 8u;\n"
" if (m < M && n < Ntot) {\n"
" uint b = dP.div(n), p = n - b * P;\n"
" mk_st(out, ((ulong)b * M + m) * P + p, T, Cs[e]);\n"
" }\n"
" }\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
"}\n"
"\n"
"// The tile shapes (BM x BN, k steps of BK, four simdgroups as (4 / WN) x WN):\n"
"// 32 x 32 and 64 x 64, and for few output channels 32 x 128 and 8 x 128.\n"
"// When the channels summed over (Ci forward, Co for the input gradient) are\n"
"// a multiple of BK, k runs over (tap, channel) instead of (channel, tap): a\n"
"// step of k is then one filter tap and BK channels, so a thread finds its\n"
"// pixel's tap position once per step and its loads differ by a channel\n"
"// stride. The weight gradient's k = (b, oy, ox) likewise takes whole steps\n"
"// within one image (and one row of it when Wo allows).\n"
"#define MK_TILE_SETUP \\\n"
" constexpr uint WM = 4u / WN, SM = BM / (WM * 8u), SN = BN / (WN * 8u); \\\n"
" uint m0 = tg.y * BM, n0 = tg.x * BN, sr = (sg / WN) * SM * 8u, sc = (sg % WN) * SN * 8u; \\\n"
" simdgroup_float8x8 acc[SM][SN]; \\\n"
" _Pragma(\"clang loop unroll(full)\") \\\n"
" for (uint i = 0; i < SM; i++) \\\n"
" _Pragma(\"clang loop unroll(full)\") \\\n"
" for (uint j = 0; j < SN; j++) acc[i][j] = simdgroup_float8x8(0.0f);\n"
"\n"
"template <uint T, uint BM, uint BN, uint BK, uint WN>\n"
"inline void mk_conv_fwd(constant DkConv& c, uint3 tg, uint tid, uint sg, threadgroup float* As, threadgroup float* Bs) {\n"
" MK_TILE_SETUP\n"
" uint KK = c.kH * c.kW, M = c.Co, K = c.Ci * KK, Po = c.Ho * c.Wo, Ntot = c.B * Po, HW = c.Hi * c.Wi;\n"
" bool cb = c.Ci % BK == 0u;\n"
" MkDiv dKK = mk_div(KK), dKW = mk_div(c.kW), dPo = mk_div(Po), dWo = mk_div(c.Wo), dCi = mk_div(c.Ci);\n"
" // this thread's column of the B tile: output pixel n = (b, oy, ox), fixed\n"
" uint n = min(n0 + tid % BN, Ntot - 1u), b = dPo.div(n), op = n - b * Po, oy = dWo.div(op), ox = op - oy * c.Wo;\n"
" int y0 = (int)(oy * c.sH) - (int)c.pH, x0 = (int)(ox * c.sW) - (int)c.pW;\n"
" bool nIn = n0 + tid % BN < Ntot;\n"
" ulong img = c.x + (ulong)b * c.Ci * HW * (T == 0u ? 4ul : 2ul);\n"
" for (uint k0 = 0; k0 < K; k0 += BK) {\n"
" if (cb) {\n"
" uint tap = dCi.div(k0), ci0 = k0 - tap * c.Ci, ky = dKW.div(tap), kx = tap - ky * c.kW;\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = e / BK, cc = e % BK, gm = m0 + r;\n"
" float v = mk_ld<T>(c.w, (min(gm, M - 1u) * c.Ci + ci0 + cc) * KK + tap);\n"
" As[r * BK + cc] = gm < M ? v : 0.0f;\n"
" }\n"
" int y = y0 + (int)(ky * c.dH), x = x0 + (int)(kx * c.dW);\n"
" bool ok = nIn && y >= 0 && y < (int)c.Hi && x >= 0 && x < (int)c.Wi;\n"
" uint pix = (uint)clamp(y, 0, (int)c.Hi - 1) * c.Wi + (uint)clamp(x, 0, (int)c.Wi - 1);\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint r = (tid + q * 128u) / BN;\n"
" float v = mk_ld<T>(img, (ci0 + r) * HW + pix);\n"
" Bs[r * BN + tid % BN] = ok ? v : 0.0f;\n"
" }\n"
" } else {\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = e / BK, cc = e % BK, gm = m0 + r, gk = k0 + cc;\n"
" float v = mk_ld<T>(c.w, min(gm, M - 1u) * K + min(gk, K - 1u));\n"
" As[r * BK + cc] = gm < M && gk < K ? v : 0.0f;\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint r = (tid + q * 128u) / BN, gk = k0 + r;\n"
" uint kc = min(gk, K - 1u), ci = dKK.div(kc), kr = kc - ci * KK, ky = dKW.div(kr), kx = kr - ky * c.kW;\n"
" int y = y0 + (int)(ky * c.dH), x = x0 + (int)(kx * c.dW);\n"
" bool ok = nIn && gk < K && y >= 0 && y < (int)c.Hi && x >= 0 && x < (int)c.Wi;\n"
" uint pix = (uint)clamp(y, 0, (int)c.Hi - 1) * c.Wi + (uint)clamp(x, 0, (int)c.Wi - 1);\n"
" float v = mk_ld<T>(img, ci * HW + pix);\n"
" Bs[r * BN + tid % BN] = ok ? v : 0.0f;\n"
" }\n"
" }\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" mk_mma<BM, BN, BK, SM, SN>(As, Bs, acc, sr, sc);\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
" mk_conv_store<T, BM, BN, SM, SN>(acc, c.y, M, Ntot, Po, dPo, m0, n0, sr, sc, tid, sg, As);\n"
"}\n"
"\n"
"// The input gradient, by stride phase (grid z, without dilation): the input\n"
"// pixels y = y0 + j sH (and likewise x) whose windows reach them through the\n"
"// taps ky = ry + i sH only, so no step of k is spent on taps that miss --\n"
"// 3 in 4 of them for a stride-2 transposed convolution's forward pass.\n"
"template <uint T, uint BM, uint BN, uint BK, uint WN>\n"
"inline void mk_conv_bwd_x(constant DkConv& c, uint3 tg, uint tid, uint sg, threadgroup float* As, threadgroup float* Bs) {\n"
" MK_TILE_SETUP\n"
" uint KK = c.kH * c.kW, M = c.Ci, Pi = c.Hi * c.Wi, Po = c.Ho * c.Wo;\n"
" bool phased = c.dH == 1u && c.dW == 1u;\n"
" // the phase: taps ky = ry + i sH, kx = rx + j sW; pixels from (yb, xb) in steps of the stride\n"
" uint sH = phased ? c.sH : 1u, sW = phased ? c.sW : 1u;\n"
" uint ry = tg.z / sW, rx = tg.z - ry * sW;\n"
" uint yb = (ry + sH - c.pH % sH) % sH, xb = (rx + sW - c.pW % sW) % sW;\n"
" uint py = yb < c.Hi ? (c.Hi - yb + sH - 1u) / sH : 0u, px = xb < c.Wi ? (c.Wi - xb + sW - 1u) / sW : 0u;\n"
" uint ty = phased ? (ry < c.kH ? (c.kH - ry + sH - 1u) / sH : 0u) : c.kH;\n"
" uint tx = phased ? (rx < c.kW ? (c.kW - rx + sW - 1u) / sW : 0u) : c.kW;\n"
" uint Pp = py * px, Ntot = c.B * Pp, K = c.Co * ty * tx;\n"
" if (n0 >= Ntot) return; // a phase with fewer pixels: the whole threadgroup\n"
" bool cb = c.Co % BK == 0u;\n"
" MkDiv dPp = mk_div(max(Pp, 1u)), dPx = mk_div(max(px, 1u)), dTT = mk_div(max(ty * tx, 1u)), dTx = mk_div(max(tx, 1u)),\n"
" dSH = mk_div(c.sH), dSW = mk_div(c.sW), dCo = mk_div(c.Co);\n"
" // this thread's column: input pixel n = (b, y, x)\n"
" uint n = min(n0 + tid % BN, Ntot - 1u), b = dPp.div(n), ip = n - b * Pp, jy = dPx.div(ip), jx = ip - jy * px;\n"
" uint y = yb + jy * sH, x = xb + jx * sW;\n"
" int ny0 = (int)(y + c.pH), nx0 = (int)(x + c.pW);\n"
" bool nIn = n0 + tid % BN < Ntot;\n"
" ulong img = c.y + (ulong)b * c.Co * Po * (T == 0u ? 4ul : 2ul);\n"
" for (uint k0 = 0; k0 < K; k0 += BK) {\n"
" if (cb) {\n"
" // A (m = ci, k = (tap, co)) = W[co, ci, ky, kx]: neighbouring threads take neighbouring ci\n"
" uint t = dCo.div(k0), co0 = k0 - t * c.Co, iy = dTx.div(t), ix = t - iy * tx;\n"
" uint ky = phased ? ry + iy * sH : iy, kx = phased ? rx + ix * sW : ix;\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = e % BM, cc = e / BM, gm = m0 + r;\n"
" float v = mk_ld<T>(c.w, ((co0 + cc) * c.Ci + min(gm, M - 1u)) * KK + ky * c.kW + kx);\n"
" As[cc * BM + r] = gm < M ? v : 0.0f; // k-major\n"
" }\n"
" // B (k = (tap, co), n = (b, y, x)) = O[b, co, oy, ox], the window at\n"
" // (oy, ox) holding (y, x) through this tap\n"
" int ny = ny0 - (int)(ky * c.dH), nx = nx0 - (int)(kx * c.dW);\n"
" uint oy = dSH.div((uint)max(ny, 0)), ox = dSW.div((uint)max(nx, 0));\n"
" bool ok = nIn && ny >= 0 && nx >= 0 && oy * c.sH == (uint)ny && ox * c.sW == (uint)nx && oy < c.Ho && ox < c.Wo;\n"
" uint pix = min(oy, c.Ho - 1u) * c.Wo + min(ox, c.Wo - 1u);\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint r = (tid + q * 128u) / BN;\n"
" float v = mk_ld<T>(img, (co0 + r) * Po + pix);\n"
" Bs[r * BN + tid % BN] = ok ? v : 0.0f;\n"
" }\n"
" } else {\n"
" // k = (co, tap)\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = e % BM, cc = e / BM, gm = m0 + r, gk = k0 + cc;\n"
" uint kc = min(gk, K - 1u), co = dTT.div(kc), t = kc - co * ty * tx, iy = dTx.div(t), ix = t - iy * tx;\n"
" uint ky = phased ? ry + iy * sH : iy, kx = phased ? rx + ix * sW : ix;\n"
" float v = mk_ld<T>(c.w, (co * c.Ci + min(gm, M - 1u)) * KK + ky * c.kW + kx);\n"
" As[cc * BM + r] = gm < M && gk < K ? v : 0.0f;\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint r = (tid + q * 128u) / BN, gk = k0 + r;\n"
" uint kc = min(gk, K - 1u), co = dTT.div(kc), t = kc - co * ty * tx, iy = dTx.div(t), ix = t - iy * tx;\n"
" uint ky = phased ? ry + iy * sH : iy, kx = phased ? rx + ix * sW : ix;\n"
" int ny = ny0 - (int)(ky * c.dH), nx = nx0 - (int)(kx * c.dW);\n"
" uint oy = dSH.div((uint)max(ny, 0)), ox = dSW.div((uint)max(nx, 0));\n"
" bool ok = nIn && gk < K && ny >= 0 && nx >= 0 && oy * c.sH == (uint)ny && ox * c.sW == (uint)nx && oy < c.Ho && ox < c.Wo;\n"
" float v = mk_ld<T>(img, co * Po + min(oy, c.Ho - 1u) * c.Wo + min(ox, c.Wo - 1u));\n"
" Bs[r * BN + tid % BN] = ok ? v : 0.0f;\n"
" }\n"
" }\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" mk_mma<BM, BN, BK, SM, SN, true>(As, Bs, acc, sr, sc);\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
" if (sH == 1u && sW == 1u) {\n"
" mk_conv_store<T, BM, BN, SM, SN>(acc, c.x, M, Ntot, Pi, dPp, m0, n0, sr, sc, tid, sg, As);\n"
" return;\n"
" }\n"
" // a phase's pixels are strided: element by element\n"
" threadgroup float* Cs = As + sg * 64u;\n"
" uint lane = tid % 32u;\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" simdgroup_store(acc[i][j], Cs, 8);\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < 2u; q++) {\n"
" uint e = lane + q * 32u, m = m0 + sr + i * 8u + e / 8u, nn = n0 + sc + j * 8u + e % 8u;\n"
" if (m < M && nn < Ntot) {\n"
" uint bb = dPp.div(nn), pp = nn - bb * Pp, qy = dPx.div(pp), qx = pp - qy * px;\n"
" mk_st(c.x, ((ulong)bb * M + m) * Pi + (yb + qy * sH) * c.Wi + xb + qx * sW, T, Cs[e]);\n"
" }\n"
" }\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
"}\n"
"\n"
"// fp32 only (atomics); W must be zeroed first\n"
"template <uint BM, uint BN, uint BK, uint WN>\n"
"inline void mk_conv_bwd_w(constant DkConv& c, uint3 tg, uint tid, uint sg, threadgroup float* As, threadgroup float* Bs) {\n"
" MK_TILE_SETUP\n"
" uint KK = c.kH * c.kW, M = c.Co, Ntot = c.Ci * KK, Po = c.Ho * c.Wo, K = c.B * Po, HW = c.Hi * c.Wi;\n"
" uint kBeg = tg.z * c.kSplit, kEnd = min(K, kBeg + c.kSplit);\n"
" bool img1 = Po % BK == 0u && c.kSplit % BK == 0u, row1 = img1 && c.Wo % BK == 0u; // a step within one image / row\n"
" MkDiv dPo = mk_div(Po), dWo = mk_div(c.Wo), dKK = mk_div(KK), dKW = mk_div(c.kW);\n"
" // this thread's column of B: n = (ci, ky, kx)\n"
" uint n = min(n0 + tid % BN, Ntot - 1u), ci = dKK.div(n), kr = n - ci * KK, ky = dKW.div(kr), kx = kr - ky * c.kW;\n"
" int yo = (int)(ky * c.dH) - (int)c.pH, xo = (int)(kx * c.dW) - (int)c.pW;\n"
" bool nIn = n0 + tid % BN < Ntot;\n"
" const device float* O = (const device float*)c.y;\n"
" const device float* I = (const device float*)c.x;\n"
" for (uint k0 = kBeg; k0 < kEnd; k0 += BK) {\n"
" if (img1) {\n"
" uint b = dPo.div(k0), p0 = k0 - b * Po;\n"
" // A (m = co, k = (b, p)) = O[b, co, p]\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = e / BK, cc = e % BK, gm = m0 + r;\n"
" float v = O[((ulong)b * c.Co + min(gm, M - 1u)) * Po + p0 + cc];\n"
" As[r * BK + cc] = gm < M ? v : 0.0f;\n"
" }\n"
" const device float* src = I + ((ulong)b * c.Ci + ci) * HW;\n"
" uint oy0 = dWo.div(p0), ox0 = p0 - oy0 * c.Wo;\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint r = (tid + q * 128u) / BN;\n"
" uint oy = oy0, ox = ox0 + r;\n"
" if (!row1) { uint p = p0 + r; oy = dWo.div(p); ox = p - oy * c.Wo; }\n"
" int y = (int)(oy * c.sH) + yo, x = (int)(ox * c.sW) + xo;\n"
" bool ok = nIn && y >= 0 && y < (int)c.Hi && x >= 0 && x < (int)c.Wi;\n"
" float v = src[(uint)clamp(y, 0, (int)c.Hi - 1) * c.Wi + (uint)clamp(x, 0, (int)c.Wi - 1)];\n"
" Bs[r * BN + tid % BN] = ok ? v : 0.0f;\n"
" }\n"
" } else {\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BM * BK / 128u; q++) {\n"
" uint e = tid + q * 128u, r = e / BK, cc = e % BK, gm = m0 + r, gk = k0 + cc;\n"
" uint kc = min(gk, kEnd - 1u), b = dPo.div(kc), op = kc - b * Po;\n"
" float v = O[((ulong)b * c.Co + min(gm, M - 1u)) * Po + op];\n"
" As[r * BK + cc] = gm < M && gk < kEnd ? v : 0.0f;\n"
" }\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < BK * BN / 128u; q++) {\n"
" uint r = (tid + q * 128u) / BN, gk = k0 + r;\n"
" uint kc = min(gk, kEnd - 1u), b = dPo.div(kc), op = kc - b * Po, oy = dWo.div(op), ox = op - oy * c.Wo;\n"
" int y = (int)(oy * c.sH) + yo, x = (int)(ox * c.sW) + xo;\n"
" bool ok = nIn && gk < kEnd && y >= 0 && y < (int)c.Hi && x >= 0 && x < (int)c.Wi;\n"
" float v = I[(((ulong)b * c.Ci + ci) * c.Hi + (uint)clamp(y, 0, (int)c.Hi - 1)) * c.Wi + (uint)clamp(x, 0, (int)c.Wi - 1)];\n"
" Bs[r * BN + tid % BN] = ok ? v : 0.0f;\n"
" }\n"
" }\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" mk_mma<BM, BN, BK, SM, SN>(As, Bs, acc, sr, sc);\n"
" threadgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
" threadgroup float* Cs = As + sg * 64u;\n"
" uint lane = tid % 32u;\n"
" #pragma clang loop unroll(full)\n"
" for (uint i = 0; i < SM; i++)\n"
" #pragma clang loop unroll(full)\n"
" for (uint j = 0; j < SN; j++) {\n"
" simdgroup_store(acc[i][j], Cs, 8);\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" #pragma clang loop unroll(full)\n"
" for (uint q = 0; q < 2u; q++) {\n"
" uint e = lane + q * 32u, m = m0 + sr + i * 8u + e / 8u, nn = n0 + sc + j * 8u + e % 8u;\n"
" if (m < M && nn < Ntot) atomicAdd((device float*)c.w + (ulong)m * Ntot + nn, Cs[e]);\n"
" }\n"
" simdgroup_barrier(mem_flags::mem_threadgroup);\n"
" }\n"
"}\n"
"\n"
"// mk_conv_<op><shape>_<type>: shapes 32 (32 x 32), 64 (64 x 64), w (32 x 128), s (8 x 128)\n"
"#if defined(DK_VULKAN)\n"
"#define MK_CONV(name, call, BM, BN) \\\n"
"[shader(\"compute\")] [numthreads(128, 1, 1)] \\\n"
"void name(uniform DkConv c, uint3 tg : SV_GroupID, uint tid : SV_GroupIndex) { \\\n"
" call(c, tg, tid, tid / 32u, TGPtr(0u), TGPtr(BM * 16u > 8u * 64u ? BM * 16u : 8u * 64u)); \\\n"
"}\n"
"#else\n"
"#define MK_CONV(name, call, BM, BN) \\\n"
"kernel void name(constant DkConv& c [[buffer(0)]], uint3 tg [[threadgroup_position_in_grid]], \\\n"
" uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { \\\n"
" threadgroup float As[BM * 16u > 8u * 64u ? BM * 16u : 8u * 64u], Bs[16u * BN]; \\\n"
" call(c, tg, tid, sg, As, Bs); \\\n"
"}\n"
"#endif\n"
"#define MK_CONV_SHAPES(op, T, sfx) \\\n"
" MK_CONV(mk_conv_##op##32_##sfx, (mk_conv_##op<T, 32u, 32u, 16u, 2u>), 32u, 32u) \\\n"
" MK_CONV(mk_conv_##op##64_##sfx, (mk_conv_##op<T, 64u, 64u, 16u, 2u>), 64u, 64u) \\\n"
" MK_CONV(mk_conv_##op##w_##sfx, (mk_conv_##op<T, 32u, 128u, 16u, 4u>), 32u, 128u) \\\n"
" MK_CONV(mk_conv_##op##s_##sfx, (mk_conv_##op<T, 8u, 128u, 16u, 4u>), 8u, 128u)\n"
"MK_CONV_SHAPES(fwd, 0u, f32)\n"
"MK_CONV_SHAPES(fwd, 1u, f16)\n"
"MK_CONV_SHAPES(fwd, 2u, bf16)\n"
"MK_CONV_SHAPES(bwd_x, 0u, f32)\n"
"MK_CONV_SHAPES(bwd_x, 1u, f16)\n"
"MK_CONV_SHAPES(bwd_x, 2u, bf16)\n"
"MK_CONV(mk_conv_bwd_w32_f32, (mk_conv_bwd_w<32u, 32u, 16u, 2u>), 32u, 32u)\n"
"MK_CONV(mk_conv_bwd_w64_f32, (mk_conv_bwd_w<64u, 64u, 16u, 2u>), 64u, 64u)\n"
"MK_CONV(mk_conv_bwd_ww_f32, (mk_conv_bwd_w<32u, 128u, 16u, 4u>), 32u, 128u)\n"
"MK_CONV(mk_conv_bwd_ws_f32, (mk_conv_bwd_w<8u, 128u, 16u, 4u>), 8u, 128u)\n"
"#if defined(DK_VULKAN)\n"
"// dst[i] = src[i] for n 32-bit words: p0 src, p2 dst, u7 = n (VulkanBackend's\n"
"// copies and staged uploads: a dispatch, where MoltenVK would switch to a\n"
"// blit encoder for vkCmdCopyBuffer)\n"
"DK_KERNEL(mk_copy) { GS(i, N) ((uint*)a.p[2])[i] = ((uint*)a.p[0])[i]; }\n"
"#endif\n"
"#endif\n"
"#if defined(__METAL_VERSION__)\n"
"// dst[i] = src[i] for n 32-bit words: p0 src, p2 dst, u7 = n (bit-exact,\n"
"// unlike a float copy of words holding two 16-bit values)\n"
"kernel void mk_copy(constant DevArgs_& a [[buffer(0)]], uint i [[thread_position_in_grid]]) {\n"
" if (i < a.u[7]) ((device uint*)a.p[2])[i] = ((const device uint*)a.p[0])[i];\n"
"}\n"
"// dst[i] = src[i] for n words of host data that came with the dispatch\n"
"// (setBytes): an upload ordered with the queued work. p2 dst, u7 = n\n"
"kernel void mk_upload(constant DevArgs_& a [[buffer(0)]], constant uint* src [[buffer(1)]],\n"
" uint i [[thread_position_in_grid]]) {\n"
" if (i < a.u[7]) ((device uint*)a.p[2])[i] = src[i];\n"
"}\n"
"#endif\n"
"\n";
}
} // namespace std