Feasibility of Adding a New Backend Library to oneDNN

[avinashcpandey]

I’m exploring the possibility of integrating an additional backend library into oneDNN to leverage specialized kernels for certain workloads. Currently, oneDNN uses a single backend model with JIT-generated kernels.

My questions are:

1. Is there any existing mechanism or roadmap for supporting multiple backends or runtime backend selection in oneDNN?
2. If not, what would be the recommended approach for adding a new backend library (e.g., via plugin architecture or extending the current design)?
3. Are there any constraints or guidelines for maintaining compatibility with oneDNN’s API and memory format system?

Any insights or references to prior discussions would be greatly appreciated.
Thank you!

cc @vpirogov

[vpirogov]

oneDNN supports integration with several libraries:

• Wide selection of BLAS API compatible libraries including ArmPL, BLIS, OpenBLAS and Apple Accelerate. See FindBLAS.cmake and DNNL_BLAS_VENDOR build knob.
• Arm Compute Library. See DNNL_AARCH64_USE_ACL build knob. See RFC document and related discussion in rfc: prototype integration with ArmCL #795.
• cuBLAS, cuDNN, rocBLAS, and MIOpen for SYCL backend.

These are used to improve performance of the functionality available in the library and any implementation relying on external libraries is required to conform with API and numerical behavior defined by oneDNN. This is essential to ensure predictable behavior for oneDNN users across platforms.

Note, that using external libraries adds complexity for distribution, limits support matrix, and increases binary size. These considerations may become critical for upstream projects maintainers, like PyTorch or Tensorflow.

[avinashcpandey]

Thanks @vpirogov for the prompt response!

Few follow up question:

1. In workloads like FlashAttention, each thread executes many tiny GEMMs, and the current primitive-based API introduces overheads (descriptor creation, initialization, format checks) that can outweigh compute time.
   Small models like DIEN and Wide & Deep exhibit similar behavior for the Linear operator, where small problem sizes lead to the same performance challenges.
   Question:
   Does oneDNN provide or plan to provide a lightweight MatMul API that minimizes per-call overhead and repeatedly invokes optimized kernels (e.g., BRGEMM or external BLAS) without full primitive setup? If not, would adding such an API align with oneDNN’s design philosophy?
   Native BRGEMM or specialized backends (e.g., AOCL-DLP, LibXSMM) can be integrated in that low overhead API without going through the full primitive layer?

2. For improving the performance of the MatMul implementation on a given architecture, There are tunable knobs—such as MC/NC blocking factors, BRGEMM block sizes, register level tile sizes (MR/NR), thread balancing strategies, or similar parameters—that can be adjusted for the particular hardware?
   Question:
   If such tuning is possible, is there a recommended approach to read hardware characteristics (e.g., cache sizes, core count, ISA) and then select optimal values for these knobs?
   Additionally, what is the feasibility of defining an architecture-specific tuning module where these parameters (e.g., blocking factors, BRGEMM block sizes, JIT heuristics/intelligence, thread balancing strategies) can be maintained and evolved as the architecture changes—without requiring fundamental design changes to oneDNN core kernel generation/implementation?
   Could you also point me to the relevant parts of the codebase where such tuning logic currently resides or could be experimented with?

[vpirogov]

Does oneDNN provide or plan to provide a lightweight MatMul API that minimizes per-call overhead and repeatedly invokes optimized kernels (e.g., BRGEMM or external BLAS) without full primitive setup?

For repeated primitive calls full setup is typically not required:
• Primitive can be reused as long as the operation definition (i.e. data types and input shapes) does not change.
• Memory objects allow updating data buffer using dnnl::memory:set_data_handle() method.
In addition oneDNN has caching mechanism that reduces the cost of repeated primitive creation.

For complex functionality like FlashAttention oneDNN offers couple more optimization approaches:

• Using uKernel API as a basic building block for tight control over thread decomposition and blocking that works best for particular case. This approach is used in PyTorch for CPUs.
• Using Graph API to express larger operation and allow oneDNN manage all the implementation details. This approach is used in PyTorch for Intel GPUs. See Scaled Dot Product Attention (SDPA) for details on Graph API.

If you see use cases not covered by existing mechanisms we could definitely explore other options. A good starting point would be sharing a oneDNN-based implementation that demonstrates usage model and quantifies the overheads using some kind of profiling.

If such tuning is possible, is there a recommended approach to read hardware characteristics (e.g., cache sizes, core count, ISA) and then select optimal values for these knobs?

oneDNN supports wide range of x64-based platforms ranging from client to server systems. It also extensively specializes for hardware characteristics when generating target-specific code. Hardware characteristics are detected using Xbyak:

• ISA
• cache sizes
• number of cores

Decisions related to tiling, loop order and isa selection are handled inside specific implementations based on the hardware information and known problem sizes. Heuristics are defined based on performance benchmarking and actual approach depends on specifics of target architecture. You may want to study BRGEMM-based matmul implementation to understand the approach we are currently using. BRGEMM-based implementation is specifically designed to abstract out ISA and focus on tunable parameters of the implementation.

[avinashcpandey]

Thank you @vpirogov for sharing all this detailed information—it’s extremely helpful for us in shaping the solution.
We plan to pursue both paths in parallel:

1. Adding ZenDNN as a backend, https://github.com/amd/ZenDNN/tree/zendnnl
In order to create an RFC for adding ZenDNN as a oneDNN backend, we wanted to get your perspective on this approach.

High-level ZenDNN info:
ZenDNN is AMD’s deep learning library optimized for Zen architectures. It provides highly tuned kernels for operations like MatMul and Embedding Bag, leveraging Zen-specific features and cache-aware blocking strategies.
ZenDNN currently supports multiple backends: AOCL-DLP, oneDNN, and LibXSMM.
Additionally, ZenDNN offers two intelligent mechanisms for automatic kernel selection based on workload characteristics:

AutoTuner – Dynamically benchmarks and selects the fastest algorithm at runtime.
ML-based Decision Tree – Uses precomputed logic derived from machine learning to pick the optimal kernel deterministically.

To get the best from backend microkernel, ZenDNN has Parallel Primitive high-performance execution layer optimized for AMD architectures. It delivers cache-aware, scalable compute for AI workloads through dynamic partitioning and ML-driven dispatch. Currently its enabled for BMM. MM support with this layer is in development.
Key Highlights:

• Dynamic Partitioning & Tiling
  Splits MatMul and Batched MatMul operations into smaller, cache-friendly tiles based on matrix dimensions, thread/core count, and L1/L2/L3 cache size. This ensures optimal data locality and minimizes memory bandwidth pressure.

• Dynamic Dispatch to Optimized Microkernels
  Each thread invokes highly tuned microkernels from AOCL-DLP and LibXSMM. The dispatch mechanism selects the most efficient kernel variant for the given tile size and hardware configuration.

• ML-Driven Heuristics for Kernel Selection
  A machine learning model predicts the best-performing kernel and tiling strategy based on workload characteristics, enabling adaptive optimization without manual tuning.

Integrating ZenDNN as an optional backend in oneDNN will allow us to immediately benefit from these specialized kernels for all workloads using ZenDNN currently.

Before proceeding, we want to ensure that this approach aligns with oneDNN’s design philosophy of adding backend. If the community agrees that adding ZenDNN as an optional backend fits within the overall architecture and goals, we’ll move forward with an RFC and prototype. We will follow the ACL RFC process as a reference for integration steps and design considerations.

2. Exploring tuning within oneDNN
In parallel, we’ll experiment with the existing BRGEMM-based implementation and adjust parameters such as blocking factors, tile sizes, and thread balancing strategies based on architecture-specific characteristics. This will help us understand the heuristics and identify opportunities for improvement without major design changes to oneDNN’s core. This will enable us to get the acceleration for other ops as well which are not part of ZenDNN.

[vpirogov]

Adding ZenDNN as a backend, https://github.com/amd/ZenDNN/tree/zendnnl

A good starting point would be a proposal that includes a clear problem statement, describes proposed solutions, and includes supporting performance data as well as distribution and support matrix considerations. I would suggest separating questions of improving oneDNN performance on Zen architecture and addressing API overheads into separate proposals to keep discussions focused.

Here are a few thoughts based on the limited information:

• Using "ZenDNN" moniker when referring to ZenDNN/tree/zendnnl branch is misleading. Looking at the code it's clear that this is a new library with API and implementation that have little in common with ZenDNN v5.1. Consider using a name that clearly differentiates these two.
• libxsmm is a research project not intended for production use. oneDNN offers production implementation of the approach invented by Intel Labs and pioneered in libxsmm. As you can see libxsmm is not adopted in PyTorch or Tensorflow.
• AOCL-DLP looks like a fairly new project. It would be great to learn more about used optimization approach and see performance comparisons.
• It would be great to review the code that implements approaches listed as ZenDNN highlights. Could you please share relevant pointers?

[vpirogov]

Before proceeding, we want to ensure that this approach aligns with oneDNN’s design philosophy of adding backend. If the community agrees that adding ZenDNN as an optional backend fits within the overall architecture and goals, we’ll move forward with an RFC and prototype. We will follow the ACL RFC process as a reference for integration steps and design considerations.

Improving oneDNN performance on AMD processors, as well as improving oneDNN performance on x64 processors in general, are well aligned with the project goals. I would encourage you to consider contributing to existing x64 implementation as the primary path forward.

There's no philosophical considerations that block using third party libraries as an optimization tool. Specific proposal will be considered on technical merits. Please refer to contributing guidelines for details.

[avinashcpandey]

Thank you @vpirogov for the additional information.

The branch ZenDNN/zendnnl is expected to become the main branch this week. ZenDNN has recently undergone a strategic re-architecture and refactoring aimed at improving performance, maintainability, and scalability.

Additionally, AOCL-DLP is inspired by AOCL-BLIS and is specifically designed to target deep learning MatMul primitives.

I will share code pointers and documentation as we open-source latest bits this week.

We will explore both possibilities—integrating ZenDNN as an optional backend and contributing to oneDNN performance improvements.

[vpirogov]

Thanks for sharing the details, @avinashcpandey. Looking forward to future updates.