Add XPU enabled jagged tensor operators - #129
Conversation
Port fbgemm_gpu/src/jagged_tensor_ops/common.cuh and common.h to SYCL/XPU as a header-only jagged_common.h, providing the kernels and launchers that the individual jagged tensor operators build on. Kernels, with CUDA functions expressed as SYCL functor classes: JaggedDenseElementwiseDenseOutputKernel JaggedDenseDenseElementwiseJaggedOutputKernel JaggedDenseDenseElementwiseJaggedOutputOptSearchKernel JaggedDenseDenseElementwiseJaggedOutputOptGatherKernel Host launchers: jagged_dense_elementwise_dense_output_, jagged_dense_elementwise_jagged_output_, jagged_dense_elementwise_jagged_output_opt_, and the jagged_dense_dense_elementwise_jagged_output_matches_opt gate that chooses between the vectorized and generic jagged-output paths.
| JaggedLaunchConfig cfg = \ | ||
| check_shape_and_partition_(x_values, x_offsets, y); \ | ||
| cfg.blocks = std::max<int64_t>( \ | ||
| 1, div_round_up(x_values.size(0), cfg.threads_y)); \ |
There was a problem hiding this comment.
Critical — generic launch throws away the grid cap
This overwrites the SYCL-safe grid cap computed by
check_shape_and_partition_(). I reproduced the consequence on B60 with
BF16, nnz=67,108,833, D=16: the flattened range becomes exactly 2^31
work-items, and both add-jagged-output and
jagged_to_padded_dense_backward fail with:
Provided range and/or offset does not fit in int
The kernel already has a group-stride loop, so could we cap this replacement
with xpu_cap_grid_dim_x(requested_blocks, threads_x * threads_y) and add
large-grid tests for add-jagged-output, backward, and dense-to-jagged? The
hardware reproduction exercised the first two; dense-to-jagged uses this same
macro for non-FP16 inputs. The comment above currently says both consumers
preserve the cap, but this macro discards it.
| const IdxT outer_begin = static_cast<IdxT>( | ||
| item.get_group(0) * item.get_local_range(0) + item.get_local_id(0)); | ||
| // CUDA: gridDim.x * blockDim.y | ||
| const IdxT outer_stride = static_cast<IdxT>( | ||
| item.get_group_range(0) * item.get_local_range(0)); | ||
| const IdxT total_outer = outer_dense_size * jagged_folded_size; | ||
|
|
||
| // CUDA: threadIdx.x / blockDim.x | ||
| const IdxT inner_begin = static_cast<IdxT>(item.get_local_id(1)); | ||
| const IdxT inner_stride = static_cast<IdxT>(item.get_local_range(1)); | ||
|
|
||
| for (IdxT outer = outer_begin; outer < total_outer; | ||
| outer += outer_stride) { |
There was a problem hiding this comment.
Critical — grid-stride induction can overflow int32
The per-tensor numel < INT_MAX check does not make this int32 induction safe.
For total_outer=2^30+1, D=1, the capped launch gives
outer_stride=1,073,741,856; the final active work-item then evaluates
1,073,741,824 + 1,073,741,856 = 2,147,483,680, overflowing int32 before the
next loop condition. The generic jagged-output int offset loop has the same
pattern.
The paired inner loops also evaluate iidx * 2 + 1 in int32. With
inner_dense_size=INT_MAX-1, that expression eventually overflows even though
the tensor-size gate selected the int32 path.
Because signed overflow is undefined behavior, this can become an
out-of-bounds device iteration rather than a clean error. Could we keep
grid-stride and paired-inner counters/arithmetic in int64 (while retaining
validated 32-bit accessors if desired), or explicitly guard the final
increment/multiply, and add large-outer and large-inner regressions?
| int count = B_ - 1; | ||
| int first = 1; | ||
| while (count > 0) { | ||
| int idx = first; | ||
| int step = count / 2; | ||
| idx += step; | ||
| if (offsets_sh[idx] <= row) { | ||
| first = ++idx; | ||
| count -= step + 1; | ||
| } else { | ||
| count = step; | ||
| } | ||
| } | ||
| --first; | ||
|
|
||
| const int dense_row = first; | ||
| const int offset = static_cast<int>(offsets_sh[dense_row]); | ||
| const int dense_col = row - offset; | ||
| rows_[row] = dense_row; |
There was a problem hiding this comment.
High — FP16 search misses the final sentinel
The generic search covers offsets.size()-1 == B entries, but this FP16 path
uses B-1 and never examines the final offsets[B] sentinel.
I reproduced wrong output on B60 for both int32 and int64 offsets with
dense=[2,1,8], offsets=[0,1,1], total_L=2: the second jagged row should
be zero, but this path returns dense row 1 ([9,...,16]).
Could we search the complete sentinel range and add this regression for both
offset dtypes? After correcting the search, an out-of-range row can produce
dense_row == B, so y0_ptr/y1_ptr must also be formed only after the
existing bounds check; currently the invalid references are constructed first.
| if (y.numel() == 0 || x_values.numel() == 0) { | ||
| return; | ||
| } |
There was a problem hiding this comment.
High — zero max length leaves outputs uninitialized
This returns before writing a non-empty output when the dense max length is
zero. Both callers allocate with at::empty, so the result is uninitialized.
With deterministic uninitialized-memory filling enabled on B60:
- dense-to-jagged with
dense=[1,0,8],offsets=[0,2],total_L=2returned
[2,8]NaNs instead of zeros; - add-jagged-output with
x=[2,8],y=[1,0,8]returned NaNs instead of
unchangedx.
These differ from the CPU semantics. Could we initialize these zero-max-length
cases explicitly and add regressions for both operators?
| constexpr int64_t kWarpSize = static_cast<int64_t>(kThreadGroupSize); | ||
|
|
||
| // CUDA: inner_dense_size >= kWarpSize / 2 ? kWarpSize : inner_dense_size. | ||
| // Clamped to at least 1 so a zero-width inner dim cannot produce an empty | ||
| // nd_range (CUDA tolerates a zero-sized block dim by never launching). | ||
| const int64_t threads_x = std::max<int64_t>( | ||
| 1, inner_dense_size >= kWarpSize / 2 ? kWarpSize : inner_dense_size); | ||
| const int64_t threads_y = static_cast<int64_t>(kMaxThreads) / kWarpSize; |
There was a problem hiding this comment.
High — CUDA geometry is up to 3.8x slower on BMG
Each work-item below handles two adjacent inner elements, but this copied CUDA
rule jumps to threads_x=32 at D=16. Together with threads_y=32, that
creates a 1024-item work-group while only 8 inner lanes do useful work.
I A/B-built a variant changing only this expression to use
min(32, ceil(D/2)). The same 24-case parity/autograd matrix passed. On B60,
BF16 D=16 improved:
- jagged-to-padded-dense:
0.568 -> 0.159 ms(3.57x); - dense-to-jagged:
1.337 -> 0.350 ms(3.82x); - add-jagged-output:
1.349 -> 0.359 ms(3.76x).
D=32 improved about 1.8-2.0x and D>=64 was unchanged. Could we derive the
inner range from the two-element chunks and validate the full D sweep? The
inner axis should remain SYCL's fastest-varying dimension.
| TENSOR_ON_SYCL_XPU(x_values); | ||
| for (const auto& x_offset : x_offsets) { | ||
| TENSOR_ON_SYCL_XPU(x_offset); | ||
| } |
There was a problem hiding this comment.
High — mixed-device pointers are not rejected
These checks validate only that values and each offset are some XPU tensor.
They do not validate y/output placement or require every tensor to use the
same XPU index. The add wrapper selects its guard from dense, so a CPU dense
tensor or a different XPU index can select a queue and pass foreign pointers
to the kernel.
Could we validate values, dense/output, and every offset are on the same XPU
before selecting the stream, and add mixed CPU/XPU plus XPU:0/XPU:1 argument
tests? This should fail cleanly at the API boundary rather than risking invalid
device-memory access.
| if (jagged_dense_dense_elementwise_jagged_output_matches_opt( | ||
| num_jagged_dim, | ||
| x_values, | ||
| x_offsets, | ||
| y_reshaped, | ||
| y_reshaped, | ||
| output_values)) { | ||
| AT_DISPATCH_INDEX_TYPES( | ||
| x_offsets[0].scalar_type(), "jagged_indices_fast_path", [&] { | ||
| const auto nnz = output_values.size(0); | ||
| const auto B = y_reshaped.size(0); | ||
| const auto E = y_reshaped.size(2); |
There was a problem hiding this comment.
High — FP16 fast path skips shape validation
This branch bypasses check_shape_and_partition_(), so unlike the generic
path it never validates the inner widths or offsets[0].numel() == B+1.
I reproduced both omissions on B60. FP16 add with x_values=[2,16],
dense=[2,1,8], offsets [0,1,2] succeeded and left the final eight output
values per row as NaNs; CPU rejects the size mismatch. Dense-to-jagged with
dense=[2,1,8], offsets [0,1,2,2], total_L=2 also silently succeeded
while CPU rejects the extra interval. Reversing the width mismatch makes the
vector gather store the larger dense width into shorter output rows.
Could we run the common shape validation before selecting the optimized path
and add FP16 regressions for both width directions and short/long offset
lists?
| // Each row is aligned to 128-bit | ||
| matches &= (x_values.stride(-2) % 8 == 0); | ||
| matches &= (output_values.stride(-2) % 8 == 0); | ||
| matches &= (y_0_reshaped.stride(-2) % 8 == 0); | ||
| matches &= (y_1_reshaped.stride(-2) % 8 == 0); | ||
|
|
||
| // Base addresses aligned to 128-bit | ||
| matches &= (reinterpret_cast<uint64_t>(x_values.data_ptr()) % 16 == 0); | ||
| matches &= (reinterpret_cast<uint64_t>(output_values.data_ptr()) % 16 == 0); | ||
| matches &= (reinterpret_cast<uint64_t>(y_0_reshaped.data_ptr()) % 16 == 0); | ||
| matches &= (reinterpret_cast<uint64_t>(y_1_reshaped.data_ptr()) % 16 == 0); |
There was a problem hiding this comment.
High — FP16 matcher misses dense batch alignment
These checks cover the base and stride(-2), but not the reshaped dense
tensor's outer batch stride. A later batch can therefore be misaligned even
though the gather dereferences it as HalfVec8*.
On B60, a valid FP16 dense tensor with shape [2,1,8] and strides (17,8,1)
entered this path. Adding contiguous x_values=[2,8] should have produced row
1 [66,68,70,72,74,76,78,80]; it returned
[8,67,69,71,73,75,77,79].
Could we require every dense row used by the vector loads to be 16-byte
aligned (including stride(0)), or fall back to unaligned/generic loads, and
add this strided regression?
2 participants
The operators enable conversion between dense and jagged representations and support element-wise operations with jagged output:
dense_to_jagged: Convert dense padded tensors to compact jagged formatjagged_to_padded_dense: Convert jagged tensors back to padded dense formatjagged_dense_elementwise_add_jagged_output: Element-wise addition with jagged outputjagged_2d_to_dense: Convert 2D jagged tensors to dense formatFour backend operators are registered under the XPU dispatch key:
dense_to_jagged_forward,jagged_to_padded_dense_forward,jagged_to_padded_dense_backwardandjagged_dense_elementwise_add_jagged_output. The four public operators listed above are composite or Autograd-key entry points owned by the upstreamfbgemm-gpu-cpupackage and come up automatically once the backends are registered, so noAutogradXPUblock is needed here.Each file is a per-file port of its FBGEMM CUDA counterpart at v1.8.0, with the CUDA source path and symbol mapping recorded in a banner comment at the top of every header.
Changes
Shared SYCL Infrastructure
•
sycl_kernels/jagged_common.h: SYCL counterpart ofjagged_tensor_ops/common.cuhandcommon.h.Operator Implementation
One
.cpp/.hpair per operator, each mapping to a single CUDA source file:•
sycl_kernels/dense_to_jagged_forward.{cpp,h}→jagged_tensor_ops/dense_to_jagged_forward.cu•
sycl_kernels/jagged_to_padded_dense_forward.{cpp,h}→jagged_tensor_ops/jagged_to_padded_dense_forward.cu•
sycl_kernels/jagged_to_padded_dense_backward.{cpp,h}→jagged_tensor_ops/jagged_to_padded_dense_backward.cu•
sycl_kernels/jagged_dense_elementwise_add_jagged_output.{cpp,h}→jagged_tensor_ops/jagged_to_padded_dense_forward.cu, including theJaggedDenseAddJaggedOutputXPUOpautograd Function that mirrors CUDA'sJaggedDenseAddJaggedOutputGPUOpone-for-one.Support Helpers
•
fbgemm_utils/dispatch_macros.h: AddsFBGEMM_DISPATCH_FLOATING_TYPES,FBGEMM_DISPATCH_FLOAT_AND_BFLOAT16_CASEandFBGEMM_DISPATCH_ALL_TYPES_BUT_HALF_CASE, byte-identical to upstream.•
fbgemm_utils/tensor_utils.h: AddsTENSOR_ON_SYCL_XPU, the XPU counterpart of upstreamTENSOR_ON_CUDA_GPU.•
fbgemm_utils/utils.h: Addsdiv_round_upandround_down, widened from upstream'sint32_ttoint64_tbecause the jagged launchers derive work-group counts from int64 tensor extents.Build System Integration
•
CMakeLists.txt: Adds the four operator sources tohost_sources.•
ops_registry.cpp: AddsschemaExists-guarded schema declarations for the four backend operators. The schema strings are copied verbatim from upstreamjagged_tensor_ops_cpu.cpp, so a standalonefbgemm-xpuinstall exposes identical signatures, and each guard becomes a no-op whenfbgemm-gpu-cpuis present.Testing
Instead of a standalone test file, this reuses FBGEMM's own jagged test suite through the existing test patch:
•
patches/0001-Add-XPU-support-to-fbgemm-tests.patch: Adds XPU support todense_to_jagged_test.py,2d_to_dense_test.py,1d_to_dense_test.py,elementwise_binary_test.py, and the sharedjagged/common.py. CUDA-specific helpers give way to device-agnostic accelerator probes (accelerator_unavailable,accelerator_memory_lt_gb,current_accelerator).•
.github/workflows/ci.yml: Runs the five jagged test modules.Documentation
•
README.md: Lists the newly reachable jagged operators.