Skip to content

Commit 3eb956e

Browse files
oleksandr-pavlykoleksandr-pavlyk
committed
Inclusive scan iter chunk update kernel (generic and 1d) improved
The chunk update kernels processed consecutive elements in contiguous memory, hence sub-group memory access pattern was sub-optimal (no coalescing). This PR changes these kernels to process n_wi elements which are sub-group size apart, improving memory access patern. Running a micro-benchmark based on code from gh-1249 (for shape =(n, n,) where n = 4096) with this change: ``` (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=cuda:gpu python index.py 0.010703916665753004 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=cuda:gpu python index.py 0.01079747307597211 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=cuda:gpu python index.py 0.010864820314088353 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ python index.py 0.023878061203975922 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ python index.py 0.023666468500677083 ``` while before: ``` (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=cuda:gpu python index.py 0.011415911812542213 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=cuda:gpu python index.py 0.011722088705196424 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=level_zero:gpu python index.py 0.030126182353813893 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=level_zero:gpu python index.py 0.030459783371986338 ``` Running the same code using NumPy (same size): ``` (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ python index_np.py 0.01416253090698134 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ python index_np.py 0.014979530811413296 ``` The reason Level-Zero device is slower has to do with slow allocation/deallocation bug. OpenCL device has better timing. With this change: ``` (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=opencl:gpu python index.py 0.015038836885381627 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=opencl:gpu python index.py 0.01527448468496678 ``` before: ``` (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=opencl:gpu python index.py 0.01758851639115838 (dev_dpctl) opavlyk@mtl-world:~/repos/dpctl$ ONEAPI_DEVICE_SELECTOR=opencl:gpu python index.py 0.017089676241286926 ```
1 parent ec924c3 commit 3eb956e

File tree

1 file changed

+66
-22
lines changed

1 file changed

+66
-22
lines changed

‎dpctl/tensor/libtensor/include/kernels/accumulators.hpp‎

