Metal: tensor-op chunk-state probe, layer bench with before/after, A/B driver

metal/kernels/gdn_tprobe.metal

@@ -0,0 +1,75 @@
+// metal/kernels/gdn_tprobe.metal - go / no-go probe for a tensor-op (MPP matmul2d) chunk-state kernel: the five products of one GDN
+// chunk step (W S, Q S, M delta, K^T delta' accumulated into S) on fp16 operands, S kept in fp32 threadgroup memory, 64 chunk steps
+// a threadgroup, grid h_v * 4 (the shape of gdn_cstate2).  No numerics: timing only.
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+using namespace metal;
+using namespace mpp::tensor_ops;
+
+constant int PROBE [[function_constant(0)]];
+constant int probe = is_function_constant_defined(PROBE) ? PROBE : 0;   // 1 = skip the products (SIMT parts only)
+
+kernel void gdn_tprobe(device float* out [[buffer(0)]], device half* Wh [[buffer(1)]], device half* Qh [[buffer(2)]],
+                       device half* Kh [[buffer(3)]], device half* Mh [[buffer(4)]], constant int& NC [[buffer(5)]],
+                       uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
+    threadgroup half Sh[128 * 32];
+    threadgroup half Dh[32 * 32];
+    threadgroup float X[32 * 32], Y[32 * 32], Z[32 * 32];
+    tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tSh(Sh, dextents<int32_t, 2>(32, 128));
+    tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline> tDh(Dh, dextents<int32_t, 2>(32, 32));
+    tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tX(X, dextents<int32_t, 2>(32, 32));
+    tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tY(Y, dextents<int32_t, 2>(32, 32));
+    tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tZ(Z, dextents<int32_t, 2>(32, 32));
+    constexpr auto d_ws = matmul2d_descriptor(32, 32, 128, false, false, false);   // [32 x 128] . [128 x 32]
+    constexpr auto d_md = matmul2d_descriptor(32, 32, 32, false, false, false);    // [32 x 32] . [32 x 32]
+    constexpr auto d_up = matmul2d_descriptor(128, 32, 32, true, false, false, matmul2d_descriptor::mode::multiply_accumulate);   // K^T delta'
+    matmul2d<d_ws, execution_simdgroups<4>> op_ws;
+    matmul2d<d_md, execution_simdgroups<4>> op_md;
+    matmul2d<d_up, execution_simdgroups<4>> op_up;
+    const ulong hb = ulong(tg.x) / 4;
+    using KT = tensor<device half, dextents<int32_t, 2>, tensor_inline>;
+    using DT = tensor<threadgroup half, dextents<int32_t, 2>, tensor_inline>;
+    auto cS = op_up.get_destination_cooperative_tensor<KT, DT, float>();
+#pragma clang loop unroll(full)
+    for (uint16_t i = 0; i < cS.get_capacity(); ++i)
+        if (cS.is_valid_element(i)) cS[i] = 0.001f * float(i & 15);
+    for (int c = 0; c < NC; ++c) {
+        const ulong cb = (hb * NC + c);
+#pragma clang loop unroll(full)
+        for (uint16_t i = 0; i < cS.get_capacity(); ++i)
+            if (cS.is_valid_element(i)) {
+                auto ids = cS.get_multidimensional_index(i);
+                Sh[ids[1] * 32 + ids[0]] = half(cS[i]);
+            }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (probe != 1) {
+            tensor<device half, dextents<int32_t, 2>, tensor_inline> tW(Wh + cb * 32 * 128, dextents<int32_t, 2>(128, 32));
+            tensor<device half, dextents<int32_t, 2>, tensor_inline> tQ(Qh + cb * 32 * 128, dextents<int32_t, 2>(128, 32));
+            op_ws.run(tW, tSh, tX);
+            op_ws.run(tQ, tSh, tY);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (int i = tid; i < 32 * 32; i += 128) Dh[i] = half(0.01f * X[i]);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (probe != 1) {
+            tensor<device half, dextents<int32_t, 2>, tensor_inline> tM(Mh + cb * 32 * 32, dextents<int32_t, 2>(32, 32));
+            op_md.run(tM, tDh, tZ);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (int i = tid; i < 32 * 32; i += 128) Dh[i] = half(Z[i] + 0.01f * Y[i]);
+#pragma clang loop unroll(full)
+        for (uint16_t i = 0; i < cS.get_capacity(); ++i)
+            if (cS.is_valid_element(i)) cS[i] *= 0.99f;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (probe != 1) {
+            tensor<device half, dextents<int32_t, 2>, tensor_inline> tK(Kh + cb * 32 * 128, dextents<int32_t, 2>(128, 32));
+            op_up.run(tK, tDh, cS);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+#pragma clang loop unroll(full)
+    for (uint16_t i = 0; i < cS.get_capacity(); ++i)
+        if (cS.is_valid_element(i)) {
+            auto ids = cS.get_multidimensional_index(i);
+            out[ulong(tg.x) * 4096 + ids[1] * 32 + ids[0]] = cS[i];
+        }
+}

metal/tests/gdn_layer_bench.cpp

@@ -146,6 +146,7 @@ int main(int argc, char** argv) {
     dpost.grid[0] = HV; dpost.grid[1] = T; dpost.tg[0] = 128;
     const double tpre = timeit({dpre}), tscan = timeit({dscan}), tpost = timeit({dpost});
     std::printf("== gdn_pre %.3f, scan %.3f, gdn_post %.3f ms (stage %.3f)\n", tpre, tscan, tpost, tpre + tscan + tpost);
+    double tF = 0;
     {   // the fused-stride kernels (function constant 7) against the separate ones: bit for bit
         for (int t = 0; t < T; ++t) {
             std::memcpy(YF.as<float>() + (size_t)t * NF, qkv.as<float>() + (size_t)t * CH, 4 * CH);
@@ -172,7 +173,7 @@ int main(int argc, char** argv) {
                     std::memcmp(qn.ptr, rqn.data(), rqn.size()) == 0 ? "equal" : "DIFFER", std::memcmp(vc.ptr, rvc.data(), rvc.size()) == 0 ? "equal" : "DIFFER",
                     std::memcmp(gb.ptr, rgb.data(), rgb.size()) == 0 ? "equal" : "DIFFER", std::memcmp(h1.ptr, rh.data(), rh.size()) == 0 ? "equal" : "DIFFER",
                     std::memcmp(yh.ptr, ryh.data(), ryh.size()) == 0 ? "equal" : "DIFFER");
-        const double tF = timeit({fpre}) + timeit({fpost});
+        tF = timeit({fpre}) + timeit({fpost});
         std::printf("fused-stride pre + post %.3f ms (separate %.3f)\n", tF, tpre + tpost);
     }
     const double tout = timeit({gemm_d(ygh, Wh[4], Y[4], T, N, ZD, 64, 64)});
@@ -180,5 +181,14 @@ int main(int argc, char** argv) {
     for (int i = 0; i < 5; ++i) dqall += timeit({dq_d(i)});
     std::printf("== layer in engine order (GPU ms): dq %.3f + input GEMMs %.3f + pre/scan/post %.3f + out GEMM %.3f = %.3f (sum of parts)\n",
                 dqall, sep, tpre + tscan + tpost, tout, dqall + sep + tpre + tscan + tpost + tout);
+    {   // the layer as the engine now runs it on the M5 family: fused input GEMM, fused-stride pre/post, tiled out GEMM
+        double tout_best = 1e9;
+        for (auto [tm, tn] : std::vector<std::pair<int, int>>{{64, 64}, {128, 64}, {64, 128}})
+            tout_best = std::min(tout_best, timeit({gemm_d(ygh, Wh[4], Y[4], T, N, ZD, tm, tn)}));
+        const double base_l = dqall + sep + tpre + tscan + tpost + tout;
+        const double new_l = dqall + best_f + tF + tscan + tout_best;
+        std::printf("== layer: before %.3f ms, after (fused input GEMM %.3f, fused-stride pre+post %.3f, scan %.3f, out GEMM %.3f) %.3f ms = x%.3f\n",
+                    base_l, best_f, tF, tscan, tout_best, new_l, base_l / new_l);
+    }
     return 0;
 }

metal/tests/gdn_tprobe.cpp

@@ -0,0 +1,36 @@
+// metal/tests/gdn_tprobe.cpp - times kernels/gdn_tprobe.metal: the five MPP matmul2d products of a GDN chunk step, NC chunk steps a
+// threadgroup, 192 threadgroups (h_v * 4).  Compare with gdn_scanr_4's stage (chunks of 32 -> T = 32 * NC).
+//   gdn_tprobe [NC]
+#include "mtl_rt.h"
+#include "ref_quant.h"
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+int main(int argc, char** argv) {
+    const int NC = argc > 1 ? std::atoi(argv[1]) : 64;
+    mtl::Context cx;
+    const int lib = cx.compile({"gdn_tprobe.metal"}, "#define NO_GGML_COMMON 1\n");
+    std::printf("device %s, NC %d (T %d)\n", cx.device_name().c_str(), NC, NC * 32);
+    const size_t NB = (size_t)48 * NC;
+    mtl::Buffer out = cx.alloc(4 * 192 * 4096), W = cx.alloc(2 * NB * 32 * 128), Q = cx.alloc(2 * NB * 32 * 128),
+                K = cx.alloc(2 * NB * 32 * 128), M = cx.alloc(2 * NB * 32 * 32);
+    ref::Rng r(3);
+    for (auto* b : {&W, &Q, &K}) for (size_t i = 0; i < NB * 32 * 128; ++i) b->as<uint16_t>()[i] = ref::f32_to_f16(r.sym() * 0.1f);
+    for (size_t i = 0; i < NB * 32 * 32; ++i) M.as<uint16_t>()[i] = ref::f32_to_f16(r.sym() * 0.1f);
+    for (int pr : {0, 1}) {
+        mtl::Pipeline p = cx.pipeline(lib, "gdn_tprobe", pr ? mtl::Constants().i(0, pr) : mtl::Constants());
+        mtl::Dispatch d; d.p = &p;
+        d.args = {mtl::Arg::B(out), mtl::Arg::B(W), mtl::Arg::B(Q), mtl::Arg::B(K), mtl::Arg::B(M), mtl::Arg::V(NC)};
+        d.grid[0] = 192; d.tg[0] = 128;
+        cx.warm({d}, 0.3);
+        std::vector<double> t;
+        for (int i = 0; i < 7; ++i) t.push_back(cx.run({d}) * 1e3);
+        std::sort(t.begin(), t.end());
+        std::printf("probe %d (%s): %.3f ms for %d chunk steps x 192 threadgroups  (max threads %d, tg mem %d)\n", pr,
+                    pr ? "SIMT parts only" : "with the five products", t[3], NC, p.max_threads, p.static_tg_mem);
+    }
+    return 0;
+}