SYCLUtils.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2026 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
 
#pragma once
 
#include <cstdint>
#include <vector>
 
#include "open3d/core/Device.h"
#include "open3d/core/SYCLContext.h"
 
#ifdef SYCL_LANGUAGE_VERSION
#include <algorithm>
#include <sycl/sycl.hpp>
#endif
 
namespace open3d {
namespace core {
namespace sy {
 
void PrintSYCLDevices(bool print_all = false);
 
bool IsAvailable();
 
bool IsDeviceAvailable(const Device& device);
 
SYCLDevice GetSYCLDeviceProperties(const Device& device);
 
bool IsCPUDevice(const Device& device);
 
std::vector<Device> GetAvailableSYCLDevices();
 
inline size_t GetDeviceCount() { return GetAvailableSYCLDevices().size(); }
 
void enablePersistentJITCache();
 
#if defined(SYCL_LANGUAGE_VERSION) && defined(BUILD_SYCL_MODULE)
 
inline size_t PreferredWorkGroupSize(const sycl::device& sycl_device) {
    size_t max_wg =
            sycl_device.get_info<sycl::info::device::max_work_group_size>();
    size_t wg = std::min<size_t>(256, max_wg);
    auto sg_sizes = sycl_device.get_info<sycl::info::device::sub_group_sizes>();
    size_t sg = sg_sizes.empty()
                        ? 1
                        : *std::max_element(sg_sizes.begin(), sg_sizes.end());
    if (sg > 1 && wg >= sg) {
        wg -= wg % sg;
    }
    return std::max<size_t>(wg, 1);
}
 
inline size_t PreferredWorkGroupSize(const Device& device) {
    return PreferredWorkGroupSize(GetQueue(device).get_device());
}
 
inline size_t MaxWorkGroupSizeForSLM(const sycl::device& sycl_device,
                                     size_t slm_bytes_per_wi,
                                     size_t sg_size) {
    size_t wg = sycl_device.get_info<sycl::info::device::max_work_group_size>();
    if (slm_bytes_per_wi > 0) {
        size_t local_mem_size =
                sycl_device.get_info<sycl::info::device::local_mem_size>();
        wg = std::min<size_t>(wg, local_mem_size / slm_bytes_per_wi);
    }
    size_t sg = std::max<size_t>(sg_size, 1);
    if (sg > 1 && wg >= sg) {
        wg -= wg % sg;
    }
    return std::max<size_t>(wg, sg);
}
 
inline size_t MaxWorkGroupSizeForSLM(const Device& device,
                                     size_t slm_bytes_per_wi,
                                     size_t sg_size) {
    return MaxWorkGroupSizeForSLM(GetQueue(device).get_device(),
                                  slm_bytes_per_wi, sg_size);
}
 
inline size_t MaxWorkGroupSizeForSLM(const Device& device,
                                     size_t slm_bytes_per_wi) {
    const SYCLDevice props = GetSYCLDeviceProperties(device);
    const size_t sg_size =
            props.sub_group_sizes.empty()
                    ? 1
                    : *std::max_element(props.sub_group_sizes.begin(),
                                        props.sub_group_sizes.end());
    return MaxWorkGroupSizeForSLM(device, slm_bytes_per_wi, sg_size);
}
 
template <int N, typename scalar_t, typename Func>
inline void PersistentReduce(sycl::queue& queue,
                             int64_t n,
                             size_t wgs,
                             scalar_t* global_sum_ptr,
                             Func&& compute_local_sum) {
    const size_t compute_units =
            SYCLContext::GetInstance().GetComputeUnits(queue.get_device());
    const int64_t natural_num_groups =
            std::max<int64_t>(1, (n + int64_t(wgs) - 1) / int64_t(wgs));
    // One work-group per compute unit is enough to saturate the device for
    // this persistent-kernel pattern; never launch more groups than there is
    // work for.
    const size_t num_groups = static_cast<size_t>(std::min<int64_t>(
            std::max<int64_t>(1, int64_t(compute_units)), natural_num_groups));
 
    sycl::buffer<scalar_t, 1> partial_sum_buf(sycl::range<1>(num_groups * N));
    sycl::buffer<int, 1> ticket_buf(sycl::range<1>(1));
    queue.submit([&](sycl::handler& cgh) {
        auto ticket_acc =
                ticket_buf
                        .template get_access<sycl::access::mode::discard_write>(
                                cgh);
        cgh.single_task([=]() { ticket_acc[0] = 0; });
    });
 
    queue.submit([&](sycl::handler& cgh) {
             // partial_sum layout is [k * num_groups + group_id] (item 17):
             // transposed from the original [group_id * N + k] so the merge
             // phase's per-lane reads (below) are contiguous across groups.
             auto partial_sum_acc = partial_sum_buf.template get_access<
                     sycl::access::mode::read_write>(cgh);
             auto ticket_acc = ticket_buf.template get_access<
                     sycl::access::mode::read_write>(cgh);
             sycl::local_accessor<int, 1> is_last(sycl::range<1>(1), cgh);
             cgh.parallel_for(
                     sycl::nd_range<1>{num_groups * wgs, wgs},
                     [=](sycl::nd_item<1> item) {
                         const size_t group_id = item.get_group(0);
                         const size_t lid = item.get_local_id(0);
                         const int64_t global_stride =
                                 int64_t(num_groups) * int64_t(wgs);
 
                         scalar_t local_sum[N] = {};
                         for (int64_t gid = int64_t(group_id * wgs + lid);
                              gid < n; gid += global_stride) {
                             compute_local_sum(gid, local_sum);
                         }
 
                         auto grp = item.get_group();
                         for (int k = 0; k < N; ++k) {
                             scalar_t v = sycl::reduce_over_group(
                                     grp, local_sum[k], sycl::plus<scalar_t>{});
                             if (lid == 0) {
                                 partial_sum_acc[size_t(k) * num_groups +
                                                 group_id] = v;
                             }
                         }
 
                         if (lid == 0) {
                             sycl::atomic_ref<int, sycl::memory_order::acq_rel,
                                              sycl::memory_scope::device>
                                     tick_ref(ticket_acc[0]);
                             int my_ticket = tick_ref.fetch_add(1);
                             is_last[0] = (my_ticket ==
                                           static_cast<int>(num_groups) - 1);
                         }
                         item.barrier(sycl::access::fence_space::local_space);
 
                         if (is_last[0]) {
                             // Merge phase (item 17): each lane privately
                             // sums its share of one output k's num_groups
                             // partial sums (no contention -- each lane
                             // touches distinct memory), then
                             // reduce_over_group combines the wgs per-lane
                             // partials. Replaces the old SLM-atomic
                             // accumulation of all num_groups * N partials
                             // into only N addresses.
                             for (int k = 0; k < N; ++k) {
                                 scalar_t partial = 0;
                                 for (size_t g = lid; g < num_groups;
                                      g += wgs) {
                                     partial += partial_sum_acc
                                             [size_t(k) * num_groups + g];
                                 }
                                 scalar_t v = sycl::reduce_over_group(
                                         grp, partial, sycl::plus<scalar_t>{});
                                 if (lid == 0) {
                                     global_sum_ptr[k] = v;
                                 }
                             }
                         }
                     });
         }).wait_and_throw();
}
 
#endif  // SYCL_LANGUAGE_VERSION && BUILD_SYCL_MODULE
 
}  // namespace sy
}  // namespace core
}  // namespace open3d