reductions: frobenius/cosine/triplet (CUDA) + Triton parity for the reduction family

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<float> predictions(64 * 64);
-    tensor::Buffer<float> targets(64 * 64);
+    tensor::Buffer<float> predictions(M * N);
+    tensor::Buffer<float> targets(M * N);
     tensor::Buffer<float> output(1);
     tensor::Buffer<size_t> shape(ndim);
 
     predictions.fill_random();
     targets.fill_random();
-    shape.set({64, 64});
+    shape.set({M, N});
 
     BENCHMARK(solution(predictions, targets, output, shape, ndim));
 

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

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 <cuda_runtime.h>
-#include <cuda_fp16.h>
-#include <cuda_fp8.h>
-#include <cstdint>
+#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<size_t>(row) * d;
+    const float* t = T + static_cast<size_t>(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<BLOCK, reduce_sum, id_zero>::apply(dot);
+    np  = SmemTreeReduce<BLOCK, reduce_sum, id_zero>::apply(np);
+    nt  = SmemTreeReduce<BLOCK, reduce_sum, id_zero>::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<<<n, BLOCK>>>(predictions, targets, output, static_cast<int>(d));
 }

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 <cuda_runtime.h>
-#include <cuda_fp16.h>
-#include <cuda_fp8.h>
-#include <cstdint>
+#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<size_t>(blockIdx.x * blockDim.x + threadIdx.x) * 4;
+    float r = 0.0f;
+    if (base + 3 < n) {
+        float4 v = *reinterpret_cast<const float4*>(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<BLOCK, reduce_sum, id_zero>::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<size_t>(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<<<g1, BLOCK>>>(X, scratch, size);
+    int g2 = (size + BLOCK - 1) / BLOCK;
+    normalize_kernel<<<g2, BLOCK>>>(X, Y, scratch, size);
 }

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 <cuda_runtime.h>
-#include <cuda_fp16.h>
-#include <cuda_fp8.h>
-#include <cstdint>
+#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<size_t>(b) * E;
+    const float* pp = P + static_cast<size_t>(b) * E;
+    const float* nn = N + static_cast<size_t>(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<BLOCK, reduce_sum, id_zero>::apply(dp);
+    dn = SmemTreeReduce<BLOCK, reduce_sum, id_zero>::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<<<B, BLOCK>>>(anchor, positive, negative, loss,
+                                 static_cast<int>(E), margin, 1.0f / static_cast<float>(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)

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)

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)

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)

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)