metal: OUT_PROD with a quantized src0 — the backward through a resident quantized base on Apple silicon (card 55f314b8, S3b)

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -334,6 +334,41 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu(ggml_metal_l
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const ggml_tensor * op) {
+    assert(op->op == GGML_OP_CROSS_ENTROPY_LOSS || op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_cross_entropy_loss%s_f32", op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK ? "_back" : "");
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
+// RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK: one f32 kernel each, named after the op.
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back(ggml_metal_library_t lib, const ggml_tensor * op) {
+    const char * kernel = nullptr;
+    switch (op->op) {
+        case GGML_OP_RMS_NORM_BACK: kernel = "kernel_rms_norm_back_f32"; break;
+        case GGML_OP_SOFT_MAX_BACK: kernel = "kernel_soft_max_back_f32"; break;
+        case GGML_OP_REPEAT_BACK:   kernel = "kernel_repeat_back_f32";   break;
+        default: GGML_ABORT("fatal error");
+    }
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, kernel);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, kernel, kernel, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_SUM);
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -120,6 +120,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu               (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum               (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk        (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add        (ggml_metal_library_t lib, const struct ggml_tensor * op);

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1,4 +1,5 @@
 #import "ggml-metal-device.h"
+#include "ggml-metal-ops.h"
 
 #import "ggml-impl.h"
 #import "ggml-backend-impl.h"
@@ -1467,6 +1468,27 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                 default:
                     return false;
             }
+        case GGML_OP_RMS_NORM_BACK:
+        case GGML_OP_SOFT_MAX_BACK:
+            // src0 the incoming gradient, src1 the forward input/output; rows contiguous.
+            // SOFT_MAX_BACK carries no max_bias (the CPU reference asserts the same).
+            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+                   ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op) &&
+                   (op->op != GGML_OP_SOFT_MAX_BACK || ggml_get_op_params_f32(op, 1) == 0.0f);
+        case GGML_OP_REPEAT_BACK:
+            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op);
+        case GGML_OP_CROSS_ENTROPY_LOSS:
+            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+                   ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
+        case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+            return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 &&
+                   ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op);
+        case GGML_OP_OUT_PROD:
+            // The backward through a frozen weight: out_prod(W, grad^T). Composed on this
+            // backend from a dequantizing copy, two transposing copies and the F32 simdgroup
+            // matmul (see ggml_metal_op_out_prod). 2D, src1 F32 with contiguous rows, src0 F32
+            // or a type with a kernel_cpy_<type>_f32.
+            return ggml_metal_op_out_prod_supported(op) && ggml_metal_device_get_props(dev)->has_simdgroup_mm;
         case GGML_OP_SILU_BACK:
             return (op->src[0]->type == GGML_TYPE_F32) &&
                 (op->src[1]->type == GGML_TYPE_F32) &&
@@ -1782,6 +1804,17 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                             default:
                                 return false;
                         }
+                    case GGML_TYPE_Q2_K:
+                    case GGML_TYPE_Q3_K:
+                    case GGML_TYPE_Q4_K:
+                    case GGML_TYPE_Q5_K:
+                    case GGML_TYPE_Q6_K:
+                        switch (op->type) {
+                            case GGML_TYPE_F32:
+                                return true;
+                            default:
+                                return false;
+                        }
                     case GGML_TYPE_I32:
                         return op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_I32;
                     default:

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -880,6 +880,23 @@ typedef struct {
     int32_t  n_head_log2;
 } ggml_metal_kargs_soft_max;
 
+typedef struct {
+    int32_t  ne00;
+    int32_t  nr;
+    uint64_t nb01;
+    uint64_t nb11;
+    uint64_t nb1;
+} ggml_metal_kargs_cross_entropy_loss;
+
+// a per-row backward with one scalar parameter (eps for RMS_NORM_BACK, scale for SOFT_MAX_BACK)
+typedef struct {
+    int32_t  ne00;
+    uint64_t nb01;
+    uint64_t nb11;
+    uint64_t nb1;
+    float    param;
+} ggml_metal_kargs_row_back;
+
 typedef struct {
     int64_t  ne00;
     int64_t  ne01;