Add Jinja2 code generation infrastructure for SYCL embedding kernels - #80
Draft
aagalleg wants to merge 42 commits into
Draft
Add Jinja2 code generation infrastructure for SYCL embedding kernels#80aagalleg wants to merge 42 commits into
aagalleg wants to merge 42 commits into
Conversation
aagalleg
force-pushed
the
feat/add_generator_scripts
branch
from
July 17, 2026 23:43
efba2dc to
e5d3f39
Compare
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.
…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.
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]).
Delete the accidentally tracked submodule reference to FBGEMM-v1.7.0.
Rename asynchronous_complete_cumsum files to sparse_async_cumsum.
Update 0001-Add-XPU-support-to-fbgemm-tests.patch to enable XPU testing for asynchronous cumsum operators in cumsum_test.py: - Add XPU device to test_cumsum (tests exclusive, inclusive, and complete cumsum) - Add XPU device to test_asynchronous_complete_cumsum_2d - Skip test_batched_complete_cumsum (operator not implemented on XPU)
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)
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.
…mpatible Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
…rators Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
…tests Integrate block_bucketize_sparse_features SYCL kernel implementation and comprehensive test suite from experimentation_prs_integration branch. - Add SYCL kernel implementation for block_bucketize_sparse_features - Add block_bucketize_sparse_features_inference variant - Add populate_bucketized_permute helper function - Include comprehensive test suite with 18 test cases covering: * Variable bucket sizes and batch sizes * Long indices and keep_orig_idx modes * Total num blocks variations * Float64 weights support * Edge cases and error handling All tests pass successfully on XPU hardware.
…t in rebase Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
… coverage Update block_bucketize_test.py hunk to use the accelerator_unavailable, xpu_available, and fbgemm_xpu registration helpers introduced by earlier patches (misc_ops_test.py, permute_sparse_features_test.py, test_utils.py).
Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com>
… 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.
Set explicit Jinja autoescape policy in generator environments using select_autoescape with HTML/XML extensions only and non-HTML default. This keeps C++ template rendering behavior unchanged while satisfying Bandit B701 in codegen scripts. Replace regex assert checks in PTA parsing helpers with explicit validation that raises ValueError on malformed patterns. This preserves runtime safety under optimization and resolves Bandit B101 findings.
flezaalv
force-pushed
the
feat/add_generator_scripts
branch
from
August 10, 2026 21:22
93bcb95 to
9a31109
Compare
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 introduces the Python-based code generation infrastructure
(
genscript/) used to produce SYCL forward and backward embeddingkernel source files from Jinja2 templates. This is the first step
toward a fully automated build pipeline for the XPU training kernels.
Depends on #76
Changes
Code Generation Framework (
src/codegen/genscript/)jinja_environment.py: Instantiates the Jinja2 environmentpointing to the template root, registers global variables
(
max_embedding_dim,items_per_warp,fixed_max_vecs_per_thread),and provides helper functions for kernel dispatch code generation
including
get_max_vecs_template_configs,dispatch_non_vec_blocking_kernel,dispatch_vec_blocking_kernel,and
dispatch_optimal_kernelcommon.py:CodeTemplateclass that wraps Jinja2 templateloading and rendering. Adds auto-generated file headers (with
__TEMPLATE_SOURCE_FILE__macro / variable) and writes renderedoutput to the configured install directory
torch_type_utils.py:ArgTypeenum andTensorTypedataclassfor mapping between PyTorch argument types and their C++
primitive/scalar type representations
scripts_argsparse.py: Sharedargparseconfiguration for allgeneration scripts — exposes
--install_dir,--opensource, and--is_rocmflagsgenerate_forward_split.py: Generates dense and split forwardkernel variants (host dispatch, SYCL small kernel header, PT2
wrapper)
generate_backward_split.py: Generates dense and rowwise-Adagradbackward kernel variants (warp-per-row and CTA-per-row headers, host
dispatch) for both the dense gradient accumulation and split optimizer
paths
cc: @flezaalv