[SYCL] Adding SYCL backend to support for Intel XPUs - #215
Open
abagusetty wants to merge 9 commits into
Open
Conversation
Adds a third kernel backend alongside CUDA and HIP, targeting Intel GPUs
through PyTorch's `xpu` device. Kernels are generated as SYCL free
functions and compiled at runtime via the oneAPI kernel_compiler
extension (`create_kernel_bundle_from_source`).
Rather than fork the kernel templates, a compatibility shim
(`templates/sycl_compat.cuh`) maps the CUDA spellings the templates
already use onto their SYCL equivalents: `__global__` becomes a free
function with `nd_range_kernel` and a fixed sub-group size, `threadIdx`
and friends resolve through `this_work_item`, `__shfl_down_sync` becomes
`shift_group_left`, and `atomicAdd` becomes an `atomic_ref`. The shim is
included only for SYCL, so CUDA and HIP render byte-identical kernels.
Shared memory is the one place the templates had to change. CUDA sizes
`extern __shared__` at launch, but SYCL runtime compilation has no
dynamic local memory for free-function kernels, so a new `declare_smem`
macro emits a `work_group_static` buffer sized from the schedule. The
CUDA and HIP branches still emit `extern __shared__`.
Unifying changes:
- The backend is a string ("cuda"/"hip"/"sycl") rather than an `is_hip`
boolean, in the Jinja environment and the LoopUnroll constructors.
- `JITKernel::execute` also takes argument sizes, which SYCL needs to
launch with raw arguments; CUDA and HIP ignore them. A small
`make_kernel_args` helper packs pointers and sizes at each call site.
- Hardcoded `device="cuda"` strings now resolve through
`extlib.DEVICE_TYPE`, and the tests take the device from a conftest
helper, so the suite runs unmodified on any backend.
Two fixes that were required to make SPIR-V work, both no-ops elsewhere:
- float64 Clebsch-Gordon coefficients no longer carry an `L` (long
double) suffix. A hex float literal is already exactly a double, so the
emitted values are bitwise unchanged, but SPIR-V has no 128-bit float
and rejected the suffix.
- SYCL kernel bundles are released via an atexit handler registered
after the first compile. libsycl-jit is dlopened at that point and
registers its own teardown; since atexit runs in reverse order, a
handler registered any earlier still ran too late and segfaulted at
exit.
`group_gemm` uses oneMKL's pointer-array `gemm_batch` on SYCL. The
strided form requires `stride_c >= ldc * n`, which this interleaved
layout deliberately violates.
Verified on Aurora (PVC Max 1550, oneAPI 2026.1.0, PyTorch 2.13):
mace-large batch tests 40 passed; import and input-validation tests 23
passed with 89 subtests; forward, backward, and double-backward match a
NumPy/CPU reference to 1e-16 in fp64; atomic and deterministic
convolutions agree to 4e-16 and the deterministic path is bitwise
repeatable; group_gemm matches a NumPy reference in both branches;
results are identical across 4 tiles, on a non-default stream, and under
torch.compile. CUDA/HIP template rendering is unchanged (verified by
diffing generated kernels against main).
JAX keeps its existing CUDA and HIP backends; XPU is PyTorch-only for
now. CI gains no SYCL job as there is no Intel GPU runner.
Adds a `verify_sycl_extension` job mirroring the existing CUDA one: it installs the oneAPI DPC++ compiler and oneMKL from Intel's apt repository, builds the extension with `CXX=icpx`, and runs the import tests both with and without `OEQ_JIT_EXTENSION`. `torch` is pinned to the 2.10 XPU wheel. The job only exercises the build, not kernel execution. Running the kernels on the OpenCL CPU device instead of an Intel GPU is not viable for two independent reasons: - The backend launches kernels with `raw_kernel_arg`, the analogue of the untyped `void*[]` that `cuLaunchKernel` takes. OpenCL requires `clSetKernelArgSVMPointer` for USM pointers, and an untyped argument carries no indication of which arguments are pointers, so the same kernel that runs on Level Zero fails on the CPU device with `UR_RESULT_ERROR_INVALID_MEM_OBJECT`. Passing the pointer as a typed argument works, but the generated kernels cannot be launched that way. - PyTorch's XPU device binds to Level Zero. Under `ONEAPI_DEVICE_SELECTOR=opencl:cpu`, `torch.xpu.is_available()` is False and no tensor can be placed on the device. Backend detection is also corrected to use `torch.version.xpu` rather than `torch.xpu.is_available()`. The CUDA and HIP branches both test build-time attributes, so they import fine with no accelerator attached; the SYCL branch tested for a live device and so raised an AssertionError at import on any machine without an Intel GPU, including a CI builder.
The SYCL backend needs PyTorch >= 2.7. 2.6 introduces the XPU device,
`shim_xpu.h` and the SYCL support in `torch.utils.cpp_extension`, and 2.7
adds `torch.library.register_autocast`, which the operators register
with. Every XPU symbol the backend touches was checked against the pinned
release: `aoti_torch_get_current_sycl_queue`, `aoti_torch_device_type_xpu`,
`c10::xpu::getCurrentXPUStream`, `library_paths("xpu")`,
`include_paths("xpu")`, `torch.version.xpu`, `ProfilerActivity.XPU` and
the `torch.xpu` module functions the tests use.
The CI job now pins torch 2.12.1+xpu, the newest XPU wheel, rather than
2.10. An import-time check raises a clear error on an older PyTorch
instead of failing later inside the extension build, and the installation
docs state the 2.7 floor alongside the existing 2.4 and 2.8 ones.
PyTorch locates the SYCL toolchain from `icpx` on PATH, which the job
already provides through setvars.sh; this is now mentioned in the docs
since it is the reason `icpx` must be on PATH rather than merely set
as CXX.
…quirement The SYCL backend's own minimum is 2.7 (torch.library.register_autocast; 2.6 for the XPU device and cpp_extension's SYCL support), but it is only tested against the 2.8 floor the project already documents for AOTI and export. Enforcing 2.8 rather than introducing a third version number keeps the requirement list to one number per feature. Also switches the Aurora environment snippet to `module restore`, which resets to the site default set rather than layering a specific oneAPI release on top of whatever is already loaded.
… to it The documented ">= 2.8 for AOTI and export" predates the move to the stable ABI. The headers libtorch_tp_jit_stable.cpp includes are not all available that early: torch/csrc/stable/tensor_struct.h and torch/csrc/stable/accelerator.h first appear in 2.9, and torch/headeronly/core/DeviceType.h in 2.10. The precompiled-extension check in extlib has required 2.10 all along, so the documentation was the part that was out of date. The SYCL backend's import guard now enforces the same 2.10 rather than a version of its own, keeping one floor for every feature that needs a recent PyTorch.
Member
|
Great feature, thanks! This will take some time to review and test - so bear with us, but with some back and forth we'll see where we land. |
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.
Major changes: