allow both operation and source tensors to be repacked [no ci]

This a a work in progress see how I can refactor this but until then
this allows both operation tensors and source tensor to be repacked.

Before this commit I had broken the repacking for the source tensor
as I'd only been running test-backend-ops.cpp. Using a model that has
tensors the support repacking would fails which is the reason for this
fix.

I also want to see if this fixes an issue I ran into on macos before
reworking the testing side of repack/extra buffer types.

Author: danbev

ggml/src/ggml-backend.cpp

@@ -1993,14 +1993,14 @@ struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, s
 
         // Again just here to see if I can get the repacking to work.
         if (extra_buffer && !ggml_op_is_empty(node->op) && node->src[0]) {
-            auto dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(extra_buffer));
-            if (ggml_backend_dev_supports_op(dev, node)) {
-                size_t id = ggml_hash_find(&hash_set, node->src[0]);
-                ggml_status status = ggml_backend_buffer_init_tensor(extra_buffer, node_copies[id]);
-                if (status != GGML_STATUS_SUCCESS) {
-                    GGML_LOG_ERROR("%s: failed to initialize tensor in extra buffer for graph copy\n", __func__);
+            size_t id = ggml_hash_find(&hash_set, node);
+            ggml_status status = ggml_backend_buffer_init_tensor(extra_buffer, node_copies[id]);
+            if (status == GGML_STATUS_SUCCESS) {
+                if (node_copies[id]->extra != nullptr) {
+                    node_copies[id]->buffer = extra_buffer;
                 }
-                node_copies[id]->buffer = extra_buffer;
+            } else {
+                GGML_LOG_ERROR("%s: failed to initialize tensor in extra buffer for graph copy\n", __func__);
             }
         }
 

ggml/src/ggml-cpu/repack.cpp

@@ -1869,8 +1869,45 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons
     return nullptr;
 }
 
+static bool supports_tensor(const struct ggml_tensor * op) {
+    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;
+        }
+
+    } 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;
+        }
+    }
+    return false;
+}
+
 static enum ggml_status ggml_backend_cpu_repack_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) {
-    tensor->extra = (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_repack_get_optimal_repack_type(tensor));
+    if (tensor->op == GGML_OP_NONE) {
+        tensor->extra = (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_repack_get_optimal_repack_type(tensor));
+    }
+
+    if (supports_tensor(tensor)) {
+        tensor->src[0]->extra = (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_repack_get_optimal_repack_type(tensor->src[0]));
+    }
 
     GGML_UNUSED(buffer);
     return GGML_STATUS_SUCCESS;
@@ -1918,39 +1955,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 supports_tensor(op);
     }
 
     ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {