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

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

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

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

Add comprehensive test suite for asynchronous_complete_cumsum operator covering: - Basic functionality with int32 and int64 dtypes - Empty tensor handling - Random input validation with numpy reference

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

aagalleg added 10 commits June 18, 2026 22:37

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

081cc88

Remove .gitkeep placeholders

feat(fbgemm-xpu): register invert_permute operation

07bfb4b

- 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

d6d6bad

Add SYCL/XPU kernel implementation for invert_permute operation.

feat(fbgemm-xpu): register permute_1D_sparse_data operator

3056382

- 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

b7f3bbf

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

fd70ab5

feat(fbgemm-xpu): add asynchronous_complete_cumsum XPU operator

50f9982

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

79272c9

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): add tests for asynchronous_complete_cumsum

389845d

Add comprehensive test suite for asynchronous_complete_cumsum operator covering: - Basic functionality with int32 and int64 dtypes - Empty tensor handling - Random input validation with numpy reference

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

Open

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

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 10 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!

2 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!

2 participants

aagalleg commented Jun 19, 2026 •

edited

Loading