trunk/d288c4204390786fb6ecadcba8a45c395b91da6e: Add oneDNN LSTM primitive support for XPU inference (#185531)

Summary

Use dnnl::lstm_forward primitive on XPU for LSTM inference, replacing the per-timestep fused cell approach. The oneDNN primitive processes the entire sequence internally, eliminating O(T) kernel launches.

Changes

• New file aten/src/ATen/native/mkldnn/xpu/RNN.cpp: Implements lstm_onednn_xpu using dnnl::lstm_forward, registered via REGISTER_XPU_DISPATCH(lstm_mkldnn_stub, ...). Handles weight layout transformation (PyTorch [4*H, I] → oneDNN ldgoi), weight reorder, and scratchpad management.


• aten/src/ATen/native/RNN.cpp:
  
  
  
  • Extended use_mkldnn() to return true for XPU in inference mode (float/bf16/fp16)
  
  
  • Added packed sequence unwrap: when batch_sizes are uniform (common in batch=1 inference), reshape to regular 3D tensor and route to oneDNN
  
  
  • Added XPU-specific LSTMCell path with pre-computed input gates as fallback for training
  
  
  
  

Correctness

• oneDNN and PyTorch use identical gate order: i, f, c̃, o


• oneDNN formula: gates = W·x + U·h + B where B = b_ih + b_hh (we sum before passing)


• Verified max diff < 1e-6 vs CPU reference across multiple configurations

Performance (Intel PVC, Kokoro TTS, bidirectional LSTM, hidden=256)

Dependencies

• Depends on
  
  Approved by: , https://github.com/atalman
  Vollständiger Original-Bericht

  Ausführliche Details, Code-Beispiele & Hersteller-Stellungnahme auf github.com.

  ↗ Original-Artikel auf github.com lesen