Add reorder_batched_ad_lengths and reorder_batched_ad_indices XPU compatible operators - #87
Draft
aagalleg wants to merge 98 commits into
Draft
Add reorder_batched_ad_lengths and reorder_batched_ad_indices XPU compatible operators#87aagalleg wants to merge 98 commits into
aagalleg wants to merge 98 commits into
Conversation
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.
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>.
Replaced test for upstream patched FBGEMM tests that enables testing XPU. This file is no longer needed.
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
Delete the accidentally tracked submodule reference to FBGEMM-v1.7.0.
Add SYCL infrastructure headers from intel/torch-xpu-ops/ to support advanced kernel implementations: - DeviceProperties.h: Device capability queries and work group sizing - SYCLContext.h: SYCL context management and namespace aliases - SYCLHelpers.h: SYCL kernel submission and utility functions - TensorInfo.h: Tensor metadata and dimension handling structures - TensorOptions.h: Tensor configuration and options management - Runtime.h: SYCL runtime utilities - Macros.h: Common macro definitions - Scalar.h: Scalar type conversion utilities These headers provide the foundation for implementing 2D sparse data permutation and other complex SYCL operations on XPU devices.
Add foundational utility headers and implementations to support complex SYCL kernel operations: - utils.h/cpp: Core constants, type definitions, kernel launch helpers, and device property queries - dispatch_macros.h: Type dispatch macros for handling multiple data types (int32, int64, float, etc.) - tensor_utils.h: Tensor manipulation and metadata utilities - function_types.h: Symbol visibility definitions for shared library exports These utilities provide essential infrastructure for implementing 2D sparse data permutation and other advanced operators on XPU devices, including work group sizing, kernel launch helpers, and type-safe dispatching mechanisms.
Add SYCL port of FBGEMM's permute_2D_sparse_data operator for Intel XPU devices. This operator permutes 2D sparse data including lengths [T, B], indices, and optional weights according to a permutation vector, commonly used for reordering embedding table features. Implementation includes: - SYCL kernels: permute_2D_lengths_kernel and permute_2D_data_kernel - Host function: permute_2D_sparse_data_xpu
Integrate permute_2D_sparse_data operator into fbgemm-xpu: - Add Python wrapper with type hints and documentation - Register operator schema in torch library - Include implementation files in CMake build (utils.cpp, SYCL kernels, and operator implementation)
Add comprehensive test suite for permute_2D_sparse_data operator covering: - Basic functionality with int32 and int64 data types - Sparse data with and without weights - Permutations with repeated indices - Exact value validation - CPU-XPU consistency verification
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
The permute_2d_sparse_data_op.cpp file was incorrectly emptied. Restore the SYCL implementation.
Register operator schemas and add Python bindings for split and dense embedding lookup functions. Changes span ops_registry.cpp and ops.py.
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
Integrate permute_2D_sparse_data operator into fbgemm-xpu: - Add Python wrapper with type hints and documentation - Register operator schema in torch library - Include implementation files in CMake build (utils.cpp, SYCL kernels, and operator implementation)
Fixes CMake configuration and import ordering to properly build and load the block_bucketize_sparse_features XPU operator. - Configure CMake for XPU-only PyTorch builds - Import torch before _C extension to load libtorch.so dependencies - Adjust test imports for consistency All 18 tests passing.
Register operator schemas and add Python bindings for split and dense embedding lookup functions. Changes span ops_registry.cpp and ops.py.
Add src/codegen/CMakeLists.txt to drive code generation and build the _C_training Python extension module, and wire it into the top-level build.
NOTE: This files are going to be remove in the future and replaced by FBGEMM tests. Add four test modules covering forward and backward passes for both dense and split (rowwise Adagrad) embedding codegen operators on XPU: test_dense_embedding_codegen_forward.py: - Forward correctness, shape, dtype, and device validation - Kernel dispatch verification (small kernel D<=32 vs general kernel) - NaN/Inf checks and deterministic behavior - Comparison against PyTorch reference implementation test_dense_embedding_codegen_backward.py: - CtaPerRow kernel (long segments, SL >= 32) in isolation - WarpPerRow kernel (short segments, SL < 32) in isolation - Both kernels working together across segment boundaries - Numerical gradient correctness test_split_lookup_operator_forward_pass_xpu.py: - Split embedding nobag forward pass with rowwise Adagrad - DEVICE and MANAGED placement types - Multiple tables, index types, and output dtypes test_split_lookup_operator_backward_pass_xpu.py: - Split embedding nobag backward pass with rowwise Adagrad - Momentum state updates and learning rate application - Weight decay and stochastic rounding validation - Integration with fbgemm_gpu mock for standalone execution
…nsion Define FBGEMM_XPU_TRAINING_BUILD for _C_training and guard fbgemm::asynchronous_complete_cumsum XPU TORCH_LIBRARY_IMPL in asynchronous_complete_cumsum.cpp. This keeps asynchronous_complete_cumsum_xpu() available as a C++ helper for backward_utils while ensuring dispatcher registration is emitted only once from the C extension.
aagalleg
force-pushed
the
feat/reorder_batched_ad
branch
from
July 20, 2026 17:41
ecd2a4d to
bfe2040
Compare
- Add packages/fbgemm-xpu/test/patches/0002-Add-XPU-support-to-fbgemm-tbe- training-tests.patch - The patch is independent of 0001-Add-XPU-support-to-fbgemm-tests.patch - Removed previous tests for lookup operators.
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
Integrate permute_2D_sparse_data operator into fbgemm-xpu: - Add Python wrapper with type hints and documentation - Register operator schema in torch library - Include implementation files in CMake build (utils.cpp, SYCL kernels, and operator implementation)
Fixes CMake configuration and import ordering to properly build and load the block_bucketize_sparse_features XPU operator. - Configure CMake for XPU-only PyTorch builds - Import torch before _C extension to load libtorch.so dependencies - Adjust test imports for consistency All 18 tests passing.
Register operator schemas and add Python bindings for split and dense embedding lookup functions. Changes span ops_registry.cpp and ops.py.
Add src/codegen/CMakeLists.txt to drive code generation and build the _C_training Python extension module, and wire it into the top-level build.
Implement reorder_batched_ad_lengths and reorder_batched_ad_indices operators with SYCL kernels for Intel XPU devices.
Add build system integration and Python API for reorder_batched_ad operators introduced in the previous commit.
Add comprehensive test suite for reorder_batched_ad_lengths and reorder_batched_ad_indices operators, covering both broadcast and non-broadcast modes.
aagalleg
force-pushed
the
feat/reorder_batched_ad
branch
from
July 21, 2026 18:39
bfe2040 to
0cc129f
Compare
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
…ng standards Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
Contributor
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR implements
reorder_batched_ad_lengthsandreorder_batched_ad_indicesoperators for Intel XPU devices, porting FBGEMM's batched advertisement reordering functionality from CUDA to SYCL. These operators transform sparse embedding data from ragged[B x T x #num_ads_b]layout to table-first[T][B][#num_ads_b]layout for efficient embedding lookups in advertising and recommendation systems.Depends on #85
Changes
Operator Implementation
reorder_batched_adOp.cpp: Top-level operator implementations for XPU.SYCL Kernel Implementation
sycl_kernels/reorder_batched_ad.h: Public API header with comprehensive documentation.sycl_kernels/reorder_batched_ad.cpp: SYCL kernel implementations.Build System Integration
CMakeLists.txt: Addsreorder_batched_ad.cppandreorder_batched_adOp.cppto thehost_sourceslist, ensuring SYCL compilation with theicpxcompiler and AOT targetsTesting
test_reorder_batched.py: Comprehensive test suite with 4 test cases:test_reorder_ad_lengths_no_broadcast: Validates identity reordering when layout already matches (B=3, T=2, A=2, L=5). Verifies output matches input for non-broadcast casetest_reorder_ad_lengths_broadcast: Tests broadcast mode where[B x T]lengths are expanded to[B x T x A]via tiling. Compares against reference implementation usingtorch.tiletest_reorder_ad_indices: End-to-end test for indices reordering withasynchronous_complete_cumsumto compute offsets. Uses random indices and validates reordering via tensor reshaping and permutation (view(T, B, A, L).permute(1, 0, 2, 3))test_reorder_ad_indices_broadcast: Tests broadcast mode where first batch's indices are replicated across all batches. Validates output against tiled reference with shape[B, T, A, L]from input[B, T, 1, L]cc: @flezaalv