viable/strict/1780997283: 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 intel/torch-xpu-ops#3770 for the fused cell fallback path (bias fix)


• Only affects XPU. CUDA and CPU paths unchanged.

Pull Request resolved: #185531

Approved by: https://github.com/EikanWang, https://github.com/guangyey, https://github.com/atalman