• A
    Reuse OneDNN handler for SGD and SUM for SelectedRows input tensors. (#35510) · 799f3861
    Adam Osewski 提交于
    * Create stateful OneDNNAXPYHandler object.
    
    This makes it possible to call it multiple times without recreating the
    oneDNN primitives every time.
    
    * Prepare SGDOpKernel to reuse its implementation from OneDNN kernel.
    
    * OneDNN SGD kernel.
    
    * Update call to use new OneDNNAXPYHandler object api.
    
    * Setup seed in proper place.
    
    * Enable OneDNN kernel only for single case.
    
    * For dense param and sparse grad.
    
    * Small refactor.
    
    * Enable oneDNN by op attr or by cmd line flag.
    
    * Use int64_t type for number of elements.
    
    * Support dense param and grad from OneDNN kernel.
    
    * Enable SGD OneDNN kernel when use MP BF16 optimizer.
    
    * Force non-copyable/movable OneDNNAXPYHandler.
    
    * Reuse OneDNNAXPYHandler for spare tensors in SUM op.
    
    * Fix SFINAE rules.
    
    * Remove recording event inside AXPY.
    
    * Get rid of internal primitive caching.
    
    * Stop use PP cache mechanims to store mem and primitive obj.
    * Handler obj store and reuse needed desc & prim
    
    * Do not derive from MKLDNNHandlerT
    799f3861
sgd_op.h 12.5 KB