diff --git a/kernel-harnesses/mse-loss.cu b/kernel-harnesses/mse-loss.cu index c7d673e..299ac1f 100644 --- a/kernel-harnesses/mse-loss.cu +++ b/kernel-harnesses/mse-loss.cu @@ -6,15 +6,17 @@ int main() { tensor::begin("mse-loss"); size_t ndim = 2; + size_t M = tensor::bench_size("M", 64); // rows + size_t N = tensor::bench_size("N", 64); // cols - tensor::Buffer predictions(64 * 64); - tensor::Buffer targets(64 * 64); + tensor::Buffer predictions(M * N); + tensor::Buffer targets(M * N); tensor::Buffer output(1); tensor::Buffer shape(ndim); predictions.fill_random(); targets.fill_random(); - shape.set({64, 64}); + shape.set({M, N}); BENCHMARK(solution(predictions, targets, output, shape, ndim)); diff --git a/run-bench.sh b/run-bench.sh index 657ad72..394fb4a 100755 --- a/run-bench.sh +++ b/run-bench.sh @@ -99,6 +99,9 @@ bench_profile() { # softmax family (rows M × reduced dim N). dim-reduce kernels are stubs # and still size via a fixed shape array — they'll need the same wiring. softmax|log-softmax) echo "TENSOR_M=65536 TENSOR_N=1024" ;; + # mse-loss reduces over M*N (rows*cols) to a scalar; set both dims so the + # default (N only) doesn't leave M=64 or blow up memory. + mse-loss) echo "TENSOR_M=8192 TENSOR_N=8192" ;; # distances / margin losses over rows×cols triplet-margin) echo "TENSOR_B=8192 TENSOR_E=8192" ;; # graphs (O(n^2)/O(n^3)) — keep moderate diff --git a/solutions-cuda/cosine-similarity.cu b/solutions-cuda/cosine-similarity.cu index 9c2d6ed..b57bd79 100644 --- a/solutions-cuda/cosine-similarity.cu +++ b/solutions-cuda/cosine-similarity.cu @@ -1,16 +1,34 @@ -// Solution stub for "cosine-similarity". -// TODO: implement the body below. The signature is derived from -// kernel-harnesses/cosine-similarity.cu and must stay in sync with it. -// -// Build it (the build system auto-picks this file): -// make build/bin/cosine-similarity.exe -// ./build/bin/cosine-similarity.exe +// cosine-similarity per row: output[i] = dot(p_i, t_i) / (||p_i|| * ||t_i||), +// for n rows of length d. A PER-ROW, MULTI-ACCUMULATOR reduction: one block per +// row reduces three sums at once (dot, ||p||^2, ||t||^2) with the shared +// SmemTreeReduce, then thread 0 writes the row's result. #include -#include -#include -#include +#include "../kernel-implementation/reduction.cuh" // SmemTreeReduce, reduce_sum, id_zero + +#define BLOCK 256 + +__global__ void cosine_kernel(const float* __restrict__ P, const float* __restrict__ T, + float* __restrict__ out, int d) { + const int row = blockIdx.x; + const float* p = P + static_cast(row) * d; + const float* t = T + static_cast(row) * d; + + float dot = 0.0f, np = 0.0f, nt = 0.0f; + for (int j = threadIdx.x; j < d; j += BLOCK) { + float a = p[j], b = t[j]; + dot += a * b; np += a * a; nt += b * b; + } + // Three block reductions (all threads reach each one uniformly). + dot = SmemTreeReduce::apply(dot); + np = SmemTreeReduce::apply(np); + nt = SmemTreeReduce::apply(nt); + + if (threadIdx.x == 0) + out[row] = dot * rsqrtf(np) * rsqrtf(nt); +} // Note: all pointer arguments are device pointers. -extern "C" void solution(const float* predictions, const float* targets, float* output, size_t n, size_t d) { - // TODO: implement cosine-similarity +extern "C" void solution(const float* predictions, const float* targets, float* output, + size_t n, size_t d) { + cosine_kernel<<>>(predictions, targets, output, static_cast(d)); } diff --git a/solutions-cuda/frobenius-norm.cu b/solutions-cuda/frobenius-norm.cu index fb69c1f..a10c305 100644 --- a/solutions-cuda/frobenius-norm.cu +++ b/solutions-cuda/frobenius-norm.cu @@ -1,16 +1,41 @@ -// Solution stub for "frobenius-norm". -// TODO: implement the body below. The signature is derived from -// kernel-harnesses/frobenius-norm.cu and must stay in sync with it. -// -// Build it (the build system auto-picks this file): -// make build/bin/frobenius-norm.exe -// ./build/bin/frobenius-norm.exe +// frobenius-norm: Y = X / ||X||_F, ||X||_F = sqrt(sum_i X[i]^2). +// A GLOBAL reduction (sum of squares) followed by a broadcast/normalize pass — +// the two-pass shape. Pass 1 reduces to a scalar with the shared SmemTreeReduce +// + one atomic per block; pass 2 rescales every element. #include -#include -#include -#include +#include "../kernel-implementation/reduction.cuh" // SmemTreeReduce, reduce_sum, id_zero + +#define BLOCK 512 + +// Pass 1: scratch[0] = sum_i X[i]^2 (float4-vectorized + scalar tail). +__global__ void sumsq_kernel(const float* __restrict__ X, float* scratch, size_t n) { + size_t base = static_cast(blockIdx.x * blockDim.x + threadIdx.x) * 4; + float r = 0.0f; + if (base + 3 < n) { + float4 v = *reinterpret_cast(X + base); + r = v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w; + } else { + for (size_t i = base; i < n; ++i) r += X[i] * X[i]; + } + r = SmemTreeReduce::apply(r); // all threads call uniformly + if (threadIdx.x == 0) atomicAdd(scratch, r); +} + +// Pass 2: Y[i] = X[i] * rsqrt(sum_sq) = X[i] / ||X||_F. +__global__ void normalize_kernel(const float* __restrict__ X, float* __restrict__ Y, + const float* scratch, size_t n) { + size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < n) Y[i] = X[i] * rsqrtf(scratch[0]); +} // Note: all pointer arguments are device pointers. extern "C" void solution(const float* X, float* Y, size_t size) { - // TODO: implement frobenius-norm + static float* scratch = nullptr; // one-float device accumulator (reused) + if (!scratch) cudaMalloc(&scratch, sizeof(float)); + cudaMemset(scratch, 0, sizeof(float)); + + int g1 = ((size + 3) / 4 + BLOCK - 1) / BLOCK; + sumsq_kernel<<>>(X, scratch, size); + int g2 = (size + BLOCK - 1) / BLOCK; + normalize_kernel<<>>(X, Y, scratch, size); } diff --git a/solutions-cuda/triplet-margin.cu b/solutions-cuda/triplet-margin.cu index 60365a9..37a1a55 100644 --- a/solutions-cuda/triplet-margin.cu +++ b/solutions-cuda/triplet-margin.cu @@ -1,16 +1,39 @@ -// Solution stub for "triplet-margin". -// TODO: implement the body below. The signature is derived from -// kernel-harnesses/triplet-margin.cu and must stay in sync with it. -// -// Build it (the build system auto-picks this file): -// make build/bin/triplet-margin.exe -// ./build/bin/triplet-margin.exe +// triplet-margin loss: mean_b max(0, ||a_b - p_b|| - ||a_b - n_b|| + margin), +// Euclidean distances, B triplets of dim E. A TWO-LEVEL reduction: one block per +// triplet reduces two squared distances over E (shared SmemTreeReduce), thread 0 +// forms the per-triplet loss and atomically adds loss/B to the global scalar. #include -#include -#include -#include +#include "../kernel-implementation/reduction.cuh" // SmemTreeReduce, reduce_sum, id_zero + +#define BLOCK 256 + +__global__ void triplet_kernel(const float* __restrict__ A, const float* __restrict__ P, + const float* __restrict__ N, float* loss, + int E, float margin, float inv_B) { + const int b = blockIdx.x; + const float* a = A + static_cast(b) * E; + const float* pp = P + static_cast(b) * E; + const float* nn = N + static_cast(b) * E; + + float dp = 0.0f, dn = 0.0f; // ||a-p||^2, ||a-n||^2 + for (int e = threadIdx.x; e < E; e += BLOCK) { + float ap = a[e] - pp[e]; + float an = a[e] - nn[e]; + dp += ap * ap; dn += an * an; + } + dp = SmemTreeReduce::apply(dp); + dn = SmemTreeReduce::apply(dn); + + if (threadIdx.x == 0) { + float v = fmaxf(0.0f, sqrtf(dp) - sqrtf(dn) + margin); + atomicAdd(loss, v * inv_B); + } +} // Note: all pointer arguments are device pointers. -extern "C" void solution(const float* anchor, const float* positive, const float* negative, float* loss, size_t B, size_t E, float margin) { - // TODO: implement triplet-margin +extern "C" void solution(const float* anchor, const float* positive, const float* negative, + float* loss, size_t B, size_t E, float margin) { + cudaMemset(loss, 0, sizeof(float)); // scalar accumulator + triplet_kernel<<>>(anchor, positive, negative, loss, + static_cast(E), margin, 1.0f / static_cast(B)); } diff --git a/solutions-triton/cosine-similarity.py b/solutions-triton/cosine-similarity.py new file mode 100644 index 0000000..9c4e0ec --- /dev/null +++ b/solutions-triton/cosine-similarity.py @@ -0,0 +1,30 @@ +"""Triton solution for `cosine-similarity` — mirrors solutions-cuda/cosine-similarity.cu.""" +import os, sys, torch, triton +import triton.language as tl +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tensor-lib")) +import triton_bench as tb + +@triton.jit +def _k(p_ptr, t_ptr, out_ptr, d, BLOCK: tl.constexpr): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + mask = offs < d + p = tl.load(p_ptr + row * d + offs, mask=mask, other=0.0) + t = tl.load(t_ptr + row * d + offs, mask=mask, other=0.0) + dot = tl.sum(p * t); npp = tl.sum(p * p); ntt = tl.sum(t * t) + tl.store(out_ptr + row, dot / (tl.sqrt(npp) * tl.sqrt(ntt))) + +def solution(p, t, out, n, d): + _k[(n,)](p, t, out, d, BLOCK=triton.next_power_of_2(d)) + +def main(do_check): + n = tb.bench_size("N", 64); d = tb.bench_size("D", 128) + p = tb.rand(n * d); t = tb.rand(n * d); out = torch.empty(n, device="cuda") + tb.benchmark(lambda: solution(p, t, out, n, d)) + tb.preview(out, "output") + if do_check: + pv = p.view(n, d); tv = t.view(n, d) + ref = (pv * tv).sum(1) / (pv.norm(dim=1) * tv.norm(dim=1)) + return tb.check("cosine-similarity", out, ref, rtol=1e-3, atol=1e-4) + return 0 +if __name__ == "__main__": tb.run("cosine-similarity", main) diff --git a/solutions-triton/frobenius-norm.py b/solutions-triton/frobenius-norm.py new file mode 100644 index 0000000..0ccb244 --- /dev/null +++ b/solutions-triton/frobenius-norm.py @@ -0,0 +1,36 @@ +"""Triton solution for `frobenius-norm` — mirrors solutions-cuda/frobenius-norm.cu.""" +import os, sys, torch, triton +import triton.language as tl +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tensor-lib")) +import triton_bench as tb + +@triton.jit +def _sumsq(x_ptr, s_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask, other=0.0) + tl.atomic_add(s_ptr, tl.sum(tl.where(mask, x * x, 0.0))) + +@triton.jit +def _norm(x_ptr, y_ptr, s_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask) + inv = 1.0 / tl.sqrt(tl.load(s_ptr)) + tl.store(y_ptr + offs, x * inv, mask=mask) + +def solution(x, y, n, s): + s.zero_() + BLOCK = 1024; g = (triton.cdiv(n, BLOCK),) + _sumsq[g](x, s, n, BLOCK=BLOCK) + _norm[g](x, y, s, n, BLOCK=BLOCK) + +def main(do_check): + n = tb.bench_size("SIZE", 4096) + x = tb.rand(n); y = torch.empty_like(x); s = torch.zeros(1, device="cuda") + tb.benchmark(lambda: solution(x, y, n, s)) + tb.preview(y, "Y") + if do_check: + return tb.check("frobenius-norm", y, x / x.norm(), rtol=1e-3, atol=1e-4) + return 0 +if __name__ == "__main__": tb.run("frobenius-norm", main) diff --git a/solutions-triton/kl-loss.py b/solutions-triton/kl-loss.py new file mode 100644 index 0000000..b8eb791 --- /dev/null +++ b/solutions-triton/kl-loss.py @@ -0,0 +1,31 @@ +"""Triton solution for `kl-loss` — mirrors solutions-cuda/kl-loss.cu.""" +import os, sys, torch, triton +import triton.language as tl +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tensor-lib")) +import triton_bench as tb + +@triton.jit +def _k(p_ptr, t_ptr, out_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + p = tl.load(p_ptr + offs, mask=mask, other=1.0) + t = tl.load(t_ptr + offs, mask=mask, other=1.0) + r = t * (tl.log(t) - tl.log(p)) + tl.atomic_add(out_ptr, tl.sum(tl.where(mask, r, 0.0)) / n) + +def solution(p, t, out, n): + out.zero_() + BLOCK = 1024 + _k[(triton.cdiv(n, BLOCK),)](p, t, out, n, BLOCK=BLOCK) + +def main(do_check): + n = tb.bench_size("N", 4096) + p = torch.rand(n, device="cuda") * 0.999 + 1e-3 # (0,1] + t = torch.rand(n, device="cuda") * 0.999 + 1e-3 + out = torch.zeros(1, device="cuda") + tb.benchmark(lambda: solution(p, t, out, n)) + tb.preview(out, "output") + if do_check: + return tb.check("kl-loss", out, (t * (t.log() - p.log())).mean().view(1), rtol=1e-2, atol=1e-3) + return 0 +if __name__ == "__main__": tb.run("kl-loss", main) diff --git a/solutions-triton/mse-loss.py b/solutions-triton/mse-loss.py index 0b9770c..0df3a12 100644 --- a/solutions-triton/mse-loss.py +++ b/solutions-triton/mse-loss.py @@ -1,45 +1,30 @@ -"""Triton solution for `huber-loss` — mirrors solutions-cuda/huber-loss.cu.""" -import os -import sys -import torch -import triton +"""Triton solution for `mse-loss` — mirrors solutions-cuda/mse-loss.cu.""" +import os, sys, torch, triton import triton.language as tl - sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tensor-lib")) import triton_bench as tb - @triton.jit -def _kernel(p_ptr, t_ptr, out_ptr, n, BLOCK: tl.constexpr): +def _k(p_ptr, t_ptr, out_ptr, n, BLOCK: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) mask = offs < n p = tl.load(p_ptr + offs, mask=mask, other=0.0) t = tl.load(t_ptr + offs, mask=mask, other=0.0) - a = tl.abs(p - t) - h = tl.where(a < 1.0, 0.5 * a * a, a - 0.5) - h = tl.where(mask, h, 0.0) - tl.atomic_add(out_ptr, tl.sum(h) / n) - + d = p - t + tl.atomic_add(out_ptr, tl.sum(tl.where(mask, d * d, 0.0)) / n) def solution(p, t, out, n): out.zero_() BLOCK = 1024 - _kernel[(triton.cdiv(n, BLOCK),)](p, t, out, n, BLOCK=BLOCK) - + _k[(triton.cdiv(n, BLOCK),)](p, t, out, n, BLOCK=BLOCK) def main(do_check): - n = tb.bench_size("N", 1024) - p = tb.rand(n) - t = tb.rand(n) - out = torch.zeros(1, device="cuda") + # mse-loss reduces over M*N (matches the CUDA harness's 2-D shape). + n = tb.bench_size("M", 64) * tb.bench_size("N", 64) + p = tb.rand(n); t = tb.rand(n); out = torch.zeros(1, device="cuda") tb.benchmark(lambda: solution(p, t, out, n)) tb.preview(out, "output") if do_check: - d = (p - t).abs() - ref = torch.where(d < 1.0, 0.5 * d * d, d - 0.5).mean().view(1) - return tb.check("huber-loss", out, ref, rtol=1e-2, atol=1e-3) + return tb.check("mse-loss", out, ((p - t) ** 2).mean().view(1), rtol=1e-2, atol=1e-3) return 0 - - -if __name__ == "__main__": - tb.run("huber-loss", main) +if __name__ == "__main__": tb.run("mse-loss", main) diff --git a/solutions-triton/triplet-margin.py b/solutions-triton/triplet-margin.py new file mode 100644 index 0000000..56f9021 --- /dev/null +++ b/solutions-triton/triplet-margin.py @@ -0,0 +1,33 @@ +"""Triton solution for `triplet-margin` — mirrors solutions-cuda/triplet-margin.cu.""" +import os, sys, torch, triton +import triton.language as tl +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tensor-lib")) +import triton_bench as tb + +@triton.jit +def _k(a_ptr, p_ptr, n_ptr, loss_ptr, E, margin, inv_B, BLOCK: tl.constexpr): + b = tl.program_id(0) + offs = tl.arange(0, BLOCK) + mask = offs < E + a = tl.load(a_ptr + b * E + offs, mask=mask, other=0.0) + p = tl.load(p_ptr + b * E + offs, mask=mask, other=0.0) + nn = tl.load(n_ptr + b * E + offs, mask=mask, other=0.0) + dp = tl.sum((a - p) * (a - p)); dn = tl.sum((a - nn) * (a - nn)) + v = tl.maximum(0.0, tl.sqrt(dp) - tl.sqrt(dn) + margin) + tl.atomic_add(loss_ptr, v * inv_B) + +def solution(a, p, nn, loss, B, E, margin): + loss.zero_() + _k[(B,)](a, p, nn, loss, E, margin, 1.0 / B, BLOCK=triton.next_power_of_2(E)) + +def main(do_check): + B = tb.bench_size("B", 8); E = tb.bench_size("E", 128); margin = 1.0 + a = tb.rand(B * E); p = tb.rand(B * E); nn = tb.rand(B * E); loss = torch.zeros(1, device="cuda") + tb.benchmark(lambda: solution(a, p, nn, loss, B, E, margin)) + tb.preview(loss, "loss") + if do_check: + av = a.view(B, E); pv = p.view(B, E); nv = nn.view(B, E) + d = torch.clamp((av - pv).norm(dim=1) - (av - nv).norm(dim=1) + margin, min=0) + return tb.check("triplet-margin", loss, d.mean().view(1), rtol=1e-2, atol=1e-3) + return 0 +if __name__ == "__main__": tb.run("triplet-margin", main)