// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Embedding gather + scatter-grad. ids are i32. Forward: one thread per
// (token, channel). Backward: one thread per (row, channel) accumulating
// with device atomic_float — repeated tokens collide, so order is
// nondeterministic (sum of floats); deterministic mode routes embedding
// backward through the CPU reference instead.
#include <metal_stdlib>
using namespace metal;

kernel void embedding_fwd_f32(device const float* W   [[buffer(0)]],
                              device const int*   ids [[buffer(1)]],
                              device float*       out [[buffer(2)]],
                              constant uint&      C   [[buffer(3)]],
                              uint2 gid [[thread_position_in_grid]]) {
    // gid.x = channel, gid.y = token index
    out[ulong(gid.y) * C + gid.x] = W[ulong(ids[gid.y]) * C + gid.x];
}

kernel void embedding_bwd_f32(device const int*     ids  [[buffer(0)]],
                              device const float*   dout [[buffer(1)]],
                              device atomic_float*  dW   [[buffer(2)]],
                              constant uint&        C    [[buffer(3)]],
                              uint2 gid [[thread_position_in_grid]]) {
    const float g = dout[ulong(gid.y) * C + gid.x];
    atomic_fetch_add_explicit(&dW[ulong(ids[gid.y]) * C + gid.x], g,
                              memory_order_relaxed);
}