Add invert_permute operator with SYCL implementation for XPU

.github/workflows/ci.yml

@@ -393,13 +393,11 @@ jobs:
     - name: 'Patch FBGEMM'
       working-directory: fbgemm
       run: |
-        # TODO: Apply XPU patches here
-        echo "Patches will be applied here"
+        git am ../torchlib-xpu/packages/fbgemm-xpu/patches/0001-Add-XPU-support-to-fbgemm-tests.patch
     - name: 'Test fbgemm-xpu'
       working-directory: fbgemm
       run: |
-        # TODO: Run pytest commands to run the XPU tests here
-        echo "Instructions to run the tests will be added here"
+        pytest -rsf fbgemm_gpu/test/sparse/misc_ops_test.py
 
   publish:
     needs: [ compose, test-torchcodec, test-fbgemm ]

packages/fbgemm-xpu/patches/0001-Add-XPU-support-to-fbgemm-tests.patch

@@ -0,0 +1,92 @@
+From 1ff5ad5b0c0a780081a8d6f90748c9c9df08e26d Mon Sep 17 00:00:00 2001
+From: Alberto Gallegos <alberto.gallegos.muro@intel.com>
+Date: Thu, 9 Jul 2026 18:28:28 +0000
+Subject: [PATCH] Add XPU support to misc ops tests
+
+Signed-off-by: Alberto Gallegos <alberto.gallegos.muro@intel.com>
+---
+ fbgemm_gpu/test/sparse/misc_ops_test.py | 18 +++++++++++++++---
+ 1 file changed, 15 insertions(+), 3 deletions(-)
+
+diff --git a/fbgemm_gpu/test/sparse/misc_ops_test.py b/fbgemm_gpu/test/sparse/misc_ops_test.py
+index edb0d3615..f825ce27c 100644
+--- a/fbgemm_gpu/test/sparse/misc_ops_test.py
++++ b/fbgemm_gpu/test/sparse/misc_ops_test.py
+@@ -33,6 +33,11 @@ else:
+         gpu_unavailable,
+     )
+
++# Registers the XPU (Intel GPU) FBGEMM operator implementations.
++import fbgemm_xpu  # noqa: F401, E402
++
++xpu_available: bool = torch.xpu.is_available()
++
+
+ class MiscOpsTest(unittest.TestCase):
+     @given(
+@@ -57,10 +62,11 @@ class MiscOpsTest(unittest.TestCase):
+         inverse_permute_cpu = torch.ops.fbgemm.invert_permute(permute)
+         torch.testing.assert_close(inverse_permute_cpu, inverse_permute_ref)
+
+-        if gpu_available:
+-            inverse_permute_gpu = torch.ops.fbgemm.invert_permute(permute.cuda())
+-            torch.testing.assert_close(inverse_permute_gpu.cpu(), inverse_permute_cpu)
++        if xpu_available:
++            inverse_permute_xpu = torch.ops.fbgemm.invert_permute(permute.xpu())
++            torch.testing.assert_close(inverse_permute_xpu.cpu(), inverse_permute_cpu)
+
++    @unittest.skip("Operator not implemented on XPU")
+     @given(
+         N=st.integers(min_value=1, max_value=20),
+         offsets_type=st.sampled_from([torch.int32, torch.int64]),
+@@ -90,6 +96,7 @@ class MiscOpsTest(unittest.TestCase):
+             range_ref = range_ref.to(range_gpu.dtype)
+             torch.testing.assert_close(range_gpu.cpu(), range_ref, rtol=0, atol=0)
+
++    @unittest.skip("Operator not implemented on XPU")
+     @given(
+         index_type=st.sampled_from([torch.int, torch.long]),
+         has_weight=st.booleans(),
+@@ -160,6 +167,7 @@ class MiscOpsTest(unittest.TestCase):
+             if bucketize_pos:
+                 torch.testing.assert_close(new_pos_gpu.cpu(), new_pos_cpu)
+
++    @unittest.skip("Operator not implemented on XPU")
+     def test_segment_sum_csr(self) -> None:
+         segment_sum_cpu = torch.ops.fbgemm.segment_sum_csr(
+             2,
+@@ -181,6 +189,7 @@ class MiscOpsTest(unittest.TestCase):
+                 segment_sum_cuda.cpu(), torch.Tensor([10.0, 11.0, 34.0]), rtol=0, atol=0
+             )
+
++    @unittest.skip("Operator not implemented on XPU")
+     def test_segment_sum_csr_empty_input(self) -> None:
+         segment_sum_cpu = torch.ops.fbgemm.segment_sum_csr(
+             0,
+@@ -254,6 +263,7 @@ class MiscOpsTest(unittest.TestCase):
+
+         torch.testing.assert_close(output_gpu.cpu(), output_cpu)
+
++    @unittest.skip("Operator not implemented on XPU")
+     @given(
+         batch_size=st.just(2),
+         m=st.just(3),
+@@ -290,6 +300,7 @@ class MiscOpsTest(unittest.TestCase):
+         C = torch.ops.fbgemm.permute102_baddbmm_permute102(bias, A, B)
+         torch.testing.assert_close(C.cpu(), C_ref.cpu())
+
++    @unittest.skip("Operator not implemented on XPU")
+     @given(
+         T=st.integers(1, 5),
+         B=st.integers(1, 5),
+@@ -321,6 +332,7 @@ class MiscOpsTest(unittest.TestCase):
+         all_indices_deduped_ref = torch.as_tensor(all_indices[:, :, :L])
+         torch.testing.assert_close(all_indices_deduped, all_indices_deduped_ref)
+
++    @unittest.skip("Operator not implemented on XPU")
+     def test_lengths_range(self) -> None:
+         # When 'output_shape' is None, the function will return a tensor with dynamic shape.
+         with self.assertRaisesRegex(
+-- 
+2.43.0
+

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})

packages/fbgemm-xpu/src/fbgemm_xpu/__init__.py

@@ -2,6 +2,10 @@
 # Copyright (c) 2026 Intel Corporation. All Rights Reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+# Import fbgemm_gpu first so that all "fbgemm" operator schemas are registered
+# before _C loads the XPU implementations via TORCH_LIBRARY_IMPL.
+import fbgemm_gpu  # noqa: F401, E402
+
 # Import the compiled C extension (_C) which contains the registered operators.
 # If native dependencies (for example libtorch.so) are unavailable, keep import
 # working so metadata like __version__ remains accessible.
@@ -10,9 +14,7 @@
 except ImportError:
     _C = None
 
-from . import ops as ops
-
-__all__ = ["_C", "ops", "__version__"]
+__all__ = ["_C", "__version__"]
 
 try:
     from ._version import __version__

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 <ATen/core/dispatch/Dispatcher.h>
+
+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

packages/fbgemm-xpu/src/fbgemm_xpu/ops_registry.cpp

@@ -9,6 +9,9 @@
 #include <ATen/core/Tensor.h>
 #include <torch/library.h>
 
+#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"
+        );
+    }
 }

packages/fbgemm-xpu/src/fbgemm_xpu/sycl_kernels/invert_permute_kernel.cpp

@@ -0,0 +1,99 @@
+/*
+ * 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
+// ============================================================================
+
+template <typename index_t>
+void InvertPermuteKernel<index_t>::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 index_t target_idx = permute_[i];
+        inversed_permute_[target_idx] = static_cast<index_t>(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;
+
+    // Dispatch over the index type (int32_t or int64_t), mirroring the
+    // reference CUDA kernel invert_permute_kernel<index_t>.
+    AT_DISPATCH_INDEX_TYPES(
+        permute_contig.scalar_type(), "invert_permute_kernel", [&] {
+            const index_t* permute_ptr = permute_contig.data_ptr<index_t>();
+            index_t* inversed_ptr = inversed_permute.data_ptr<index_t>();
+
+            queue.submit([&](sycl::handler& cgh) {
+                cgh.parallel_for<InvertPermuteKernel<index_t>>(
+                    sycl::nd_range<1>(
+                        sycl::range<1>(global_size),  // Global range
+                        sycl::range<1>(local_size)    // Local range
+                    ),
+                    InvertPermuteKernel<index_t>(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