ggml : add fused SINKHORN_NORM op (Sinkhorn-Knopp normalization)

Author: jadenmach2

ggml/include/ggml.h

@@ -509,6 +509,7 @@ extern "C" {
         GGML_OP_RMS_NORM_BACK,
         GGML_OP_GROUP_NORM,
         GGML_OP_L2_NORM,
+        GGML_OP_SINKHORN_NORM,
 
         GGML_OP_MUL_MAT,
         GGML_OP_MUL_MAT_ID,
@@ -1406,6 +1407,12 @@ extern "C" {
             struct ggml_tensor  * a,
             float                 eps);
 
+    GGML_API struct ggml_tensor * ggml_sinkhorn_norm(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   n_iters,
+            float                 eps);
+
     // a - x
     // b - dy
     GGML_API struct ggml_tensor * ggml_rms_norm_back(

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1833,6 +1833,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_l2_norm(params, tensor);
             } break;
+        case GGML_OP_SINKHORN_NORM:
+            {
+                ggml_compute_forward_sinkhorn_norm(params, tensor);
+            } break;
         case GGML_OP_MUL_MAT:
             {
                 ggml_compute_forward_mul_mat(params, tensor);
@@ -2306,6 +2310,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_RMS_NORM:
         case GGML_OP_RMS_NORM_BACK:
         case GGML_OP_L2_NORM:
+        case GGML_OP_SINKHORN_NORM:
         case GGML_OP_GROUP_NORM:
         case GGML_OP_CONCAT:
         case GGML_OP_MUL_MAT:

ggml/src/ggml-cpu/ops.cpp

@@ -4237,6 +4237,125 @@ void ggml_compute_forward_l2_norm(
     }
 }
 
+// ggml_compute_forward_sinkhorn_norm
+
+#define GGML_SINKHORN_NORM_MAX_N 32
+
+static void ggml_compute_forward_sinkhorn_norm_f32(
+    const ggml_compute_params * params,
+    ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];
+
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
+    GGML_ASSERT(ggml_is_contiguous(src0));
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    GGML_TENSOR_UNARY_OP_LOCALS
+
+    const int32_t n_iters = ggml_get_op_params_i32(dst, 0);
+    const float   eps     = ggml_get_op_params_f32(dst, 1);
+
+    const int64_t n = ne00;
+    GGML_ASSERT(ne01 == n);
+    GGML_ASSERT(n <= GGML_SINKHORN_NORM_MAX_N);
+    GGML_ASSERT(n_iters >= 1);
+
+    const int64_t n_slices = ne02 * ne03;
+
+    float m[GGML_SINKHORN_NORM_MAX_N * GGML_SINKHORN_NORM_MAX_N];
+
+    for (int64_t s = ith; s < n_slices; s += nth) {
+        const int64_t i02 = s % ne02;
+        const int64_t i03 = s / ne02;
+
+        const float * x = (const float *) ((const char *) src0->data + i02*nb02 + i03*nb03);
+        float *       y = (float *)       ((      char *) dst->data  + i02*nb2  + i03*nb3);
+
+        for (int64_t b = 0; b < n; ++b) {
+            for (int64_t a = 0; a < n; ++a) {
+                m[b*n + a] = x[b*n + a];
+            }
+        }
+
+        for (int64_t b = 0; b < n; ++b) {
+            float mx = -INFINITY;
+            for (int64_t a = 0; a < n; ++a) {
+                mx = fmaxf(mx, m[b*n + a]);
+            }
+            float sum = 0.0f;
+            for (int64_t a = 0; a < n; ++a) {
+                const float e = expf(m[b*n + a] - mx);
+                m[b*n + a] = e;
+                sum += e;
+            }
+            for (int64_t a = 0; a < n; ++a) {
+                m[b*n + a] /= sum;
+            }
+        }
+
+        for (int64_t i = 0; i < n*n; ++i) {
+            m[i] += eps;
+        }
+
+        auto norm_cols = [&]() {
+            for (int64_t a = 0; a < n; ++a) {
+                float r = 0.0f;
+                for (int64_t b = 0; b < n; ++b) {
+                    r += m[b*n + a];
+                }
+                r += eps;
+                for (int64_t b = 0; b < n; ++b) {
+                    m[b*n + a] /= r;
+                }
+            }
+        };
+
+        auto norm_rows = [&]() {
+            for (int64_t b = 0; b < n; ++b) {
+                float c = 0.0f;
+                for (int64_t a = 0; a < n; ++a) {
+                    c += m[b*n + a];
+                }
+                c += eps;
+                for (int64_t a = 0; a < n; ++a) {
+                    m[b*n + a] /= c;
+                }
+            }
+        };
+
+        norm_cols();
+        for (int32_t it = 1; it < n_iters; ++it) {
+            norm_rows();
+            norm_cols();
+        }
+
+        for (int64_t i = 0; i < n*n; ++i) {
+            y[i] = m[i];
+        }
+    }
+}
+
+void ggml_compute_forward_sinkhorn_norm(
+    const ggml_compute_params * params,
+    ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];
+
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_sinkhorn_norm_f32(params, dst);
+            } break;
+        default:
+            {
+                GGML_ABORT("fatal error");
+            }
+    }
+}
+
 // ggml_compute_forward_out_prod
 
 static void ggml_compute_forward_out_prod_f32(

ggml/src/ggml-cpu/ops.h

@@ -48,6 +48,7 @@ void ggml_compute_forward_rms_norm_mul_fused(const struct ggml_compute_params *
 void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_sinkhorn_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -33,6 +33,7 @@
 #include "ggml-cuda/mmvf.cuh"
 #include "ggml-cuda/mmvq.cuh"
 #include "ggml-cuda/norm.cuh"
+#include "ggml-cuda/sinkhorn-norm.cuh"
 #include "ggml-cuda/opt-step-adamw.cuh"
 #include "ggml-cuda/opt-step-sgd.cuh"
 #include "ggml-cuda/out-prod.cuh"
@@ -2083,6 +2084,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_GROUP_NORM:
             ggml_cuda_op_group_norm(ctx, dst);
             break;
+        case GGML_OP_SINKHORN_NORM:
+            ggml_cuda_op_sinkhorn_norm(ctx, dst);
+            break;
         case GGML_OP_L2_NORM:
             ggml_cuda_op_l2_norm(ctx, dst);
             break;
@@ -4850,6 +4854,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_SILU_BACK:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
             break;
+        case GGML_OP_SINKHORN_NORM:
+            return op->src[0]->type == GGML_TYPE_F32 &&
+                   ggml_is_contiguous(op->src[0]) &&
+                   op->src[0]->ne[0] == op->src[0]->ne[1] &&
+                   op->src[0]->ne[0] <= 8;
         case GGML_OP_NORM:
         case GGML_OP_RMS_NORM:
         case GGML_OP_L2_NORM:

ggml/src/ggml-cuda/sinkhorn-norm.cu

@@ -0,0 +1,192 @@
+#include "sinkhorn-norm.cuh"
+
+#define SINKHORN_NORM_MAX_N 8
+
+template <int MAXN>
+static __global__ void sinkhorn_norm_f32(
+        const float * __restrict__ x,
+        float       * __restrict__ y,
+        const int     n,
+        const int     n_iters,
+        const float   eps,
+        const int64_t n_slices,
+        const int64_t slice_stride) {
+
+    const int64_t s = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (s >= n_slices) {
+        return;
+    }
+
+    const float * xs = x + s * slice_stride;
+    float       * ys = y + s * slice_stride;
+
+    float m[MAXN * MAXN];
+    const int nn = n * n;
+
+    for (int i = 0; i < nn; ++i) {
+        m[i] = xs[i];
+    }
+
+    for (int b = 0; b < n; ++b) {
+        float mx = -INFINITY;
+        for (int a = 0; a < n; ++a) {
+            mx = fmaxf(mx, m[b*n + a]);
+        }
+        float sum = 0.0f;
+        for (int a = 0; a < n; ++a) {
+            const float e = expf(m[b*n + a] - mx);
+            m[b*n + a] = e;
+            sum += e;
+        }
+        for (int a = 0; a < n; ++a) {
+            m[b*n + a] /= sum;
+        }
+    }
+
+    for (int i = 0; i < nn; ++i) {
+        m[i] += eps;
+    }
+
+    #define NORM_COLS() do {                             \
+        for (int a = 0; a < n; ++a) {                    \
+            float r = 0.0f;                              \
+            for (int b = 0; b < n; ++b) r += m[b*n + a]; \
+            r += eps;                                    \
+            for (int b = 0; b < n; ++b) m[b*n + a] /= r; \
+        }                                                \
+    } while (0)
+
+    #define NORM_ROWS() do {                             \
+        for (int b = 0; b < n; ++b) {                    \
+            float c = 0.0f;                              \
+            for (int a = 0; a < n; ++a) c += m[b*n + a]; \
+            c += eps;                                    \
+            for (int a = 0; a < n; ++a) m[b*n + a] /= c; \
+        }                                                \
+    } while (0)
+
+    NORM_COLS();
+    for (int it = 1; it < n_iters; ++it) {
+        NORM_ROWS();
+        NORM_COLS();
+    }
+
+    #undef NORM_COLS
+    #undef NORM_ROWS
+
+    for (int i = 0; i < nn; ++i) {
+        ys[i] = m[i];
+    }
+}
+
+template <int N>
+static __global__ void sinkhorn_norm_warp_f32(
+        const float * __restrict__ x,
+        float       * __restrict__ y,
+        const int     n_iters,
+        const float   eps,
+        const int64_t n_slices) {
+
+    constexpr int NN = N * N;
+
+    const int64_t g     = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    const int64_t token = g / NN;
+    if (token >= n_slices) {
+        return;
+    }
+
+    const int e = threadIdx.x % NN;
+
+    float v = x[token * NN + e];
+
+    float mx = v;
+#pragma unroll
+    for (int off = 1; off < N; off <<= 1) {
+        mx = fmaxf(mx, __shfl_xor_sync(0xffffffff, mx, off, NN));
+    }
+    float ex  = expf(v - mx);
+    float sum = ex;
+#pragma unroll
+    for (int off = 1; off < N; off <<= 1) {
+        sum += __shfl_xor_sync(0xffffffff, sum, off, NN);
+    }
+    v = ex / sum + eps;
+
+    {
+        float r = v;
+#pragma unroll
+        for (int off = N; off < NN; off <<= 1) {
+            r += __shfl_xor_sync(0xffffffff, r, off, NN);
+        }
+        v /= (r + eps);
+    }
+
+    for (int it = 1; it < n_iters; ++it) {
+        {
+            float c = v;
+#pragma unroll
+            for (int off = 1; off < N; off <<= 1) {
+                c += __shfl_xor_sync(0xffffffff, c, off, NN);
+            }
+            v /= (c + eps);
+        }
+        {
+            float r = v;
+#pragma unroll
+            for (int off = N; off < NN; off <<= 1) {
+                r += __shfl_xor_sync(0xffffffff, r, off, NN);
+            }
+            v /= (r + eps);
+        }
+    }
+
+    y[token * NN + e] = v;
+}
+
+static bool sinkhorn_norm_is_pow2(int n) {
+    return n > 0 && (n & (n - 1)) == 0;
+}
+
+void ggml_cuda_op_sinkhorn_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const float * src0_d = (const float *) src0->data;
+    float       * dst_d  = (float *)       dst->data;
+    cudaStream_t  stream = ctx.stream();
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+
+    GGML_TENSOR_UNARY_OP_LOCALS;
+
+    const int32_t n_iters = ggml_get_op_params_i32(dst, 0);
+    const float   eps     = ggml_get_op_params_f32(dst, 1);
+
+    const int n = (int) ne00;
+    GGML_ASSERT(ne01 == ne00);
+    GGML_ASSERT(n <= SINKHORN_NORM_MAX_N);
+    GGML_ASSERT(n_iters >= 1);
+
+    const int64_t n_slices     = ne02 * ne03;
+    const int64_t slice_stride = (int64_t) n * n;
+
+    const int block_size = 256;
+
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+
+    if (sinkhorn_norm_is_pow2(n) && n*n <= warp_size) {
+        const int64_t total_threads = n_slices * (int64_t) (n*n);
+        const int64_t num_blocks    = (total_threads + block_size - 1) / block_size;
+        switch (n) {
+            case 2: sinkhorn_norm_warp_f32<2><<<num_blocks, block_size, 0, stream>>>(src0_d, dst_d, n_iters, eps, n_slices); break;
+            case 4: sinkhorn_norm_warp_f32<4><<<num_blocks, block_size, 0, stream>>>(src0_d, dst_d, n_iters, eps, n_slices); break;
+            case 8: sinkhorn_norm_warp_f32<8><<<num_blocks, block_size, 0, stream>>>(src0_d, dst_d, n_iters, eps, n_slices); break;
+            default: GGML_ABORT("sinkhorn_norm: unsupported warp N=%d", n);
+        }
+        return;
+    }
+
+    const int64_t num_blocks = (n_slices + block_size - 1) / block_size;
+    sinkhorn_norm_f32<SINKHORN_NORM_MAX_N><<<num_blocks, block_size, 0, stream>>>(
+        src0_d, dst_d, n, n_iters, eps, n_slices, slice_stride);
+}