sycl: contiguous fast path + 32-bit index math for unary elementwise ops (#25946)

* sycl: contiguous fast path + 32-bit index math for unary elementwise ops

* sycl: use fastdiv for elementwise index math

diff --git a/ggml/src/ggml-sycl/element_wise.cpp b/ggml/src/ggml-sycl/element_wise.cpp
index b2406e11b5afc73e1bb65e520b15c45cdbfe7812..3cd055494ecf056f5648ea525929d65fec1305e6 100644 (file)
--- a/ggml/src/ggml-sycl/element_wise.cpp
+++ b/ggml/src/ggml-sycl/element_wise.cpp
@@ -306,29 +306,43 @@ static __dpct_inline__ T op_trunc(T x) {
     }
 }
 
+template<typename T, typename F>
+static void unary_op_flat_kernel(const T * x, T * dst, const int k, const sycl::nd_item<1> & item_ct1, F func) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        dst[i] = func(x[i]);
+    }
+}
+
 template<typename T, typename F>
 static void unary_op_generic_kernel(
         const T * x,
         T * dst,
         const int k,
-        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3,
+        const sycl::uint3 ne0_fd, const sycl::uint3 ne1_fd, const sycl::uint3 ne2_fd,
         const size_t nb0,  const size_t nb1,  const size_t nb2,  const size_t nb3,
         const size_t nbd0, const size_t nbd1, const size_t nbd2, const size_t nbd3,
         const sycl::nd_item<1> & item_ct1,
         F func) {
 
-        (void) ne3;
+    // 32-bit index math: k is int, so every logical index fits u32. 64-bit integer div/mod is
+    // emulated on Xe and dominates this kernel otherwise, and even the 32-bit divide is worth
+    // avoiding -- the divisors are launch-invariant, so the magic numbers are precomputed
+    // host-side and each division becomes a multiply-high plus a shift.
+    // Byte offsets are widened back to size_t only for the final address math.
     SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
-        const int64_t i0 =  i % ne0;
-        const int64_t i1 = (i / ne0)        % ne1;
-        const int64_t i2 = (i / (ne0*ne1))  % ne2;
-        const int64_t i3 =  i / (ne0*ne1*ne2);
+        sycl::uint2 dm = fast_div_modulo((uint32_t) i, ne0_fd);
+        const uint32_t i0 = dm.y();
+        dm = fast_div_modulo(dm.x(), ne1_fd);
+        const uint32_t i1 = dm.y();
+        dm = fast_div_modulo(dm.x(), ne2_fd);
+        const uint32_t i2 = dm.y();
+        const uint32_t i3 = dm.x();
 
         const char * src_base = (const char *) x;
         char       * dst_base = (char *) dst;
 
-        const T * srcp = (const T *)(src_base + i0*nb0  + i1*nb1  + i2*nb2  + i3*nb3 );
-        T *       dstp = (T *)(dst_base + i0*nbd0 + i1*nbd1 + i2*nbd2 + i3*nbd3);
+        const T * srcp = (const T *)(src_base + (size_t) i0*nb0  + (size_t) i1*nb1  + (size_t) i2*nb2  + (size_t) i3*nb3 );
+        T *       dstp = (T *)(dst_base + (size_t) i0*nbd0 + (size_t) i1*nbd1 + (size_t) i2*nbd2 + (size_t) i3*nbd3);
 
         *dstp = func(*srcp);
     }
@@ -407,46 +421,51 @@ static void clamp(const T * x, T * dst, const float min, const float max, const
 }
 
 template<typename T>
-static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
     SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
-        const int64_t j0 = (i / n) * o0 + (i % n);
-        const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
         dst[i] = op_gelu(x[j0]) * g[j1];
     }
 }
 
 template<typename T>
-static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
     SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
-        const int64_t j0 = (i / n) * o0 + (i % n);
-        const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
         dst[i] = op_relu(x[j0]) * g[j1];
     }
 }
 
 template<typename T>
-static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
     SYCL_GLOBAL_ID_LOOP(k, item_ct1)  {
-        const int64_t j0 = (i / n) * o0 + (i % n);
-        const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
         dst[i] = op_silu(x[j0]) * g[j1];
     }
 }
 
 template<typename T>
-static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
     SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
-        const int64_t j0 = (i / n) * o0 + (i % n);
-        const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
         dst[i] = op_gelu_erf(x[j0]) * g[j1];
     }
 }
 
 template<typename T>
-static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
     SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
-        const int64_t j0 = (i / n) * o0 + (i % n);
-        const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
         dst[i] = op_gelu_quick(x[j0]) * g[j1];
     }
 }
