From 081cc885df27d353895063f04232bcc1c07b4b6d Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 20:02:22 +0000 Subject: [PATCH 01/15] chore(fbgemm-xpu): remove .gitkeep files from populated directories Remove .gitkeep placeholders --- packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/.gitkeep | 0 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/.gitkeep | 0 packages/fbgemm-xpu/test/.gitkeep | 0 3 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/.gitkeep delete mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/.gitkeep delete mode 100644 packages/fbgemm-xpu/test/.gitkeep 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/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/test/.gitkeep b/packages/fbgemm-xpu/test/.gitkeep deleted file mode 100644 index e69de29..0000000 From 07bfb4b66be300d4cc0ad38b2d39927712c84641 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 20:05:07 +0000 Subject: [PATCH 02/15] feat(fbgemm-xpu): register invert_permute operation - Add invert_permute kernel to CMake build - Implement invert_permute Python wrapper in ops.py - Register invert_permute operator with schema existence check - Add torch_library.h utility for schema validation --- .../fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt | 1 + .../fbgemm_xpu/fbgemm_utils/torch_library.h | 17 +++++++++++ packages/fbgemm-xpu/src/fbgemm_xpu/ops.py | 13 +++++---- .../src/fbgemm_xpu/ops_registry.cpp | 29 +++++-------------- 4 files changed, 33 insertions(+), 27 deletions(-) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/torch_library.h diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt index 551ef95..abbc6d2 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt @@ -43,6 +43,7 @@ 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 ) set(all_sources ${host_sources}) 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/ops.py b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py index 0d10647..284f889 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py @@ -5,12 +5,13 @@ # Python wrapper functions for all custom operators under the fbgemm namespace # This module provides user-friendly interfaces to the C++ operators +import torch +from torch import Tensor + __all__ = [ - "dense_embedding_codegen_lookup_function", + "invert_permute", ] -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) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp index f4692b5..87424f0 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,9 @@ 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" + ); + } } From d6d6badef2cbd91285f44544d3fdc635cf33e30e Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 20:06:47 +0000 Subject: [PATCH 03/15] feat(fbgemm-xpu): add SYCL kernel implementation for invert_permute Add SYCL/XPU kernel implementation for invert_permute operation. --- .../sycl_kernels/invert_permute_kernel.cpp | 123 ++++++++++++++ .../sycl_kernels/invert_permute_kernel.h | 152 ++++++++++++++++++ 2 files changed, 275 insertions(+) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.cpp create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.h 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 From 14fe23cd62f3d6446b717cfa1e4a8370db055c47 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 20:11:14 +0000 Subject: [PATCH 04/15] test(fbgemm-xpu): add comprehensive test suite for invert_permute Add complete test coverage for invert_permute operator on XPU devices, covering correctness, validation, parity, and performance. Test coverage includes: - Correctness tests for int32/int64 with edge cases (empty, single element, identity, reverse, random permutations) - Input validation tests for invalid dimensions and dtypes - Meta function tests for torch.compile compatibility - PyTorch opcheck validation for operator conventions - Parametric tests with varying sizes (1 to 1M elements) - CPU-XPU parity tests to ensure consistent results - Performance benchmarks measuring execution time and bandwidth --- .../fbgemm-xpu/test/test_invert_permute.py | 374 ++++++++++++++++++ 1 file changed, 374 insertions(+) create mode 100644 packages/fbgemm-xpu/test/test_invert_permute.py 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() From 3056382db13ace437722d427bc1b23d640ff651f Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 21:30:13 +0000 Subject: [PATCH 05/15] feat(fbgemm-xpu): register permute_1D_sparse_data operator - CMakeLists: add permute_1d_sparse_data.cpp to build sources - ops.py: add Python wrapper with type hints - ops_registry.cpp: register operator schema in fbgemm namespace --- .../fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt | 1 + packages/fbgemm-xpu/src/fbgemm_xpu/ops.py | 16 ++++++++++++++++ .../fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp | 13 +++++++++++++ 3 files changed, 30 insertions(+) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt index abbc6d2..ca5dc30 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt @@ -44,6 +44,7 @@ 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 ) set(all_sources ${host_sources}) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py index 284f889..2551866 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py @@ -5,13 +5,29 @@ # 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__ = [ "invert_permute", + "permute_1D_sparse_data", ] 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 + ) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp index 87424f0..f052d15 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp @@ -54,4 +54,17 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) "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?)" + ); + } } From b7f3bbfeeb990b4b1f58a24425fbe426fd7c52a8 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 21:32:54 +0000 Subject: [PATCH 06/15] feat(fbgemm-xpu): add SYCL kernels for permute_1D_sparse_data Implement SYCL/XPU kernel implementation of permute_1D_sparse_data operator for sparse jagged/1D format data permutation. --- .../sycl_kernels/permute_1d_sparse_data.cpp | 430 ++++++++++++++++++ .../sycl_kernels/permute_1d_sparse_data.h | 240 ++++++++++ 2 files changed, 670 insertions(+) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.cpp create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_1d_sparse_data.h 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 From fd70ab5b036944ee5df2f28fae7ab2dc3146f489 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Thu, 18 Jun 2026 21:46:11 +0000 Subject: [PATCH 07/15] test(fbgemm-xpu): add comprehensive tests for permute_1D_sparse_data --- .../test/test_permute_1d_sparse_data.py | 460 ++++++++++++++++++ 1 file changed, 460 insertions(+) create mode 100644 packages/fbgemm-xpu/test/test_permute_1d_sparse_data.py 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() From 50f998298e4bef6084701c4494b4224a242d5fe7 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Fri, 19 Jun 2026 23:36:39 +0000 Subject: [PATCH 08/15] feat(fbgemm-xpu): add asynchronous_complete_cumsum XPU operator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add SYCL port of FBGEMM's asynchronous_complete_cumsum operator for Intel XPU devices. The operator computes a complete cumulative sum with a leading zero (e.g., [a, b, c] → [0, a, a+b, a+b+c]). --- .../asynchronous_complete_cumsum.cpp | 45 ++++++++++++++++ .../fbgemm_xpu/asynchronous_complete_cumsum.h | 52 +++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.cpp create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/asynchronous_complete_cumsum.h 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 From 79272c9777f10a4b8ca5fb95bb58025a2bf6a244 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Fri, 19 Jun 2026 23:37:57 +0000 Subject: [PATCH 09/15] feat(fbgemm-xpu): register asynchronous_complete_cumsum operator Integrate asynchronous_complete_cumsum operator into fbgemm-xpu: - Add Python wrapper with complete cumsum documentation - Register operator schema in torch library - Include implementation in CMake build --- packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt | 1 + packages/fbgemm-xpu/src/fbgemm_xpu/ops.py | 5 +++++ packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp | 6 ++++++ 3 files changed, 12 insertions(+) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt index ca5dc30..0b3b400 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt @@ -45,6 +45,7 @@ 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 ) set(all_sources ${host_sources}) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py index 2551866..2bc4ce1 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py @@ -13,6 +13,7 @@ __all__ = [ "invert_permute", "permute_1D_sparse_data", + "asynchronous_complete_cumsum", ] def invert_permute(permute: Tensor) -> Tensor: @@ -31,3 +32,7 @@ def permute_1D_sparse_data( 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) \ No newline at end of file diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp index f052d15..acc2b60 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp @@ -67,4 +67,10 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) ") -> (Tensor, Tensor, Tensor?)" ); } + + if (!utils::torch::schemaExists("fbgemm::asynchronous_complete_cumsum")) { + m.def( + "asynchronous_complete_cumsum(Tensor t_in) -> Tensor" + ); + } } From 389845d4af10370a4e98315405e3d267c7ea9ea4 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Fri, 19 Jun 2026 23:39:03 +0000 Subject: [PATCH 10/15] test(fbgemm-xpu): add tests for asynchronous_complete_cumsum Add comprehensive test suite for asynchronous_complete_cumsum operator covering: - Basic functionality with int32 and int64 dtypes - Empty tensor handling - Random input validation with numpy reference --- .../test/test_asynchronous_complete_cumsum.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 packages/fbgemm-xpu/test/test_asynchronous_complete_cumsum.py 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() From 26e18adfe3fe73be3766450bd8033e31b0ce0042 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Sat, 20 Jun 2026 01:06:18 +0000 Subject: [PATCH 11/15] feat(fbgemm-xpu): add SYCL communication utilities form torch-xpu-ops Add SYCL infrastructure headers from intel/torch-xpu-ops/ to support advanced kernel implementations: - DeviceProperties.h: Device capability queries and work group sizing - SYCLContext.h: SYCL context management and namespace aliases - SYCLHelpers.h: SYCL kernel submission and utility functions - TensorInfo.h: Tensor metadata and dimension handling structures - TensorOptions.h: Tensor configuration and options management - Runtime.h: SYCL runtime utilities - Macros.h: Common macro definitions - Scalar.h: Scalar type conversion utilities These headers provide the foundation for implementing 2D sparse data permutation and other complex SYCL operations on XPU devices. --- .../fbgemm_utils/comm/DeviceProperties.h | 220 ++++++++++++++++++ .../src/fbgemm_xpu/fbgemm_utils/comm/Macros.h | 17 ++ .../fbgemm_xpu/fbgemm_utils/comm/Runtime.h | 21 ++ .../fbgemm_utils/comm/SYCLContext.h | 19 ++ .../fbgemm_utils/comm/SYCLHelpers.h | 216 +++++++++++++++++ .../src/fbgemm_xpu/fbgemm_utils/comm/Scalar.h | 85 +++++++ .../fbgemm_xpu/fbgemm_utils/comm/TensorInfo.h | 220 ++++++++++++++++++ .../fbgemm_utils/comm/TensorOptions.h | 72 ++++++ 8 files changed, 870 insertions(+) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/DeviceProperties.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Macros.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Runtime.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLContext.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/SYCLHelpers.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/Scalar.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorInfo.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/comm/TensorOptions.h 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 From 3b438ad1a1d0b8b7d35f2d53b619c5f69ea094e5 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Sat, 20 Jun 2026 01:08:40 +0000 Subject: [PATCH 12/15] feat(fbgemm-xpu): add core utility infrastructure for SYCL kernels Add foundational utility headers and implementations to support complex SYCL kernel operations: - utils.h/cpp: Core constants, type definitions, kernel launch helpers, and device property queries - dispatch_macros.h: Type dispatch macros for handling multiple data types (int32, int64, float, etc.) - tensor_utils.h: Tensor manipulation and metadata utilities - function_types.h: Symbol visibility definitions for shared library exports These utilities provide essential infrastructure for implementing 2D sparse data permutation and other advanced operators on XPU devices, including work group sizing, kernel launch helpers, and type-safe dispatching mechanisms. --- .../fbgemm_xpu/fbgemm_utils/dispatch_macros.h | 262 ++++++++++++ .../fbgemm_xpu/fbgemm_utils/function_types.h | 17 + .../fbgemm_xpu/fbgemm_utils/tensor_utils.h | 192 +++++++++ .../src/fbgemm_xpu/fbgemm_utils/utils.cpp | 65 +++ .../src/fbgemm_xpu/fbgemm_utils/utils.h | 399 ++++++++++++++++++ 5 files changed, 935 insertions(+) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/dispatch_macros.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/function_types.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/tensor_utils.h create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.cpp create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/utils.h 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/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 From d32c585aec2719341710f332371f357a9882a961 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Sat, 20 Jun 2026 01:10:44 +0000 Subject: [PATCH 13/15] feat(fbgemm-xpu): add permute_2D_sparse_data XPU operator Add SYCL port of FBGEMM's permute_2D_sparse_data operator for Intel XPU devices. This operator permutes 2D sparse data including lengths [T, B], indices, and optional weights according to a permutation vector, commonly used for reordering embedding table features. Implementation includes: - SYCL kernels: permute_2D_lengths_kernel and permute_2D_data_kernel - Host function: permute_2D_sparse_data_xpu --- .../fbgemm_xpu/permute_2d_sparse_dataOp.cpp | 133 ++++++++++ .../sycl_kernels/permute_2d_sparse_data.cpp | 250 ++++++++++++++++++ .../sycl_kernels/permute_2d_sparse_data.h | 134 ++++++++++ 3 files changed, 517 insertions(+) create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/permute_2d_sparse_dataOp.cpp create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.cpp create mode 100644 packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/permute_2d_sparse_data.h 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/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 From a11d9a10693dfe27ee35d4a44772d4a807b8cf73 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Sat, 20 Jun 2026 01:12:07 +0000 Subject: [PATCH 14/15] feat(fbgemm-xpu): register permute_2D_sparse_data operator Integrate permute_2D_sparse_data operator into fbgemm-xpu: - Add Python wrapper with type hints and documentation - Register operator schema in torch library - Include implementation files in CMake build (utils.cpp, SYCL kernels, and operator implementation) --- packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt | 3 +++ packages/fbgemm-xpu/src/fbgemm_xpu/ops.py | 15 ++++++++++++++- .../fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp | 13 +++++++++++++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt index 0b3b400..b31f9cf 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt @@ -46,6 +46,9 @@ set(host_sources ${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/ops.py b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py index 2bc4ce1..a710817 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops.py @@ -14,6 +14,7 @@ "invert_permute", "permute_1D_sparse_data", "asynchronous_complete_cumsum", + "permute_2D_sparse_data", ] def invert_permute(permute: Tensor) -> Tensor: @@ -35,4 +36,16 @@ def permute_1D_sparse_data( 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) \ No newline at end of file + 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 acc2b60..e6755f4 100644 --- a/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp +++ b/packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp @@ -73,4 +73,17 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) "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?)" + ); + } + } From f9097cb68203b9cb408db1e633414e883fa058f5 Mon Sep 17 00:00:00 2001 From: Alberto Gallegos Date: Sat, 20 Jun 2026 01:13:54 +0000 Subject: [PATCH 15/15] test(fbgemm-xpu): add tests for permute_2D_sparse_data Add comprehensive test suite for permute_2D_sparse_data operator covering: - Basic functionality with int32 and int64 data types - Sparse data with and without weights - Permutations with repeated indices - Exact value validation - CPU-XPU consistency verification --- .../test/test_permute_2d_sparse_data.py | 190 ++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 packages/fbgemm-xpu/test/test_permute_2d_sparse_data.py 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()