Add permute_1D_sparse_data Operation for Intel XPU by aagalleg · Pull Request #72 · intel/torchlib-xpu

aagalleg · 2026-06-18T23:36:09Z

Depends on #71

Add `permute_1D_sparse_data` Operation for Intel XPU

This PR introduces the permute_1D_sparse_data operator to fbgemm-xpu, enabling permutation of sparse data in jagged/1D format on Intel XPU devices.

Changes

Core Implementation

SYCL Kernels (permute_1d_sparse_data.cpp/h): SYCL implementation with three specialized functors for permuting lengths, indices, and weighted data, plus optimized cumsum helpers for offset computation
Operator Registration (ops_registry.cpp): Registers the operator with PyTorch's dispatch system using conditional schema registration to avoid conflicts
Python API (ops.py): Clean Python wrapper function with type hints supporting optional weights and length sum parameters

Infrastructure

Build System (CMakeLists.txt): Added SYCL kernel to build configuration

Testing

Test Suite (test_permute_1d_sparse_data.py):
- Correctness validation with and without weights
- Various segment configurations (uniform, mixed lengths, empty segments)
- Edge cases (empty tensors, single segments, identity permutations)
- Random permutations with validation
- CPU-XPU consistency verification

cc: @manuelhsantana, @flezaalv

flezaalv · 2026-07-16T22:22:46Z

Depends on #71

Add permute_1D_sparse_data Operation for Intel XPU

This PR introduces the permute_1D_sparse_data operator to fbgemm-xpu, enabling permutation of sparse data in jagged/1D format on Intel XPU devices.

Changes

Core Implementation

SYCL Kernels (permute_1d_sparse_data.cpp/h): SYCL implementation with three specialized functors for permuting lengths, indices, and weighted data, plus optimized cumsum helpers for offset computation

Operator Registration (ops_registry.cpp): Registers the operator with PyTorch's dispatch system using conditional schema registration to avoid conflicts

Python API (ops.py): Clean Python wrapper function with type hints supporting optional weights and length sum parameters

Infrastructure

Build System (CMakeLists.txt): Added SYCL kernel to build configuration

Testing

Test Suite (test_permute_1d_sparse_data.py):

Correctness validation with and without weights

Various segment configurations (uniform, mixed lengths, empty segments)

Edge cases (empty tensors, single segments, identity permutations)

Random permutations with validation

CPU-XPU consistency verification

cc: @manuelhsantana, @flezaalv

The test suite has been replaced with the patch 0001-Add-XPU-support-to-fbgemm-tests.patch for running upstream FBGEMM tests for permute_2d_sparse_data.
Removed ops.py file (no longer necessary).

Remove .gitkeep placeholders

- 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

Add SYCL/XPU kernel implementation for invert_permute operation.

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

Replace the custom standalone test_invert_permute.py with a git-am patch applied to upstream FBGEMM v1.7.0 misc_ops_test.py, following the torchcodec-xpu convention. The patch makes test_invert_permute run on XPU (permute.xpu(), gated on torch.xpu.is_available()) and skips the remaining operator tests that are not implemented on XPU.

Collapse the two type-specific kernel functors (InvertPermuteKernelInt32, InvertPermuteKernelInt64) into a single templated functor InvertPermuteKernel<index_t>, mirroring the reference CUDA kernel invert_permute_kernel<index_t>.

… order

…cture - Fix test patches to match FBGEMM v1.8.0 tests. - Move test patches from test/patches/ subdirectory to patches/ at the package root level for better organization. Remove now-unnecessary .gitkeep file and update patch with correct base commit reference.

Replace TODO placeholders with actual commands to apply the FBGEMM XPU test patch and execute pytest for misc_ops_test.py in the CI pipeline.

- 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

Implement SYCL/XPU kernel implementation of permute_1D_sparse_data operator for sparse jagged/1D format data permutation.

Replaced test for upstream patched FBGEMM tests that enables testing XPU. This file is no longer needed.

Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>

dvrogozh · 2026-07-29T23:50:04Z

+             low=1,
+             high=L,
+             size=(T,),  # 1D tensor with T elements
+-            device=torch.accelerator.current_accelerator(),


This should not be needed as torch.accelerator.current_accelerator() should return xpu as in trivial:

$ python -c "import torch; print(torch.xpu.is_available()); print(torch.accelerator.current_accelerator())" True xpu

That line will be deleted by the patch and the device="xpu" will be added, the result code will be:

