Add CMake build system and tests for the lookup operators#85
Draft
aagalleg wants to merge 41 commits into
Draft
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.
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
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
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.
… generation Add jinja_environment.py module to support template-based code generation for FBGEMM-XPU kernels.
Add common.py module with CodeTemplate class that provides functionality for loading Jinja2 templates, rendering them with context variables, and writing generated files with appropriate headers.
Add torch_type_utils.py module with utilities for handling PyTorch data types in template-based code generation.
Add generate_forward_split.py and generate_backward_split.py scripts for generating SYCL embedding kernels from Jinja2 templates.
Add backward_utils.cpp and backward_utils.h with SYCL ports of FBGEMM backward pass utilities for embedding operations.
…mization Add vec4.h header implementing 4-element vectorized data structures for efficient memory access in embedding operations.
Add feature_gates module to enable/disable features at runtime via environment variables.
Add pt2_arg_utils.h header defining argument index enumerations for PyTorch 2 compiled embedding operations.
Add split_embeddings_cache_xpu.h header defining indices for UVM (Unified Virtual Memory) cache performance statistics.
…version Add stochastic_rounding.h header implementing stochastic rounding algorithms for float-to-half precision conversions in embedding operations.
Add weight_row.h header implementing abstractions for efficient access to embedding table weight rows with support for both direct table access and cache-resident data.
Add Jinja2 template for generating optimized SYCL forward kernels for small embedding dimensions (D <= 32).
Add embedding_forward_split_kernel_template.h as a Jinja2 template for generating the main SYCL forward kernel for no-bag (sequence) embeddings.
Add embedding_forward_nobag_unweighted_host_template.cpp as a Jinja2 template for generating the host-side dispatch function for no-bag (sequence) embedding forward passes.
…plit Add embedding_backward_split_kernel_templates.h as a Jinja2 template for generating SYCL backward kernels for no-bag unweighted embedding lookups. Generates both dense (gradient-only) and split (rowwise Adagrad optimizer) variants.
…ghted Add embedding_backward_nobag_unweighted_host_template.cpp as a Jinja2 template for generating the host-side backward dispatch function for no-bag (sequence) embedding passes. Generates both dense (gradient accumulation) and split (rowwise Adagrad) variants.
…ding Add embedding_forward_nobag_unweighted_pt2_wrapper_template.cpp as a Jinja2 template for generating PyTorch 2.0 compilation wrapper functions for split embedding nobag unweighted operations. The single template generates both forward and backward wrappers (controlled by the is_forward flag).
Add fbgemm_dense_lookups_ops.cpp implementing the dense embedding lookup operator for XPU with full autograd support. NOTE: Since the code in this file and the code in te split operator counterpart are from different templates, the code wasn't fused into a single template.
…ise Adagrad Add fbgemm_split_lookups_ops.cpp implementing the PT2-compatible autograd function for split embedding lookups with rowwise Adagrad optimizer on XPU. NOTE: Since the code in this file and the code in the dense operator counterpart are from different templates, the code wasn't fused into a single template.
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
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 wires the code generation scripts, Jinja2 kernel templates,
static operator files, and utility sources into a fully automated CMake
build pipeline. The output is a new
fbgemm_xpu._C_trainingPythonextension module that registers all XPU and AutogradXPU embedding
operator implementations at import time.
Depends on #84
Changes
Build System
src/codegen/CMakeLists.txt: The central builddefinition for the
_C_trainingextension module:Toolchain validation: enforces Intel oneAPI
icpxcompiler;fails fast with a descriptive error if SYCL support is unavailable
XPU architecture selection: reads
TORCH_XPU_ARCH_LISTfromthe environment (defaults to
pvc) and sets the AOT compilationtargets (
-fsycl-targets=spir64_gen,spir64) and device-specificflags (
-cl-poison-unsupported-fp64-kernels)Forward code generation (
add_custom_command): runsgenerate_forward_split.py --opensourceat build time to producehost dispatch files, SYCL kernel headers, and the PT2 wrapper for
the dense and split forward passes. Output is tracked as explicit
CMake outputs so incremental rebuilds are dependency-correct
Backward code generation (
add_custom_command): runsgenerate_backward_split.pyat build time to produce the dense androwwise Adagrad backward host dispatch files, kernel headers, and
PT2 wrapper. Both generators share the same output directory
(
${CMAKE_CURRENT_BINARY_DIR}/generated/sycl_kernels/)generate_lookup_kernelstarget: convenience phony target thatdepends on all generated outputs, useful for manual incremental
regeneration without triggering a full compile
_C_trainingextension module: built from the generatedsources, static ops files (
fbgemm_dense_lookups_ops.cpp,fbgemm_split_lookups_ops.cpp), and the shared utility sources(
backward_utils.cpp,asynchronous_complete_cumsum.cpp,feature_gates.cpp,utils.cpp). ThePyInit__C_trainingstub is inlined via
file(GENERATE ...)at CMake-configure timeso the
.socan be imported without a separate compilation stepInclude paths: adds
fbgemm_xpu/,fbgemm_utils/, and thegenerated headers directory so that all
#includedirectives inboth static and generated sources resolve correctly
CMakeLists.txt: addsadd_subdirectory(src/codegen)toinclude the new sub-project in the top-level build
src/fbgemm_xpu/__init__.py: imports_C_trainingafter_Cso that
TORCH_LIBRARY_IMPLstatic initialisers forXPUandAutogradXPUrun only after all operator schemas have been declaredby
_C. Exports_C_trainingin__all__Testing
NOTE: The following tests will be replaced in the future by FBGEMM tests.
test_dense_embedding_codegen_forward.py: forward correctness,small/large kernel dispatch, output dtype, shape, NaN/Inf, and
reference comparison against
torch.nn.Embeddingtest_dense_embedding_codegen_backward.py: WarpPerRow andCtaPerRow kernels in isolation, mixed-segment dispatch, boundary at
SL 31/32, and numerical gradient correctness
test_split_lookup_operator_forward_pass_xpu.py: split forwardpass across DEVICE, MANAGED, and MANAGED_CACHING placements with
FP32/FP16/BF16 weights
test_split_lookup_operator_backward_pass_xpu.py: split backwardpass with rowwise Adagrad, momentum state validation
cc: @flezaalv