sycl: XMX path for PQ2_0 and PTQ1_0 on 16-wide DPAS devices

ggml/src/ggml-sycl/CMakeLists.txt

@@ -215,4 +215,22 @@ if (GGML_SYCL_DEVICE_ARCH)
         "SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\""
         -fsycl-max-parallel-link-jobs=${GGML_SYCL_MAX_PARALLEL_LINK_JOBS}
     )
+
+    # The PQ2_0/PTQ1_0 XMX kernels (pq2_xmx.cpp) need 16-wide DPAS and 2D block loads, and AOT compiles them for
+    # every listed device, so they are only built when all of them are Xe-HPC, Xe2 or later. Otherwise those
+    # types keep the existing paths.
+    string(REPLACE "," ";" _ggml_sycl_aot_devices "${GGML_SYCL_DEVICE_ARCH}")
+    set(_ggml_sycl_pq2_xmx ON)
+    foreach (_ggml_sycl_dev IN LISTS _ggml_sycl_aot_devices)
+        string(STRIP "${_ggml_sycl_dev}" _ggml_sycl_dev)
+        string(TOLOWER "${_ggml_sycl_dev}" _ggml_sycl_dev)
+        # the -vg parts of Xe-HPC have no XMX
+        if (NOT _ggml_sycl_dev MATCHES "^(pvc|bmg|lnl|ptl|wcl|nvl|cri|xe-hpc|xe2|xe3)" OR _ggml_sycl_dev MATCHES "-vg")
+            set(_ggml_sycl_pq2_xmx OFF)
+        endif()
+    endforeach()
+    if (NOT _ggml_sycl_pq2_xmx)
+        message(STATUS "GGML_SYCL_DEVICE_ARCH includes a device without 16-wide DPAS, not building the PQ2_0/PTQ1_0 XMX path")
+        set_property(SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/pq2_xmx.cpp APPEND PROPERTY COMPILE_DEFINITIONS GGML_SYCL_NO_PQ2_XMX)
+    endif()
 endif()

ggml/src/ggml-sycl/common.hpp