lengths = torch.randint( low=1, high=L, size=(T,), # 1D tensor with T elements device="xpu", ).type(index_dtype)

I understand what this patch is currently doing and I am arguing that this should not be needed as upstream code should already do what you are trying to achieve - torch.accelerator.current_accelerator() should return xpu.

Reverted, using the original torch.accelerator.current_accelerator() instead.

Did you verify that it works?

Yes, I did. The pytest command is already integrated into the Test fbgemm-xpu step: https://github.com/intel/torchlib-xpu/actions/runs/30579440544/job/90997602510#step:13:3

…st patch Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>

dvrogozh · 2026-07-31T03:32:59Z

+-    @unittest.skipIf(*running_in_oss)
+-    @unittest.skipIf(*gpu_unavailable)
+    @unittest.skipIf(
+        gpu_unavailable[0] and not torch.xpu.is_available(),


Ok for now, but I suspect that going forward we will need to mimic gpu_unavailable.

dvrogozh · 2026-07-31T03:43:47Z

+/**
+ * @brief Permute1DLengthsKernel operator implementation
+ */
+template <typename index_t>


Why we don't parameterize template for permute as in a reference: https://github.com/pytorch/FBGEMM/blob/3bc4f702d9e60c21d6b97150946e2e339625b73a/fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu#L16?

Done at commit ce10815

dvrogozh · 2026-07-31T03:46:53Z

+ * output[n] = sum(all input)
+ */
+template <typename T>
+at::Tensor complete_cumsum_xpu(const at::Tensor& input) {


Same as above. This seems to be defined elsewhere in the cuda reference. Why do we define it here?

@manuelhsantana why does exclusive_cumsum_xpu and complete_cumsum_xpu have the same code?
Also, the original code of permute_1D_sparse_data uses asynchronous_complete_cumsum that is an operator that we did implement to enable TorchRec, why wasn't that implementation used?

complete_cumsum_xpu can be changed in PR #73 for asynchronouse_complete_cumsum if they provide the same results. But there is still the mystery of why exclusive_cumsum_xpu and complete_cumsum_xpu have the same code.

Introduce a permute_t template parameter on Permute1DLengthsKernel (defaulting to int32_t).

dvrogozh · 2026-08-03T05:16:33Z

@aagalleg , this PR requires manual rebase for whatever reason (I am not sure what's wrong). Please, try to rebase. Once done I will attempt to merge.

aagalleg requested review from dvrogozh, eromomon and luis-real as code owners June 18, 2026 23:36

aagalleg mentioned this pull request Jun 19, 2026

Add asynchronous_complete_cumsum operator for XPU #73

Draft

dvrogozh marked this pull request as draft June 22, 2026 18:36

flezaalv added a commit to aagalleg/torchlib-xpu that referenced this pull request Jul 13, 2026

fix(fbgemm-xpu): address PR intel#72 review comments

b71efbe

flezaalv added a commit to aagalleg/torchlib-xpu that referenced this pull request Jul 28, 2026

fix(fbgemm-xpu): address PR intel#72 review comments

152e470

flezaalv force-pushed the feat/permute_1d_sparse_data branch from 33197a5 to ea4b40f Compare July 28, 2026 19:01

aagalleg added 10 commits July 29, 2026 16:24

chore(fbgemm-xpu): remove .gitkeep files from populated directories

f7eea23

Remove .gitkeep placeholders

feat(fbgemm-xpu): register invert_permute operation

c050b11

- 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

feat(fbgemm-xpu): add SYCL kernel implementation for invert_permute

cce1fee

Add SYCL/XPU kernel implementation for invert_permute operation.

style(fbgemm-xpu): add missing newline at end of torch_library.h

d280078

fix(fbgemm-xpu): remove redundant ops.py and ensure fbgemm_gpu import…

da64394

… order

ci(fbgemm-xpu): apply XPU patch and run tests in workflow

bacbaec

Replace TODO placeholders with actual commands to apply the FBGEMM XPU test patch and execute pytest for misc_ops_test.py in the CI pipeline.

flezaalv added a commit to aagalleg/torchlib-xpu that referenced this pull request Jul 29, 2026

fix(fbgemm-xpu): address PR intel#72 review comments

be7e845

flezaalv force-pushed the feat/permute_1d_sparse_data branch from e15792b to a31b216 Compare July 29, 2026 17:36

aagalleg and others added 7 commits July 29, 2026 17:42

feat(fbgemm-xpu): register permute_1D_sparse_data operator

0ba5f24

- 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

feat(fbgemm-xpu): add SYCL kernels for permute_1D_sparse_data

831a650

Implement SYCL/XPU kernel implementation of permute_1D_sparse_data operator for sparse jagged/1D format data permutation.

test(fbgemm-xpu): add comprehensive tests for permute_1D_sparse_data

a857133

fix(fbgemm-xpu): address PR intel#72 review comments

ee6c379

Consolidate FBGEMM XPU test patches into single unified patch

d0c278b

Fix permute_1D_sparse_data to support 2D weights tensor

8c82a1a

chore(fbgemm-xpu): drop standalone test_permute_1d_sparse_data.py

936e378

Replaced test for upstream patched FBGEMM tests that enables testing XPU. This file is no longer needed.

flezaalv force-pushed the feat/permute_1d_sparse_data branch from a31b216 to 936e378 Compare July 29, 2026 17:43

ci(fbgemm-xpu): run permute_sparse_features tests in CI

f90f4b1

Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>

Merge branch 'main' into feat/permute_1d_sparse_data

4e249eb

aagalleg marked this pull request as ready for review July 29, 2026 21:55

dvrogozh requested changes Jul 29, 2026

View reviewed changes

fix(fbgemm-xpu): extend GPU skip gate instead of hard XPU check in te…

c27e3df

…st patch Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>

dvrogozh reviewed Jul 31, 2026

View reviewed changes

fix(fbgemm-xpu): templatize permute_t in Permute1DLengthsKernel

ce10815

Introduce a permute_t template parameter on Permute1DLengthsKernel (defaulting to int32_t).

dvrogozh approved these changes Aug 1, 2026

View reviewed changes

Uh oh!

Conversation

aagalleg commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Add permute_1D_sparse_data Operation for Intel XPU

Changes

Core Implementation

Infrastructure

Testing

Uh oh!

flezaalv commented Jul 16, 2026

Add permute_1D_sparse_data Operation for Intel XPU

Changes

Core Implementation

Infrastructure

Testing

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

flezaalv Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dvrogozh commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aagalleg commented Jun 18, 2026 •

edited

Loading

Add `permute_1D_sparse_data` Operation for Intel XPU

Add `permute_1D_sparse_data` Operation for Intel XPU

flezaalv Jul 31, 2026 •

edited

Loading

dvrogozh commented Aug 3, 2026 •

edited

Loading