Add asynchronous_complete_cumsum operator for XPU by aagalleg · Pull Request #73 · intel/torchlib-xpu

aagalleg · 2026-06-19T23:50:16Z

Depends on: #72

This PR introduces the asynchronous_complete_cumsum operator to fbgemm-xpu, enabling complete cumulative sum computation with a leading zero on Intel XPU devices.

Changes

Core Implementation

XPU Implementation (asynchronous_complete_cumsum.cpp/h): SYCL port of FBGEMM's asynchronous complete cumsum operator that computes cumulative sum with prepended zero (e.g., [a, b, c] → [0, a, a+b, a+b+c])
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 and documentation for easy integration

Infrastructure

Build System (CMakeLists.txt): Added implementation to build configuration
Tensor Support: Handles both 1D and 2D tensors (cumsum along last dimension for 2D)
Data Types: Supports int32 and int64 tensors

Testing

Comprehensive Test Suite (test_asynchronous_complete_cumsum.py):
- Correctness validation for int32/int64 data types
- Edge cases (empty tensors)
- Random input validation with numpy reference
- XPU device testing with proper error handling

cc: @flezaalv, @manuelhsantana

dvrogozh

In addition to inline comments:

Repository has lint CI - please, address identified issues.
What we implement and merge must be tested under CI. And IMHO doing that by engineers who implemented the code is the best variant rather than offloading that to dedicated CI engineers. Thus, please, modify .github/workflows/ci.yml and add a test section for FBGEMM following the pattern outlined for torchcodec (see link below). The test job for FBGEMM will be much simpler. Target only bmg for now is fine. Note that fbgemm wheel is already getting build by CI.

torchlib-xpu/.github/workflows/ci.yml

Line 199 in d4230aa

test-torchcodec:

dvrogozh · 2026-06-22T18:06:11Z

+      .has_value();
+}
+
+} // namespace fbgemm_xpu::utils::torch


add new line in end of file

dvrogozh · 2026-06-22T18:10:21Z

+    }
+}
+
+void InvertPermuteKernelInt64::operator()(const sycl::nd_item<1>& item) const {


This seems templetable kernel over the operand type. Why don't we do that especially considering that the reference CUDA kernel is doing exactly that [1]?

[1] https://github.com/pytorch/FBGEMM/blob/8744183e4bd7390a674a0d8736bf4d7ef2d0516e/fbgemm_gpu/src/sparse_ops/sparse_invert_permute.cu#L15

dvrogozh · 2026-06-22T18:13:39Z

+ * @brief Permute1DDataWithWeightsKernel operator implementation
+ */
+template <typename offsets_t, typename indices_t, typename weights_t>
+void Permute1DDataWithWeightsKernel<offsets_t, indices_t, weights_t>::operator()(


Why not to follow CUDA reference [1] and additionally parameterize with has_weight?

[1] https://github.com/pytorch/FBGEMM/blob/8744183e4bd7390a674a0d8736bf4d7ef2d0516e/fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu#L30

dvrogozh · 2026-06-22T18:19:36Z

+
 __all__ = [
-	"dense_embedding_codegen_lookup_function",
+    "invert_permute",


Here I am confused. Why do we expose these XPU specific APIs from fbgemm xpu? who is supposed to use them? My understanding was that we are just register low level entrypoints for XPU key so that standard fbgemm APIs start to work on XPU device. Shouldn't that be that user calls fbgemm::invert_permute for example?

dvrogozh · 2026-06-22T18:19:50Z

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


add empty line in end of file

dvrogozh · 2026-06-22T18:20:46Z

@@ -0,0 +1,71 @@
+"""


Should also have a copyright.

dvrogozh · 2026-06-22T18:24:37Z

+For 1D input [a, b, c], returns [0, a, a+b, a+b+c].
+For 2D input, cumsum along dimension 1 (columns).
+"""
+


Tests are eventually good, but our side tests won't guarantee that what we've implemented correctly integrates with the upstream FBGEMM. We must explore if upstream tests cover what we are implementing and use them if they are. If not, then we need to raise PRs with upstream to add test coverage. That's critical concern - we need to resolve that before proceeding.

dvrogozh · 2026-06-22T18:32:11Z

+        torch.manual_seed(SEED)
+
+    def test_basic_int32(self):
+        if not torch.xpu.is_available():


This check is inconsistent with checks in other test file: @unittest.skipIf(not torch.xpu.is_available(), "XPU not available").

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>

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

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

…complete_cumsum

Delete the accidentally tracked submodule reference to FBGEMM-v1.7.0.

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

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

aagalleg mentioned this pull request Jun 20, 2026

Add permute_2D_sparse_data operator with SYCL implementation for XPU #74

Draft

dvrogozh requested changes Jun 22, 2026

View reviewed changes

dvrogozh mentioned this pull request Jun 22, 2026

Add invert_permute operator with SYCL implementation for XPU #71

Merged

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

aagalleg force-pushed the feat/asynchronous_complete_cumsum branch from 389845d to dcdd4df Compare July 15, 2026 17:03

aagalleg and others added 22 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.

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.

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

feat(fbgemm-xpu): add asynchronous_complete_cumsum XPU operator

f276807

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

feat(fbgemm-xpu): register asynchronous_complete_cumsum operator

aaa1fe9

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

test(fbgemm-xpu): use patched upstream FBGEMM tests for asynchronous_…

87c0c80

…complete_cumsum

chore(repo): remove stray FBGEMM-v1.7.0 gitlink

bc7c7a9

Delete the accidentally tracked submodule reference to FBGEMM-v1.7.0.

flezaalv force-pushed the feat/asynchronous_complete_cumsum branch from c458218 to bc7c7a9 Compare July 29, 2026 22:39

flezaalv added 2 commits July 29, 2026 23:37

fix(fbgemm-xpu): update cumsum_test hunk for FBGEMM v1.8.0

5900fcf

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

ci(fbgemm-xpu): run cumsum tests

b4ab572

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add asynchronous_complete_cumsum operator for XPU - #73

Add asynchronous_complete_cumsum operator for XPU#73
aagalleg wants to merge 25 commits into
intel:mainfrom
aagalleg:feat/asynchronous_complete_cumsum

aagalleg commented Jun 19, 2026 •

edited

Loading

Uh oh!

dvrogozh left a comment

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

dvrogozh Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

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

Uh oh!

Changes

Core Implementation

Infrastructure

Testing

Uh oh!

dvrogozh left a comment

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aagalleg commented Jun 19, 2026 •

edited

Loading