@@ -224,6 +224,7 @@ inline dpct::err0 ggml_sycl_set_device(const int device) try {
 //////////////////////
 struct optimize_feature {
     bool reorder=false;
+    bool xmx_pq2=false; // PQ2_0 rewritten into the XMX layout (pq2_xmx.hpp); only that path can read it
 };
 
 struct sycl_device_info {
@@ -243,6 +244,7 @@ struct sycl_device_info {
     sycl_hw_info hw_info;
     optimize_feature opt_feature;
     bool    usm_system_support; // support for USM system allocations
+    int     dpas_exec_size;     // XMX DPAS width (8 or 16) for int8, 0 without XMX
 };
 
 

ggml/src/ggml-sycl/getrows.cpp

@@ -277,10 +277,16 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_PTQ1_0:
+            GGML_ASSERT(!(dst->src[0]->extra &&
+                          ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.xmx_pq2) &&
+                        "PTQ1_0 in the XMX layout reached get_rows");
             get_rows_sycl<QK_PTQ1_0, 1, dequantize_ptq1_0>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_PQ2_0:
+            GGML_ASSERT(!(dst->src[0]->extra &&
+                          ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.xmx_pq2) &&
+                        "PQ2_0 in the XMX layout reached get_rows");
             get_rows_sycl<QK_PQ2_0, 1, dequantize_pq2_0>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -56,6 +56,7 @@
 
 #include "ggml-sycl/add-id.hpp"
 #include "ggml-sycl/backend.hpp"
+#include "ggml-sycl/pq2_xmx.hpp"
 #include "ggml-sycl/common.hpp"
 #include "ggml-sycl/element_wise.hpp"
 #include "ggml-sycl/fwht.hpp"
@@ -106,6 +107,23 @@ int g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL;
 int g_ggml_sycl_usm_system = 0;
 int g_ggml_sycl_enable_host_pinned_mem = 1;
 
+// int8 DPAS execution size of the device's XMX units (8 or 16), as the runtime reports it; 0 without XMX
+static int ggml_sycl_dpas_exec_size(const sycl::device & device) {
+    namespace matrix = syclex::matrix;
+    if (!device.has(sycl::aspect::ext_intel_matrix)) {
+        return 0;
+    }
+    try {
+        for (const matrix::combination & c : device.get_info<syclex::info::device::matrix_combinations>()) {
+            if (c.atype == matrix::matrix_type::sint8 && c.btype == matrix::matrix_type::sint8) {
+                return (int) c.nsize;
+            }
+        }
+    } catch (const sycl::exception &) {
+    }
+    return 0;
+}
+
 static ggml_sycl_device_info ggml_sycl_init() {
     ggml_sycl_device_info info = {};
 
@@ -170,6 +188,7 @@ static ggml_sycl_device_info ggml_sycl_init() {
         info.max_work_group_sizes[i] = prop.get_max_work_group_size();
         info.devices[i].max_wg_per_cu = info.max_work_group_sizes[i] / prop.get_max_compute_units();
         info.devices[i].hw_info = get_device_hw_info(&device);
+        info.devices[i].dpas_exec_size = ggml_sycl_dpas_exec_size(device);
 
         // Only check GPU devices; CPU devices use OpenCL and would otherwise
         // disable Level Zero for the GPUs on systems without ONEAPI_DEVICE_SELECTOR set.
@@ -593,7 +612,9 @@ ggml_backend_sycl_buffer_init_tensor(ggml_backend_buffer_t buffer,
             case GGML_TYPE_Q3_K:
             case GGML_TYPE_Q4_K:
             case GGML_TYPE_Q5_K:
-            case GGML_TYPE_Q6_K:{
+            case GGML_TYPE_Q6_K:
+            case GGML_TYPE_PQ2_0:
+            case GGML_TYPE_PTQ1_0:{
                 ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{};
                 tensor->extra                 = extra;
                 ctx->tensor_extras.push_back(extra);
@@ -954,6 +975,26 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ
     GGML_UNUSED(buft);
 }
 
+static bool ggml_sycl_device_has_dpas16(int device);
+
+// GGML_SYCL_DISABLE_XMX=1 keeps PQ2_0/PTQ1_0 off the XMX path (and its layout)
+static bool ggml_sycl_xmx_disabled() {
+    static const bool disabled = ggml_sycl_get_env("GGML_SYCL_DISABLE_XMX", 0);
+    return disabled;
+}
+
+// PTQ1_0 weights are expanded into the 34-byte PQ2_0 XMX blocks on first use (pq2_xmx.hpp), so on devices
+// that run that path their allocation reserves room for the expanded form
+static bool ggml_sycl_ptq1_xmx_expands(int device, const ggml_tensor * tensor) {
+    return tensor->type == GGML_TYPE_PTQ1_0 && g_ggml_sycl_enable_optimize && !ggml_sycl_xmx_disabled() &&
+           tensor->ne[2] == 1 && tensor->ne[3] == 1 && ggml_sycl_pq2_xmx_supports_ne0(tensor->ne[0]) &&
+           ggml_sycl_device_has_dpas16(device);
+}
+
+static size_t ggml_sycl_ptq1_xmx_bytes(const ggml_tensor * tensor) {
+    return (size_t) (ggml_nelements(tensor) / QK_PTQ1_0) * sizeof(block_pq2_0);
+}
+
 static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
     size_t size = ggml_nbytes(tensor);
     int64_t ne0 = tensor->ne[0];
@@ -964,9 +1005,12 @@ static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_t
         }
     }
 
-    return size;
+    const auto * buft_ctx = (const ggml_backend_sycl_buffer_type_context *) buft->context;
+    if (ggml_sycl_ptq1_xmx_expands(buft_ctx->device, tensor)) {
+        size = std::max(size, ggml_sycl_ptq1_xmx_bytes(tensor));
+    }
 
-    GGML_UNUSED(buft);
+    return size;
 }
 
 static const ggml_backend_buffer_type_i ggml_backend_sycl_buffer_type_interface = {
@@ -3764,6 +3808,13 @@ inline bool ggml_sycl_supports_mmq(enum ggml_type type) {
     return false;
 }
 
+// The PQ2_0/PTQ1_0 XMX path feeds 2-bit weights to ESIMD DPAS at execution size 16 through 2D block loads, which
+// every XMX device with 16-wide DPAS has (Xe-HPC, Xe2 and later). 8-wide XMX (Xe-HPG, Arrow Lake-H) keeps the
+// existing paths.
+static bool ggml_sycl_device_has_dpas16(int device) {
+    return ggml_sycl_info().devices[device].dpas_exec_size == 16;
+}
+
 inline bool ggml_sycl_supports_reorder_mul_mat_sycl(enum ggml_type type) {
     switch (type) {
         case GGML_TYPE_Q1_0:
@@ -4529,6 +4580,50 @@ static bool can_use_mul_mat_vec_q(const ggml_tensor * src0, const ggml_tensor *
            src1->ne[1] <= MMVQ_MAX_BATCH_SIZE;
 }
 
+// PQ2_0/PTQ1_0 weights on 16-wide DPAS devices are rewritten into the XMX layout on first use. From then on every
+// mul_mat on them has to take that path, so the layout flag alone decides once it is set.
+static bool ggml_sycl_pq2_xmx_use(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1,
+                                  const ggml_tensor * dst) {
+    if (src0->type != GGML_TYPE_PQ2_0 && src0->type != GGML_TYPE_PTQ1_0) {
+        return false;
+    }
+    // MUL_MAT_ID passes each expert as a 2D copy of the 3D weight that shares its extra, and a view shares
+    // its parent's data: rewriting either in place would corrupt the rest of the tensor
+    if (dst->op != GGML_OP_MUL_MAT || src0->view_src != nullptr) {
+        return false;
+    }
+    ggml_tensor_extra_gpu * extra = static_cast<ggml_tensor_extra_gpu *>(src0->extra);
+    if (extra && extra->optimized_feature.xmx_pq2) {
+        return true;
+    }
+    if (!g_ggml_sycl_enable_optimize || ggml_sycl_xmx_disabled() || !ggml_sycl_device_has_dpas16(ctx.device)) {
+        return false;
+    }
+    // op offload refills COMPUTE buffers from host memory every time, so an in-place layout there goes stale
+    if (!extra || ggml_backend_buffer_is_sycl_split(src0->buffer) ||
+        src0->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+        return false;
+    }
+    if (src0->ne[2] != 1 || src0->ne[3] != 1 || !ggml_is_contiguous(src0) ||
+        !ggml_sycl_pq2_xmx_supports_ne0(src0->ne[0]) || (uintptr_t) src0->data % 64 != 0) {
+        return false;
+    }
+    if (src1->type != GGML_TYPE_F32 || src1->nb[0] != sizeof(float) || dst->type != GGML_TYPE_F32 ||
+        !ggml_is_contiguous(dst)) {
+        return false;
+    }
+    // PTQ1_0 expands to 34 bytes a block, which only fits where the buffer reserved room for it
+    if (src0->type == GGML_TYPE_PTQ1_0 &&
+        ggml_backend_buft_get_alloc_size(src0->buffer->buft, src0) < ggml_sycl_ptq1_xmx_bytes(src0)) {
+        return false;
+    }
+    if (!ggml_sycl_pq2_xmx_reorder(const_cast<ggml_tensor *>(src0), ctx.stream())) {
+        return false;
+    }
+    extra->optimized_feature.xmx_pq2 = true;
+    return true;
+}
+
 static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
 
@@ -4543,6 +4638,11 @@ static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor
         return;
     }
 
+    if (ggml_sycl_pq2_xmx_use(ctx, src0, src1, dst)) {
+        ggml_sycl_pq2_xmx_mul_mat(ctx, src0, src1, dst);
+        return;
+    }
+
     const bool split = ggml_backend_buffer_is_sycl_split(src0->buffer);
     int64_t min_compute_capability = INT_MAX;