ggml-cpu : check src[1] dims are 2 for repack

This commit adds an additional check for src[1] dimensions to be 2
when determining if a tensor supports repacking.

The motivation for this change is to ensure that both source tensors
are strictly 2D before using repack. The repack implementation does
not support broadcasting in dimensions 2 and 3, which occurs when
src1 has more dimensions than src0 (like when nr != [1,1] in
test-backend-ops.cpp).

Without this check, operations with broadcasting would use repack and
produce incorrect results because repack assumes.
With this check, broadcasting operations fall back to the standard
CPU implementation which correctly handles the index mapping
(i02 = i12/r2, i03 = i13/r3).

This fixes test failures like:
```console
MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[2,1])
MUL_MAT(type_a=q4_K,type_b=f32,m=16,n=16,k=256,bs=[1,1],nr=[1,2])
```
which were consistently failing across all architectures (x86, ARM,
macOS) with high NMSE values (~0.4-0.7).

Author: danbev

ggml/src/ggml-cpu/repack.cpp

@@ -1869,10 +1869,13 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons
     return nullptr;
 }
 
-static bool supports_tensor(const struct ggml_tensor * op) {
+static bool repack_supports_op(const struct ggml_tensor * op) {
     if (op->op == GGML_OP_MUL_MAT &&
         op->src[0]->buffer &&
-        (ggml_n_dims(op->src[0]) == 2) && ggml_repack_get_optimal_repack_type(op->src[0])) {
+        (ggml_n_dims(op->src[0]) == 2) &&
+        (ggml_n_dims(op->src[1]) == 2) &&
+        op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() &&
+        ggml_repack_get_optimal_repack_type(op->src[0])) {
 
         if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
             return false;
@@ -1882,8 +1885,12 @@ static bool supports_tensor(const struct ggml_tensor * op) {
             return true;
         }
 
-    } else if (op->op == GGML_OP_MUL_MAT_ID && op->src[0]->buffer &&
-            (ggml_n_dims(op->src[0]) == 3) && ggml_repack_get_optimal_repack_type(op->src[0])) {
+    } else if (op->op == GGML_OP_MUL_MAT_ID &&
+               op->src[0]->buffer &&
+               (ggml_n_dims(op->src[0]) == 3) &&
+               (ggml_n_dims(op->src[1]) == 2) &&
+               op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() &&
+               ggml_repack_get_optimal_repack_type(op->src[0])) {
 
         if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
             return false;
@@ -1902,7 +1909,7 @@ static enum ggml_status ggml_backend_cpu_repack_buffer_init_tensor(ggml_backend_
         tensor->buffer = buffer;
     }
 
-    if (supports_tensor(tensor)) {
+    if (repack_supports_op(tensor)) {
         tensor->src[0]->extra = (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_repack_get_optimal_repack_type(tensor->src[0]));
         tensor->src[0]->buffer = buffer;
     }
@@ -1953,39 +1960,7 @@ static size_t ggml_backend_cpu_repack_buffer_type_get_alignment(ggml_backend_buf
 namespace ggml::cpu::repack {
 class extra_buffer_type : ggml::cpu::extra_buffer_type {
     bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override {
-        if (    op->op == GGML_OP_MUL_MAT &&
-                op->src[0]->buffer &&
-                (ggml_n_dims(op->src[0]) == 2) &&
-                op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() &&
-                ggml_repack_get_optimal_repack_type(op->src[0])
-                ) {
-            if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
-                return false;
-            }
-            if (op->src[1]->type == GGML_TYPE_F32) {
-                return true;
-            }
-            //if (op->src[1]->type == GGML_TYPE_Q8_0) {
-            //    return true;
-            //}
-            // may be possible if Q8_0 packed...
-        } else if (op->op == GGML_OP_MUL_MAT_ID
-                && op->src[0]->buffer
-                && (ggml_n_dims(op->src[0]) == 3)
-                && op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type()
-                && ggml_repack_get_optimal_repack_type(op->src[0])
-                ) {
-            if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
-                return false;
-            }
-            if (op->src[1]->type == GGML_TYPE_F32) {
-                return true;
-            }
-            //if (op->src[1]->type == GGML_TYPE_Q8_0) {
-            //    return true;
-            //}
-        }
-        return false;
+        return repack_supports_op(op);
     }
 
     ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {