complete rebase against fused adds - multi_add shader can also compute partial sums

Author: jeffbolznv

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -102,9 +102,9 @@ static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
 
 struct ggml_backend_vk_context;
 
-#define MAX_PARAMETER_COUNT 8
+#define MAX_PARAMETER_COUNT 12
 // Max number of adds that can be fused without exceeding MAX_PARAMETER_COUNT.
-#define MAX_FUSED_ADDS (MAX_PARAMETER_COUNT - 2)
+#define MAX_FUSED_ADDS (MAX_PARAMETER_COUNT - 3)
 
 struct vk_pipeline_struct {
     std::string name;
@@ -459,6 +459,7 @@ struct vk_device_struct {
 
     // indexed by num_additional_fused_ops == num_adds - 1
     vk_pipeline pipeline_multi_add[MAX_FUSED_ADDS];
+    vk_pipeline pipeline_multi_add_rms[MAX_FUSED_ADDS];
 
     vk_pipeline pipeline_add_id_f32;
 
@@ -819,8 +820,13 @@ struct vk_op_multi_add_push_constants {
     uint32_t ne20; uint32_t ne21; uint32_t ne22; uint32_t ne23;
 
     // strides for srcs+dst
-    uint32_t nb[8][4];
+    uint32_t nb[MAX_PARAMETER_COUNT][4];
+
+    uint32_t rms_partials;
 };
+// update multi_add.comp if this changes
+static_assert(MAX_PARAMETER_COUNT == 12);
+static_assert(sizeof(vk_op_multi_add_push_constants) <= 256);
 
 struct vk_op_add_id_push_constants {
     uint32_t ne0;
@@ -3032,7 +3038,8 @@ static void ggml_vk_load_shaders(vk_device& device) {
 
     if (device->multi_add) {
         for (uint32_t i = 0; i < MAX_FUSED_ADDS; ++i) {
-            ggml_vk_create_pipeline(device, device->pipeline_multi_add[i], "multi_add_f32_" + std::to_string(i+1), multi_add_f32_len, multi_add_f32_data, "main", MAX_PARAMETER_COUNT, sizeof(vk_op_multi_add_push_constants), {512, 1, 1}, {i+2}, 1);
+            ggml_vk_create_pipeline(device, device->pipeline_multi_add[i],     "multi_add_f32_"     + std::to_string(i+1), multi_add_f32_len,     multi_add_f32_data,     "main", MAX_PARAMETER_COUNT, sizeof(vk_op_multi_add_push_constants), {512, 1, 1}, {i+2}, 1);
+            ggml_vk_create_pipeline(device, device->pipeline_multi_add_rms[i], "multi_add_rms_f32_" + std::to_string(i+1), multi_add_rms_f32_len, multi_add_rms_f32_data, "main", MAX_PARAMETER_COUNT, sizeof(vk_op_multi_add_push_constants), {512, 1, 1}, {i+2}, 1);
         }
     }
 
@@ -6912,7 +6919,7 @@ static std::array<uint32_t, 3> ggml_vk_get_conv_elements(const ggml_tensor *dst)
     return elements;
 }
 
-static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, ggml_op op) {
+static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) {
     switch (op) {
     case GGML_OP_GET_ROWS:
         GGML_ASSERT(src1->type == GGML_TYPE_I32);
@@ -7836,7 +7843,7 @@ static void ggml_vk_multi_add(ggml_backend_vk_context * ctx, vk_context& subctx,
     const ggml_tensor *tensors[MAX_PARAMETER_COUNT];
     uint32_t num_srcs = ctx->num_additional_fused_ops + 2;
     uint32_t num_tensors = num_srcs + 1;
-    GGML_ASSERT(num_tensors <= MAX_PARAMETER_COUNT);
+    GGML_ASSERT(num_tensors + ctx->do_add_rms_partials <= MAX_PARAMETER_COUNT);
 
     tensors[0] = first_node->src[0];
     tensors[1] = first_node->src[1];
@@ -7863,8 +7870,9 @@ static void ggml_vk_multi_add(ggml_backend_vk_context * ctx, vk_context& subctx,
         pc.nb[i][2] = (uint32_t)t->nb[2] / sizeof(float);
         pc.nb[i][3] = (uint32_t)t->nb[3] / sizeof(float);
     }
+    pc.rms_partials = ctx->do_add_rms_partials;
 
-    vk_pipeline pipeline = ctx->device->pipeline_multi_add[ctx->num_additional_fused_ops];
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, tensors[0], tensors[1], nullptr, dst, dst->op);
 
     if (pipeline == nullptr) {
         std::cerr << "ggml_vulkan: Error: Missing multi_add";
@@ -7902,6 +7910,10 @@ static void ggml_vk_multi_add(ggml_backend_vk_context * ctx, vk_context& subctx,
         buf[i] = buf[0];
         offset[i] = 0;
     }
+    if (ctx->do_add_rms_partials) {
+        buf[num_tensors] = ctx->prealloc_add_rms_partials;
+        offset[num_tensors] = ctx->prealloc_size_add_rms_partials_offset;
+    }
 
     std::array<uint32_t, 3> elements;
 
@@ -7915,6 +7927,7 @@ static void ggml_vk_multi_add(ggml_backend_vk_context * ctx, vk_context& subctx,
     }
 
     ggml_vk_sync_buffers(subctx);
+    static_assert(MAX_PARAMETER_COUNT == 12);
     ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
         {
             vk_subbuffer{ buf[0], offset[0], VK_WHOLE_SIZE },
@@ -7925,6 +7938,10 @@ static void ggml_vk_multi_add(ggml_backend_vk_context * ctx, vk_context& subctx,
             vk_subbuffer{ buf[5], offset[5], VK_WHOLE_SIZE },
             vk_subbuffer{ buf[6], offset[6], VK_WHOLE_SIZE },
             vk_subbuffer{ buf[7], offset[7], VK_WHOLE_SIZE },
+            vk_subbuffer{ buf[8], offset[8], VK_WHOLE_SIZE },
+            vk_subbuffer{ buf[9], offset[9], VK_WHOLE_SIZE },
+            vk_subbuffer{ buf[10], offset[10], VK_WHOLE_SIZE },
+            vk_subbuffer{ buf[11], offset[11], VK_WHOLE_SIZE },
         }, pc, elements);
 }
 
@@ -9771,17 +9788,19 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
         }
         break;
     case GGML_OP_ADD:
-        if (node_idx + 1 < cgraph->n_nodes &&
-            cgraph->nodes[node_idx + 1]->op == GGML_OP_RMS_NORM &&
-            cgraph->nodes[node_idx + 1]->src[0] == cgraph->nodes[node_idx] &&
-            ggml_nrows(cgraph->nodes[node_idx + 1]) == 1 &&
-            ctx->device->add_rms_fusion) {
-            if (dryrun) {
-                ctx->prealloc_size_add_rms_partials += ggml_vk_rms_partials_size(ctx, cgraph->nodes[node_idx]);
+        {
+            int next_node_idx = node_idx + 1 + ctx->num_additional_fused_ops;
+            if (next_node_idx < cgraph->n_nodes &&
+                cgraph->nodes[next_node_idx]->op == GGML_OP_RMS_NORM &&
+                cgraph->nodes[next_node_idx]->src[0] == cgraph->nodes[next_node_idx - 1] &&
+                ggml_nrows(cgraph->nodes[next_node_idx]) == 1 &&
+                ctx->device->add_rms_fusion) {
+                if (dryrun) {
+                    ctx->prealloc_size_add_rms_partials += ggml_vk_rms_partials_size(ctx, cgraph->nodes[node_idx]);
+                }
+                ctx->do_add_rms_partials = true;
             }
-            ctx->do_add_rms_partials = true;
-        }
-        break;
+        } break;
     case GGML_OP_REPEAT:
     case GGML_OP_REPEAT_BACK:
     case GGML_OP_GET_ROWS:

ggml/src/ggml-vulkan/vulkan-shaders/multi_add.comp

@@ -3,6 +3,10 @@
 #extension GL_EXT_shader_16bit_storage : require
 #extension GL_EXT_nonuniform_qualifier : enable
 #extension GL_EXT_control_flow_attributes : require
+#if ADD_RMS
+#extension GL_KHR_shader_subgroup_arithmetic : enable
+#extension GL_KHR_shader_subgroup_basic : enable
+#endif
 
 #include "rte.comp"
 #include "types.comp"
@@ -14,12 +18,16 @@ layout (push_constant) uniform parameter2
     uint ne20; uint ne21; uint ne22; uint ne23;
 
     // strides for srcs+dst
-    uint nb[8][4];
+    uint nb[12][4];
+
+    uint rms_partials;
 } p;
 
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];} a[];
 layout (binding = 0) writeonly buffer D {D_TYPE data_d[];} d[];
 
+layout (binding = 0, std430) buffer PartialBuf {float partial_sums[];} partials[];
+
 layout(constant_id = 0) const uint num_srcs = 2;
 
 uint src_idx(uint s, uint i00, uint i01, uint i02, uint i03) {
@@ -42,14 +50,22 @@ const uint num_threads = 256;
 
 layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in;
 
+#if ADD_RMS
+// XXX TODO this could be sized based on number of subgroups, but that't not considered a constant
+shared FLOAT_TYPE sumsh[num_threads];
+#endif
+
 void main() {
     uint idx = get_idx();
+    uint orig_idx = idx;
 
     uint ne = p.ne20 * p.ne21 * p.ne22 * p.ne23;
 
     // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation
     const uint num_iter = 2;
 
+    FLOAT_TYPE sum_sq = 0;
+
     [[unroll]] for (uint i = 0; i < num_iter; ++i) {
         if (idx >= ne) {
             continue;
@@ -61,8 +77,32 @@ void main() {
         [[unroll]] for (uint s = 0; s < num_srcs; ++s) {
             sum += FLOAT_TYPE(a[s].data_a[src_idx(s, i00, i01, i02, i03)]);
         }
+        sum_sq += sum*sum;
         d[num_srcs].data_d[dst_idx(i00, i01, i02, i03)] = D_TYPE(sum);
 
         idx += num_threads;
     }
+
+#if ADD_RMS
+    if (p.rms_partials != 0) {
+        // reduce the sum within each subgroup, then across subgroups
+        const uint NumSubgroups = num_threads / gl_SubgroupSize;
+        sum_sq = subgroupAdd(sum_sq);
+        if (gl_SubgroupInvocationID == 0) {
+            sumsh[gl_SubgroupID] = sum_sq;
+        }
+        barrier();
+        [[unroll]] for (uint s = NumSubgroups / 2; s > 0; s >>= 1) {
+            if (gl_SubgroupID < s && gl_SubgroupInvocationID == 0) {
+                sum_sq += sumsh[gl_SubgroupID + s];
+                sumsh[gl_SubgroupID] = sum_sq;
+            }
+            barrier();
+        }
+
+        if (gl_SubgroupID == 0 && gl_SubgroupInvocationID == 0) {
+            partials[num_srcs + 1].partial_sums[orig_idx / (num_iter * num_threads)] = sum_sq;
+        }
+    }
+#endif
 }

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -680,7 +680,8 @@ void process_shaders() {
 
     string_to_spv("add_id_f32", "add_id.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
 
-    string_to_spv("multi_add_f32", "multi_add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"RTE16", "1"}});
+    string_to_spv("multi_add_f32", "multi_add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"RTE16", "1"}, {"ADD_RMS" , "0"}});
+    string_to_spv("multi_add_rms_f32", "multi_add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"RTE16", "1"}, {"ADD_RMS" , "1"}});
 
     for (auto &c : compiles) {
         c.wait();

tests/test-backend-ops.cpp

@@ -2858,6 +2858,7 @@ struct test_rms_norm_mul_add : public test_case {
     const std::array<int64_t, 4> ne;
     const float eps;
     const bool broadcast;
+    const bool multi_add; // test a sequence of adds feeding into rms_norm
 
     std::string op_desc(ggml_tensor * t) override {
         GGML_UNUSED(t);
@@ -2867,13 +2868,13 @@ struct test_rms_norm_mul_add : public test_case {
     bool run_whole_graph() override { return true; }
 
     std::string vars() override {
-        return VARS_TO_STR4(type, ne, eps, broadcast);
+        return VARS_TO_STR5(type, ne, eps, broadcast, multi_add);
     }
 
     test_rms_norm_mul_add(ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne = {64, 5, 4, 3},
-            float eps = 1e-6f, bool broadcast = false)
-        : type(type), ne(ne), eps(eps), broadcast(broadcast) {}
+            float eps = 1e-6f, bool broadcast = false, bool multi_add = false)
+        : type(type), ne(ne), eps(eps), broadcast(broadcast), multi_add(multi_add) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         std::array<int64_t, 4> broadcast_dims = {ne[0]*2, ne[1]*3, ne[2]*3, ne[3]*4};
@@ -2891,6 +2892,9 @@ struct test_rms_norm_mul_add : public test_case {
 
         // Use a, b and c early, so we don't end up with an OP_NONE between rms_norm and mul
         a = ggml_add(ctx, ggml_add(ctx, a, b), c);
+        if (multi_add) {
+            a = ggml_add(ctx, ggml_add(ctx, a, b), c);
+        }
         ggml_tensor * out = ggml_add(ctx, ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b), c);
         ggml_set_name(out, "out");
 
@@ -5679,7 +5683,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {64, 5, 4, 3}, eps, true));
     }
     for (uint32_t n : {1, 511, 1025, 8192, 33*512}) {
-        test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {n, 1, 1, 1}, 1e-6f));
+        for (bool multi_add : {false, true}) {
+            test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {n, 1, 1, 1}, 1e-6f, false, multi_add));
+        }
     }
 
     test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, {64, 5, 4, 3}, 1e-12f));