Lines changed: 66 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -404,25 +404,47 @@ sycl::event inclusive_scan_iter_1d(sycl::queue &exec_q,
404404
const size_t src_size = stack_elem.get_size();
405405
outputT *local_scans = stack_elem.get_local_scans_ptr();
406406

407+
using UpdateKernelName =
408+
class inclusive_scan_1d_iter_chunk_update_krn<
409+
inputT, outputT, n_wi, IndexerT, TransformerT,
410+
NoOpTransformerT, ScanOpT, include_initial>;
411+
412+
const auto &kernel_id = sycl::get_kernel_id<UpdateKernelName>();
413+
414+
auto const &ctx = exec_q.get_context();
415+
auto const &dev = exec_q.get_device();
416+
auto kb = sycl::get_kernel_bundle<sycl::bundle_state::executable>(
417+
ctx, {dev}, {kernel_id});
418+
419+
auto krn = kb.get_kernel(kernel_id);
420+
421+
const std::uint32_t sg_size = krn.template get_info<
422+
sycl::info::kernel_device_specific::max_sub_group_size>(dev);
423+
407424
// output[ chunk_size * (i + 1) + j] += temp[i]
408425
dependent_event = exec_q.submit([&](sycl::handler &cgh) {
409426
cgh.depends_on(dependent_event);
427+
cgh.use_kernel_bundle(kb);
410428

411429
constexpr nwiT updates_per_wi = n_wi;
412-
const size_t n_items = ceiling_quotient<size_t>(src_size, n_wi);
430+
const size_t n_items =
431+
ceiling_quotient<size_t>(src_size, sg_size * n_wi) *
432+
sg_size;
413433

414-
using UpdateKernelName =
415-
class inclusive_scan_1d_iter_chunk_update_krn<
416-
inputT, outputT, n_wi, IndexerT, TransformerT,
417-
NoOpTransformerT, ScanOpT, include_initial>;
434+
sycl::range<1> gRange{n_items};
435+
sycl::range<1> lRange{sg_size};
436+
sycl::nd_range<1> ndRange{gRange, lRange};
418437

419438
cgh.parallel_for<UpdateKernelName>(
420-
{n_items}, [chunk_size, src, src_size, local_scans, scan_op,
421-
identity](auto wiid) {
422-
const size_t gid = n_wi * wiid[0];
439+
ndRange, [chunk_size, src, src_size, local_scans, scan_op,
440+
identity](sycl::nd_item<1> ndit) {
441+
const std::uint32_t lws = ndit.get_local_range(0);
442+
const size_t block_offset =
443+
ndit.get_group(0) * n_wi * lws;
423444
#pragma unroll
424445
for (size_t i = 0; i < updates_per_wi; ++i) {
425-
const size_t src_id = gid + i;
446+
const size_t src_id =
447+
block_offset + ndit.get_local_id(0) + i * lws;
426448
if (src_id < src_size) {
427449
const size_t scan_id = (src_id / chunk_size);
428450
src[src_id] =
@@ -661,33 +683,55 @@ sycl::event inclusive_scan_iter(sycl::queue &exec_q,
661683
outputT *local_scans = stack_elem.get_local_scans_ptr();
662684
size_t local_stride = stack_elem.get_local_stride();
663685

686+
using UpdateKernelName = class inclusive_scan_iter_chunk_update_krn<
687+
inputT, outputT, n_wi, TransformerT, NoOpTransformerT, ScanOpT,
688+
include_initial>;
689+
690+
const auto &kernel_id = sycl::get_kernel_id<UpdateKernelName>();
691+
692+
auto const &ctx = exec_q.get_context();
693+
auto const &dev = exec_q.get_device();
694+
auto kb = sycl::get_kernel_bundle<sycl::bundle_state::executable>(
695+
ctx, {dev}, {kernel_id});
696+
697+
auto krn = kb.get_kernel(kernel_id);
698+
699+
const std::uint32_t sg_size = krn.template get_info<
700+
sycl::info::kernel_device_specific::max_sub_group_size>(dev);
701+
664702
constexpr nwiT updates_per_wi = n_wi;
665703
const size_t update_nelems =
666-
ceiling_quotient<size_t>(src_size, updates_per_wi);
704+
ceiling_quotient<size_t>(src_size, sg_size * updates_per_wi) *
705+
sg_size;
667706

668707
dependent_event = exec_q.submit([&](sycl::handler &cgh) {
669708
cgh.depends_on(dependent_event);
709+
cgh.use_kernel_bundle(kb);
670710

671-
using UpdateKernelName =
672-
class inclusive_scan_iter_chunk_update_krn<
673-
inputT, outputT, n_wi, TransformerT, NoOpTransformerT,
674-
ScanOpT, include_initial>;
711+
sycl::range<1> gRange{iter_nelems * update_nelems};
712+
sycl::range<1> lRange{sg_size};
713+
714+
sycl::nd_range<1> ndRange{gRange, lRange};
675715

676716
cgh.parallel_for<UpdateKernelName>(
677-
{iter_nelems * update_nelems},
717+
ndRange,
678718
[chunk_size, update_nelems, src_size, local_stride, src,
679-
local_scans, scan_op, identity](auto wiid) {
680-
const size_t gid = wiid[0];
719+
local_scans, scan_op, identity](sycl::nd_item<1> ndit) {
720+
const size_t gr_id = ndit.get_group(0);
681721

682-
const size_t iter_gid = gid / update_nelems;
683-
const size_t axis_gid =
684-
gid - (iter_gid * update_nelems);
722+
const size_t iter_gid = gr_id / update_nelems;
723+
const size_t axis_gr_id =
724+
gr_id - (iter_gid * update_nelems);
685725

686-
const size_t src_axis_id0 = axis_gid * updates_per_wi;
726+
const std::uint32_t lws = ndit.get_local_range(0);
727+
728+
const size_t src_axis_id0 =
729+
axis_gr_id * updates_per_wi * lws;
687730
const size_t src_iter_id = iter_gid * src_size;
688731
#pragma unroll
689732
for (nwiT i = 0; i < updates_per_wi; ++i) {
690-
const size_t src_axis_id = src_axis_id0 + i;
733+
const size_t src_axis_id =
734+
src_axis_id0 + ndit.get_local_id(0) + i * lws;
691735
const size_t src_id = src_axis_id + src_iter_id;
692736

693737
if (src_axis_id < src_size) {

0 commit comments

Comments
 (0)