diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt index 551ef95..b31f9cf 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt @@ -43,6 +43,12 @@ set(SYCL_DEVICE_LIST -Xs "-device ${TORCH_XPU_ARCH_LIST} -options -cl-poison-uns # -------------------------------------------------------------------------- set(host_sources ${CMAKE_CURRENT_SOURCE_DIR}/ops_registry.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sycl_kernels/invert_permute_kernel.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sycl_kernels/permute_1d_sparse_data.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/asynchronous_complete_cumsum.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/fbgemm_utils/utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sycl_kernels/permute_2d_sparse_data.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/permute_2d_sparse_dataOp.cpp ) set(all_sources ${host_sources}) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.cpp new file mode 100644 index 0000000..20751e0 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.cpp @@ -0,0 +1,45 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#include "asynchronous_complete_cumsum.h" + +namespace fbgemm_xpu { + +// ============================================================================ +// Host Function - XPU Implementation +// ============================================================================ + +at::Tensor asynchronous_complete_cumsum_xpu(const at::Tensor& t_in) { + // Input validation + TORCH_CHECK(t_in.is_contiguous(), "Input tensor must be contiguous"); + TORCH_CHECK( + t_in.dtype() == at::kInt || t_in.dtype() == at::kLong, + "Input tensor must have dtype int32 or int64"); + TORCH_CHECK( + t_in.dim() == 1 || t_in.dim() == 2, + "Input tensor must be 1D or 2D"); + + // Handle 1D case: input [a, b, c] → output [0, a, a+b, a+b+c] + if (t_in.dim() == 1) { + at::Tensor t_out = at::zeros({t_in.numel() + 1}, t_in.options()); + auto r_out = t_out.slice(0, 1); // View excluding the first element + at::cumsum_out(r_out, t_in, 0); // Compute cumsum into the view + return t_out; + } + + // Handle 2D case: cumsum along dimension 1 (columns) + // input shape [M, N] → output shape [M, N+1] + at::Tensor t_out = at::zeros({t_in.size(0), t_in.size(1) + 1}, t_in.options()); + auto r_out = t_out.slice(1, 1); // View excluding the first column + at::cumsum_out(r_out, t_in, 1); // Compute cumsum along dim 1 into the view + return t_out; +} + +TORCH_LIBRARY_IMPL(fbgemm, XPU, m) { + m.impl("asynchronous_complete_cumsum", &fbgemm_xpu::asynchronous_complete_cumsum_xpu); +} + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.h b/packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.h new file mode 100644 index 0000000..4103789 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +//////////////////////////////////////////////////////////////////////////////// +// SYCL PORT MAPPING TO FBGEMM CUDA SOURCE - ASYNCHRONOUS COMPLETE CUMSUM OPERATOR +//////////////////////////////////////////////////////////////////////////////// +// +// This file contains SYCL ports of FBGEMM asynchronous_complete_cumsum operator. +// +// ORIGINAL CUDA SOURCE: +// File: fbgemm_gpu/src/sparse_ops/sparse_async_cumsum.cu +// Header: fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h +// +// HOST FUNCTION MAPPING: +// asynchronous_complete_cumsum_xpu (SYCL) +// → asynchronous_complete_cumsum_gpu (CUDA) +// +// DESCRIPTION: +// Computes complete cumulative sum with a leading zero on Intel XPU devices. +// Given input [a, b, c, d], returns [0, a, a+b, a+b+c, a+b+c+d]. +// Supports both 1D and 2D tensors (cumsum along last dimension for 2D). +// Uses PyTorch's native cumsum operation for computation. +// +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include +#include + +namespace fbgemm_xpu { + +/** + * @brief Computes complete cumulative sum with a leading zero. + * + * This function computes a cumulative sum with a prepended zero element. + * For a 1D input [a, b, c], it returns [0, a, a+b, a+b+c]. + * For a 2D input, cumsum is computed along dimension 1 (columns). + * + * @param t_in Input tensor (1D or 2D) with dtype int32 or int64 + * @return Output tensor with one additional element along the cumsum dimension + * + * @pre t_in.is_contiguous() == true + * @pre t_in.dtype() == at::kInt || t_in.dtype() == at::kLong + * @pre t_in.dim() == 1 || t_in.dim() == 2 + */ +at::Tensor asynchronous_complete_cumsum_xpu(const at::Tensor& t_in); + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/.gitkeep b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/DeviceProperties.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/DeviceProperties.h new file mode 100644 index 0000000..59a668e --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/DeviceProperties.h @@ -0,0 +1,220 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#include + +#include +#include + +namespace syclext = sycl::ext::oneapi; +namespace syclexp = sycl::ext::oneapi::experimental; + +namespace xpu { +namespace sycl { + +template +static int64_t syclMaxWorkGroupSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto& ctx = c10::xpu::get_device_context(); + auto& dev = c10::xpu::get_raw_device(dev_id); + + auto kid = ::sycl::get_kernel_id(); + // The kernel won't be built for devices except for the first device. + // Launching kernel on devices except for the first device will raise + // runtime error. Here is an alternative as a temporary solution to + // provide an extra hint to SYCL runtime. + // https://github.com/intel/llvm/issues/15127 + auto kbundle = ::sycl::get_kernel_bundle<::sycl::bundle_state::executable>( + ctx, {dev}, {kid}); + + ::sycl::kernel k = kbundle.get_kernel(kid); + return k.get_info<::sycl::info::kernel_device_specific::work_group_size>(dev); +} + +template +static int64_t syclMaxWorkGroupSize( + const KernelClass& /*kfn*/, + at::DeviceIndex dev_id = at::xpu::current_device()) { + return syclMaxWorkGroupSize(dev_id); +} + +// For SYCL free function +template +static int64_t syclMaxWorkGroupSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto q = c10::xpu::getCurrentXPUStream(dev_id).queue(); + auto ctxt = q.get_context(); + auto dev = q.get_device(); + auto exe_bndl = + ::syclexp::get_kernel_bundle( + ctxt); + ::sycl::kernel k = exe_bndl.template ext_oneapi_get_kernel(); + return k.get_info<::sycl::info::kernel_device_specific::work_group_size>(dev); +} + +static inline int64_t syclDeviceMaxWorkGroupSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->max_work_group_size; +} + +static inline int64_t syclMaxSubGroupSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + const auto& subgroup_sizes = dev_prop->sub_group_sizes; + TORCH_CHECK( + !subgroup_sizes.empty(), + "The device subgroup sizes is empty, please check the device status."); + return *std::max_element(subgroup_sizes.begin(), subgroup_sizes.end()); +} + +static inline int64_t syclMinSubGroupSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + const auto& subgroup_sizes = dev_prop->sub_group_sizes; + TORCH_CHECK( + !subgroup_sizes.empty(), + "The device subgroup sizes is empty, please check the device status."); + return *std::min_element(subgroup_sizes.begin(), subgroup_sizes.end()); +} + +static inline int64_t syclMaxComputeUnitSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->max_compute_units; +} + +static inline int64_t syclGpuEuCount( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->gpu_eu_count; +} + +static inline int64_t syclGpuEuSimdWidth( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->gpu_eu_simd_width; +} + +static inline int64_t syclGpuHWThreadsPerEU( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->gpu_hw_threads_per_eu; +} + +static inline int64_t syclGpuEUCountPerSubslice( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->gpu_eu_count_per_subslice; +} + +static inline int64_t syclMaxWorkItemsPerTile( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + int64_t eu_cnt = dev_prop->gpu_eu_count; + int64_t simd_width = syclMaxSubGroupSize(dev_id); + int64_t hw_threads = dev_prop->gpu_hw_threads_per_eu; + return eu_cnt * simd_width * hw_threads; +} + +static inline int64_t syclMaxWorkItemsPerSubSlice( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + int64_t simd_width = syclMaxSubGroupSize(dev_id); + int64_t eu_count = dev_prop->gpu_eu_count_per_subslice; + return simd_width * eu_count; +} + +static inline int64_t syclMaxWorkItemsPerEU( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + int64_t simd_width = syclMaxSubGroupSize(dev_id); + int64_t hw_threads = dev_prop->gpu_hw_threads_per_eu; + return simd_width * hw_threads; +} + +static inline int64_t syclMaxNumSubGroups( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->max_num_sub_groups; +} + +static inline int64_t syclMaxDSSNum( + at::DeviceIndex dev_id = at::xpu::current_device()) { + int64_t dss_num = + syclMaxComputeUnitSize(dev_id) / syclGpuEUCountPerSubslice(dev_id); + return dss_num; +} + +static inline size_t syclGlobalMemSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->global_mem_size; +} + +static inline int64_t syclLocalMemSize( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->local_mem_size; +} + +template +uint32_t syclPrefVectorWidth( + at::DeviceIndex dev_id = at::xpu::current_device()) { + (void)dev_id; + + constexpr uint32_t vec_width = 16; + + if constexpr ( + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v) { + return vec_width / sizeof(T); + } else { + throw std::invalid_argument( + "Invalid data type to fetch preferred vector width!"); + } +} + +template +uint32_t syclNativeVectorWidth( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_char; + } else if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_short; + } else if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_int; + } else if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_long; + } else if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_float; + } else if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_double; + } else if constexpr (std::is_same_v) { + return dev_prop->native_vector_width_half; + } else { + throw std::invalid_argument( + "Invalid data type to fetch native vector width!"); + } +} + +static inline bool syclHasFloat64( + at::DeviceIndex dev_id = at::xpu::current_device()) { + auto* dev_prop = at::xpu::getDeviceProperties(dev_id); + return dev_prop->has_fp64; +} + +} // namespace sycl +} // namespace xpu \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Macros.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Macros.h new file mode 100644 index 0000000..e2b8a14 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Macros.h @@ -0,0 +1,17 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#ifdef _WIN32 +#define RESTRICT __restrict +#else +#define RESTRICT __restrict__ +#endif \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Runtime.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Runtime.h new file mode 100644 index 0000000..6fbc5de --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Runtime.h @@ -0,0 +1,21 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#include + +namespace at::xpu { + +static inline sycl::queue& getCurrentSYCLQueue() { + return c10::xpu::getCurrentXPUStream().queue(); +} + +} // namespace at::xpu \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLContext.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLContext.h new file mode 100644 index 0000000..82c764a --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLContext.h @@ -0,0 +1,19 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +using namespace at::xpu; +using namespace xpu::sycl; \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLHelpers.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLHelpers.h new file mode 100644 index 0000000..8cd7ba4 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLHelpers.h @@ -0,0 +1,216 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#include +#include // Remove it once the header is exposed by sycl.hpp +#include + +namespace syclext = sycl::ext::oneapi; +namespace syclexp = sycl::ext::oneapi::experimental; + +// sycl access address space +static constexpr auto sycl_priv_space = + sycl::access::address_space::private_space; +static constexpr auto sycl_local_space = + sycl::access::address_space::local_space; +static constexpr auto sycl_global_space = + sycl::access::address_space::global_space; + +// sycl memory ordering +static constexpr auto sycl_mem_odr_rlx = sycl::memory_order::relaxed; +static constexpr auto sycl_mem_odr_acq = sycl::memory_order::acquire; +static constexpr auto sycl_mem_odr_rel = sycl::memory_order::release; +static constexpr auto sycl_mem_odr_acq_rel = sycl::memory_order::acq_rel; +static constexpr auto sycl_mem_odr_seq_cst = sycl::memory_order::seq_cst; + +// sycl memory scope +static constexpr auto sycl_mem_scp_wi = sycl::memory_scope::work_item; +static constexpr auto sycl_mem_scp_sg = sycl::memory_scope::sub_group; +static constexpr auto sycl_mem_scp_wg = sycl::memory_scope::work_group; +static constexpr auto sycl_mem_scp_dev = sycl::memory_scope::device; +static constexpr auto sycl_mem_scp_sys = sycl::memory_scope::system; + +template +using sycl_local_acc_t = sycl::local_accessor; + +template +using sycl_local_ptr = typename sycl::local_ptr; + +template +using sycl_global_ptr = typename sycl::global_ptr; + +template +using sycl_atomic_ref_rlx_dev_global_t = + sycl::atomic_ref; + +template +static inline void sycl_kernel_submit( + ::sycl::range range, + ::sycl::queue q, + ker_t ker) { + auto cgf = [&](::sycl::handler& cgh) { cgh.parallel_for(range, ker); }; + q.submit(cgf); +} + +// Additional convention of SYCL kernel configuration. Besides construct kernel +// functor, SYCL has some additional conventions to be called during setuping +// SYCL command group handler, e.g. declaring SYCL local accessor when the +// kernel requires shared local memory usage. Helpers below help simpilfiy +// submission of SYCL kernels requiring additional conventions. + +// Defining additional convention. Can use `sycl_kernel_submit` simply to +// submit a kernel, if the kernel functor inherits from the struct below. +// Since cannot offload non-device-copyable (sycl::is_device_copyable) kernel +// functor, a structure has virtual function is non-device-copyable. +// Using an empty class, the kernel functor derived by it will be required to +// define member method `void convention(sycl::handler&)`, or fails in +// compilation. +struct __SYCL_KER_CONFIG_CONVENTION__ {}; + +template +static inline typename std::enable_if< + std::is_base_of_v<__SYCL_KER_CONFIG_CONVENTION__, ker_t>, + void>::type +sycl_kernel_submit( + ::sycl::range global_range, + ::sycl::range local_range, + ::sycl::queue q, + ker_t ker) { + auto cgf = [&](::sycl::handler& cgh) { + ker.sycl_ker_config_convention(cgh); + cgh.parallel_for( + ::sycl::nd_range(global_range, local_range), ker); + }; + q.submit(cgf); +} + +template +static inline typename std::enable_if< + !std::is_base_of_v<__SYCL_KER_CONFIG_CONVENTION__, ker_t>, + void>::type +sycl_kernel_submit( + ::sycl::range global_range, + ::sycl::range local_range, + ::sycl::queue q, + ker_t ker) { + auto cgf = [&](::sycl::handler& cgh) { + cgh.parallel_for( + ::sycl::nd_range(global_range, local_range), ker); + }; + q.submit(cgf); +} + +template +static inline typename std::enable_if< + std::is_base_of_v<__SYCL_KER_CONFIG_CONVENTION__, ker_t>, + void>::type +sycl_kernel_submit( + int64_t global_range, + int64_t local_range, + ::sycl::queue q, + ker_t ker) { + auto cgf = [&](::sycl::handler& cgh) { + ker.sycl_ker_config_convention(cgh); + cgh.parallel_for( + ::sycl::nd_range<1>( + ::sycl::range<1>(global_range), ::sycl::range<1>(local_range)), + ker); + }; + q.submit(cgf); +} + +template +static inline typename std::enable_if< + !std::is_base_of_v<__SYCL_KER_CONFIG_CONVENTION__, ker_t>, + void>::type +sycl_kernel_submit( + int64_t global_range, + int64_t local_range, + ::sycl::queue q, + ker_t ker) { + auto cgf = [&](::sycl::handler& cgh) { + cgh.parallel_for( + ::sycl::nd_range<1>( + ::sycl::range<1>(global_range), ::sycl::range<1>(local_range)), + ker); + }; + q.submit(cgf); +} + +// For SYCL free function +template +static inline void sycl_kernel_submit( + int64_t global_range, + int64_t local_range, + ::sycl::queue q, + int slm_sz, + Kargs... args) { + sycl::context ctxt = q.get_context(); + auto exe_bndl = + syclexp::get_kernel_bundle(ctxt); + sycl::kernel ker = exe_bndl.template ext_oneapi_get_kernel(); + if (slm_sz != 0) { + syclexp::launch_config cfg{ + ::sycl::nd_range<1>( + ::sycl::range<1>(global_range), ::sycl::range<1>(local_range)), + syclexp::properties{syclexp::work_group_scratch_size(slm_sz)}}; + syclexp::nd_launch(q, cfg, ker, args...); + } else { + syclexp::launch_config cfg{::sycl::nd_range<1>( + ::sycl::range<1>(global_range), ::sycl::range<1>(local_range))}; + syclexp::nd_launch(q, cfg, ker, args...); + } +} + +template +static inline void sycl_kernel_submit( + ::sycl::range global_range, + ::sycl::range local_range, + ::sycl::queue q, + int slm_sz, + Kargs... args) { + sycl::context ctxt = q.get_context(); + auto exe_bndl = + syclexp::get_kernel_bundle(ctxt); + sycl::kernel ker = exe_bndl.template ext_oneapi_get_kernel(); + if (slm_sz != 0) { + syclexp::launch_config cfg{ + ::sycl::nd_range( + ::sycl::range(global_range), ::sycl::range(local_range)), + syclexp::properties{syclexp::work_group_scratch_size(slm_sz)}}; + syclexp::nd_launch(q, cfg, ker, args...); + } else { + syclexp::launch_config cfg{::sycl::nd_range( + ::sycl::range(global_range), ::sycl::range(local_range))}; + syclexp::nd_launch(q, cfg, ker, args...); + } +} + +#ifdef __SYCL_DEVICE_ONLY__ +#define SYCL_KERNEL_STRING(var, str) \ + static const __attribute__((opencl_constant)) char var[] = str +#else +#define SYCL_KERNEL_STRING(var, str) static const char var[] = str +#endif +#define SYCL_KERNEL_PRINTF sycl::ext::oneapi::experimental::printf + +#define SYCL_PRINT(fmt_str, ...) \ + { \ + SYCL_KERNEL_STRING(fmt_var, fmt_str); \ + SYCL_KERNEL_PRINTF(fmt_var, ##__VA_ARGS__); \ + } + +#ifdef __SYCL_DEVICE_ONLY__ +#define SYCL_REQD_SUB_GROUP_SIZE(x) [[sycl::reqd_sub_group_size(x)]] +#else +#define SYCL_REQD_SUB_GROUP_SIZE(x) +#endif \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Scalar.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Scalar.h new file mode 100644 index 0000000..d2757a8 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Scalar.h @@ -0,0 +1,85 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#include + +union u32_to_f32 { + uint32_t in; + float out; +}; + +union f32_to_u32 { + float in; + uint32_t out; +}; + +union double_to_ull { + double in; + unsigned long long out; +}; + +union ull_to_double { + unsigned long long in; + double out; +}; + +union ushort_to_half { + unsigned short int in; + sycl::half out; +}; + +union double_to_u32 { + double in; + struct { + uint32_t high, low; + }; +}; + +static inline uint32_t __float_as_int(float val) { + f32_to_u32 cn; + cn.in = val; + return cn.out; +} + +static inline float __int_as_float(uint32_t val) { + u32_to_f32 cn; + cn.in = val; + return cn.out; +} + +static inline unsigned long long __double_as_long_long(double val) { + double_to_ull cn; + cn.in = val; + return cn.out; +} + +static inline double __long_long_as_double(unsigned long long val) { + ull_to_double cn; + cn.in = val; + return cn.out; +} + +static inline uint32_t __double_as_int(double val) { + double_to_u32 cn; + cn.in = val; + return cn.low; +} + +static inline float __int_as_double(uint32_t val) { + return (double)val; +} + +static inline sycl::half __ushort_as_half(unsigned short int val) { + ushort_to_half cn; + cn.in = val; + return cn.out; +} \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorInfo.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorInfo.h new file mode 100644 index 0000000..4c8bceb --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorInfo.h @@ -0,0 +1,220 @@ +// Porting from ipex +// will upstream to pytorch when in tree +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +namespace at { +namespace xpu { +namespace detail { + +#define XPU_MAX_TENSORINFO_DIMS 12 + +template +struct TensorInfo { + using scalar_t = T; + + TensorInfo(); + TensorInfo( + T* p, + int dim, + IndexType sz[XPU_MAX_TENSORINFO_DIMS], + IndexType st[XPU_MAX_TENSORINFO_DIMS]); + + void reduceDim(int dim); + + int collapseDims(const int excludeDim = -1); + + IndexType outerSize(const int dim); + + IndexType innerSize(const int dim); + + inline bool isContiguous() const { + return dims == 1 && strides[0] == 1; + } + + inline bool isContiguousCheckStrict(bool strict_contiguous) const { + if (strict_contiguous) { + return is_strict_contiguous; + } else { + return is_contiguous; + } + } + + T* data = nullptr; + IndexType sizes[XPU_MAX_TENSORINFO_DIMS]; + IndexType strides[XPU_MAX_TENSORINFO_DIMS]; + int dims = 0; + bool is_contiguous; + bool is_strict_contiguous; +}; + +template +TensorInfo::TensorInfo() { + data = nullptr; + dims = 0; + is_contiguous = true; + is_strict_contiguous = true; +} + +template +TensorInfo::TensorInfo( + T* p, + int dim, + IndexType sz[XPU_MAX_TENSORINFO_DIMS], + IndexType st[XPU_MAX_TENSORINFO_DIMS]) { + data = p; + dims = dim; + TORCH_INTERNAL_ASSERT(dims <= XPU_MAX_TENSORINFO_DIMS); + + is_contiguous = true; + IndexType z = 1; + for (int i = dim - 1; i >= 0; i--) { + sizes[i] = sz[i]; + strides[i] = st[i]; + + if (is_contiguous && strides[i] == z) { + z *= sizes[i]; + } else { + is_contiguous = false; + } + } + + is_strict_contiguous = dims == 1 && strides[0] == 1; +} + +template +void TensorInfo::reduceDim(int dim) { + TORCH_CHECK(dim < dims && dim >= 0, "expect dim between 0 and dims - 1"); + sizes[dim] = 1; +} + +template +int TensorInfo::collapseDims(const int excludeDim) { + auto result = at::collapse_dims(sizes, strides, dims, excludeDim); + dims = std::get<1>(result); + return std::get<0>(result); +} + +template +IndexType TensorInfo::innerSize(const int exclusive) { + IndexType size = 1; + for (int i = dims - 1; i > exclusive; i--) { + size *= sizes[i]; + } + return size; +} + +template +IndexType TensorInfo::outerSize(const int exclusive) { + IndexType size = 1; + for (int i = 0; i < exclusive; i++) { + size *= sizes[i]; + } + return size; +} + +template +struct IndexToOffset { + static constexpr bool STRICT_CONTIGUOUS = true; + static constexpr bool NON_STRICT_CONTIGUOUS = false; + static inline IndexType get( + IndexType linearId, + const TensorInfo& info, + bool strict_contiguous = true) { + IndexType offset = 0; + + if (info.isContiguousCheckStrict(strict_contiguous)) { + return linearId; + } + + for (int dim = info.dims - 1; dim > 0; --dim) { + IndexType curDimIndex = linearId % info.sizes[dim]; + IndexType curDimOffset = curDimIndex * info.strides[dim]; + offset += curDimOffset; + linearId /= info.sizes[dim]; + } + return offset + linearId * info.strides[0]; + } +}; + +template +struct IndexToOffset { + static constexpr bool STRICT_CONTIGUOUS = true; + static constexpr bool NON_STRICT_CONTIGUOUS = false; + static inline IndexType get( + IndexType linearId, + const TensorInfo& info, + bool strict_contiguous = true) { + return linearId; + } +}; + +template +TensorInfo getTensorInfo(const at::TensorBase& t) { + IndexType sz[XPU_MAX_TENSORINFO_DIMS]; + IndexType st[XPU_MAX_TENSORINFO_DIMS]; + + TORCH_CHECK( + t.dim() <= XPU_MAX_TENSORINFO_DIMS, + "dim:", + t.dim(), + " exceed max allowed dim:", + XPU_MAX_TENSORINFO_DIMS); + + int dims; + if (t.dim()) { + dims = t.dim(); + for (int i = 0; i < dims; ++i) { + sz[i] = t.size(i); + st[i] = t.stride(i); + } + } else { + dims = 1; + sz[0] = 1; + st[0] = 1; + } + + scalar* data_ptr = nullptr; + + if constexpr (std::is_const::value) { + data_ptr = t.const_data_ptr(); + } else { + data_ptr = t.mutable_data_ptr(); + } + + return TensorInfo(data_ptr, dims, sz, st); +} + +} // namespace detail +} // namespace xpu +} // namespace at \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorOptions.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorOptions.h new file mode 100644 index 0000000..dafab37 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorOptions.h @@ -0,0 +1,72 @@ +/* + * Copyright 2020-2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +#pragma once + +#include + +namespace at { +template +static inline TensorOptions map_options() { + if (std::is_same::value) + return at::TensorOptions().dtype(kByte).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kChar).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kShort).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kInt).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kLong).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kFloat).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kDouble).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kHalf).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kBFloat16).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kBool).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same>::value) + return at::TensorOptions() + .dtype(kComplexFloat) + .device(kXPU) + .memory_format(LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same>::value) + return at::TensorOptions() + .dtype(kComplexDouble) + .device(kXPU) + .memory_format(LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kUInt16).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kUInt32).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else if (std::is_same::value) + return at::TensorOptions().dtype(kUInt64).device(kXPU).memory_format( + LEGACY_CONTIGUOUS_MEMORY_FORMAT); + else { + AT_ERROR("PSTLFunctions: data type cannot be mapped to tensor's dtype."); + } + return at::TensorOptions(); +} +} // namespace at \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/dispatch_macros.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/dispatch_macros.h new file mode 100644 index 0000000..fdaf849 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/dispatch_macros.h @@ -0,0 +1,262 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#pragma once + +#include + +#include + +namespace fbgemm_xpu::utils { + +using source_location = std::experimental::source_location; + +//////////////////////////////////////////////////////////////////////////////// +// Source Context +// +// This is a wrapper abstraction around some source context information, +// including the source location, template filepath, and summary string. It is +// used to generate consistent descriptions in log messages around kernel +// executions. +//////////////////////////////////////////////////////////////////////////////// + +struct SourceContext { + // The source location + const source_location location; + // A summary of the context (usually the kernel name) + const std::string_view summary; + // The originating template filepath (for template-generated source files) + const std::string_view template_; + // The file descriptor for DSA error reporting (needs to be generated at + // compile-time) + const std::string_view dsa_file_descriptor_; + + constexpr inline SourceContext( + const source_location& loc_, + const std::string_view& sum_, + const std::string_view& tmpl_, + const std::string_view& dsa_) noexcept + : location(loc_), + summary(sum_), + template_(tmpl_), + dsa_file_descriptor_(dsa_) {} + + inline const std::string description() const noexcept { + // Generate and cache the description if it hasn't been generated yet + std::stringstream ss; + + // Append template source file location if it exists + if (!template_.empty()) { + ss << "[" << template_ << "] "; + } + + ss << "[" << location.file_name() << '(' << location.line() << ':' + << location.column() << ")] [" << summary << "]"; + + return ss.str(); + } + + inline SourceContext withSummary( + const std::string_view& sum_) const noexcept { + return SourceContext(location, sum_, template_, dsa_file_descriptor_); + } +}; + +} // namespace fbgemm_xpu::utils + +#define SOURCE_CONTEXT_CURRENT(LABEL) \ + fbgemm_xpu::utils::SourceContext( \ + fbgemm_gpu::utils::source_location::current(), \ + #LABEL, \ + _FBGEMM_TFILE_, \ + _FBGEMM_DSA_FILESRC_); + +#define FBGEMM_LAUNCH_KERNEL(KERNEL, GRID, BLOCK, SMEM, STREAM, ...) \ + ([&] { \ + constexpr auto context = SOURCE_CONTEXT_CURRENT(KERNEL); \ + auto& kernel = KERNEL; \ + \ + return fbgemm_gpu::utils:: \ + KernelLauncher(context) \ + .launch_kernel(kernel, GRID, BLOCK, SMEM, STREAM, __VA_ARGS__); \ + }()) + + +#define PRIVATE_CASE_TYPE_CACHE(enum_type, type, ...) \ + case enum_type: { \ + using cache_t = type; \ + return __VA_ARGS__(); \ + } + + +#define PRIVATE_CASE_TYPE_EMB(enum_type1, enum_type2, type1, NAME, ...) \ + case enum_type1: { \ + using emb_t = type1; \ + switch (enum_type2) { \ + PRIVATE_CASE_TYPE_CACHE(at::ScalarType::Float, float, __VA_ARGS__) \ + PRIVATE_CASE_TYPE_CACHE(at::ScalarType::Half, at::Half, __VA_ARGS__) \ + default: \ + AT_ERROR( \ + #NAME, \ + " not implemented for cache_t '", \ + toString(enum_type2), \ + "'"); \ + } \ + } + +#define _DISPATCH_EMB_CACHE_TYPES(emb_enum_type, cache_enum_type, NAME, ...) \ + at::ScalarType _emb_t = emb_enum_type; \ + at::ScalarType _cache_t = cache_enum_type; \ + switch (_emb_t) { \ + PRIVATE_CASE_TYPE_EMB( \ + at::ScalarType::Float, _cache_t, float, NAME, __VA_ARGS__) \ + PRIVATE_CASE_TYPE_EMB( \ + at::ScalarType::Half, _cache_t, at::Half, NAME, __VA_ARGS__) \ + PRIVATE_CASE_TYPE_EMB( \ + at::ScalarType::Float8_e4m3fnuz, \ + _cache_t, \ + at::Float8_e4m3fnuz, \ + NAME, \ + __VA_ARGS__) \ + default: \ + AT_ERROR(#NAME, " not implemented for emb_t '", toString(_emb_t), "'"); \ + } + + +#define PRIVATE_CASE_TYPE_OUTPUT( \ + output_enum_type1, \ + emb_enum_type1, \ + cache_enum_type1, \ + output_type1, \ + NAME, \ + ...) \ + case output_enum_type1: { \ + using output_t = output_type1; \ + _DISPATCH_EMB_CACHE_TYPES( \ + emb_enum_type1, cache_enum_type1, NAME, __VA_ARGS__) \ + } + +#define DISPATCH_EMB_CACHE_OUTPUT_TYPES( \ + EMB_TYPE, CACHE_TYPE, OUTPUT_TYPE, NAME, ...) \ + [&] { \ + const at::ScalarType _output_t = OUTPUT_TYPE; /* */ \ + const auto& emb_type = EMB_TYPE; \ + const auto& cache_type = CACHE_TYPE; \ + switch (_output_t) { \ + PRIVATE_CASE_TYPE_OUTPUT( \ + at::ScalarType::Half, \ + emb_type, \ + cache_type, \ + at::Half, \ + NAME, \ + __VA_ARGS__) \ + PRIVATE_CASE_TYPE_OUTPUT( \ + at::ScalarType::Float, \ + emb_type, \ + cache_type, \ + float, \ + NAME, \ + __VA_ARGS__) \ + PRIVATE_CASE_TYPE_OUTPUT( \ + at::ScalarType::BFloat16, \ + emb_type, \ + cache_type, \ + at::BFloat16, \ + NAME, \ + __VA_ARGS__) \ + default: \ + AT_ERROR( \ + #NAME, \ + " not implemented for output_t '", \ + toString(_output_t), \ + "'"); \ + } \ + }() + +#define PRIVATE_CASE_TYPE_CACHE(enum_type, type, ...) \ + case enum_type: { \ + using cache_t = type; \ + return __VA_ARGS__(); \ + } + +#define PRIVATE_CASE_TYPE_CACHE_EMB( \ + grad_enum_type, _cache_t, _emb_t, grad_cxx_type, NAME, ...) \ + case grad_enum_type: { \ + using grad_t = grad_cxx_type; \ + switch (_emb_t) { \ + PRIVATE_CASE_TYPE_EMB( \ + at::ScalarType::Float, _cache_t, float, NAME, __VA_ARGS__) \ + PRIVATE_CASE_TYPE_EMB( \ + at::ScalarType::Half, _cache_t, at::Half, NAME, __VA_ARGS__) \ + default: \ + AT_ERROR( \ + #NAME, " not implemented for emb_t '", toString(_emb_t), "'"); \ + } \ + } + + +#define DISPATCH_EMB_GRAD_CACHE_TYPES( \ + EMB_TYPE, GRAD_TYPE, CACHE_TYPE, NAME, ...) \ + [&] { \ + const auto& emb_type = EMB_TYPE; \ + const auto& grad_type = GRAD_TYPE; \ + const auto& cache_type = CACHE_TYPE; \ + at::ScalarType _emb_t = emb_type; \ + at::ScalarType _grad_t = grad_type; \ + at::ScalarType _cache_t = cache_type; \ + switch (_grad_t) { \ + PRIVATE_CASE_TYPE_CACHE_EMB( \ + at::ScalarType::Float, _cache_t, _emb_t, float, NAME, __VA_ARGS__) \ + PRIVATE_CASE_TYPE_CACHE_EMB( \ + at::ScalarType::Half, _cache_t, _emb_t, at::Half, NAME, __VA_ARGS__) \ + PRIVATE_CASE_TYPE_CACHE_EMB( \ + at::ScalarType::BFloat16, \ + _cache_t, \ + _emb_t, \ + at::BFloat16, \ + NAME, \ + __VA_ARGS__) \ + default: \ + AT_ERROR( \ + #NAME, " not implemented for grad_t '", toString(_grad_t), "'"); \ + } \ + }() + +// ============================================================================ +// Type Dispatch Macros for Sparse Operations +// ============================================================================ + +#define FBGEMM_DISPATCH_FLOATING_TYPES_CASE(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define FBGEMM_DISPATCH_INTEGRAL_TYPES_CASE(...) \ + AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) + +#define FBGEMM_DISPATCH_ALL_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, \ + NAME, \ + FBGEMM_DISPATCH_FLOATING_TYPES_CASE(__VA_ARGS__) \ + FBGEMM_DISPATCH_INTEGRAL_TYPES_CASE(__VA_ARGS__)) + +#define FBGEMM_DISPATCH_ALL_TYPES_AND_DOUBLE(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, \ + NAME, \ + FBGEMM_DISPATCH_FLOATING_TYPES_CASE(__VA_ARGS__) \ + FBGEMM_DISPATCH_INTEGRAL_TYPES_CASE(__VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__)) \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/function_types.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/function_types.h new file mode 100644 index 0000000..a9e400a --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/function_types.h @@ -0,0 +1,17 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#pragma once + +#define DLL_PUBLIC __attribute__((visibility("default"))) \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/tensor_utils.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/tensor_utils.h new file mode 100644 index 0000000..c57a133 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/tensor_utils.h @@ -0,0 +1,192 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#pragma once + +#include + +#include + +namespace fbgemm_xpu { + +inline std::optional get_device_index_from_tensor( + const at::Tensor& ten) { + return {ten.device().index()}; +} + +inline std::optional get_device_index_from_tensor( + const std::optional& ten) { + if (ten) { + return {ten->device().index()}; + } else { + return {}; + } +} + +inline bool torch_tensor_on_sycl_xpu_check(const at::Tensor& ten) { + return ten.is_xpu(); +} + +inline std::string torch_tensor_device_name(const at::Tensor& ten) { + return c10::DeviceTypeName(ten.device().type()); +} + +inline bool torch_tensor_undefined(const at::Tensor& ten) { + return ten.defined(); +} + +inline bool torch_tensor_undefined(const std::optional& ten) { + return !ten.has_value() || torch_tensor_undefined(ten.value()); +} + +inline bool torch_tensor_on_cpu_or_on_mtia_check(const at::Tensor& ten) { + return ten.is_cpu() || ten.is_mtia(); +} + +inline bool torch_tensor_on_same_device_check( + const at::Tensor& ten1, + const std::optional& ten2) { + return !ten2.has_value() || ten1.get_device() == ten2->get_device(); +} + +inline std::string torch_tensor_device_name( + const std::optional& ten) { + if (ten.has_value()) { + return torch_tensor_device_name(ten.value()); + } else { + return "N/A"; + } +} + +inline bool torch_tensor_on_sycl_xpu_check( + const std::optional& ten) { + return !ten.has_value() || torch_tensor_on_sycl_xpu_check(ten.value()); +} + +#define TENSOR_ON_CPU_OR_MTIA(x) \ + TORCH_CHECK( \ + torch_tensor_on_cpu_or_on_mtia_check(x), \ + #x " must be a CPU or MTIA tensor; it is currently on device ", \ + torch_tensor_device_name(x)) + +#define TENSORS_EMPTY_OR_ON_SAME_DEVICE(x, y) \ + TORCH_CHECK( \ + torch_tensor_on_same_device_check(x, y) || (x.numel() == 0), \ + #x " must be empty or a XPU tensor; it is currently on device ", \ + torch_tensor_device_name(x)) + +#define TENSOR_ON_SYCL_XPU(x) \ + TORCH_CHECK( \ + torch_tensor_on_sycl_xpu_check(x), \ + #x " must be a SYCL XPU tensor; it is currently on device ", \ + torch_tensor_device_name(x)) + +// Generate constexpr array of variable names to improve diagnostic output and +// raise a message if any non-empty tensor is not on a XPU or not on the same +// XPU as all the other non-empty tensors. +#define TENSORS_ON_SAME_SYCL_XPU_IF_NOT_OPTIONAL(...) \ + do { \ + const auto tensors_on_same_xpu = \ + tensor_on_same_xpu_if_not_optional_check(#__VA_ARGS__, __VA_ARGS__); \ + TORCH_CHECK(tensors_on_same_xpu.empty(), tensors_on_same_xpu); \ + } while (false) + +inline at::Tensor aligned_grad_output_tensor_for_xpu_backwards( + const at::Tensor& grad_output) { + auto aligned_grad_output = grad_output; + // FIXME: to support aligned memory access in Vec4T load/store function + // 16 for FP32 and 8 for FP16 + if (!aligned_grad_output.is_contiguous()) { + aligned_grad_output = aligned_grad_output.contiguous(); + } + if (reinterpret_cast(aligned_grad_output.data_ptr()) % 16 != 0) { + aligned_grad_output = + at::empty_like(aligned_grad_output).copy_(aligned_grad_output); + } + TORCH_CHECK(aligned_grad_output.is_contiguous()); + TORCH_CHECK( + reinterpret_cast(aligned_grad_output.data_ptr()) % 16 == 0); + return aligned_grad_output; +} + +template +std::string tensor_on_same_xpu_if_not_optional_check( + const std::string& var_names_str, + const Tensors&... tensors) { + std::optional xpu_index; + bool on_same_xpu = true; + + // Collect the GPU index of the first non-empty optional tensor and make sure + // that all tensors are on this same index. + ( + [&](const auto& tensor) { + if (!torch_tensor_undefined(tensor)) { + return; + } + if (!torch_tensor_on_sycl_xpu_check(tensor)) { + on_same_xpu = false; + return; + } + const auto my_xpu_index = get_device_index_from_tensor(tensor); + if (my_xpu_index) { + if (!xpu_index) { + xpu_index = my_xpu_index; + } else if (*xpu_index != my_xpu_index) { + on_same_xpu = false; + } + } + }(tensors), + ...); + + if (on_same_xpu) { + return ""; + } + + std::vector var_names; + { + std::string temp; + for (const auto& x : var_names_str) { + if (x == ',') { + var_names.push_back(temp); + temp = ""; + } else { + temp.push_back(x); + } + } + var_names.push_back(temp); + } + + // Not all the tensors on a XPU or on the same XPU, generate a message. + std::string msg = "Not all tensors were on the same XPU: "; + size_t current_idx = 0; + ( + [&](const auto& tensor) { + if (current_idx > 0) { + msg.append(", "); + } + msg.append( + var_names.at(current_idx++) + "(" + + torch_tensor_device_name(tensor)); + const auto xpu_device_index = get_device_index_from_tensor(tensor); + if (xpu_device_index) { + msg.append(":" + std::to_string(*xpu_device_index)); + } + msg.append(")"); + }(tensors), + ...); + + return msg; +} + +} // namespace fbgemm_xpu \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/torch_library.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/torch_library.h new file mode 100644 index 0000000..b644603 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/torch_library.h @@ -0,0 +1,17 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#include + +namespace fbgemm_xpu::utils::torch { + +inline bool schemaExists(const std::string& qualified_name) { + return c10::Dispatcher::singleton() + .findSchema({qualified_name, ""}) + .has_value(); +} + +} // namespace fbgemm_xpu::utils::torch \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.cpp new file mode 100644 index 0000000..b9a8ad0 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.cpp @@ -0,0 +1,65 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#include "utils.h" + +using Tensor = at::Tensor; + +namespace fbgemm_xpu { + DLL_PUBLIC std::tuple adjust_info_B_num_bits( + int32_t B, + int32_t T) { + int32_t info_B_num_bits = kDefaultInfoBNumBits; + uint32_t info_B_mask = kDefaultInfoBMask; + uint32_t max_T = kMaxT; + uint32_t max_B = kMaxB; + bool invalid_T = T > max_T; + bool invalid_B = B > max_B; + + TORCH_CHECK( + !(invalid_T && invalid_B), + "Not enough infos bits to accommodate T and B. Default num bits = ", + kDefaultInfoNumBits); + + if (invalid_T) { + // Reduce info_B_num_bits + while (invalid_T && !invalid_B && info_B_num_bits > 0) { + info_B_num_bits--; + max_T = ((max_T + 1) << 1) - 1; + max_B = ((max_B + 1) >> 1) - 1; + invalid_T = T > max_T; + invalid_B = B > max_B; + } + } else if (invalid_B) { + // Increase info_B_num_bits + while (!invalid_T && invalid_B && info_B_num_bits < kDefaultInfoNumBits) { + info_B_num_bits++; + max_T = ((max_T + 1) >> 1) - 1; + max_B = ((max_B + 1) << 1) - 1; + invalid_T = T > max_T; + invalid_B = B > max_B; + } + } + + TORCH_CHECK( + !invalid_T && !invalid_B, + "Not enough infos bits to accommodate T and B. Default num bits = ", + kDefaultInfoNumBits); + + // Recompute info_B_mask using new info_B_num_bits + info_B_mask = (1u << info_B_num_bits) - 1; + + return {info_B_num_bits, info_B_mask}; + } +}; \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.h b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.h new file mode 100644 index 0000000..616880e --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.h @@ -0,0 +1,399 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#pragma once + +#include + +#include + +#include +#include +#include + +#include "comm/SYCLContext.h" +#include "function_types.h" +#include "dispatch_macros.h" + +namespace fbgemm_xpu { + +constexpr int kVecWidth = 4; +constexpr size_t kThreadGroupSize = 32; +constexpr int32_t kCacheLocationMissing = -1; +constexpr size_t kMaxThreads = 1024; +constexpr size_t kForwardMaxThreads = 512; +constexpr size_t kBackwardMaxThreads = 512; + +using overflow_safe_int_t = int64_t; + +// These values are adjusted in backward based on B and T +constexpr int kDefaultInfoNumBits = 32; +constexpr int kDefaultInfoBNumBits = 26; +constexpr uint32_t kDefaultInfoBMask = (1u << kDefaultInfoBNumBits) - 1; +constexpr uint32_t kMaxT = + (1u << (kDefaultInfoNumBits - kDefaultInfoBNumBits)) - 1; +constexpr uint32_t kMaxB = (1u << kDefaultInfoBNumBits) - 1; + +enum class SparseType : uint8_t { + FP32 = 0, + FP16 = 1, + INT8 = 2, + INT4 = 3, + INT2 = 4, + BF16 = 5, + FP8 = 6, + INVALID = 7, + MX4 = 8, + NFP8 = 9, +}; + +// SYCL atomic add (equivalent to CUDA atomicAdd) +template +inline T xpuAtomicAdd(T* address, T val) { + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device, + sycl::access::address_space::global_space> + atomic_val(*address); + return atomic_val.fetch_add(val); +} + +template +inline T div_round_up(T numerator, T denominator) { + return (numerator + denominator - 1) / denominator; +} + +inline at::ScalarType getScalarType(SparseType dtype) { + switch (dtype) { + case SparseType::FP32: + return at::kFloat; + case SparseType::FP16: + return at::kHalf; + case SparseType::INT8: + return at::kByte; + case SparseType::BF16: + return at::kBFloat16; + case SparseType::INT4: + return at::kQUInt4x2; + case SparseType::INT2: + return at::kQUInt2x4; + case SparseType::NFP8: + return at::kFloat8_e4m3fn; + default: + return at::ScalarType::Undefined; + } +}; + +enum class PoolingMode : uint8_t { SUM = 0, MEAN = 1, NONE = 2 }; + +// Keep in sync with EmbeddingLocation in split_table_batched_embeddings_ops.py +enum class PlacementType : uint8_t { + DEVICE = 0, + MANAGED = 1, + MANAGED_CACHING = 2, + HOST = 3, +}; + +std::tuple adjust_info_B_num_bits(int32_t B, int32_t T); + + +// PhiloxXpuState is not available in PyTorch 2.8 +// Use native PhiloxState in future implementations +struct PhiloxXpuState { + uint64_t seed = 0; + uint64_t offset = 0; +}; + +using fint32 = union fint32 { + uint32_t I; + float F; +}; + +// ============================================================================ +// Block Count Calculation Utilities (from fbgemm_utils.h/sycl) +// ============================================================================ + +/** + * @brief Calculate the number of SYCL work-groups (blocks) needed + * + * Base function for calculating block count with overflow protection. + * + * @param num_items Total number of items to process + * @param threads_per_block Number of work-items per work-group + * @return Number of work-groups needed (capped at max_blocks) + */ +inline uint32_t xpu_calc_xblock_count_base(int num_items, int threads_per_block) { + // The number of threads can be as high as 2048 on some newer architectures, + // but this is not portable. + TORCH_CHECK( + threads_per_block <= syclDeviceMaxWorkGroupSize(), + "Number of threads must be <=1024!"); + constexpr uint64_t max_blocks = 2147483647; + const auto u_num_items = static_cast(num_items); + const auto u_threads = static_cast(threads_per_block); + // Overflow safe variant of (a + b - 1) / b + const uint64_t blocks = + u_num_items / u_threads + (u_num_items % u_threads != 0); + return static_cast(std::min(blocks, max_blocks)); +} + +/** + * @brief Calculate the number of SYCL work-groups (blocks) needed + * + * Validates input and calls xpu_calc_xblock_count_base. + * + * @param num_items Total number of items to process (must be >= 0) + * @param threads_per_block Number of work-items per work-group + * @return Number of work-groups needed + */ +inline uint32_t xpu_calc_xblock_count(int num_items, int threads_per_block) { + TORCH_CHECK( + num_items >= 0, + "When calculating block counts, the number of items must be positive!"); + return xpu_calc_xblock_count_base(num_items, threads_per_block); +} + +class FixedDivisor { + public: + explicit FixedDivisor(const int32_t d) : d_(d) { + CalcSignedMagic(); + } + + /// Calculates `q = n / d`. + int32_t Div(const int32_t n) const { + // In lieu of a mulhi instruction being available, perform the + // work in uint64 + return (int32_t)((magic_ * (uint64_t)n) >> shift_); + } + + /// Calculates `r = n % d`. + int32_t Mod(const int32_t n) const { + return n - d_ * Div(n); + } + + /// Calculates `q = n / d` and `r = n % d` together. + void DivMod(const int32_t n, int32_t* q, int32_t* r) const { + *q = Div(n); + *r = n - d_ * *q; + } + int32_t D() const { + return d_; + } + + private: + // Calculates magic multiplicative value and shift amount for calculating `q = + // n / d` for signed 32-bit integers. + // Implementation taken from Hacker's Delight section 10. + void CalcSignedMagic() { + if (d_ == 1) { + magic_ = UINT64_C(0x1) << 32; + shift_ = 32; + return; + } + + const uint32_t two31 = UINT32_C(0x80000000); + const uint32_t ad = std::abs(d_); + const uint32_t t = two31 + ((uint32_t)d_ >> 31); + const uint32_t anc = t - 1 - t % ad; // Absolute value of nc. + uint32_t p = 31; // Init. p. + uint32_t q1 = two31 / anc; // Init. q1 = 2**p/|nc|. + uint32_t r1 = two31 - q1 * anc; // Init. r1 = rem(2**p, |nc|). + uint32_t q2 = two31 / ad; // Init. q2 = 2**p/|d|. + uint32_t r2 = two31 - q2 * ad; // Init. r2 = rem(2**p, |d|). + uint32_t delta = 0; + do { + ++p; + q1 <<= 1; // Update q1 = 2**p/|nc|. + r1 <<= 1; // Update r1 = rem(2**p, |nc|). + if (r1 >= anc) { // (Must be an unsigned comparison here). + ++q1; + r1 -= anc; + } + q2 <<= 1; // Update q2 = 2**p/|d|. + r2 <<= 1; // Update r2 = rem(2**p, |d|). + if (r2 >= ad) { // (Must be an unsigned comparison here). + ++q2; + r2 -= ad; + } + delta = ad - r2; + } while (q1 < delta || (q1 == delta && r1 == 0)); + int32_t magic = q2 + 1; + if (d_ < 0) { + magic = -magic; + } + shift_ = p; + magic_ = (uint64_t)(uint32_t)magic; + } + int32_t d_ = 1; + uint64_t magic_; + int shift_; +}; + + +// Based on the empirical study, max grid size that is 64x larger than the +// number of compute units gives good performance across the board +constexpr int32_t kMaxWorkGroupsFactor = 64; + +inline int32_t get_max_work_groups_() { + auto device = c10::xpu::getCurrentXPUStream().queue().get_device(); + return kMaxWorkGroupsFactor * + device.get_info(); +} + +#define SYCL_DEVICE_GUARD(TENSOR) \ + c10::OptionalDeviceGuard device_guard; \ + device_guard.reset_device(TENSOR.device()) + +#ifdef FBGEMM_USE_SUBWARP_SHUFFLE +#define DISPATCH_OPTIMAL_KERNEL(MAX_D, ...) \ + [&] { \ + if (MAX_D <= 32) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + 1; \ + constexpr int kFixedMaxVecsPerThread = 1; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = \ + 8; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = \ + false; \ + return __VA_ARGS__(); \ + } \ + if (MAX_D <= 64) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + 1; \ + constexpr int kFixedMaxVecsPerThread = 1; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = \ + 16; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = \ + false; \ + return __VA_ARGS__(); \ + } \ + if (MAX_D <= 128) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + 1; \ + constexpr int kFixedMaxVecsPerThread = 1; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = \ + 32; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = \ + false; \ + return __VA_ARGS__(); \ + } \ + if (MAX_D <= 256) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + 2; \ + constexpr int kFixedMaxVecsPerThread = 2; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = \ + 32; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = \ + false; \ + return __VA_ARGS__(); \ + } \ + if (MAX_D > 256) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + (MAX_D + 128 - 1) / 128; \ + constexpr int kFixedMaxVecsPerThread = 2; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = fbgemm_xpu::kThreadGroupSize; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = true; \ + return __VA_ARGS__(); \ + } \ + }() + +#else +#define DISPATCH_OPTIMAL_KERNEL(MAX_D, ...) \ + [&] { \ + if (MAX_D <= 128) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + 1; \ + constexpr int kFixedMaxVecsPerThread = 1; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = \ + 32; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = \ + false; \ + return __VA_ARGS__(); \ + } \ + if (MAX_D <= 256) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + 2; \ + constexpr int kFixedMaxVecsPerThread = 2; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = \ + 32; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = \ + false; \ + return __VA_ARGS__(); \ + } \ + if (MAX_D > 256) { \ + [[ maybe_unused ]] const int max_vecs_per_thread = \ + (MAX_D + 128 - 1) / 128; \ + constexpr int kFixedMaxVecsPerThread = 2; \ + [[ maybe_unused ]] constexpr int kThreadGroupSize = fbgemm_xpu::kThreadGroupSize; \ + [[ maybe_unused ]] constexpr bool kUseVecBlocking = true; \ + return __VA_ARGS__(); \ + } \ + }() +#endif + +inline void validate_local_mem_size( + sycl::queue& q, + const int32_t local_mem_bytes) { + + auto device = q.get_device(); + auto max_local_mem = device.get_info(); + TORCH_CHECK( + local_mem_bytes <= max_local_mem, + "Attempted to allocate ", + local_mem_bytes / 1024, + " KB of local memory but only ", + max_local_mem / 1024, + " KB is available"); +} + +template +int32_t compute_num_groups_and_dynamic_smem_bytes( + int32_t* num_groups, + const func_t compute_smem_bytes_fn, + const int32_t used_shared_bytes) { + int32_t smem_bytes = 0; + while ( + (smem_bytes = compute_smem_bytes_fn(*num_groups)) + >= used_shared_bytes + ) { + *num_groups /= 2; + } + TORCH_CHECK_GE(*num_groups, 1); + + return smem_bytes; +} + +#define DISPATCH_OPTIMAL_NOBAG_FORWARD_KERNEL(DD_, ...) \ + [&] { \ + if (DD_ <= 4) { \ + constexpr int kEmbeddingSize = 4; \ + return __VA_ARGS__(); \ + } \ + if (DD_ <= 8) { \ + constexpr int kEmbeddingSize = 8; \ + return __VA_ARGS__(); \ + } \ + if (DD_ <= 16) { \ + constexpr int kEmbeddingSize = 16; \ + return __VA_ARGS__(); \ + } \ + if (DD_ <= 32) { \ + constexpr int kEmbeddingSize = 32; \ + return __VA_ARGS__(); \ + } \ + return; \ + }() + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py index 0d10647..a710817 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py @@ -5,12 +5,47 @@ # Python wrapper functions for all custom operators under the fbgemm namespace # This module provides user-friendly interfaces to the C++ operators +from typing import Optional, Tuple + +import torch +from torch import Tensor + __all__ = [ - "dense_embedding_codegen_lookup_function", + "invert_permute", + "permute_1D_sparse_data", + "asynchronous_complete_cumsum", + "permute_2D_sparse_data", ] -def dense_embedding_codegen_lookup_function(*args, **kwargs): - """Temporary stub for the planned dense embedding API.""" - raise NotImplementedError( - "dense_embedding_codegen_lookup_function is not implemented yet in src/fbgemm_xpu/ops.py" - ) +def invert_permute(permute: Tensor) -> Tensor: + """Computes the inverse of a permutation tensor.""" + return torch.ops.fbgemm.invert_permute.default(permute) + + +def permute_1D_sparse_data( + permute: Tensor, + lengths: Tensor, + indices: Tensor, + weights: Optional[Tensor] = None, + permuted_lengths_sum: Optional[int] = None, +) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """Permutes sparse data in jagged/1D format according to permutation indices.""" + return torch.ops.fbgemm.permute_1D_sparse_data.default( + permute, lengths, indices, weights, permuted_lengths_sum + ) + +def asynchronous_complete_cumsum(t_in: Tensor) -> Tensor: + """Computes complete cumulative sum: output[0] = 0, output[i] = sum(t_in[0:i]).""" + return torch.ops.fbgemm.asynchronous_complete_cumsum.default(t_in) + +def permute_2D_sparse_data( + permute: Tensor, + lengths: Tensor, + indices: Tensor, + weights: Optional[Tensor] = None, + permuted_lengths_sum: Optional[int] = None, +) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """Permutes 2D sparse data (lengths: [T, B]) according to permutation indices.""" + return torch.ops.fbgemm.permute_2D_sparse_data.default( + permute, lengths, indices, weights, permuted_lengths_sum + ) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp index f4692b5..e6755f4 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp @@ -9,6 +9,9 @@ #include #include +#include "fbgemm_utils/torch_library.h" + +using namespace fbgemm_xpu; extern "C" { /** @@ -46,25 +49,41 @@ extern "C" { */ TORCH_LIBRARY_FRAGMENT(fbgemm, m) { - m.def("dense_embedding_codegen_lookup_function(" - " Tensor dev_weights, " - " Tensor weights_offsets, " - " Tensor D_offsets, " - " SymInt total_D, " - " SymInt max_D, " - " Tensor hash_size_cumsum, " - " int total_hash_size_bits, " - " Tensor indices, " - " Tensor offsets, " - " int pooling_mode, " - " Tensor? indice_weights, " - " Tensor? feature_requires_grad, " - " int output_dtype=0, " - " Tensor? B_offsets=None, " - " Tensor? vbe_output_offsets_feature_rank=None, " - " Tensor? vbe_B_offsets_rank_per_feature=None, " - " SymInt max_B=-1, " - " SymInt max_B_feature_rank=-1, " - " SymInt vbe_output_size=-1, " - " bool mixed_D=True) -> Tensor"); + if (!utils::torch::schemaExists("fbgemm::invert_permute")) { + m.def( + "invert_permute(Tensor permute) -> Tensor" + ); + } + + + if (!utils::torch::schemaExists("fbgemm::permute_1D_sparse_data")) { + m.def( + "permute_1D_sparse_data(" + " Tensor permute, " + " Tensor lengths, " + " Tensor values, " + " Tensor? weights=None, " + " SymInt? permuted_lengths_sum=None" + ") -> (Tensor, Tensor, Tensor?)" + ); + } + + if (!utils::torch::schemaExists("fbgemm::asynchronous_complete_cumsum")) { + m.def( + "asynchronous_complete_cumsum(Tensor t_in) -> Tensor" + ); + } + + if (!utils::torch::schemaExists("fbgemm::permute_2D_sparse_data")) { + m.def( + "permute_2D_sparse_data(" + " Tensor permute, " + " Tensor lengths, " + " Tensor values, " + " Tensor? weights=None, " + " SymInt? permuted_lengths_sum=None" + ") -> (Tensor, Tensor, Tensor?)" + ); + } + } diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/permute_2d_sparse_dataOp.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/permute_2d_sparse_dataOp.cpp new file mode 100644 index 0000000..b036c61 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/permute_2d_sparse_dataOp.cpp @@ -0,0 +1,133 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#include "asynchronous_complete_cumsum.h" +#include "sycl_kernels/permute_2d_sparse_data.h" + +namespace fbgemm_xpu { + +// ============================================================================ +// Top-Level Operator Function +// ============================================================================ + +/** + * @brief Permute 2D sparse data operator + * + * Top-level function that permutes 2D sparse data including lengths, indices, + * and optional weights according to a permutation vector. + * + * @param permute Permutation indices [T] + * @param lengths Input lengths [T, B] + * @param indices Input sparse indices + * @param weights Optional input weights + * @param permuted_lengths_sum Optional pre-computed sum of permuted lengths + * @return Tuple of (permuted_lengths, permuted_indices, permuted_weights) + */ + +std::tuple> +permute_2D_sparse_data_xpu( + const at::Tensor& permute, + const at::Tensor& lengths, + const at::Tensor& indices, + const std::optional& weights, + const std::optional& permuted_lengths_sum) { + TENSORS_ON_SAME_SYCL_XPU_IF_NOT_OPTIONAL(permute, lengths, indices, weights); + TORCH_CHECK(lengths.dim() == 2); + + SYCL_DEVICE_GUARD(indices); + + const auto permute_contig = permute.contiguous(); + const auto lengths_contig = lengths.contiguous(); + const auto indices_contig = indices.contiguous(); + // the data to permute over can be less or more with or without + // repetitions + const auto T = permute.numel(); + const auto B = lengths.size(1); + + if (T == 0 || B == 0) { + // When T = 0 or B = 0, permutation will not be performed. Return the + // input tensors. + return { + lengths.clone(), + indices.clone(), + weights.has_value() ? std::make_optional(weights->clone()) + : std::nullopt}; + } + + at::Tensor permuted_lengths = at::empty({T, B}, lengths.options()); + at::Tensor permuted_indices; + at::Tensor permuted_weights; + + permute_2D_lengths_kernel_xpu( + T, B, lengths_contig, permute_contig, permuted_lengths); + + // convert lengths to offsets + const auto input_offsets = asynchronous_exclusive_cumsum(lengths_contig); + const auto output_offsets = + asynchronous_complete_cumsum_xpu(permuted_lengths.flatten()); + int64_t permuted_indices_size = 0; + if (permuted_lengths_sum.has_value()) { + permuted_indices_size = permuted_lengths_sum.value(); + } else { + permuted_indices_size = output_offsets[-1].item(); + } + + permuted_indices = at::empty(permuted_indices_size, indices.options()); + + if (weights.has_value()) { + const at::Tensor weights_value = weights.value(); + int32_t weights_columns = 1; + if (weights_value.dense_dim() > 1) { + weights_columns = weights_value.size(1); + permuted_weights = at::empty( + {permuted_indices_size, weights_columns}, weights_value.options()); + } else { + permuted_weights = + at::empty(permuted_indices_size, weights_value.options()); + } + permute_2D_data_kernel_xpu( + permuted_indices_size, + T, + B, + indices_contig, + std::optional{weights_value}, + weights_columns, + permute_contig, + input_offsets, + output_offsets, + permuted_indices, + std::optional{permuted_weights}); + } else { + permute_2D_data_kernel_xpu( + permuted_indices_size, + T, + B, + indices_contig, + std::nullopt, + 0, + permute_contig, + input_offsets, + output_offsets, + permuted_indices, + std::nullopt); + } + + return {permuted_lengths, permuted_indices, permuted_weights}; +} + +TORCH_LIBRARY_IMPL(fbgemm, XPU, m) { + m.impl("permute_2D_sparse_data", &fbgemm_xpu::permute_2D_sparse_data_xpu); +} + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/.gitkeep b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.cpp new file mode 100644 index 0000000..5026e7f --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.cpp @@ -0,0 +1,123 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#include "invert_permute_kernel.h" + +namespace fbgemm_xpu { + +// ============================================================================ +// SYCL Kernel Functor Implementations +// ============================================================================ + +void InvertPermuteKernelInt32::operator()(const sycl::nd_item<1>& item) const { + const int64_t global_id = item.get_global_id(0); + const int64_t global_range = item.get_global_range(0); + + // Grid-stride loop pattern for scalability + for (int64_t i = global_id; i < numel_; i += global_range) { + const int32_t target_idx = permute_[i]; + inversed_permute_[target_idx] = static_cast(i); + } +} + +void InvertPermuteKernelInt64::operator()(const sycl::nd_item<1>& item) const { + const int64_t global_id = item.get_global_id(0); + const int64_t global_range = item.get_global_range(0); + + // Grid-stride loop pattern for scalability + for (int64_t i = global_id; i < numel_; i += global_range) { + const int64_t target_idx = permute_[i]; + inversed_permute_[target_idx] = i; + } +} + +// ============================================================================ +// Host Function - XPU Implementation +// ============================================================================ +at::Tensor invert_permute_forward_xpu(const at::Tensor& permute) { + // Input validation + TORCH_CHECK(permute.dim() == 1, + "invert_permute: input must be 1-dimensional, got ", permute.dim(), "D"); + TORCH_CHECK(permute.dtype() == at::kInt || permute.dtype() == at::kLong, + "invert_permute: input must be int32 or int64, got ", permute.dtype()); + TORCH_INTERNAL_ASSERT(permute.device().type() == at::DeviceType::XPU, + "invert_permute_forward_xpu: input must be on XPU device"); + + // Get input size + const int64_t N = permute.size(0); + + // Handle empty tensor + if (N == 0) { + return at::empty_like(permute); + } + + // Ensure input is contiguous for efficient memory access + at::Tensor permute_contig = permute.contiguous(); + + // Allocate output tensor on the same device + at::Tensor inversed_permute = at::empty_like(permute_contig); + + // Get SYCL queue from current XPU stream + sycl::queue& queue = c10::xpu::getCurrentXPUStream().queue(); + + // Kernel configuration + // Work-group size (local size): number of work-items per work-group + constexpr int threads = 256; + // Global size: total number of work-items across all work-groups + // Rounded up to multiple of work-group size + const int blocks = (N + threads - 1) / threads; + const size_t global_size = blocks * threads; + const size_t local_size = threads; + + // Launch kernel based on dtype + if (permute.dtype() == at::kInt) { + // int32 path + const int32_t* permute_ptr = permute_contig.data_ptr(); + int32_t* inversed_ptr = inversed_permute.data_ptr(); + + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>( + sycl::range<1>(global_size), // Global range + sycl::range<1>(local_size) // Local range (work-group size) + ), + InvertPermuteKernelInt32(N, permute_ptr, inversed_ptr) + ); + }); + + } else { + // int64 path + const int64_t* permute_ptr = permute_contig.data_ptr(); + int64_t* inversed_ptr = inversed_permute.data_ptr(); + + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>( + sycl::range<1>(global_size), + sycl::range<1>(local_size) + ), + InvertPermuteKernelInt64(N, permute_ptr, inversed_ptr) + ); + }); + } + + // Note: We don't call queue.wait() here for asynchronous execution. + // PyTorch's stream synchronization will handle proper ordering. + + return inversed_permute; +} + +/** + * Register XPU implementation with PyTorch dispatch system + * + * This binds the SYCL/XPU implementation to the operator schema defined in ops_registry.cpp + * When an XPU tensor is passed to the operator, this implementation will be called. + */ +TORCH_LIBRARY_IMPL(fbgemm, XPU, m) { + m.impl("invert_permute", &invert_permute_forward_xpu); +} + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.h b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.h new file mode 100644 index 0000000..5d5c9b8 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.h @@ -0,0 +1,152 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +//////////////////////////////////////////////////////////////////////////////// +// SYCL PORT MAPPING TO FBGEMM CUDA SOURCE - INVERT PERMUTE OPERATOR +//////////////////////////////////////////////////////////////////////////////// +// +// This file contains SYCL ports of FBGEMM invert_permute operator. +// +// ORIGINAL CUDA SOURCE: +// File: fbgemm_gpu/src/sparse_ops/sparse_invert_permute.cu +// +// KERNEL MAPPING: +// InvertPermuteKernelInt32 (SYCL) +// → invert_permute_kernel (CUDA) +// InvertPermuteKernelInt64 (SYCL) +// → invert_permute_kernel (CUDA) +// +// HOST FUNCTION MAPPING: +// invert_permute_forward_xpu (SYCL) +// → invert_permute_cuda (CUDA) +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_invert_permute.cu +// +// DESCRIPTION: +// Computes the inverse of a permutation tensor on Intel XPU devices. +// If permute[i] = j, then inversed_permute[j] = i. +// The SYCL kernel uses a grid-stride loop pattern for scalability, +// allowing efficient handling of permutations of any size. +// +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include + +#include + +#include + +#include +#include + +namespace fbgemm_xpu { + +// ============================================================================ +// SYCL Kernel Functors +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// InvertPermuteKernelInt32 - Device Kernel +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Kernel: invert_permute_kernel +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_invert_permute.cu +// +// DESCRIPTION: +// Main kernel for inverting a permutation tensor (int32 version). +// Implements the core logic: inversed_permute[permute[i]] = i +// using a grid-stride loop pattern for scalability. +// +// The grid-stride loop allows each work-item to process multiple elements, +// ensuring correct behavior regardless of the relationship between N +// (number of elements) and the number of work-items launched. +// +//////////////////////////////////////////////////////////////////////////////// +class InvertPermuteKernelInt32 { +public: + InvertPermuteKernelInt32( + int64_t numel, + const int32_t* permute, + int32_t* inversed_permute) + : numel_(numel), + permute_(permute), + inversed_permute_(inversed_permute) {} + + void operator()(const sycl::nd_item<1>& item) const; + +private: + int64_t numel_; + const int32_t* permute_; + int32_t* inversed_permute_; +}; + +//////////////////////////////////////////////////////////////////////////////// +// InvertPermuteKernelInt64 - Device Kernel +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Kernel: invert_permute_kernel +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_invert_permute.cu +// +// DESCRIPTION: +// Main kernel for inverting a permutation tensor (int64 version). +// Same logic as Int32 version but for 64-bit integers. +// Uses grid-stride loop pattern for scalability. +// +//////////////////////////////////////////////////////////////////////////////// +class InvertPermuteKernelInt64 { +public: + InvertPermuteKernelInt64( + int64_t numel, + const int64_t* permute, + int64_t* inversed_permute) + : numel_(numel), + permute_(permute), + inversed_permute_(inversed_permute) {} + + void operator()(const sycl::nd_item<1>& item) const; + +private: + int64_t numel_; + const int64_t* permute_; + int64_t* inversed_permute_; +}; + +// ============================================================================ +// Host Function Declaration +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// invert_permute_forward_xpu - Host Function +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Function: invert_permute_cuda +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_invert_permute.cu +// +// DESCRIPTION: +// XPU implementation of invert_permute forward pass. +// Computes the inverse of a permutation tensor using SYCL on Intel GPUs. +// +// Input validation: +// - Ensures input is 1D tensor +// - Verifies dtype is int32 or int64 +// - Checks tensor is on XPU device +// +// Performance: +// - Uses grid-stride loop for scalability +// - Work-group size: 256 threads +// - Asynchronous execution via PyTorch XPU stream +// +// @param permute Input permutation tensor (1D, int32 or int64, on XPU device) +// @return at::Tensor Output inverse permutation tensor (same shape, dtype, device) +// +//////////////////////////////////////////////////////////////////////////////// +at::Tensor invert_permute_forward_xpu(const at::Tensor& permute); + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.cpp new file mode 100644 index 0000000..cbd5b50 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.cpp @@ -0,0 +1,430 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +/* + * SYCL/XPU Implementation of permute_1D_sparse_data operator + * + * This operator permutes sparse data represented in jagged/1D format. + * It reorders (permutes) segments of data according to a permutation array. + * + * The operation is commonly used in: + * - Deep Learning Recommendation Models (DLRMs) for reordering embedding features + * - Sparse tensor operations in graph neural networks + * - Distributed training for all-to-all communication preparation + * + * Input format (jagged/CSR-like): + * lengths: [L0, L1, L2, ...] - length of each segment + * values: [seg0 data][seg1 data][seg2 data]... - concatenated values + * permute: [P0, P1, P2, ...] - new segment order + * + * Output: + * permuted_lengths[j] = lengths[permute[j]] + * permuted_values contains data from segment permute[j] at position j + */ + +//////////////////////////////////////////////////////////////////////////////// +// SYCL PORT MAPPING TO FBGEMM CUDA SOURCE +//////////////////////////////////////////////////////////////////////////////// +// +// This file contains SYCL implementations of FBGEMM sparse data permutation +// kernels and host functions. +// +// ORIGINAL CUDA SOURCE: +// File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +// KERNEL IMPLEMENTATIONS: +// Permute1DLengthsKernel::operator() +// → permute_1D_lengths_kernel (CUDA) +// +// Permute1DDataKernel::operator() +// → permute_1D_data_kernel (CUDA) +// +// Permute1DDataWithWeightsKernel::operator() +// → permute_1D_data_kernel (CUDA) +// +// HOST FUNCTION: +// permute_1D_sparse_data_xpu +// → permute_1D_sparse_data_cuda (CUDA) +// +// HELPER FUNCTIONS: +// exclusive_cumsum_xpu +// → asynchronous_exclusive_cumsum_gpu (CUDA) +// +// complete_cumsum_xpu +// → asynchronous_complete_cumsum_gpu (CUDA) +// +//////////////////////////////////////////////////////////////////////////////// + +#include "permute_1d_sparse_data.h" + +namespace fbgemm_xpu { + +// ============================================================================ +// SYCL Kernel Functors - Operator Implementations +// ============================================================================ + +/** + * @brief Permute1DLengthsKernel operator implementation + */ +template +void Permute1DLengthsKernel::operator()( + const sycl::nd_item<1>& item) const { + const int64_t i = item.get_global_id(0); + if (i < permuted_lengths_size_) { + permuted_lengths_[i] = lengths_[permute_[i]]; + } +} + +/** + * @brief Permute1DDataKernel operator implementation + */ +template +void Permute1DDataKernel::operator()( + const sycl::nd_item<2>& item) const { + // Get segment ID and thread ID within segment + const int32_t segment_base = item.get_group(0) * item.get_local_range(1); + const int32_t segment_local = item.get_local_id(1); + const int32_t segment_id = segment_base + segment_local; + const int32_t tid = item.get_local_id(0); + const int32_t threads_per_segment = item.get_local_range(0); + + if (segment_id >= permuted_lengths_size_) { + return; + } + + // Calculate segment boundaries + const offsets_t output_start = output_offsets_[segment_id]; + const offsets_t output_end = (segment_id == permuted_lengths_size_ - 1) + ? static_cast(permuted_indices_size_) + : output_offsets_[segment_id + 1]; + const int32_t segment_length = static_cast(output_end - output_start); + const offsets_t input_start = input_offsets_[permute_[segment_id]]; + + // Copy data using strided access pattern + for (int32_t i = tid; i < segment_length; i += threads_per_segment) { + permuted_indices_[output_start + i] = indices_[input_start + i]; + } +} + +/** + * @brief Permute1DDataWithWeightsKernel operator implementation + */ +template +void Permute1DDataWithWeightsKernel::operator()( + const sycl::nd_item<2>& item) const { + const int32_t segment_base = item.get_group(0) * item.get_local_range(1); + const int32_t segment_local = item.get_local_id(1); + const int32_t segment_id = segment_base + segment_local; + const int32_t tid = item.get_local_id(0); + const int32_t threads_per_segment = item.get_local_range(0); + + if (segment_id >= permuted_lengths_size_) { + return; + } + + const offsets_t output_start = output_offsets_[segment_id]; + const offsets_t output_end = (segment_id == permuted_lengths_size_ - 1) + ? static_cast(permuted_indices_size_) + : output_offsets_[segment_id + 1]; + const int32_t segment_length = static_cast(output_end - output_start); + const offsets_t input_start = input_offsets_[permute_[segment_id]]; + + // Copy both indices and weights + for (int32_t i = tid; i < segment_length; i += threads_per_segment) { + permuted_indices_[output_start + i] = indices_[input_start + i]; + permuted_weights_[output_start + i] = weights_[input_start + i]; + } +} + +// ============================================================================ +// Helper Functions +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// exclusive_cumsum_xpu - Helper Function +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Function: asynchronous_exclusive_cumsum_gpu +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_ops_gpu.cu +// +// DESCRIPTION: +// Computes exclusive prefix sum (cumulative sum with 0 prepended). +// Used to convert lengths array to offsets array for jagged tensor access. +// +// Example: input = [3, 2, 4] → output = [0, 3, 5, 9] +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief Compute exclusive cumsum (prefix sum) for offsets + * + * output[0] = 0 + * output[i] = sum(input[0:i]) for i > 0 + */ +template +at::Tensor exclusive_cumsum_xpu(const at::Tensor& input) { + auto output = at::empty({input.numel() + 1}, input.options().dtype(at::kLong)); + output[0] = 0; + if (input.numel() > 0) { + auto cumsum = input.cumsum(0, at::kLong); + output.slice(0, 1).copy_(cumsum); + } + return output; +} + +//////////////////////////////////////////////////////////////////////////////// +// complete_cumsum_xpu - Helper Function +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Function: asynchronous_complete_cumsum_gpu +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_ops_gpu.cu +// +// DESCRIPTION: +// Computes complete cumulative sum (inclusive prefix sum with 0 prepended). +// Used to convert permuted lengths to output offsets. +// +// Identical to exclusive_cumsum_xpu in current implementation. +// Example: input = [4, 3, 2] → output = [0, 4, 7, 9] +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief Compute complete cumsum (inclusive + final element) + * + * output[0] = 0 + * output[i] = sum(input[0:i]) for i > 0 + * output[n] = sum(all input) + */ +template +at::Tensor complete_cumsum_xpu(const at::Tensor& input) { + auto output = at::empty({input.numel() + 1}, input.options().dtype(at::kLong)); + output[0] = 0; + if (input.numel() > 0) { + auto cumsum = input.cumsum(0, at::kLong); + output.slice(0, 1).copy_(cumsum); + } + return output; +} + +// ============================================================================ +// Host Function - XPU Implementation +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// permute_1D_sparse_data_xpu - Host Function +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Function: permute_1D_sparse_data_cuda +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +// DESCRIPTION: +// Main host function that orchestrates permutation of sparse jagged data. +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief XPU implementation of permute_1D_sparse_data + * + * Permutes sparse data represented in jagged format according to permutation indices. + * + * @param permute Permutation indices tensor [P] - int32 + * @param lengths Segment lengths tensor [L] - int32/int64 + * @param indices Concatenated values tensor [V] - any type + * @param weights Optional weights tensor [V] - float/double + * @param permuted_lengths_sum Optional precomputed sum of permuted lengths + * @return Tuple of (permuted_lengths, permuted_indices, permuted_weights) + */ +std::tuple> +permute_1D_sparse_data_xpu( + const at::Tensor& permute, + const at::Tensor& lengths, + const at::Tensor& indices, + const std::optional& weights, + const std::optional& permuted_lengths_sum) { + + // Device validation + TORCH_INTERNAL_ASSERT(permute.device().type() == at::DeviceType::XPU, + "permute must be on XPU device"); + TORCH_INTERNAL_ASSERT(lengths.device().type() == at::DeviceType::XPU, + "lengths must be on XPU device"); + TORCH_INTERNAL_ASSERT(indices.device().type() == at::DeviceType::XPU, + "indices must be on XPU device"); + + // Input validation + TORCH_CHECK(permute.dim() == 1, "permute must be 1D"); + TORCH_CHECK(lengths.dim() == 1, "lengths must be 1D"); + TORCH_CHECK(indices.dim() == 1, "indices must be 1D"); + TORCH_CHECK(permute.dtype() == at::kInt, "permute must be int32"); + + // Ensure contiguous + const auto permute_contig = permute.contiguous(); + const auto lengths_contig = lengths.contiguous(); + const auto indices_contig = indices.contiguous(); + + const int64_t permuted_lengths_size = permute.numel(); + const int64_t lengths_size = lengths.numel(); + + // Handle empty input + if (permuted_lengths_size == 0) { + // Empty permutation returns empty outputs + return { + at::empty({0}, lengths.options()), + at::empty({0}, indices.options()), + weights.has_value() ? std::make_optional(at::empty({0}, weights->options())) : std::nullopt + }; + } + + if (lengths_size == 0) { + // Empty lengths but non-empty permute - return zeros for lengths and empty indices + // Warning: if permute contains valid indices they would be out of bounds, + // but typically 0 lengths implies 0 segments, so permute should be empty too. + // If we really want to support "selecting from 0 segments" it implies permute indices are invalid + // unless we assume they map to implicit 0 length segments? + // Following CPU logic: + return { + at::zeros({permuted_lengths_size}, lengths.options()), + at::empty({0}, indices.options()), + weights.has_value() ? std::make_optional(at::empty({0}, weights->options())) : std::nullopt + }; + } + + // Get SYCL queue + sycl::queue& queue = c10::xpu::getCurrentXPUStream().queue(); + + // Phase 1: Allocate and compute permuted_lengths + at::Tensor permuted_lengths = at::empty({permuted_lengths_size}, lengths.options()); + + // Launch lengths permutation kernel + { + constexpr int kLocalSize = 256; + const int64_t global_size = ((permuted_lengths_size + kLocalSize - 1) / kLocalSize) * kLocalSize; + + AT_DISPATCH_INDEX_TYPES( + lengths.scalar_type(), "permute_1D_lengths_xpu", [&] { + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for>( + sycl::nd_range<1>( + sycl::range<1>(global_size), + sycl::range<1>(kLocalSize) + ), + Permute1DLengthsKernel( + permuted_lengths_size, + lengths_contig.data_ptr(), + permute_contig.data_ptr(), + permuted_lengths.data_ptr() + ) + ); + }); + } + ); + } + + // Phase 2: Compute input and output offsets + at::Tensor input_offsets = exclusive_cumsum_xpu(lengths_contig); + at::Tensor output_offsets = complete_cumsum_xpu(permuted_lengths); + + // Phase 3: Determine output size + int64_t permuted_indices_size = 0; + if (permuted_lengths_sum.has_value()) { + permuted_indices_size = permuted_lengths_sum.value(); + } else { + // Need to sync to get the value from device + permuted_indices_size = output_offsets[permuted_lengths_size].item(); + } + + // Phase 4: Allocate output tensors + at::Tensor permuted_indices = at::empty(permuted_indices_size, indices.options()); + std::optional permuted_weights = std::nullopt; + + // Phase 5: Launch data permutation kernel + constexpr int32_t kThreadsPerSegment = 64; + constexpr int32_t kSegmentsPerBlock = 16; + + const int64_t num_blocks = (permuted_lengths_size + kSegmentsPerBlock - 1) / kSegmentsPerBlock; + sycl::range<2> global_range{static_cast(num_blocks * kThreadsPerSegment), + static_cast(kSegmentsPerBlock)}; + sycl::range<2> local_range{kThreadsPerSegment, kSegmentsPerBlock}; + + if (weights.has_value()) { + // With weights + const auto weights_contig = weights->contiguous(); + permuted_weights = at::empty(permuted_indices_size, weights->options()); + + AT_DISPATCH_INDEX_TYPES( + input_offsets.scalar_type(), "permute_1D_data_xpu_offsets", [&] { + using offsets_t = index_t; + AT_DISPATCH_ALL_TYPES_AND2( + at::kHalf, at::kBFloat16, + indices.scalar_type(), "permute_1D_data_xpu_indices", [&] { + using indices_t = scalar_t; + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + weights->scalar_type(), "permute_1D_data_xpu_weights", [&] { + using weights_t = scalar_t; + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for>( + sycl::nd_range<2>(global_range, local_range), + Permute1DDataWithWeightsKernel( + permuted_indices_size, + permuted_lengths_size, + indices_contig.data_ptr(), + weights_contig.data_ptr(), + permute_contig.data_ptr(), + input_offsets.data_ptr(), + output_offsets.data_ptr(), + permuted_indices.data_ptr(), + permuted_weights->data_ptr() + ) + ); + }); + } + ); + } + ); + } + ); + } else { + // Without weights + AT_DISPATCH_INDEX_TYPES( + input_offsets.scalar_type(), "permute_1D_data_xpu_offsets", [&] { + using offsets_t = index_t; + AT_DISPATCH_ALL_TYPES_AND2( + at::kHalf, at::kBFloat16, + indices.scalar_type(), "permute_1D_data_xpu_indices", [&] { + using indices_t = scalar_t; + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for>( + sycl::nd_range<2>(global_range, local_range), + Permute1DDataKernel( + permuted_indices_size, + permuted_lengths_size, + indices_contig.data_ptr(), + permute_contig.data_ptr(), + input_offsets.data_ptr(), + output_offsets.data_ptr(), + permuted_indices.data_ptr() + ) + ); + }); + } + ); + } + ); + } + + return {permuted_lengths, permuted_indices, permuted_weights}; +} + +/** + * Register XPU implementation with PyTorch dispatch system + */ +TORCH_LIBRARY_IMPL(fbgemm, XPU, m) { + m.impl("permute_1D_sparse_data", &permute_1D_sparse_data_xpu); +} + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.h b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.h new file mode 100644 index 0000000..fb9a4db --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.h @@ -0,0 +1,240 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + * Copyright (c) 2026 Intel Corporation. All Rights Reserved. + * SPDX-License-Identifier: BSD-3-Clause + */ + +//////////////////////////////////////////////////////////////////////////////// +// SYCL PORT MAPPING TO FBGEMM CUDA SOURCE +//////////////////////////////////////////////////////////////////////////////// +// +// This file contains SYCL ports of FBGEMM sparse data permutation kernels. +// +// ORIGINAL CUDA SOURCE: +// File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +// KERNEL MAPPING: +// Permute1DLengthsKernel (SYCL) +// → permute_1D_lengths_kernel (CUDA) +// +// Permute1DDataKernel (SYCL) +// → permute_1D_data_kernel (CUDA - no weights variant) +// +// Permute1DDataWithWeightsKernel (SYCL) +// → permute_1D_data_kernel (CUDA - with weights variant) +// +// HOST FUNCTION MAPPING: +// permute_1D_sparse_data_xpu (SYCL) +// → permute_1D_sparse_data_cuda (CUDA) +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include +#include + +#include +#include +#include + + +namespace fbgemm_xpu { + +// ============================================================================ +// SYCL Kernel Functors - Lengths Permutation +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// Permute1DLengthsKernel - Device Kernel +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Kernel: permute_1D_lengths_kernel +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +// DESCRIPTION: +// Simple element-parallel kernel for permuting lengths array according to +// a permutation index. Each work-item handles one output element: +// permuted_lengths[i] = lengths[permute[i]] +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief SYCL kernel functor for permuting lengths + * + * Simple element-parallel kernel: permuted_lengths[i] = lengths[permute[i]] + */ +template +class Permute1DLengthsKernel { +public: + Permute1DLengthsKernel( + int64_t permuted_lengths_size, + const index_t* lengths, + const int32_t* permute, + index_t* permuted_lengths) + : permuted_lengths_size_(permuted_lengths_size), + lengths_(lengths), + permute_(permute), + permuted_lengths_(permuted_lengths) {} + + void operator()(const sycl::nd_item<1>& item) const; + +private: + int64_t permuted_lengths_size_; + const index_t* lengths_; + const int32_t* permute_; + index_t* permuted_lengths_; +}; + +// ============================================================================ +// SYCL Kernel Functors - Data Permutation (without weights) +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// Permute1DDataKernel - Device Kernel +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Kernel: permute_1D_data_kernel +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +// DESCRIPTION: +// Permutes sparse indices according to segment permutation. Uses 2D parallel +// decomposition where each row of work-items processes one segment. +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief SYCL kernel functor for permuting data without weights + * + * 2D parallel decomposition: + * - Dimension 0 (y): segments (one row of work-items per segment) + * - Dimension 1 (x): threads cooperating on one segment + */ +template +class Permute1DDataKernel { +public: + Permute1DDataKernel( + int64_t permuted_indices_size, + int64_t permuted_lengths_size, + const indices_t* indices, + const int32_t* permute, + const offsets_t* input_offsets, + const offsets_t* output_offsets, + indices_t* permuted_indices) + : permuted_indices_size_(permuted_indices_size), + permuted_lengths_size_(permuted_lengths_size), + indices_(indices), + permute_(permute), + input_offsets_(input_offsets), + output_offsets_(output_offsets), + permuted_indices_(permuted_indices) {} + + void operator()(const sycl::nd_item<2>& item) const; + +private: + int64_t permuted_indices_size_; + int64_t permuted_lengths_size_; + const indices_t* indices_; + const int32_t* permute_; + const offsets_t* input_offsets_; + const offsets_t* output_offsets_; + indices_t* permuted_indices_; +}; + +// ============================================================================ +// SYCL Kernel Functors - Data Permutation (with weights) +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// Permute1DDataWithWeightsKernel - Device Kernel +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Kernel: permute_1D_data_kernel +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// +// DESCRIPTION: +// Same as Permute1DDataKernel but also copies weight values alongside +// indices. +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief SYCL kernel functor for permuting data with weights + * + * Same structure as Permute1DDataKernel but also copies weight values. + */ +template +class Permute1DDataWithWeightsKernel { +public: + Permute1DDataWithWeightsKernel( + int64_t permuted_indices_size, + int64_t permuted_lengths_size, + const indices_t* indices, + const weights_t* weights, + const int32_t* permute, + const offsets_t* input_offsets, + const offsets_t* output_offsets, + indices_t* permuted_indices, + weights_t* permuted_weights) + : permuted_indices_size_(permuted_indices_size), + permuted_lengths_size_(permuted_lengths_size), + indices_(indices), + weights_(weights), + permute_(permute), + input_offsets_(input_offsets), + output_offsets_(output_offsets), + permuted_indices_(permuted_indices), + permuted_weights_(permuted_weights) {} + + void operator()(const sycl::nd_item<2>& item) const; + +private: + int64_t permuted_indices_size_; + int64_t permuted_lengths_size_; + const indices_t* indices_; + const weights_t* weights_; + const int32_t* permute_; + const offsets_t* input_offsets_; + const offsets_t* output_offsets_; + indices_t* permuted_indices_; + weights_t* permuted_weights_; +}; + +// ============================================================================ +// Host Function Declaration +// ============================================================================ + +//////////////////////////////////////////////////////////////////////////////// +// permute_1D_sparse_data_xpu - Host Function +//////////////////////////////////////////////////////////////////////////////// +// +// CUDA SOURCE MAPPING: +// CUDA Function: permute_1D_sparse_data_cuda +// CUDA File: fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu +// CUDA Header: fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h +// +// DESCRIPTION: +// Host function that orchestrates the permutation of sparse data in jagged +// format. +// +//////////////////////////////////////////////////////////////////////////////// + +/** + * @brief XPU implementation of permute_1D_sparse_data + * + * Permutes sparse data represented in jagged format according to permutation + * indices. + */ +std::tuple> +permute_1D_sparse_data_xpu( + const at::Tensor& permute, + const at::Tensor& lengths, + const at::Tensor& indices, + const std::optional& weights, + const std::optional& permuted_lengths_sum); + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.cpp new file mode 100644 index 0000000..5a619d1 --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.cpp @@ -0,0 +1,250 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#include "permute_2d_sparse_data.h" + +namespace syclext = sycl::ext::oneapi; +namespace syclexp = sycl::ext::oneapi::experimental; + +// ============================================================================ +// fbgemm_cumsum_kernel implementation +// This is from the original FbgemmKernels.sycl +// ============================================================================ + +namespace at { +namespace native { +namespace xpu { + +void fbgemm_cumsum_kernel( + const Tensor& result, + const Tensor& self, + int64_t dim) { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2( + ScalarType::Half, + ScalarType::BFloat16, + self.scalar_type(), + "fbgemm_cumsum_xpu", + [&]() { + scalar_t init = 0; + scan( + result, self, dim, init, std::plus()); + }); +} + +} // namespace xpu +} // namespace native +} // namespace at + +// ============================================================================ +// fbgemm_xpu namespace implementation +// ============================================================================ + +namespace fbgemm_xpu { + +// ============================================================================ +// Asynchronous Exclusive Cumsum Helper +// ============================================================================ + +at::Tensor asynchronous_exclusive_cumsum(const at::Tensor& t_in) { + TENSOR_ON_SYCL_XPU(t_in); + SYCL_DEVICE_GUARD(t_in); + + if (t_in.numel() == 0) { + return at::empty_like(t_in); + } + + TORCH_CHECK(t_in.is_contiguous()); + TORCH_CHECK(t_in.dtype() == at::kInt || t_in.dtype() == at::kLong); + // only handles up to INT_MAX elements. + TORCH_CHECK(t_in.numel() < std::numeric_limits::max()); + auto t_in_flatten = t_in.flatten(); + auto t_out = at::empty_like(t_in_flatten); + + at::native::xpu::fbgemm_cumsum_kernel(t_out, t_in_flatten, 0); + + // make it exclusive + t_out = t_out.roll(1, 0); + // set all first elemnts 0 + t_out[0] = 0; + return t_out.view_as(t_in); +} + +// ============================================================================ +// SYCL Kernel Implementations +// ============================================================================ + +// Kernel for permuting the lengths. Used for permutation of sparse features. +template +SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>)) +void permute_2D_lengths_kernel_( + int32_t T, + int32_t B, + const index_t* __restrict__ lengths, + const int32_t* __restrict__ permute, + index_t* __restrict__ permuted_lengths) { + auto item = syclext::this_work_item::get_nd_item<1>(); + XPU_KERNEL_LOOP(item, b_t, B * T) { + int32_t b = b_t % B; + int32_t t = b_t / B; + permuted_lengths[b_t] = lengths[permute[t] * B + b]; + } +} + +void permute_2D_lengths_kernel_xpu( + int32_t T, + int32_t B, + const at::Tensor& lengths_contig, + const at::Tensor& permute_contig, + at::Tensor& permuted_lengths) { + constexpr int32_t threads_1 = 256; + const auto blocks_1 = xpu_calc_xblock_count(B * T, threads_1); + AT_DISPATCH_INDEX_TYPES( + lengths_contig.scalar_type(), "permute_2D_lengths_kernel", [&] { + sycl_kernel_submit>( + sycl::range<1>(blocks_1 * threads_1), + sycl::range<1>(threads_1), + getCurrentSYCLQueue(), + 0, + T, + B, + lengths_contig.data_ptr(), + permute_contig.data_ptr(), + permuted_lengths.data_ptr()); + }); +} + +template < + bool has_weight, + typename offsets_t, + typename indices_t, + typename weights_t> +SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<2>)) +void permute_2D_data_kernel_( + int32_t len, + int32_t T, + int32_t B, + const indices_t* __restrict__ indices, + const weights_t* __restrict__ weights, + const int32_t weights_columns, + const int32_t* __restrict__ permute, + const offsets_t* __restrict__ input_offsets, + const offsets_t* __restrict__ output_offsets, + indices_t* __restrict__ permuted_indices, + weights_t* __restrict__ permuted_weights) { + auto item = syclext::this_work_item::get_nd_item<2>(); + auto b_t_start = + item.get_group(0) * item.get_local_range(1) + item.get_local_id(1); + const auto stride = item.get_group_range(0) * item.get_local_range(1); + for (int b_t = b_t_start; b_t < B * T; b_t += stride) { + int32_t b = b_t % B; + int32_t t = b_t / B; + offsets_t output_start = output_offsets[b_t]; + offsets_t segment_length; + if (b_t == B * T - 1) { + segment_length = len - output_offsets[b_t]; + } else { + segment_length = output_offsets[b_t + 1] - output_offsets[b_t]; + } + offsets_t input_start = input_offsets[permute[t] * B + b]; + for (auto i = item.get_local_id(0); i < segment_length; + i += item.get_local_range(0)) { + permuted_indices[output_start + i] = indices[input_start + i]; + if (has_weight) { + for (auto w_col = 0; w_col < weights_columns; ++w_col) { + permuted_weights[(output_start + i) * weights_columns + w_col] = + weights[(input_start + i) * weights_columns + w_col]; + } + } + } + } +} + +void permute_2D_data_kernel_xpu( + int32_t permuted_indices_size, + int32_t T, + int32_t B, + const at::Tensor& indices_contig, + const std::optional& weights, + const int32_t weights_columns, + const at::Tensor& permute_contig, + const at::Tensor& input_offsets, + const at::Tensor& output_offsets, + at::Tensor& permuted_indices, + const std::optional& permuted_weights) { + constexpr int32_t BT_blocks = 32; + const auto blocks_2 = xpu_calc_xblock_count(B * T, BT_blocks); + AT_DISPATCH_INDEX_TYPES( + input_offsets.scalar_type(), "permute_2D_data_kernel_1", [&] { + using offsets_t = index_t; + FBGEMM_DISPATCH_ALL_TYPES( + indices_contig.scalar_type(), "permute_2D_data_kernel_2", [&] { + using indices_t = scalar_t; + if (weights.has_value()) { + const auto weights_value_contig = weights.value().contiguous(); + FBGEMM_DISPATCH_ALL_TYPES_AND_DOUBLE( + weights_value_contig.scalar_type(), + "permute_2D_data_kernel_3", + [&] { + using weights_t = scalar_t; + sycl_kernel_submit>( + sycl::range<2>(blocks_2 * 32, BT_blocks), + sycl::range<2>(32, BT_blocks), + getCurrentSYCLQueue(), + 0, + permuted_indices_size, + T, + B, + indices_contig.data_ptr(), + weights_value_contig.data_ptr(), + weights_columns, + permute_contig.data_ptr(), + input_offsets.data_ptr(), + output_offsets.data_ptr(), + permuted_indices.data_ptr(), + permuted_weights.value().data_ptr()); + }); // for each weights_t + } else { + sycl_kernel_submit>( // false float type here as wa since + // std::nullptr_t cannot be + // supported in free function + // kernel + sycl::range<2>(blocks_2 * 32, BT_blocks), + sycl::range<2>(32, BT_blocks), + getCurrentSYCLQueue(), + 0, + permuted_indices_size, + T, + B, + indices_contig.data_ptr(), + nullptr, + 0, + permute_contig.data_ptr(), + input_offsets.data_ptr(), + output_offsets.data_ptr(), + permuted_indices.data_ptr(), + nullptr); + } + }); // for each indices_t + }); // for each offsets_t +} + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.h b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.h new file mode 100644 index 0000000..34c126d --- /dev/null +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.h @@ -0,0 +1,134 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Portions of this file are derived from FBGEMM + * Copyright (c) Meta Platforms, Inc. and affiliates. + * SPDX-License-Identifier: BSD-3-Clause + */ + +//////////////////////////////////////////////////////////////////////////////// +// SYCL PORT MAPPING TO FBGEMM CUDA SOURCE - PERMUTE 2D SPARSE DATA OPERATOR +//////////////////////////////////////////////////////////////////////////////// +// +// This file contains SYCL ports of FBGEMM 2D sparse data permutation kernels. +// +// ORIGINAL CUDA SOURCE: +// File: fbgemm_gpu/src/sparse_ops/sparse_permute_embeddings.cu +// +// KERNEL MAPPING: +// permute_2D_lengths_kernel_ (SYCL) +// → permute_2D_lengths_kernel (CUDA) +// +// permute_2D_data_kernel_ (SYCL) +// → permute_2D_data_kernel (CUDA) +// +// HOST FUNCTION MAPPING: +// permute_2D_sparse_data_xpu (SYCL) +// → permute_2D_sparse_data_cuda (CUDA) +// +// DESCRIPTION: +// Permutes 2D sparse data (indices, lengths, optional weights) according to +// a permutation vector. Used for reordering embedding table features. +// +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "../fbgemm_utils/utils.h" +#include "../fbgemm_utils/tensor_utils.h" + +// Forward declaration for cumsum kernel +// Implementation is in permute_2d_sparse_data.sycl +namespace at { namespace native { namespace xpu { + void fbgemm_cumsum_kernel( + const at::Tensor& result, + const at::Tensor& self, + int64_t dim); +}}} // namespace at::native::xpu + +namespace fbgemm_xpu { + +// ============================================================================ +// Helper Functions +// ============================================================================ + +/** + * @brief Compute exclusive prefix sum (cumulative sum shifted by one) + * + * Computes the exclusive cumulative sum of a 1D int32 or int64 tensor. + * The output has the same shape as the input, where output[i] = sum(input[0..i-1]) + * and output[0] = 0. Used to convert a lengths tensor into an offsets tensor. + * + * @param t_in Input tensor (int32 or int64, contiguous, size < INT_MAX) + * @return Exclusive cumsum tensor with the same shape as t_in + */ + +at::Tensor asynchronous_exclusive_cumsum(const at::Tensor& t_in); + +// ============================================================================ +// Kernel Functions +// ============================================================================ + +/** + * @brief Permute 2D lengths array + * + * Kernel function to permute lengths tensor according to permutation indices. + * + * @param T Number of tables/features + * @param B Batch size + * @param lengths_contig Input lengths tensor [T, B] + * @param permute_contig Permutation indices [T] + * @param permuted_lengths Output permuted lengths tensor [T, B] + */ +void permute_2D_lengths_kernel_xpu( + int32_t T, + int32_t B, + const at::Tensor& lengths_contig, + const at::Tensor& permute_contig, + at::Tensor& permuted_lengths); + +/** + * @brief Permute 2D sparse data (indices and optional weights) + * + * Kernel function to permute sparse indices and weights according to + * permutation indices and offset information. + * + * @param permuted_indices_size Total size of output indices + * @param T Number of tables/features + * @param B Batch size + * @param indices_contig Input indices + * @param weights Optional weights tensor + * @param weights_columns Number of weight columns (for 2D weights) + * @param permute_contig Permutation indices + * @param input_offsets Input offset tensor + * @param output_offsets Output offset tensor + * @param permuted_indices Output permuted indices + * @param permuted_weights Optional output permuted weights + */ +void permute_2D_data_kernel_xpu( + int32_t permuted_indices_size, + int32_t T, + int32_t B, + const at::Tensor& indices_contig, + const std::optional& weights, + const int32_t weights_columns, + const at::Tensor& permute_contig, + const at::Tensor& input_offsets, + const at::Tensor& output_offsets, + at::Tensor& permuted_indices, + const std::optional& permuted_weights); + +} // namespace fbgemm_xpu diff --git a/packages/fbgemm-xpu/test/.gitkeep b/packages/fbgemm-xpu/test/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/packages/fbgemm-xpu/test/test_asynchronous_complete_cumsum.py b/packages/fbgemm-xpu/test/test_asynchronous_complete_cumsum.py new file mode 100644 index 0000000..09024c4 --- /dev/null +++ b/packages/fbgemm-xpu/test/test_asynchronous_complete_cumsum.py @@ -0,0 +1,71 @@ +""" +Test suite for asynchronous_complete_cumsum operator + +Tests complete cumulative sum computation with leading zero. +For 1D input [a, b, c], returns [0, a, a+b, a+b+c]. +For 2D input, cumsum along dimension 1 (columns). +""" + +import random + +import numpy as np +import torch +import unittest + +import fbgemm_gpu +import fbgemm_xpu +from torch.testing._internal.common_utils import TestCase, run_tests + +SEED = 42 + + +class TestAsynchronousCompleteCumsum(TestCase): + def setUp(self): + super().setUp() + torch.manual_seed(SEED) + + def test_basic_int32(self): + if not torch.xpu.is_available(): + self.skipTest("XPU not available") + + x = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device="xpu") + result = torch.ops.fbgemm.asynchronous_complete_cumsum(x) + expected = torch.tensor([0, 1, 3, 6, 10], dtype=torch.int32, device="xpu") + torch.testing.assert_close(result, expected) + + def test_basic_int64(self): + if not torch.xpu.is_available(): + self.skipTest("XPU not available") + + x = torch.tensor([1, 2, 3, 4], dtype=torch.int64, device="xpu") + result = torch.ops.fbgemm.asynchronous_complete_cumsum(x) + expected = torch.tensor([0, 1, 3, 6, 10], dtype=torch.int64, device="xpu") + torch.testing.assert_close(result, expected) + + def test_empty(self): + if not torch.xpu.is_available(): + self.skipTest("XPU not available") + + x = torch.tensor([], dtype=torch.int64, device="xpu") + result = torch.ops.fbgemm.asynchronous_complete_cumsum(x) + self.assertEqual(result.numel(), 1) + self.assertEqual(result[0].item(), 0) + + def test_random(self): + if not torch.xpu.is_available(): + self.skipTest("XPU not available") + + for _ in range(5): + n = random.randint(0, 50) + for dtype in [torch.int32, torch.int64]: + x = torch.randint(0, 100, (n,), dtype=dtype, device="xpu") + result = torch.ops.fbgemm.asynchronous_complete_cumsum(x) + expected_np = np.cumsum([0] + x.cpu().numpy().tolist()) + expected = torch.from_numpy(expected_np.astype( + np.int32 if dtype == torch.int32 else np.int64 + )) + torch.testing.assert_close(result.cpu(), expected) + + +if __name__ == "__main__": + run_tests() diff --git a/packages/fbgemm-xpu/test/test_invert_permute.py b/packages/fbgemm-xpu/test/test_invert_permute.py new file mode 100644 index 0000000..e7be0af --- /dev/null +++ b/packages/fbgemm-xpu/test/test_invert_permute.py @@ -0,0 +1,374 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. +# Copyright (c) 2026 Intel Corporation. All Rights Reserved. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Test suite for invert_permute custom operator + +This module tests the correctness of the invert_permute operator +on XPU (Intel GPU) devices. + +Tests are designed to match TorchRec/DLRM sparse embedding scenarios. +""" + +import torch +import unittest +import fbgemm_gpu +import fbgemm_xpu +from torch.testing._internal.common_utils import TestCase, run_tests +from torch.testing._internal.optests import opcheck + +# Reproducible seed for all random tests +SEED = 42 + + +def generate_random_permutation(n, device="cpu", dtype=torch.int32, seed=None): + """ + Generate a random permutation of [0, 1, ..., n-1] + + Args: + n: Size of the permutation + device: Device to place the tensor on + dtype: Data type (torch.int32 or torch.int64) + seed: Random seed for reproducibility + + Returns: + Random permutation tensor + """ + if seed is not None: + torch.manual_seed(seed) + + permute = torch.arange(n, dtype=dtype, device=device) + # Generate random permutation using randperm + indices = torch.randperm(n, device=device) + return permute[indices] + + +def compute_inverse_reference(permute): + """ + Reference CPU implementation of inverse permutation + + This is the ground truth implementation used for testing. + + Args: + permute: Input permutation tensor + + Returns: + Inverse permutation tensor + """ + n = permute.size(0) + inversed = torch.empty_like(permute) + + permute_cpu = permute.cpu() + inversed_cpu = inversed.cpu() + + for i in range(n): + inversed_cpu[permute_cpu[i]] = i + + return inversed_cpu.to(permute.device) + + +def verify_inverse(permute, inversed): + """ + Verify that inversed is truly the inverse of permute + + Checks that: inversed[permute[i]] == i for all i + + Args: + permute: Original permutation + inversed: Claimed inverse permutation + + Returns: + True if inversed is correct, False otherwise + """ + n = permute.size(0) + for i in range(n): + if inversed[permute[i]] != i: + return False + return True + + +class TestInvertPermute(TestCase): + """Test cases for invert_permute operator""" + + def _test_correctness(self, device, dtype): + """ + Helper method to test correctness on a specific device and dtype + + Args: + device: Device string ("cpu" or "xpu") + dtype: torch dtype (torch.int32 or torch.int64) + """ + # Test case 1: Empty tensor + permute_empty = torch.tensor([], dtype=dtype, device=device) + result_empty = torch.ops.fbgemm.invert_permute(permute_empty) + self.assertEqual(result_empty.shape, permute_empty.shape) + self.assertEqual(result_empty.dtype, dtype) + + # Test case 2: Single element [0] + permute_single = torch.tensor([0], dtype=dtype, device=device) + result_single = torch.ops.fbgemm.invert_permute(permute_single) + expected_single = torch.tensor([0], dtype=dtype, device=device) + torch.testing.assert_close(result_single, expected_single) + self.assertTrue(verify_inverse(permute_single, result_single)) + + # Test case 3: Small permutation [2, 0, 1] + permute_small = torch.tensor([2, 0, 1], dtype=dtype, device=device) + result_small = torch.ops.fbgemm.invert_permute(permute_small) + expected_small = torch.tensor([1, 2, 0], dtype=dtype, device=device) + torch.testing.assert_close(result_small, expected_small) + self.assertTrue(verify_inverse(permute_small, result_small)) + + # Test case 4: Identity permutation [0, 1, 2, ..., 99] + n = 100 + permute_identity = torch.arange(n, dtype=dtype, device=device) + result_identity = torch.ops.fbgemm.invert_permute(permute_identity) + torch.testing.assert_close(result_identity, permute_identity) + self.assertTrue(verify_inverse(permute_identity, result_identity)) + + # Test case 5: Reverse permutation [99, 98, ..., 1, 0] + permute_reverse = torch.arange(n - 1, -1, -1, dtype=dtype, device=device) + result_reverse = torch.ops.fbgemm.invert_permute(permute_reverse) + torch.testing.assert_close(result_reverse, permute_reverse) + self.assertTrue(verify_inverse(permute_reverse, result_reverse)) + + # Test case 6: Medium random permutation (N=1000) + n = 1000 + permute_medium = generate_random_permutation(n, device=device, dtype=dtype) + result_medium = torch.ops.fbgemm.invert_permute(permute_medium) + expected_medium = compute_inverse_reference(permute_medium) + torch.testing.assert_close(result_medium, expected_medium) + self.assertTrue(verify_inverse(permute_medium, result_medium)) + + # Test case 7: Large random permutation (N=100000) + n = 100000 + permute_large = generate_random_permutation(n, device=device, dtype=dtype) + result_large = torch.ops.fbgemm.invert_permute(permute_large) + expected_large = compute_inverse_reference(permute_large) + torch.testing.assert_close(result_large, expected_large) + self.assertTrue(verify_inverse(permute_large, result_large)) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_correctness_xpu_int32(self): + """Test correctness on XPU (Intel GPU) with int32 dtype""" + self._test_correctness("xpu", torch.int32) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_correctness_xpu_int64(self): + """Test correctness on XPU (Intel GPU) with int64 dtype""" + self._test_correctness("xpu", torch.int64) + + def test_invalid_input_2d(self): + """Test that 2D input raises an error""" + permute_2d = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32, device="xpu") + with self.assertRaisesRegex(RuntimeError, "must be 1-dimensional"): + torch.ops.fbgemm.invert_permute(permute_2d) + + def test_invalid_dtype_float(self): + """Test that float dtype raises an error""" + permute_float = torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32, device="xpu") + with self.assertRaisesRegex(RuntimeError, "must be int32 or int64"): + torch.ops.fbgemm.invert_permute(permute_float) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_meta_function(self): + """Test that meta function works correctly for torch.compile""" + device = "xpu" + dtype = torch.int32 + n = 100 + + permute = generate_random_permutation(n, device=device, dtype=dtype) + + # Create a FakeTensor to test meta function + from torch._subclasses.fake_tensor import FakeTensorMode + + with FakeTensorMode(): + fake_permute = torch.empty(n, dtype=dtype, device=device) + # This should use the meta function without actual computation + fake_result = torch.ops.fbgemm.invert_permute(fake_permute) + + # Check output properties + self.assertEqual(fake_result.shape, permute.shape) + self.assertEqual(fake_result.dtype, dtype) + self.assertEqual(fake_result.device.type, device) + + def _test_opcheck(self, device, dtype): + """ + Test operator with PyTorch's opcheck utility + + opcheck validates that the operator follows PyTorch's conventions + for autograd, vmap, and other advanced features. + """ + n = 100 + permute = generate_random_permutation(n, device=device, dtype=dtype) + + # opcheck validates the operator implementation + # It checks schema, device handling, and other PyTorch conventions + opcheck(torch.ops.fbgemm.invert_permute, (permute,)) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_opcheck_xpu_int32(self): + """Run opcheck validation on XPU with int32""" + self._test_opcheck("xpu", torch.int32) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_opcheck_xpu_int64(self): + """Run opcheck validation on XPU with int64""" + self._test_opcheck("xpu", torch.int64) + + +class TestInvertPermuteParametric(TestCase): + """ + Parametric tests with varying sizes and edge cases + + These tests cover TorchRec/DLRM-like scenarios with reproducible seeds. + """ + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_varying_sizes_xpu(self): + """Test various tensor sizes on XPU""" + torch.manual_seed(SEED) + + test_sizes = [1, 10, 128, 1024, 8192, 65536, 100000] + + for n in test_sizes: + with self.subTest(size=n): + permute = generate_random_permutation(n, device="xpu", dtype=torch.int32, seed=SEED + n) + result = torch.ops.fbgemm.invert_permute(permute) + expected = compute_inverse_reference(permute) + torch.testing.assert_close(result, expected) + self.assertTrue(verify_inverse(permute, result)) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_torchrec_batch_shapes_xpu(self): + """Test with TorchRec shapes on XPU""" + torch.manual_seed(SEED) + + torchrec_shapes = [100, 1000, 10000, 100000] + + for num_embeddings in torchrec_shapes: + with self.subTest(num_embeddings=num_embeddings): + permute = generate_random_permutation( + num_embeddings, device="xpu", dtype=torch.int64, seed=SEED + ) + result = torch.ops.fbgemm.invert_permute(permute) + expected = compute_inverse_reference(permute) + torch.testing.assert_close(result, expected) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_power_of_two_sizes_xpu(self): + """Test power-of-2 sizes on XPU""" + torch.manual_seed(SEED) + + for power in [4, 8, 10, 12, 14, 16]: + n = 2 ** power + with self.subTest(size=n, power=power): + permute = generate_random_permutation(n, device="xpu", dtype=torch.int32, seed=SEED) + result = torch.ops.fbgemm.invert_permute(permute) + self.assertTrue(verify_inverse(permute, result)) + + +class TestCPUXPUParity(TestCase): + """Test that CPU and XPU implementations produce identical results""" + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_parity_small(self): + """Test CPU-XPU parity for small inputs""" + torch.manual_seed(SEED) + + for n in [10, 100, 1000]: + with self.subTest(size=n): + permute_cpu = generate_random_permutation(n, device="cpu", dtype=torch.int32, seed=SEED) + permute_xpu = permute_cpu.to("xpu") + + result_cpu = torch.ops.fbgemm.invert_permute(permute_cpu) + result_xpu = torch.ops.fbgemm.invert_permute(permute_xpu) + + # CPU and XPU should produce identical results + torch.testing.assert_close(result_cpu, result_xpu.cpu()) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_parity_large(self): + """Test CPU-XPU parity for large inputs""" + torch.manual_seed(SEED) + + for n in [10000, 100000]: + with self.subTest(size=n): + permute_cpu = generate_random_permutation(n, device="cpu", dtype=torch.int64, seed=SEED) + permute_xpu = permute_cpu.to("xpu") + + result_cpu = torch.ops.fbgemm.invert_permute(permute_cpu) + result_xpu = torch.ops.fbgemm.invert_permute(permute_xpu) + + torch.testing.assert_close(result_cpu, result_xpu.cpu()) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_parity_edge_cases(self): + """Test CPU-XPU parity for edge cases""" + torch.manual_seed(SEED) + + # Empty tensor + permute_empty = torch.tensor([], dtype=torch.int32) + result_cpu_empty = torch.ops.fbgemm.invert_permute(permute_empty) + result_xpu_empty = torch.ops.fbgemm.invert_permute(permute_empty.to("xpu")) + torch.testing.assert_close(result_cpu_empty, result_xpu_empty.cpu()) + + # Single element + permute_single = torch.tensor([0], dtype=torch.int32) + result_cpu_single = torch.ops.fbgemm.invert_permute(permute_single) + result_xpu_single = torch.ops.fbgemm.invert_permute(permute_single.to("xpu")) + torch.testing.assert_close(result_cpu_single, result_xpu_single.cpu()) + + # Identity + permute_identity = torch.arange(100, dtype=torch.int32) + result_cpu_identity = torch.ops.fbgemm.invert_permute(permute_identity) + result_xpu_identity = torch.ops.fbgemm.invert_permute(permute_identity.to("xpu")) + torch.testing.assert_close(result_cpu_identity, result_xpu_identity.cpu()) + + +class TestInvertPermutePerformance(TestCase): + """Performance benchmarking for invert_permute operator""" + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_benchmark_xpu(self): + """ + Benchmark invert_permute on XPU + + This test measures execution time for different input sizes. + It's marked as a test but primarily serves as a benchmarking tool. + """ + import time + + print("\n=== XPU Performance Benchmark ===") + sizes = [1000, 10000, 100000, 1000000] + num_iterations = 100 + + for n in sizes: + for dtype in [torch.int32, torch.int64]: + permute = generate_random_permutation(n, device="xpu", dtype=dtype, seed=SEED) + + # Warmup + for _ in range(10): + _ = torch.ops.fbgemm.invert_permute(permute) + + # Benchmark + torch.xpu.synchronize() + start = time.perf_counter() + + for _ in range(num_iterations): + result = torch.ops.fbgemm.invert_permute(permute) + + torch.xpu.synchronize() + end = time.perf_counter() + + avg_time_ms = (end - start) * 1000 / num_iterations + bytes_transferred = 2 * n * (4 if dtype == torch.int32 else 8) + bandwidth_gb_s = (bytes_transferred / (avg_time_ms / 1000)) / (1024**3) + + dtype_str = "int32" if dtype == torch.int32 else "int64" + print(f"N={n:>7} ({dtype_str}): {avg_time_ms:.4f} ms, " + f"Bandwidth: {bandwidth_gb_s:.2f} GB/s") + + +if __name__ == "__main__": + run_tests() diff --git a/packages/fbgemm-xpu/test/test_permute_1d_sparse_data.py b/packages/fbgemm-xpu/test/test_permute_1d_sparse_data.py new file mode 100644 index 0000000..7cbb01d --- /dev/null +++ b/packages/fbgemm-xpu/test/test_permute_1d_sparse_data.py @@ -0,0 +1,460 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. +# Copyright (c) 2026 Intel Corporation. All Rights Reserved. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Test suite for permute_1D_sparse_data custom operator + +This module tests the correctness of the permute_1D_sparse_data operator +on XPU (Intel GPU) devices. + +Tests are designed to match TorchRec/DLRM sparse embedding scenarios. +""" + +import torch +import unittest +import fbgemm_gpu # Provides fallback implementations +import fbgemm_xpu +from torch.testing._internal.common_utils import TestCase, run_tests +from torch.testing._internal.optests import opcheck +from typing import Optional, Tuple + +# Reproducible seed for all random tests +SEED = 42 + + +def compute_reference_permute_1d_sparse_data( + permute: torch.Tensor, + lengths: torch.Tensor, + indices: torch.Tensor, + weights: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Reference implementation of permute_1D_sparse_data using pure Python/PyTorch + + This is the ground truth implementation used for testing. + + Args: + permute: Permutation indices tensor + lengths: Segment lengths tensor + indices: Concatenated values tensor + weights: Optional weights tensor + + Returns: + Tuple of (permuted_lengths, permuted_indices, permuted_weights) + """ + # Move to CPU for reference computation + permute_cpu = permute.cpu() + lengths_cpu = lengths.cpu() + indices_cpu = indices.cpu() + weights_cpu = weights.cpu() if weights is not None else None + + permuted_lengths_size = permute_cpu.numel() + + # Compute permuted lengths + permuted_lengths = torch.empty(permuted_lengths_size, dtype=lengths_cpu.dtype) + for i in range(permuted_lengths_size): + permuted_lengths[i] = lengths_cpu[permute_cpu[i]] + + # Compute input offsets (exclusive cumsum) + input_offsets = torch.zeros(lengths_cpu.numel() + 1, dtype=torch.long) + for i in range(lengths_cpu.numel()): + input_offsets[i + 1] = input_offsets[i] + lengths_cpu[i] + + # Compute output offsets (exclusive cumsum) + output_offsets = torch.zeros(permuted_lengths_size + 1, dtype=torch.long) + for i in range(permuted_lengths_size): + output_offsets[i + 1] = output_offsets[i] + permuted_lengths[i] + + # Compute permuted indices size + permuted_indices_size = output_offsets[permuted_lengths_size].item() + + # Allocate output tensors + permuted_indices = torch.empty(permuted_indices_size, dtype=indices_cpu.dtype) + permuted_weights = torch.empty(permuted_indices_size, dtype=weights_cpu.dtype) if weights_cpu is not None else None + + # Copy data + for j in range(permuted_lengths_size): + segment_length = permuted_lengths[j].item() + input_start = input_offsets[permute_cpu[j]].item() + output_start = output_offsets[j].item() + + for i in range(segment_length): + permuted_indices[output_start + i] = indices_cpu[input_start + i] + if permuted_weights is not None: + permuted_weights[output_start + i] = weights_cpu[input_start + i] + + # Move back to original device + device = permute.device + permuted_lengths = permuted_lengths.to(device) + permuted_indices = permuted_indices.to(device) + if permuted_weights is not None: + permuted_weights = permuted_weights.to(device) + + return permuted_lengths, permuted_indices, permuted_weights + + +def generate_sparse_data( + num_segments: int, + avg_segment_length: int, + device: str = "cpu", + lengths_dtype: torch.dtype = torch.int32, + indices_dtype: torch.dtype = torch.int64, + weights_dtype: Optional[torch.dtype] = None, + seed: Optional[int] = None +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Generate random sparse data for testing + + Args: + num_segments: Number of segments + avg_segment_length: Average length of each segment + device: Device to place tensors on + lengths_dtype: Data type for lengths tensor + indices_dtype: Data type for indices tensor + weights_dtype: Data type for weights tensor (None to skip weights) + seed: Random seed for reproducibility + + Returns: + Tuple of (permute, lengths, indices, weights) + """ + if seed is not None: + torch.manual_seed(seed) + + # Generate random lengths (Poisson-like distribution around avg_segment_length) + lengths = torch.randint( + max(1, avg_segment_length - 5), + avg_segment_length + 6, + (num_segments,), + dtype=lengths_dtype, + device=device + ) + + # Total number of values + total_values = lengths.sum().item() + + # Generate random indices + indices = torch.randint(0, 10000, (total_values,), dtype=indices_dtype, device=device) + + # Generate random permutation + permute = torch.randperm(num_segments, dtype=torch.int32, device=device) + + # Generate weights if requested + weights = None + if weights_dtype is not None: + weights = torch.randn(total_values, dtype=weights_dtype, device=device) + + return permute, lengths, indices, weights + + +class TestPermute1DSparseData(TestCase): + """Test cases for permute_1D_sparse_data operator""" + + def setUp(self): + super().setUp() + torch.manual_seed(SEED) + + def _test_correctness(self, device: str, lengths_dtype: torch.dtype, + indices_dtype: torch.dtype, with_weights: bool = False): + """ + Helper method to test correctness on a specific device and dtype combination + """ + weights_dtype = torch.float32 if with_weights else None + + # Test case 1: Basic permutation [2, 0, 1] + permute = torch.tensor([2, 0, 1], dtype=torch.int32, device=device) + lengths = torch.tensor([3, 2, 4], dtype=lengths_dtype, device=device) + indices = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=indices_dtype, device=device) + weights = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + dtype=torch.float32, device=device) if with_weights else None + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, indices, weights) + expected = compute_reference_permute_1d_sparse_data(permute, lengths, indices, weights) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + if with_weights: + torch.testing.assert_close(result[2], expected[2]) + + # Verify expected values manually + expected_lengths = torch.tensor([4, 3, 2], dtype=lengths_dtype, device=device) + expected_indices = torch.tensor([5, 6, 7, 8, 0, 1, 2, 3, 4], dtype=indices_dtype, device=device) + torch.testing.assert_close(result[0], expected_lengths) + torch.testing.assert_close(result[1], expected_indices) + + # Test case 2: Identity permutation + n_segments = 10 + permute_id = torch.arange(n_segments, dtype=torch.int32, device=device) + lengths_id = torch.randint(1, 5, (n_segments,), dtype=lengths_dtype, device=device) + total = lengths_id.sum().item() + indices_id = torch.arange(total, dtype=indices_dtype, device=device) + weights_id = torch.randn(total, dtype=torch.float32, device=device) if with_weights else None + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute_id, lengths_id, indices_id, weights_id) + + torch.testing.assert_close(result[0], lengths_id) + torch.testing.assert_close(result[1], indices_id) + if with_weights: + torch.testing.assert_close(result[2], weights_id) + + # Test case 3: Reverse permutation + n_segments = 5 + permute_rev = torch.arange(n_segments - 1, -1, -1, dtype=torch.int32, device=device) + lengths_rev = torch.tensor([2, 3, 1, 4, 2], dtype=lengths_dtype, device=device) + total = lengths_rev.sum().item() + indices_rev = torch.arange(total, dtype=indices_dtype, device=device) + weights_rev = torch.randn(total, dtype=torch.float32, device=device) if with_weights else None + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute_rev, lengths_rev, indices_rev, weights_rev) + expected = compute_reference_permute_1d_sparse_data(permute_rev, lengths_rev, indices_rev, weights_rev) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + if with_weights: + torch.testing.assert_close(result[2], expected[2]) + + # Test case 4: Random medium-sized data + permute, lengths, indices, weights = generate_sparse_data( + num_segments=100, + avg_segment_length=20, + device=device, + lengths_dtype=lengths_dtype, + indices_dtype=indices_dtype, + weights_dtype=weights_dtype, + seed=SEED + ) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, indices, weights) + expected = compute_reference_permute_1d_sparse_data(permute, lengths, indices, weights) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + if with_weights: + torch.testing.assert_close(result[2], expected[2]) + + # Test case 5: Large-scale data (DLRM-like) + permute, lengths, indices, weights = generate_sparse_data( + num_segments=1000, + avg_segment_length=40, + device=device, + lengths_dtype=lengths_dtype, + indices_dtype=indices_dtype, + weights_dtype=weights_dtype, + seed=SEED + 1 + ) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, indices, weights) + expected = compute_reference_permute_1d_sparse_data(permute, lengths, indices, weights) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + if with_weights: + torch.testing.assert_close(result[2], expected[2]) + + def _test_edge_cases(self, device: str): + """Test edge cases""" + # Test case 1: Empty permute + permute_empty = torch.tensor([], dtype=torch.int32, device=device) + lengths = torch.tensor([3, 2], dtype=torch.int32, device=device) + indices = torch.arange(5, dtype=torch.int64, device=device) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute_empty, lengths, indices) + self.assertEqual(result[0].numel(), 0) + + # Test case 2: Single segment + permute_single = torch.tensor([0], dtype=torch.int32, device=device) + lengths_single = torch.tensor([5], dtype=torch.int32, device=device) + indices_single = torch.tensor([10, 20, 30, 40, 50], dtype=torch.int64, device=device) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute_single, lengths_single, indices_single) + torch.testing.assert_close(result[0], lengths_single) + torch.testing.assert_close(result[1], indices_single) + + # Test case 3: Zero-length segments + permute_zero = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) + lengths_zero = torch.tensor([0, 3, 0], dtype=torch.int32, device=device) + indices_zero = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute_zero, lengths_zero, indices_zero) + expected = compute_reference_permute_1d_sparse_data(permute_zero, lengths_zero, indices_zero) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + + # Test case 4: Permutation with duplicates (replication) + permute_dup = torch.tensor([0, 0, 1], dtype=torch.int32, device=device) + lengths_dup = torch.tensor([2, 3], dtype=torch.int32, device=device) + indices_dup = torch.tensor([10, 20, 30, 40, 50], dtype=torch.int64, device=device) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute_dup, lengths_dup, indices_dup) + expected_lengths = torch.tensor([2, 2, 3], dtype=torch.int32, device=device) + expected_indices = torch.tensor([10, 20, 10, 20, 30, 40, 50], dtype=torch.int64, device=device) + + torch.testing.assert_close(result[0], expected_lengths) + torch.testing.assert_close(result[1], expected_indices) + + def _test_with_permuted_lengths_sum(self, device: str): + """Test with precomputed permuted_lengths_sum to avoid sync""" + permute = torch.tensor([2, 0, 1], dtype=torch.int32, device=device) + lengths = torch.tensor([3, 2, 4], dtype=torch.int32, device=device) + indices = torch.arange(9, dtype=torch.int64, device=device) + + # Precompute the sum: permuted_lengths = [4, 3, 2], sum = 9 + permuted_lengths_sum = 9 + + result = torch.ops.fbgemm.permute_1D_sparse_data( + permute, lengths, indices, None, permuted_lengths_sum + ) + expected = compute_reference_permute_1d_sparse_data(permute, lengths, indices) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + + # ======================== + # XPU Tests + # ======================== + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_correctness_xpu_int32_int64(self): + """Test XPU with int32 lengths and int64 indices""" + self._test_correctness("xpu", torch.int32, torch.int64) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_correctness_xpu_int64_int64(self): + """Test XPU with int64 lengths and int64 indices""" + self._test_correctness("xpu", torch.int64, torch.int64) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_correctness_xpu_int32_int32(self): + """Test XPU with int32 lengths and int32 indices""" + self._test_correctness("xpu", torch.int32, torch.int32) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_correctness_xpu_with_weights(self): + """Test XPU with weights""" + self._test_correctness("xpu", torch.int32, torch.int64, with_weights=True) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_edge_cases_xpu(self): + """Test edge cases on XPU""" + self._test_edge_cases("xpu") + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_with_permuted_lengths_sum_xpu(self): + """Test with precomputed permuted_lengths_sum on XPU""" + self._test_with_permuted_lengths_sum("xpu") + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_opcheck_xpu(self): + """Test PT2/torch.compile compliance using opcheck""" + permute = torch.tensor([2, 0, 1], dtype=torch.int32, device="xpu") + lengths = torch.tensor([3, 2, 4], dtype=torch.int32, device="xpu") + indices = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int64, device="xpu") + + # Test without weights + opcheck( + torch.ops.fbgemm.permute_1D_sparse_data, + (permute, lengths, indices, None, None), + test_utils=["test_schema", "test_autograd_registration", "test_faketensor"] + ) + + # Test with weights + weights = torch.randn(9, dtype=torch.float32, device="xpu") + opcheck( + torch.ops.fbgemm.permute_1D_sparse_data, + (permute, lengths, indices, weights, None), + test_utils=["test_schema", "test_autograd_registration", "test_faketensor"] + ) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_cpu_xpu_consistency(self): + """Test that CPU and XPU produce same results""" + permute, lengths, indices, weights = generate_sparse_data( + num_segments=500, + avg_segment_length=30, + device="cpu", + lengths_dtype=torch.int32, + indices_dtype=torch.int64, + weights_dtype=torch.float32, + seed=SEED + ) + + # Run on CPU + result_cpu = torch.ops.fbgemm.permute_1D_sparse_data( + permute, lengths, indices, weights + ) + + # Run on XPU + permute_xpu = permute.to("xpu") + lengths_xpu = lengths.to("xpu") + indices_xpu = indices.to("xpu") + weights_xpu = weights.to("xpu") + + result_xpu = torch.ops.fbgemm.permute_1D_sparse_data( + permute_xpu, lengths_xpu, indices_xpu, weights_xpu + ) + + # Compare results + torch.testing.assert_close(result_cpu[0], result_xpu[0].cpu()) + torch.testing.assert_close(result_cpu[1], result_xpu[1].cpu()) + torch.testing.assert_close(result_cpu[2], result_xpu[2].cpu()) + + +class TestPermute1DSparseDataDLRMScale(TestCase): + """ + Performance-oriented tests at DLRM production scale + """ + + def setUp(self): + super().setUp() + torch.manual_seed(SEED) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_dlrm_v2_scale(self): + """Test at DLRMv2 typical scale""" + # DLRMv2 configuration + num_features = 26 + batch_size = 512 + avg_pooling_factor = 40 + + total_segments = num_features * batch_size # 13,312 + + permute, lengths, indices, weights = generate_sparse_data( + num_segments=total_segments, + avg_segment_length=avg_pooling_factor, + device="xpu", + lengths_dtype=torch.int32, + indices_dtype=torch.int64, + weights_dtype=torch.float32, + seed=SEED + ) + + # Run and verify + result = torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, indices, weights) + expected = compute_reference_permute_1d_sparse_data(permute, lengths, indices, weights) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + torch.testing.assert_close(result[2], expected[2]) + + @unittest.skipIf(not torch.xpu.is_available(), "XPU not available") + def test_large_segments(self): + """Test with large segment sizes""" + permute, lengths, indices, _ = generate_sparse_data( + num_segments=100, + avg_segment_length=1000, # Large segments + device="xpu", + lengths_dtype=torch.int32, + indices_dtype=torch.int64, + seed=SEED + ) + + result = torch.ops.fbgemm.permute_1D_sparse_data(permute, lengths, indices) + expected = compute_reference_permute_1d_sparse_data(permute, lengths, indices) + + torch.testing.assert_close(result[0], expected[0]) + torch.testing.assert_close(result[1], expected[1]) + + +if __name__ == "__main__": + run_tests() diff --git a/packages/fbgemm-xpu/test/test_permute_2d_sparse_data.py b/packages/fbgemm-xpu/test/test_permute_2d_sparse_data.py new file mode 100644 index 0000000..e69f5f9 --- /dev/null +++ b/packages/fbgemm-xpu/test/test_permute_2d_sparse_data.py @@ -0,0 +1,190 @@ +""" +Test suite for permute_2D_sparse_data operator + +Ported from test_fbgemm_ops.py +""" + +import random +from itertools import accumulate +from typing import Optional, Tuple + +import torch +import unittest +import fbgemm_gpu +import fbgemm_xpu +from torch.testing._internal.common_utils import TestCase, run_tests + +DEVICE = "xpu" + + +# ============================================================================= +# Reference / helper functions +# ============================================================================= + +def permute_indices_ref_( + lengths: torch.Tensor, + indices: torch.Tensor, + weights: Optional[torch.Tensor], + permute: torch.LongTensor, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + T = lengths.size(0) + B = lengths.size(1) + if T == 0 or B == 0: + return lengths, indices, weights + + permuted_lengths = torch.index_select(lengths.view(T, -1), 0, permute) + original_segment_lengths = lengths.view(T, -1).sum(dim=1, dtype=torch.int32) + original_segment_start = [0] + list(accumulate(original_segment_lengths.view(-1))) + + permuted_indices = [] + permuted_weights = [] + for i in range(permute.size(0)): + start = original_segment_start[permute[i]] + end = start + original_segment_lengths[permute[i]] + permuted_indices.append(indices[start:end]) + if weights is not None: + permuted_weights.append(weights[start:end]) + + permuted_indices = torch.cat(permuted_indices, dim=0).flatten() + + if weights is None: + permuted_weights = None + else: + permuted_weights = torch.cat(permuted_weights, dim=0).flatten() + + return permuted_lengths, permuted_indices, permuted_weights + + +# ============================================================================= +# Tests +# ============================================================================= + +class TestPermute2DSparseData(TestCase): + def test_basic(self): + T, B, L = 3, 4, 5 + index_dtype = torch.int32 + lengths = torch.randint(1, L, (T, B), dtype=index_dtype) + indices = torch.randint(1, int(1e5), (lengths.sum().item(),), dtype=index_dtype) + weights = torch.rand(lengths.sum().item()).float() + + permute_list = list(range(T)) + random.shuffle(permute_list) + permute = torch.IntTensor(permute_list) + + permuted_lengths_ref, permuted_indices_ref, permuted_weights_ref = ( + permute_indices_ref_(lengths, indices, weights, permute.long()) + ) + + permuted_lengths_xpu, permuted_indices_xpu, permuted_weights_xpu = ( + torch.ops.fbgemm.permute_2D_sparse_data( + permute.to(DEVICE), lengths.to(DEVICE), indices.to(DEVICE), + weights.to(DEVICE), None, + ) + ) + + torch.testing.assert_close(permuted_lengths_xpu.cpu(), permuted_lengths_ref) + torch.testing.assert_close(permuted_indices_xpu.cpu(), permuted_indices_ref) + torch.testing.assert_close(permuted_weights_xpu.cpu(), permuted_weights_ref) + + def test_no_weights(self): + T, B, L = 3, 4, 5 + index_dtype = torch.int32 + lengths = torch.randint(1, L, (T, B), dtype=index_dtype) + indices = torch.randint(1, int(1e5), (lengths.sum().item(),), dtype=index_dtype) + + permute_list = list(range(T)) + random.shuffle(permute_list) + permute = torch.IntTensor(permute_list) + + permuted_lengths_ref, permuted_indices_ref, _ = ( + permute_indices_ref_(lengths, indices, None, permute.long()) + ) + + permuted_lengths_xpu, permuted_indices_xpu, permuted_weights_xpu = ( + torch.ops.fbgemm.permute_2D_sparse_data( + permute.to(DEVICE), lengths.to(DEVICE), indices.to(DEVICE), + None, None, + ) + ) + + torch.testing.assert_close(permuted_lengths_xpu.cpu(), permuted_lengths_ref) + torch.testing.assert_close(permuted_indices_xpu.cpu(), permuted_indices_ref) + self.assertIsNone(permuted_weights_xpu) + + def test_int64(self): + T, B, L = 4, 3, 6 + index_dtype = torch.int64 + lengths = torch.randint(1, L, (T, B), dtype=index_dtype) + indices = torch.randint(1, int(1e5), (lengths.sum().item(),), dtype=index_dtype) + + permute_list = list(range(T)) + random.shuffle(permute_list) + permute = torch.IntTensor(permute_list) + + permuted_lengths_ref, permuted_indices_ref, _ = ( + permute_indices_ref_(lengths, indices, None, permute.long()) + ) + + permuted_lengths_xpu, permuted_indices_xpu, permuted_weights_xpu = ( + torch.ops.fbgemm.permute_2D_sparse_data( + permute.to(DEVICE), lengths.to(DEVICE), indices.to(DEVICE), + None, None, + ) + ) + + torch.testing.assert_close(permuted_lengths_xpu.cpu(), permuted_lengths_ref) + torch.testing.assert_close(permuted_indices_xpu.cpu(), permuted_indices_ref) + + def test_with_repeats(self): + T, B, L = 4, 3, 5 + index_dtype = torch.int32 + lengths = torch.randint(1, L, (T, B), dtype=index_dtype) + indices = torch.randint(1, int(1e5), (lengths.sum().item(),), dtype=index_dtype) + weights = torch.rand(lengths.sum().item()).float() + + permute_list = list(range(T)) + for _ in range(random.randint(0, T)): + permute_list.append(random.randint(0, T - 1)) + random.shuffle(permute_list) + permute = torch.IntTensor(permute_list) + + permuted_lengths_ref, permuted_indices_ref, permuted_weights_ref = ( + permute_indices_ref_(lengths, indices, weights, permute.long()) + ) + + permuted_lengths_xpu, permuted_indices_xpu, permuted_weights_xpu = ( + torch.ops.fbgemm.permute_2D_sparse_data( + permute.to(DEVICE), lengths.to(DEVICE), indices.to(DEVICE), + weights.to(DEVICE), None, + ) + ) + + torch.testing.assert_close(permuted_lengths_xpu.cpu(), permuted_lengths_ref) + torch.testing.assert_close(permuted_indices_xpu.cpu(), permuted_indices_ref) + torch.testing.assert_close(permuted_weights_xpu.cpu(), permuted_weights_ref) + + def test_exact_values(self): + """Test from torch-xpu-ops: known input/output pairs.""" + lengths = torch.tensor( + [[0, 0, 1], [0, 1, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 1]], + dtype=torch.int32, device=DEVICE, + ) + indices = torch.tensor([500, 1000, 1999], dtype=torch.int32, device=DEVICE) + permute = torch.tensor([0, 3, 1, 4, 2, 5], dtype=torch.int32, device=DEVICE) + weights = torch.rand((3, 64), device=DEVICE) + + lengths_out, values_out, weights_out = torch.ops.fbgemm.permute_2D_sparse_data( + permute, lengths, indices, weights, indices.numel() + ) + + expected_lengths = torch.tensor( + [[0, 0, 1], [0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 0], [0, 0, 1]], + dtype=torch.int32, device=DEVICE, + ) + self.assertTrue(torch.equal(lengths_out, expected_lengths)) + self.assertTrue(torch.equal(values_out, indices)) + self.assertTrue(torch.equal(weights_out, weights)) + + +if __name__ == "__main__": + run_tests()