Add permute_2D_sparse_data operator with SYCL implementation for XPU

packages/fbgemm-xpu/src/fbgemm_xpu/CMakeLists.txt

@@ -43,6 +43,12 @@ set(SYCL_DEVICE_LIST -Xs "-device ${TORCH_XPU_ARCH_LIST} -options -cl-poison-uns
 # --------------------------------------------------------------------------
 set(host_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/ops_registry.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/sycl_kernels/invert_permute_kernel.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/sycl_kernels/permute_1d_sparse_data.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/asynchronous_complete_cumsum.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/fbgemm_utils/utils.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/sycl_kernels/permute_2d_sparse_data.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/permute_2d_sparse_dataOp.cpp
 )
 
 set(all_sources ${host_sources})

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

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 <ATen/ATen.h>
+#include <torch/library.h>
+
+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

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 <ATen/xpu/XPUContext.h>
+
+#include <comm/Runtime.h>
+#include <sycl/sycl.hpp>
+
+namespace syclext = sycl::ext::oneapi;
+namespace syclexp = sycl::ext::oneapi::experimental;
+
+namespace xpu {
+namespace sycl {
+
+template <class KernelClass>
+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<KernelClass>();
+  // 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 <class KernelClass>
+static int64_t syclMaxWorkGroupSize(
+    const KernelClass& /*kfn*/,
+    at::DeviceIndex dev_id = at::xpu::current_device()) {
+  return syclMaxWorkGroupSize<KernelClass>(dev_id);
+}
+
+// For SYCL free function
+template <auto* kptr>
+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<kptr, ::sycl::bundle_state::executable>(
+          ctxt);
+  ::sycl::kernel k = exe_bndl.template ext_oneapi_get_kernel<kptr>();
+  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 <typename T>
+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<T, char> || std::is_same_v<T, short> ||
+      std::is_same_v<T, int> || std::is_same_v<T, int64_t> ||
+      std::is_same_v<T, float> || std::is_same_v<T, double> ||
+      std::is_same_v<T, ::sycl::half>) {
+    return vec_width / sizeof(T);
+  } else {
+    throw std::invalid_argument(
+        "Invalid data type to fetch preferred vector width!");
+  }
+}
+
+template <typename T>
+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<T, char>) {
+    return dev_prop->native_vector_width_char;
+  } else if constexpr (std::is_same_v<T, short>) {
+    return dev_prop->native_vector_width_short;
+  } else if constexpr (std::is_same_v<T, int>) {
+    return dev_prop->native_vector_width_int;
+  } else if constexpr (std::is_same_v<T, int64_t>) {
+    return dev_prop->native_vector_width_long;
+  } else if constexpr (std::is_same_v<T, float>) {
+    return dev_prop->native_vector_width_float;
+  } else if constexpr (std::is_same_v<T, double>) {
+    return dev_prop->native_vector_width_double;
+  } else if constexpr (std::is_same_v<T, ::sycl::half>) {
+    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

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

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 <c10/xpu/XPUStream.h>
+
+namespace at::xpu {
+
+static inline sycl::queue& getCurrentSYCLQueue() {
+  return c10::xpu::getCurrentXPUStream().queue();
+}
+
+} // namespace at::xpu

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 <comm/DeviceProperties.h>
+#include <comm/Macros.h>
+#include <comm/Runtime.h>
+#include <comm/SYCLHelpers.h>
+
+using namespace at::xpu;
+using namespace xpu::sycl;