Skip to content

Add permute_1D_sparse_data Operation for Intel XPU - #72

Open
aagalleg wants to merge 19 commits into
intel:mainfrom
aagalleg:feat/permute_1d_sparse_data
Open

Add permute_1D_sparse_data Operation for Intel XPU#72
aagalleg wants to merge 19 commits into
intel:mainfrom
aagalleg:feat/permute_1d_sparse_data

Conversation

@aagalleg

@aagalleg aagalleg commented Jun 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

@aagalleg aagalleg mentioned this pull request Jun 19, 2026
Draft
@dvrogozh
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
@flezaalv
fix(fbgemm-xpu): address PR intel#72 review comments
b71efbe
@flezaalv

flezaalv commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

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

flezaalv added a commit to aagalleg/torchlib-xpu that referenced this pull request Jul 28, 2026
@flezaalv
fix(fbgemm-xpu): address PR intel#72 review comments
152e470
@flezaalv
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
@aagalleg
chore(fbgemm-xpu): remove .gitkeep files from populated directories
f7eea23 
Remove .gitkeep placeholders
@aagalleg
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
@aagalleg
feat(fbgemm-xpu): add SYCL kernel implementation for invert_permute
cce1fee 
Add SYCL/XPU kernel implementation for invert_permute operation.
@aagalleg
test(fbgemm-xpu): add comprehensive test suite for invert_permute
4b0e034 
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
@aagalleg
test(fbgemm-xpu): use patched upstream FBGEMM tests for invert_permute
eb23224 
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.
@aagalleg
style(fbgemm-xpu): add missing newline at end of torch_library.h
d280078
@aagalleg
refactor(fbgemm-xpu): templatize invert_permute SYCL kernel over index_t
3e582b8 
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>.
@aagalleg
fix(fbgemm-xpu): remove redundant ops.py and ensure fbgemm_gpu import…
da64394 
… order
@aagalleg
fix(fbgemm-xpu): fix test patch and reorganize patches directory stru…
f70e5d5 
…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.
@aagalleg
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
@flezaalv
fix(fbgemm-xpu): address PR intel#72 review comments
be7e845
@flezaalv
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
@aagalleg @flezaalv
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
@aagalleg @flezaalv
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.
@aagalleg @flezaalv
test(fbgemm-xpu): add comprehensive tests for permute_1D_sparse_data
a857133
@flezaalv
fix(fbgemm-xpu): address PR intel#72 review comments
ee6c379
@flezaalv
Consolidate FBGEMM XPU test patches into single unified patch
d0c278b
@flezaalv
Fix permute_1D_sparse_data to support 2D weights tensor
8c82a1a
@aagalleg @flezaalv
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
flezaalv force-pushed the feat/permute_1d_sparse_data branch from a31b216 to 936e378 Compare July 29, 2026 17:43
@flezaalv
ci(fbgemm-xpu): run permute_sparse_features tests in CI
f90f4b1 
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
@aagalleg
aagalleg marked this pull request as ready for review July 29, 2026 21:55
dvrogozh
dvrogozh requested changes Jul 29, 2026
View reviewed changes
Comment thread packages/fbgemm-xpu/patches/0001-Add-XPU-support-to-fbgemm-tests.patch
low=1,
high=L,
size=(T,), # 1D tensor with T elements
- device=torch.accelerator.current_accelerator(),

@dvrogozh dvrogozh Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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


class Permute1DSparseFeaturesTest(unittest.TestCase):
- @unittest.skipIf(*running_in_oss)

@dvrogozh dvrogozh Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is running_in_oss and why that's a problem for us?

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

class Permute1DSparseFeaturesTest(unittest.TestCase):
- @unittest.skipIf(*running_in_oss)
- @unittest.skipIf(*gpu_unavailable)

@dvrogozh dvrogozh Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's better to extend this check rather than hard XPU one, something like:

    @unittest.skipIf(*gpu_unavailable and *xpu_unavailable)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@dvrogozh dvrogozh dvrogozh requested changes
@eromomon eromomon Awaiting requested review from eromomon eromomon is a code owner
@luis-real luis-real Awaiting requested review from luis-real luis-real is a code owner

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@aagalleg @flezaalv @dvrogozh