Add reorder_batched_ad_lengths and reorder_batched_ad_indices XPU compatible operators by aagalleg · Pull Request #87 · intel/torchlib-xpu

aagalleg · 2026-07-13T17:05:22Z

This PR implements reorder_batched_ad_lengths and reorder_batched_ad_indices operators 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.

Changes

Operator Implementation

sycl_kernels/reorder_batched_ad.h: Public API header with top-level operator implementations for XPU..
sycl_kernels/reorder_batched_ad.cpp: SYCL kernel implementations.

Build System Integration

CMakeLists.txt: Adds reorder_batched_ad.cpp to the host_sources list, ensuring SYCL compilation with the icpx compiler and AOT targets

Testing

The patch adds XPU support to fbgemm_gpu/test/sparse/reorder_batched_test.py by importing fbgemm_xpu and using accelerator_unavailable / current_accelerator() instead of the CUDA-only gpu_unavailable / .cuda() helpers.

Tests enabled on XPU:

test_reorder_batched_ad_lengths — validates reorder_batched_ad_lengths in both broadcast and non-broadcast modes.
test_reorder_batched_ad_indices — validates reorder_batched_ad_indices in both broadcast and non-broadcast modes.

flezaalv · 2026-07-22T19:16:20Z

reorder_batched_adOp.cpp: Top-level operator implementations for XPU.

Changed file to reorder_batched_ad_op.cpp to follow snack case convention.

Testing

test/patches/0001-Add-XPU-support-to-fbgemm-tests.patch: Extends the existing upstream FBGEMM test patch to add XPU support to fbgemm_gpu/test/sparse/reorder_batched_test.py for the reorder_batched_ad_lengths and reorder_batched_ad_indices operators (replaces .cuda() calls with .xpu(), gates the tests on torch.xpu.is_available(), and imports fbgemm_xpu to register the XPU operator implementations).
The local test/test_reorder_batched.py file has been removed; the operators are now covered by the upstream FBGEMM test suite.

Implement reorder_batched_ad_lengths and reorder_batched_ad_indices operators with SYCL kernels for Intel XPU devices. (cherry picked from commit efbf5ca)

Add build system integration and Python API for reorder_batched_ad operators introduced in the previous commit. (cherry picked from commit 693ffcf)

Add comprehensive test suite for reorder_batched_ad_lengths and reorder_batched_ad_indices operators, covering both broadcast and non-broadcast modes. (cherry picked from commit 11faaf1)

(cherry picked from commit 0cc129f)

Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com> (cherry picked from commit 603a554)

…ng standards Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com> (cherry picked from commit 8a3c89b)

Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com> (cherry picked from commit e19631d)

The four kernels in reorder_batched_ad.cpp were SYCL free-function kernels launched through the sycl_kernel_submit helper, which builds a kernel bundle per launch and pulls the comm/ helper directory into this file. Convert them to functors and submit them with queue.submit / parallel_for, which is the pattern already used by the other kernels in this package. Kernel bodies, global/local ranges and kernel arguments are unchanged, so behaviour is identical and only the submission mechanism differs. The host side now takes the queue from c10::xpu::getCurrentXPUStream() and reads the sub-group and work-group limits from at::xpu::getDeviceProperties, so reorder_batched_ad.h no longer includes comm/SYCLContext.h. (cherry picked from commit 42f484d)

Introduce packages/fbgemm-xpu/src/fbgemm_xpu/fbgemm_utils/ with: - dispatch_macros.h: FBGEMM_DISPATCH_* macros for floating/integral scalar type dispatch, mirroring upstream fbgemm_gpu conventions. - tensor_utils.h: helpers to check that a set of tensors are defined and reside on the same XPU device, raising a descriptive TORCH_CHECK message otherwise. - utils.h: XPU work-group sizing helpers (syclDeviceMaxWorkGroupSize, xpu_calc_xblock_count) and the SYCL_DEVICE_GUARD macro, factored out for reuse across SYCL kernel implementations.

…cl_kernels Fold reorder_batched_ad_op.cpp into sycl_kernels/reorder_batched_ad.{h,cpp} so each kernel dispatcher function directly returns the output tensor and registers itself with TORCH_LIBRARY_IMPL, removing the separate host-side wrapper file and its duplicate tensor-allocation logic. Replace the file-local xpu_max_sub_group_size()/xpu_device_max_work_group_size() helpers with the shared kThreadGroupSize/kMaxThreads constants from fbgemm_utils/utils.h, and switch the fbgemm_utils includes to their new top-level path. Also fix NarrowBroadcastIndicesKernel to bail out before indexing cat_ad_offsets_ for warps launched past reordered_cat_ad_batches_, avoiding an out-of-bounds read introduced by the work-group rounding.

mkrze · 2026-09-02T10:26:24Z

+    const auto b_t =
+            item.get_group(0) * item.get_local_range(1) + item.get_local_id(1);


Critical — transposed dimensions, indices kernel

Same transposed mapping as in ReorderBatchedAdLengthsKernel: b_t comes from
get_local_id(1) on line 184, while both copy loops (lines 206 and 215) index
with get_local_id(0). Adjacent work-items copy different segments instead of
adjacent elements, so the stores to reordered_cat_ad_indices_ are scattered.

Measured on B60, B=64, T=64, A=16, sweeping segment length so the inner copy
width is A*L:

width as-is transposed speedup

128 0.0439 ms (96 GB/s) 0.0200 ms (210 GB/s) 2.20×

