opencl : fix Q8 tensor readback from generic buffers

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -3894,6 +3894,7 @@ struct ggml_tensor_extra_cl_q8_0 {
 
     size_t size_q = 0;
     size_t size_d = 0;
+    bool adreno_transposed = false;
 
     ~ggml_tensor_extra_cl_q8_0() {
         reset();
@@ -3917,6 +3918,7 @@ struct ggml_tensor_extra_cl_q8_0 {
         d_img = nullptr;
         size_q = 0;
         size_d = 0;
+        adreno_transposed = false;
     }
 };
 
@@ -4246,6 +4248,15 @@ static ggml_status ggml_backend_opencl_graph_compute(ggml_backend_t backend, ggm
     return GGML_STATUS_SUCCESS;
 }
 
+enum ggml_opencl_q8_0_layout {
+    GGML_OPENCL_Q8_0_LAYOUT_UNKNOWN,
+    GGML_OPENCL_Q8_0_LAYOUT_GENERIC,
+    GGML_OPENCL_Q8_0_LAYOUT_SOA,
+};
+
+static ggml_opencl_q8_0_layout ggml_cl_q8_0_layout(const ggml_tensor * tensor);
+static bool ggml_cl_q8_0_copy_supported(const ggml_tensor * src0, const ggml_tensor * src1);
+
 static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
     ggml_backend_opencl_device_context * dev_ctx     = (ggml_backend_opencl_device_context *)dev->context;
     ggml_backend_opencl_context *        backend_ctx = dev_ctx->backend_ctx;
@@ -4310,6 +4321,18 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                         default:
                             return false;
                     }
+                case GGML_TYPE_Q8_0: {
+                    if (!(op->op == GGML_OP_CPY &&
+                        op->type == GGML_TYPE_Q8_0 &&
+                        op->src[1] != nullptr &&
+                        op->src[1]->type == GGML_TYPE_Q8_0 &&
+                        ggml_is_contiguous(op->src[0]) &&
+                        ggml_is_contiguous(op->src[1]))) {
+                        return false;
+                    }
+
+                    return ggml_cl_q8_0_copy_supported(op->src[0], op->src[1]);
+                }
                 default:
                     return false;
             }
@@ -4810,6 +4833,44 @@ struct ggml_backend_opencl_buffer_context {
         return extra;
     }
 
