diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index dc8a17443..d58ffd00d 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -2402,7 +2402,138 @@ static void argsort_f32_i32_sycl(const float *x, int *dst, const int ncols, } } +// Scan and block merge, shared by every launch shape below so a partitioned row uses the +// same insertion order as an unpartitioned one. +// +// src_map != nullptr: report src_map[col] instead of col, so a merge pass can carry the +// original column index through. +// out_vals != nullptr: also emit the k winning values, for a later merge pass. +// swap01: emit in the output order the single-pass path uses. +static void top_k_scan_merge_f32( + const float * src_vals, + const int32_t * src_map, + const int begin, + const int end, + const int k, + const int block_size, + float * shared_vals, + int * shared_idx, + float * out_vals, + int32_t * out_idx, + const bool swap01, + const sycl::nd_item<1> & item_ct1 +) { + const int tid = item_ct1.get_local_id(0); + + // The running top-k lives in SLM (shared local memory) rather than a private array: + // an array indexed by a runtime position cannot be register-allocated, so a private + // one lands in scratch, i.e. device memory, and insertion is this kernel's dominant + // cost. + // + // Lane-strided (lv[i * block_size]) rather than lane-blocked (lv[i]) so a given i is + // contiguous across lanes; a k-strided layout would put every lane of a shift step in + // the same SLM bank. + float * lv = shared_vals + tid; + int * li = shared_idx + tid; + + for (int i = 0; i < k; i++) { + lv[i * block_size] = -FLT_MAX; + li[i * block_size] = -1; + } + + // The k-th best, cached in a register. The reject test is taken for the large + // majority of elements scanned, and in that case touches no memory. + float kth = -FLT_MAX; + + for (int col = begin + tid; col < end; col += block_size) { + float val = src_vals[col]; + + if (val > kth) { + int pos = k - 1; + while (pos > 0 && val > lv[(pos - 1) * block_size]) { + pos--; + } + + for (int i = k - 1; i > pos; i--) { + lv[i * block_size] = lv[(i - 1) * block_size]; + li[i * block_size] = li[(i - 1) * block_size]; + } + lv[pos * block_size] = val; + li[pos * block_size] = src_map ? src_map[col] : col; + + kth = lv[(k - 1) * block_size]; + } + } + + item_ct1.barrier(sycl::access::fence_space::local_space); + + if (tid != 0) { + return; + } + + // Same treatment for the merge accumulator, past the per-lane region. + float * fv = shared_vals + (size_t) k * block_size; + int * fi = shared_idx + (size_t) k * block_size; + + for (int i = 0; i < k; i++) { + fv[i] = -FLT_MAX; + fi[i] = -1; + } + + float fkth = -FLT_MAX; + + // Candidates are visited in the same (t, i) order as before, so tie-breaking is + // unchanged. + for (int t = 0; t < block_size; t++) { + for (int i = 0; i < k; i++) { + float val = shared_vals[i * block_size + t]; + + if (val <= fkth) { + // Lane t's list is sorted descending, so once one of its entries loses + // to the k-th best, every later entry loses too. fkth only rises, so + // that stays true for the rest of the merge. This turns the merge from + // block_size*k steps into roughly block_size plus the candidates + // accepted. + break; + } + + int idx = shared_idx[i * block_size + t]; + + int pos = k - 1; + while (pos > 0 && val > fv[pos - 1]) { + pos--; + } + + for (int j = k - 1; j > pos; j--) { + fv[j] = fv[j - 1]; + fi[j] = fi[j - 1]; + } + fv[pos] = val; + fi[pos] = idx; + + fkth = fv[k - 1]; + } + } + + if (out_vals) { + for (int i = 0; i < k; i++) { + out_vals[i] = fv[i]; + } + } + + for (int i = 0; i < k; i++) { + out_idx[i] = fi[i]; + } + + if (swap01 && k > 1) { + int32_t temp = out_idx[0]; + out_idx[0] = out_idx[1]; + out_idx[1] = temp; + } +} + static void top_k_f32_sycl( + ggml_backend_sycl_context & ctx, const float * src, int32_t * dst_indices, const int64_t ncols, @@ -2410,98 +2541,107 @@ static void top_k_f32_sycl( const int k, dpct::queue_ptr main_stream ) { + // A row is scanned by exactly one work-group, so a vocabulary-sized row leaves the + // rest of the device idle. What the scan is short of is memory requests in flight, + // not bandwidth or per-request latency, so lanes in flight is the lever: split the + // row across independent work-groups, have each emit its partition's top-k, and + // merge those nsplit*k candidates in a second launch. + // + // split_block trades parallelism against SLM residency. Its cost is + // (split_block + 1) * k * 8 bytes of SLM per group, so at the k <= 32 ceiling 128 + // lanes need about 33 KB, which leaves a single resident group per Xe-core. Revisit + // if the supported k ever grows. + constexpr int split_block = 128; + constexpr int max_splits = 128; + constexpr int min_cols = 8192; + + int nsplit = 1; + if (ncols >= min_cols) { + // A partition is then always >= split_block = 128 columns, hence always more than + // the k <= 32 ceiling, so no pass is ever padded with -FLT_MAX sentinels. + const int64_t want = ncols / split_block; + nsplit = (int) (want > max_splits ? max_splits : want); + } + + if (nsplit > 1) { + const int nchunk = (int) ((ncols + nsplit - 1) / nsplit); + const size_t ncand = (size_t) nrows * nsplit * k; + + ggml_sycl_pool_alloc part_vals(ctx.pool(), ncand); + ggml_sycl_pool_alloc part_idx(ctx.pool(), ncand); + + float * pv = part_vals.get(); + int32_t * pi = part_idx.get(); + + const sycl::range<1> block_dims(split_block); + + main_stream->submit([&](sycl::handler &cgh) { + sycl::local_accessor shared_vals(sycl::range<1>((split_block + 1) * k), cgh); + sycl::local_accessor shared_idx(sycl::range<1>((split_block + 1) * k), cgh); + + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>(nrows * nsplit) * block_dims, block_dims), + [=](sycl::nd_item<1> item_ct1) { + const int grp = item_ct1.get_group(0); + const int row = grp / nsplit; + const int part = grp % nsplit; + + const int begin = part * nchunk; + int end = begin + nchunk; + if (end > (int) ncols) { + end = (int) ncols; + } + + top_k_scan_merge_f32( + src + (int64_t) row * ncols, nullptr, begin, end, k, split_block, + shared_vals.get_multi_ptr().get(), + shared_idx.get_multi_ptr().get(), + pv + (size_t) grp * k, pi + (size_t) grp * k, false, item_ct1); + }); + }); + + main_stream->submit([&](sycl::handler &cgh) { + sycl::local_accessor shared_vals(sycl::range<1>((split_block + 1) * k), cgh); + sycl::local_accessor shared_idx(sycl::range<1>((split_block + 1) * k), cgh); + + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>(nrows) * block_dims, block_dims), + [=](sycl::nd_item<1> item_ct1) { + const int row = item_ct1.get_group(0); + const size_t off = (size_t) row * nsplit * k; + + top_k_scan_merge_f32( + pv + off, pi + off, 0, nsplit * k, k, split_block, + shared_vals.get_multi_ptr().get(), + shared_idx.get_multi_ptr().get(), + nullptr, dst_indices + (int64_t) row * k, true, item_ct1); + }); + }); + + return; + } + const int block_size = 128; const sycl::range<1> block_dims(block_size); const sycl::range<1> grid_dims(nrows); main_stream->submit([&](sycl::handler &cgh) { - sycl::local_accessor shared_vals(sycl::range<1>(block_size * k), cgh); - sycl::local_accessor shared_idx(sycl::range<1>(block_size * k), cgh); + sycl::local_accessor shared_vals(sycl::range<1>((block_size + 1) * k), cgh); + sycl::local_accessor shared_idx(sycl::range<1>((block_size + 1) * k), cgh); cgh.parallel_for( sycl::nd_range<1>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<1> item_ct1) { const int row = item_ct1.get_group(0); - const int tid = item_ct1.get_local_id(0); if (row >= nrows) return; - const float * src_row = src + row * ncols; - int32_t * dst_idx_row = dst_indices + row * k; - - float local_vals[32]; - int local_idx[32]; - - for (int i = 0; i < k; i++) { - local_vals[i] = -FLT_MAX; - local_idx[i] = -1; - } - - for (int col = tid; col < ncols; col += block_size) { - float val = src_row[col]; - - if (val > local_vals[k-1]) { - int pos = k - 1; - while (pos > 0 && val > local_vals[pos - 1]) { - pos--; - } - - for (int i = k - 1; i > pos; i--) { - local_vals[i] = local_vals[i - 1]; - local_idx[i] = local_idx[i - 1]; - } - local_vals[pos] = val; - local_idx[pos] = col; - } - } - - for (int i = 0; i < k; i++) { - shared_vals[tid * k + i] = local_vals[i]; - shared_idx[tid * k + i] = local_idx[i]; - } - item_ct1.barrier(sycl::access::fence_space::local_space); - - if (tid == 0) { - float final_vals[32]; - int final_idx[32]; - - for (int i = 0; i < k; i++) { - final_vals[i] = -FLT_MAX; - final_idx[i] = -1; - } - - for (int t = 0; t < block_size; t++) { - for (int i = 0; i < k; i++) { - float val = shared_vals[t * k + i]; - int idx = shared_idx[t * k + i]; - - if (val > final_vals[k-1]) { - int pos = k - 1; - while (pos > 0 && val > final_vals[pos - 1]) { - pos--; - } - - for (int j = k - 1; j > pos; j--) { - final_vals[j] = final_vals[j - 1]; - final_idx[j] = final_idx[j - 1]; - } - final_vals[pos] = val; - final_idx[pos] = idx; - } - } - } - - for (int i = 0; i < k; i++) { - dst_idx_row[i] = final_idx[i]; - } - - if (k > 1) { - int32_t temp = dst_idx_row[0]; - dst_idx_row[0] = dst_idx_row[1]; - dst_idx_row[1] = temp; - } - } + top_k_scan_merge_f32( + src + (int64_t) row * ncols, nullptr, 0, (int) ncols, k, block_size, + shared_vals.get_multi_ptr().get(), + shared_idx.get_multi_ptr().get(), + nullptr, dst_indices + (int64_t) row * k, true, item_ct1); }); }); } @@ -2902,7 +3042,7 @@ static void ggml_sycl_op_top_k(ggml_backend_sycl_context & ctx, ggml_tensor * ds GGML_ASSERT(k > 0 && k <= 32); GGML_ASSERT(k <= ncols); - top_k_f32_sycl(src0_dd, dst_dd, ncols, nrows, k, main_stream); + top_k_f32_sycl(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream); } inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {