diff --git a/ggml/src/ggml-sycl/fwht.cpp b/ggml/src/ggml-sycl/fwht.cpp index 39f273beaa..fb48d7fec7 100644 --- a/ggml/src/ggml-sycl/fwht.cpp +++ b/ggml/src/ggml-sycl/fwht.cpp @@ -124,6 +124,107 @@ static void launch_fwht(const float * src, float * dst, const int64_t n_rows, co }); } +// Wide blocks: one row per work-group instead of per sub-group, so each work-item +// keeps N/NT values rather than N/WARP_SIZE. Butterflies below the sub-group width +// still shuffle; those up to NT go through work-group local memory; the rest stay +// in registers. +template +static void fwht_kernel_wide(const float * __restrict__ src, + float * __restrict__ dst, + const int64_t n_rows, + const float scale, + const sycl::nd_item<2> & item, + float * smem) { + const int64_t r = item.get_global_id(0); + if (r >= n_rows) { + return; + } + + src += r * N; + dst += r * N; + + constexpr int el_w = N / NT; + static_assert(el_w >= 1 && N % NT == 0, "row must be a whole number of work-group widths"); + + const int tid = item.get_local_id(1); + + float reg[el_w]; +#pragma unroll + for (int i = 0; i < el_w; ++i) { + reg[i] = src[i * NT + tid] * scale; + } + + const sycl::sub_group sg = item.get_sub_group(); + const int lane = sg.get_local_linear_id(); + + // Butterflies inside the sub-group, same pattern as the narrow kernel. +#pragma unroll + for (int h = 1; h < WARP_SIZE; h *= 2) { +#pragma unroll + for (int j = 0; j < el_w; ++j) { + const float val = reg[j]; + const float val2 = dpct::permute_sub_group_by_xor(sg, val, h, WARP_SIZE); + + reg[j] = (lane & h) == 0 ? val + val2 : val2 - val; + } + } + + // Butterflies from the sub-group width up to NT: the partner lane is outside + // this sub-group, so it goes through work-group local memory instead of a shuffle. + for (int h = WARP_SIZE; h < NT; h *= 2) { +#pragma unroll + for (int j = 0; j < el_w; ++j) { + smem[j * NT + tid] = reg[j]; + } + item.barrier(sycl::access::fence_space::local_space); +#pragma unroll + for (int j = 0; j < el_w; ++j) { + const float val = reg[j]; + const float val2 = smem[j * NT + (tid ^ h)]; + reg[j] = (tid & h) == 0 ? val + val2 : val2 - val; + } + item.barrier(sycl::access::fence_space::local_space); + } + + // Butterflies across registers: h is a multiple of NT, so the partner of element + // i*NT + tid lives in reg[i + h/NT] on the same work-item. + for (int h = NT; h < N; h *= 2) { + const int step = h / NT; + for (int j = 0; j < el_w; j += 2 * step) { + for (int k = 0; k < step; ++k) { + const float x = reg[j + k]; + const float y = reg[j + k + step]; + + reg[j + k] = x + y; + reg[j + k + step] = x - y; + } + } + } + +#pragma unroll + for (int i = 0; i < el_w; ++i) { + dst[i * NT + tid] = reg[i]; + } +} + +template +static void launch_fwht_wide(const float * src, + float * dst, + const int64_t n_rows, + const float scale, + dpct::queue_ptr stream) { + const sycl::range<2> global(n_rows, NT); + const sycl::range<2> local(1, NT); + + stream->submit([&](sycl::handler & cgh) { + sycl::local_accessor smem(sycl::range<1>(N), cgh); + cgh.parallel_for(sycl::nd_range<2>(global, local), + [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + fwht_kernel_wide(src, dst, n_rows, scale, item, get_pointer(smem)); + }); + }); +} + template static void kronecker_kernel(const float * __restrict__ src, float * __restrict__ dst, @@ -285,6 +386,18 @@ bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src, case 1280: launch_kronecker<1280, 20>(src_d, dst_d, rows, scale, stream); return true; + case 1024: + launch_fwht_wide<1024, 256>(src_d, dst_d, rows, scale, stream); + return true; + case 2048: + launch_fwht_wide<2048, 256>(src_d, dst_d, rows, scale, stream); + return true; + case 4096: + launch_fwht_wide<4096, 256>(src_d, dst_d, rows, scale, stream); + return true; + case 8192: + launch_fwht_wide<8192, 256>(src_d, dst_d, rows, scale, stream); + return true; default: return false; }