+    template <typename T>
+    static bool contains_extra(const std::vector<T *> & extras, const void * extra) {
+        for (const T * e : extras) {
+            if (e == extra) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    bool is_soa_q4_0_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_q4_0_in_use, extra);
+    }
+
+    bool is_soa_q4_1_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_q4_1_in_use, extra);
+    }
+
+    bool is_soa_mxfp4_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_mxfp4_in_use, extra);
+    }
+
+    bool is_soa_q8_0_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_q8_0_in_use, extra);
+    }
+
+    bool is_soa_q4_K_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_q4_K_in_use, extra);
+    }
+
+    bool is_soa_q5_K_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_q5_K_in_use, extra);
+    }
+
+    bool is_soa_q6_K_extra(const void * extra) const {
+        return contains_extra(temp_tensor_extras_q6_K_in_use, extra);
+    }
+
     void reset() {
         for (ggml_tensor_extra_cl * e : temp_tensor_extras_in_use) {
             temp_tensor_extras.push_back(e);
@@ -4890,6 +4951,73 @@ struct ggml_backend_opencl_buffer_context {
     std::string name;
 };
 
+static ggml_opencl_q8_0_layout ggml_cl_q8_0_layout(const ggml_tensor * tensor) {
+    if (tensor == nullptr || tensor->type != GGML_TYPE_Q8_0 || tensor->extra == nullptr) {
+        return GGML_OPENCL_Q8_0_LAYOUT_UNKNOWN;
+    }
+
+    const ggml_tensor * base = tensor->view_src != nullptr ? tensor->view_src : tensor;
+    if (base->buffer == nullptr ||
+        base->buffer->buft == nullptr ||
+        base->buffer->buft->iface.get_name != ggml_backend_opencl_buffer_type_get_name) {
+        return GGML_OPENCL_Q8_0_LAYOUT_UNKNOWN;
+    }
+
+    ggml_backend_opencl_buffer_context * ctx = (ggml_backend_opencl_buffer_context *) base->buffer->context;
+    return ctx->is_soa_q8_0_extra(tensor->extra) ?
+        GGML_OPENCL_Q8_0_LAYOUT_SOA :
+        GGML_OPENCL_Q8_0_LAYOUT_GENERIC;
+}
+
+static bool ggml_cl_q8_0_is_adreno_transposed(const ggml_tensor * tensor) {
+    if (ggml_cl_q8_0_layout(tensor) != GGML_OPENCL_Q8_0_LAYOUT_SOA) {
+        return false;
+    }
+
+    ggml_tensor_extra_cl_q8_0 * extra = (ggml_tensor_extra_cl_q8_0 *) tensor->extra;
+    return extra->adreno_transposed;
+}
+
+static bool ggml_cl_q8_0_same_shape(const ggml_tensor * src0, const ggml_tensor * src1) {
+    for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+        if (src0->ne[i] != src1->ne[i]) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+static bool ggml_cl_q8_0_copy_supported(const ggml_tensor * src0, const ggml_tensor * src1) {
+    const ggml_opencl_q8_0_layout src0_layout = ggml_cl_q8_0_layout(src0);
+    const ggml_opencl_q8_0_layout src1_layout = ggml_cl_q8_0_layout(src1);
+    if (src0_layout == GGML_OPENCL_Q8_0_LAYOUT_UNKNOWN ||
+        src1_layout == GGML_OPENCL_Q8_0_LAYOUT_UNKNOWN ||
+        src0_layout != src1_layout) {
+        return false;
+    }
+
+    if (src0_layout == GGML_OPENCL_Q8_0_LAYOUT_GENERIC) {
+        return true;
+    }
+
+    const bool src0_transposed = ggml_cl_q8_0_is_adreno_transposed(src0);
+    const bool src1_transposed = ggml_cl_q8_0_is_adreno_transposed(src1);
+    if (src0_transposed != src1_transposed) {
+        return false;
+    }
+
+    if (!src0_transposed) {
+        return true;
+    }
+
+    return src0->view_src == nullptr &&
+           src1->view_src == nullptr &&
+           src0->view_offs == 0 &&
+           src1->view_offs == 0 &&
+           ggml_cl_q8_0_same_shape(src0, src1);
+}
+
 static void ggml_backend_opencl_buffer_free_buffer(ggml_backend_buffer_t buffer) {
     ggml_backend_opencl_buffer_context * ctx = (ggml_backend_opencl_buffer_context *) buffer->context;
     GGML_LOG_INFO("[DEBUG] opencl buf FREE  cl_mem=%p size=%.2f MiB\n",
@@ -5484,7 +5612,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
 
         // Transpose the weights and scales
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
-        if (enable_adreno_trans_weight(backend_ctx, tensor)) {
+        if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && enable_adreno_trans_weight(backend_ctx, tensor)) {
 
             int M = tensor->ne[1];   // ne01
             int K = tensor->ne[0];   // ne00
@@ -5496,6 +5624,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
 
             transpose_2d_as_32b(backend_ctx, extra->q, extra->q, size_q, K/4,  M);
             transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/32, M);
+            extra->adreno_transposed = true;
         } // end transpose
 #endif // GGML_OPENCL_USE_ADRENO_KERNELS
 
@@ -5846,6 +5975,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
     GGML_ASSERT(tensor->extra);
 
     ggml_backend_opencl_context *backend_ctx = ggml_cl2_init(buffer->buft->device);
+    ggml_backend_opencl_buffer_context * ctx = (ggml_backend_opencl_buffer_context *) buffer->context;
 
     cl_context context = backend_ctx->context;
     cl_command_queue queue = backend_ctx->queue;
@@ -5860,7 +5990,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
     // which requires reading back quantized weight tensors.
     // To properly support this, we need to restore block_q4_0 struct arrays
     // from the flattened buffers.
-    if (tensor->type == GGML_TYPE_Q4_0) {
+    if (tensor->type == GGML_TYPE_Q4_0 && ctx->is_soa_q4_0_extra(tensor->extra)) {
         ggml_tensor_extra_cl_q4_0 * extra = (ggml_tensor_extra_cl_q4_0 *)tensor->extra;
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
@@ -5971,7 +6101,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(clReleaseMemObject(data_device));
         return;
     }
-    if (tensor->type == GGML_TYPE_Q4_1) {
+    if (tensor->type == GGML_TYPE_Q4_1 && ctx->is_soa_q4_1_extra(tensor->extra)) {
         ggml_tensor_extra_cl_q4_1 * extra = (ggml_tensor_extra_cl_q4_1 *)tensor->extra;
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
@@ -6045,7 +6175,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(clReleaseMemObject(data_device));
         return;
     }
-    if (tensor->type == GGML_TYPE_MXFP4) {
+    if (tensor->type == GGML_TYPE_MXFP4 && ctx->is_soa_mxfp4_extra(tensor->extra)) {
         ggml_tensor_extra_cl_mxfp4 * extra = (ggml_tensor_extra_cl_mxfp4 *)tensor->extra;
 
         cl_int err;
@@ -6098,7 +6228,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(clReleaseMemObject(data_device));
         return;
     }
-    if (tensor->type == GGML_TYPE_Q8_0) {
+    if (tensor->type == GGML_TYPE_Q8_0 && ctx->is_soa_q8_0_extra(tensor->extra)) {
         ggml_tensor_extra_cl_q8_0 * extra = (ggml_tensor_extra_cl_q8_0 *)tensor->extra;
 
         cl_int err;
@@ -6107,7 +6237,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(err);
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
-        if (enable_adreno_trans_weight(backend_ctx, tensor)) {
+        if (extra->adreno_transposed) {
             cl_kernel kernel = backend_ctx->kernel_restore_block_q8_0_trans;
 
             int ne00 = tensor->ne[0];
@@ -6154,7 +6284,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(clReleaseMemObject(data_device));
         return;
     }
-    if (tensor->type == GGML_TYPE_Q4_K) {
+    if (tensor->type == GGML_TYPE_Q4_K && ctx->is_soa_q4_K_extra(tensor->extra)) {
         ggml_tensor_extra_cl_q4_K * extra = (ggml_tensor_extra_cl_q4_K *)tensor->extra;
 
         cl_int err;
@@ -6230,7 +6360,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(clReleaseMemObject(data_device));
         return;
     }
-    if (tensor->type == GGML_TYPE_Q5_K) {
+    if (tensor->type == GGML_TYPE_Q5_K && ctx->is_soa_q5_K_extra(tensor->extra)) {
         ggml_tensor_extra_cl_q5_K * extra = (ggml_tensor_extra_cl_q5_K *)tensor->extra;
 
         cl_int err;
@@ -6312,7 +6442,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
         CL_CHECK(clReleaseMemObject(data_device));
         return;
     }
-    if (tensor->type == GGML_TYPE_Q6_K) {
+    if (tensor->type == GGML_TYPE_Q6_K && ctx->is_soa_q6_K_extra(tensor->extra)) {
         ggml_tensor_extra_cl_q6_K * extra = (ggml_tensor_extra_cl_q6_K *)tensor->extra;
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
@@ -11368,7 +11498,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
 
     // q8_0 x fp32
     if (src0t == GGML_TYPE_Q8_0 && src1t == GGML_TYPE_F32 &&
-        enable_adreno_trans_weight(backend_ctx, src0)) {
+        ggml_cl_q8_0_is_adreno_transposed(src0)) {
             ggml_cl_mul_mat_q8_0_f32_adreno(backend, src0, src1, dst);
             return;
     }
@@ -13115,6 +13245,93 @@ static void ggml_cl_scale(ggml_backend_t backend, const ggml_tensor * src0, cons
     backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size_ptr, dst);
 }
 
+static bool ggml_cl_has_soa_q8_0_extra(const ggml_tensor * tensor) {
+    return ggml_cl_q8_0_layout(tensor) == GGML_OPENCL_Q8_0_LAYOUT_SOA;
+}
+
+static const char * ggml_cl_q8_0_layout_name(ggml_opencl_q8_0_layout layout) {
+    switch (layout) {
+        case GGML_OPENCL_Q8_0_LAYOUT_GENERIC:
+            return "generic";
+        case GGML_OPENCL_Q8_0_LAYOUT_SOA:
+            return "SOA";
+        case GGML_OPENCL_Q8_0_LAYOUT_UNKNOWN:
+            return "unknown";
+    }
+
+    return "invalid";
+}
+
+static const char * ggml_cl_q8_0_storage_name(const ggml_tensor * tensor) {
+    if (ggml_cl_q8_0_is_adreno_transposed(tensor)) {
+        return "SOA transposed";
+    }
+
+    return ggml_cl_q8_0_layout_name(ggml_cl_q8_0_layout(tensor));
+}
+
+static void ggml_cl_copy_buffer_region(ggml_backend_opencl_context * backend_ctx, cl_mem src, size_t src_offset, cl_mem dst, size_t dst_offset, size_t size) {
+    if (size == 0) {
+        return;
+    }
+
+    const bool overlaps = src == dst && src_offset < dst_offset + size && dst_offset < src_offset + size;
+    if (!overlaps) {
+        CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, src, dst, src_offset, dst_offset, size, 0, NULL, NULL));
+        return;
+    }
+
+    cl_int err;
+    cl_mem tmp = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, size, NULL, &err);
+    CL_CHECK(err);
+    CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, src, tmp, src_offset, 0, size, 0, NULL, NULL));
+    CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, tmp, dst, 0, dst_offset, size, 0, NULL, NULL));
+    CL_CHECK(clReleaseMemObject(tmp));
+}
+
+static void ggml_cl_cpy_q8_0(ggml_backend_opencl_context * backend_ctx, const ggml_tensor * src0, const ggml_tensor * src1) {
+    GGML_ASSERT(src0->type == GGML_TYPE_Q8_0);
+    GGML_ASSERT(src1->type == GGML_TYPE_Q8_0);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(src1));
+    GGML_ASSERT(ggml_nelements(src0) == ggml_nelements(src1));
+    GGML_ASSERT(src0->view_offs % ggml_type_size(src0->type) == 0);
+    GGML_ASSERT(src1->view_offs % ggml_type_size(src1->type) == 0);
+
+    if (!ggml_cl_q8_0_copy_supported(src0, src1)) {
+        GGML_LOG_ERROR("%s: unsupported Q8_0 OpenCL copy from %s layout to %s layout\n",
+            __func__, ggml_cl_q8_0_storage_name(src0), ggml_cl_q8_0_storage_name(src1));
+        GGML_ABORT("unsupported Q8_0 OpenCL copy layout");
+    }
+
+    const ggml_opencl_q8_0_layout src0_layout = ggml_cl_q8_0_layout(src0);
+    if (src0_layout == GGML_OPENCL_Q8_0_LAYOUT_SOA) {
+        ggml_tensor_extra_cl_q8_0 * extra0 = (ggml_tensor_extra_cl_q8_0 *) src0->extra;
+        ggml_tensor_extra_cl_q8_0 * extra1 = (ggml_tensor_extra_cl_q8_0 *) src1->extra;
+
+        const size_t block_count = ggml_nelements(src0) / ggml_blck_size(src0->type);
+        const size_t src0_block_offset = src0->view_offs / ggml_type_size(src0->type);
+        const size_t src1_block_offset = src1->view_offs / ggml_type_size(src1->type);
+
+        ggml_cl_copy_buffer_region(
+            backend_ctx, extra0->d, src0_block_offset * sizeof(ggml_fp16_t),
+            extra1->d, src1_block_offset * sizeof(ggml_fp16_t),
+            block_count * sizeof(ggml_fp16_t));
+        ggml_cl_copy_buffer_region(
+            backend_ctx, extra0->q, src0_block_offset * ggml_blck_size(src0->type),
+            extra1->q, src1_block_offset * ggml_blck_size(src1->type),
+            block_count * ggml_blck_size(src0->type));
+        return;
+    }
+
+    ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *) src0->extra;
+    ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *) src1->extra;
+    ggml_cl_copy_buffer_region(
+        backend_ctx, extra0->data_device, extra0->offset + src0->view_offs,
+        extra1->data_device, extra1->offset + src1->view_offs,
+        ggml_nbytes(src0));
+}
+
 static void ggml_cl_cpy(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -13135,6 +13352,11 @@ static void ggml_cl_cpy(ggml_backend_t backend, const ggml_tensor * src0, const
 
     ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
 
+    if (src0t == GGML_TYPE_Q8_0 && src1t == GGML_TYPE_Q8_0) {
+        ggml_cl_cpy_q8_0(backend_ctx, src0, src1);
+        return;
+    }
+
     ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra;
     ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra;