Decoupling oneDNN

[Zhenzhong1]

Description

Decoupling oneDNN

Changes:

• Adds ARK_DNNL CMake gating and conditional compilation to allow XPU builds without oneDNN when SYCL-TLA is enabled.
• Introduces DeviceMemoryPool and refactors multiple call sites to use it instead of DnnlContext scratch allocations.
• Extends SYCL-TLA dense GEMM and Python wrappers/tests to support FP32 (plus adds an XPU INT8-weight GEMM wrapper via SYCL).

Progress

• Moffied CmakeLists
  CPU build: ARK_DNNL=ON by default
  XPU build: ARK_XPU & ARK_DNNL=OFF & ARK_SYCL_TLA=ON by default.
• Support FP32 sycl_tla_dense_gemm
• Decoupling oneDNN
• replace joint_maxtrix

Usage

ARK_DNNL=1 python setup.py bdist_wheel

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[Copilot]

Pull request overview

This PR decouples oneDNN from the ARK kernel build (especially for XPU), introduces a shared scratch-memory pool, and extends SYCL-TLA dense GEMM to accept FP32 inputs (along with corresponding Python/test updates).

Changes:

• Adds ARK_DNNL CMake gating and conditional compilation to allow XPU builds without oneDNN when SYCL-TLA is enabled.
• Introduces DeviceMemoryPool and refactors multiple call sites to use it instead of DnnlContext scratch allocations.
• Extends SYCL-TLA dense GEMM and Python wrappers/tests to support FP32 (plus adds an XPU INT8-weight GEMM wrapper via SYCL).

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
auto_round_extension/ark/auto_round_kernel/CMakeLists.txt	Adds ARK_DNNL option + build-time constraints for XPU backend selection and conditional oneDNN FetchContent/linking.
auto_round_extension/ark/setup.py	Adds env flag parsing and passes ARK_DNNL/ARK_SYCL_TLA toggles into the XPU CMake build.
auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp	Adds DeviceMemoryPool and refactors DnnlContext scratch allocation to use it.
auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp	Makes oneDNN/SYCL-TLA includes conditional and routes scratch allocations + S8 path through new wrappers/pool.
auto_round_extension/ark/auto_round_kernel/wrapper/include/cpu_wrapper.hpp	Switches scratch allocation to DeviceMemoryPool.
auto_round_extension/ark/auto_round_kernel/sdpa.cpp	Switches varlen workspace scratch allocation to DeviceMemoryPool.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_gemm.hpp	Adds FP32 support to SYCL-TLA dense GEMM and updates error messaging.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp	New SYCL-based S8 wrapper used by XPU paths.
auto_round_extension/ark/auto_round_kernel/ark.cpp	Adds conditional XPU backend dispatch (oneDNN vs SYCL-TLA) and uses SyclS8Wrapper for XPU S8.
auto_round_extension/ark/auto_round_kernel/init.py	Updates matmul_sycl_tla Python API dtype validation to include FP32.
auto_round_extension/ark/test/test_matmul.py	Expands SYCL-TLA test parameterization to include FP32 and adjusts backend-compare tolerances.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.