@@ -529,6 +548,10 @@ static inline void dispatch_ggml_sycl_op_fused_glu(ggml_backend_sycl_context & c
     GGML_ASSERT(dst->ne[0] == nc);
     GGML_ASSERT(ggml_is_contiguous_1(dst->src[0]));
     GGML_ASSERT(ggml_is_contiguous(dst));
+    // The fused GLU kernels index with 32-bit fastdiv, which is exact only for indices below
+    // 2^31. A dst that large is ~8 GB at f32, and the grid sizing already narrows to 32 bits,
+    // so assert the bound rather than carry a second code path for it.
+    GGML_ASSERT(ggml_nelements(dst) < ((int64_t) 1 << 31));
     const int32_t swapped = ((const int32_t *) dst->op_params)[1];
     void * src0_d = src0->data;
     void * src1_d = src1 ? src1->data : src0->data;
@@ -597,7 +620,6 @@ static inline void ggml_sycl_op_unary(
     const int64_t ne0  = dst->ne[0];
     const int64_t ne1  = dst->ne[1];
     const int64_t ne2  = dst->ne[2];
-    const int64_t ne3  = dst->ne[3];
 
     const size_t  nb0  = src0->nb[0];
     const size_t  nb1  = src0->nb[1];
@@ -609,24 +631,42 @@ static inline void ggml_sycl_op_unary(
     const size_t  nbd2 = dst->nb[2];
     const size_t  nbd3 = dst->nb[3];
 
+    // Hot unary ops (FFN/GDN silu, sigmoid, ...) run on contiguous tensors;
+    // skip the strided index math entirely for them.
+    const bool contiguous = ggml_is_contiguous(src0) && ggml_is_contiguous(dst);
+
     ggml_sycl_detail::dispatch_ggml_sycl_op_unary(ctx, dst,
         [=](const auto* src, auto* dst_ptr, int k_elements, queue_ptr stream) {
 
             const int num_blocks = ceil_div(k_elements, 256);
 
-            stream->parallel_for(
-                sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256),
-                                  sycl::range<1>(256)),
-                [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                    unary_op_generic_kernel(
-                        src, dst_ptr, k_elements,
-                        ne0, ne1, ne2, ne3,
-                        nb0, nb1, nb2, nb3,
-                        nbd0, nbd1, nbd2, nbd3,
-                        item_ct1,
-                        func
-                    );
-                });
+            if (contiguous) {
+                stream->parallel_for(
+                    sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256),
+                                      sycl::range<1>(256)),
+                    [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                        unary_op_flat_kernel(src, dst_ptr, k_elements, item_ct1, func);
+                    });
+            } else {
+                // Launch-invariant divisors: compute the magic numbers once on the host so the
+                // kernel never issues an integer divide. Only the strided path needs them.
+                const sycl::uint3 ne0_fd = init_fastdiv_values((uint32_t) ne0);
+                const sycl::uint3 ne1_fd = init_fastdiv_values((uint32_t) ne1);
+                const sycl::uint3 ne2_fd = init_fastdiv_values((uint32_t) ne2);
+                stream->parallel_for(
+                    sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256),
+                                      sycl::range<1>(256)),
+                    [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                        unary_op_generic_kernel(
+                            src, dst_ptr, k_elements,
+                            ne0_fd, ne1_fd, ne2_fd,
+                            nb0, nb1, nb2, nb3,
+                            nbd0, nbd1, nbd2, nbd3,
+                            item_ct1,
+                            func
+                        );
+                    });
+            }
         });
 }
 
@@ -930,10 +970,11 @@ static inline void ggml_sycl_op_geglu(ggml_backend_sycl_context & ctx, ggml_tens
     ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
         [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
             const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
             main_stream->parallel_for(
                     sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
                     sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+                gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
             });
         });
 }
@@ -942,10 +983,11 @@ static inline void ggml_sycl_op_reglu(ggml_backend_sycl_context & ctx, ggml_tens
     ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
         [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
             const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_RELU_BLOCK_SIZE); // Using RELU block size for reglu
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
             main_stream->parallel_for(
                     sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_RELU_BLOCK_SIZE)),
                     sycl::range<1>(SYCL_RELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+                gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
             });
         });
 }
@@ -954,10 +996,11 @@ static inline void ggml_sycl_op_swiglu(ggml_backend_sycl_context & ctx, ggml_ten
     ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
         [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
             const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_SILU_BLOCK_SIZE); // Using SILU block size for swiglu
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
             main_stream->parallel_for(
                     sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_SILU_BLOCK_SIZE)),
                     sycl::range<1>(SYCL_SILU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+                gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
             });
         });
 }
@@ -1057,10 +1100,11 @@ static inline void ggml_sycl_op_geglu_erf(ggml_backend_sycl_context & ctx, ggml_
     ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
         [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
             const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
             main_stream->parallel_for(
                     sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
                     sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+                gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
             });
         });
 }
@@ -1069,10 +1113,11 @@ static inline void ggml_sycl_op_geglu_quick(ggml_backend_sycl_context & ctx, ggm
     ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
         [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
             const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
             main_stream->parallel_for(
                     sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
                     sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1);
+                gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
             });
         });
 }