256 0.0522 ms (161 GB/s) 0.0156 ms (539 GB/s) 3.36×

1024 0.1584 ms (212 GB/s) 0.0629 ms (533 GB/s) 2.52×

Same fix: b_t from dimension 0, copy index from dimension 1. Ranges at lines
474–475 already have both dimensions equal, so they stay as they are.

Done at: c7eeb5f

dvrogozh

LGTM

aagalleg force-pushed the feat/reorder_batched_ad branch from ecd2a4d to bfe2040 Compare July 20, 2026 17:41

flezaalv mentioned this pull request Jul 20, 2026

Add jagged_index_select_2d_forward XPU operator #92

Merged

aagalleg force-pushed the feat/reorder_batched_ad branch from bfe2040 to 0cc129f Compare July 21, 2026 18:39

mkrze mentioned this pull request Aug 19, 2026

fbgemm-xpu: submit reorder_batched_ad kernels directly to the SYCL queue aagalleg/torchlib-xpu#13

Merged

5 tasks done

aagalleg and others added 12 commits August 31, 2026 22:24

feat(fbgemm-xpu): add reorder_batched_ad operators for XPU

3979bfd

Implement reorder_batched_ad_lengths and reorder_batched_ad_indices operators with SYCL kernels for Intel XPU devices. (cherry picked from commit efbf5ca)

build(fbgemm-xpu): integrate reorder_batched_ad operators

778c882

Add build system integration and Python API for reorder_batched_ad operators introduced in the previous commit. (cherry picked from commit 693ffcf)

test(fbgemm-xpu): add tests for reorder_batched_ad operators

97c3baf

Add comprehensive test suite for reorder_batched_ad_lengths and reorder_batched_ad_indices operators, covering both broadcast and non-broadcast modes. (cherry picked from commit 11faaf1)

test(fbgemm-xpu): drop resurrected test file after rebase

2d7659e

(cherry picked from commit 0cc129f)

refactor(fbgemm-xpu): rename reorder_batched_adOp.cpp to snake_case

80b9a6b

Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com> (cherry picked from commit 603a554)

refactor(fbgemm-xpu): align reorder_batched_ad SYCL kernels with codi…

6f7c571

…ng standards Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com> (cherry picked from commit 8a3c89b)

test(fbgemm-xpu): move reorder_batched tests to upstream patch

4590335

Signed-off-by: Felipe Leza Alvarez <felipe.leza.alvarez@intel.com> (cherry picked from commit e19631d)

fix(fbgemm-xpu): update CMakeLists.txt, ops_registry.cpp, and ci.yml

0c641f9

test(fbgemm-xpu): add XPU support to reorder_batched_test.py patch

e1a175e

aagalleg force-pushed the feat/reorder_batched_ad branch from e93d159 to e1a175e Compare September 2, 2026 00:24

aagalleg added 2 commits September 2, 2026 00:33

style(fbgemm-xpu): normalize trailing whitespace in fbgemm_utils headers

6b533fe

docs(fbgemm-xpu): list reorder_batched_ad ops in README

a9df3b8

aagalleg marked this pull request as ready for review September 2, 2026 00:59

aagalleg requested review from dvrogozh, flezaalv, luis-real, mkrze and molintc as code owners September 2, 2026 00:59

mkrze reviewed Sep 2, 2026

View reviewed changes

aagalleg added 3 commits September 2, 2026 17:13

fix(fbgemm-xpu): use SymInt for reorder_batched_ad schema args

649d5b9

fix(fbgemm-xpu): restore grid-stride loop and fix warp axis order

c7eeb5f

feat(fbgemm-xpu): port vectorized reorder_batched_ad_indices kernel

10310e4

aagalleg added 2 commits September 2, 2026 22:29

fix(fbgemm-xpu): get indices work-group size from the device

e716f82

test(fbgemm-xpu): enable reorder_batched_ad large-grid tests on XPU

a0d6273

aagalleg requested a review from mkrze September 2, 2026 23:38

mkrze approved these changes Sep 3, 2026

View reviewed changes

dvrogozh approved these changes Sep 3, 2026

View reviewed changes

dvrogozh merged commit 099a1c9 into intel:main Sep 3, 2026
17 checks passed

Navigation Menu

Repository navigation

Add reorder_batched_ad_lengths and reorder_batched_ad_indices XPU compatible operators - #87

Add reorder_batched_ad_lengths and reorder_batched_ad_indices XPU compatible operators#87
dvrogozh merged 19 commits into
intel:mainfrom
aagalleg:feat/reorder_batched_ad

aagalleg commented Jul 13, 2026 •

edited

Loading

Uh oh!

flezaalv commented Jul 22, 2026

Uh oh!

Uh oh!

mkrze Sep 2, 2026

Uh oh!

aagalleg Sep 2, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dvrogozh left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		const auto b_t =
		item.get_group(0) * item.get_local_range(1) + item.get_local_id(1);

width	as-is	transposed	speedup
128	0.0439 ms (96 GB/s)	0.0200 ms (210 GB/s)	2.20×
256	0.0522 ms (161 GB/s)	0.0156 ms (539 GB/s)	3.36×
1024	0.1584 ms (212 GB/s)	0.0629 ms (533 GB/s)	2.52×

Conversation

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

Uh oh!

Changes

Operator Implementation

Build System Integration

Testing

Uh oh!

flezaalv commented Jul 22, 2026

Uh oh!

Uh oh!

mkrze Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

aagalleg Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dvrogozh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aagalleg commented Jul 13, 2026 •

edited

Loading