Commits on Sep 10, 2022

1. Refactor LXU cache logic in TBE fwd training (pytorch#1295) …

   Summary:
   Pull Request resolved: pytorch#1295
   
   The LXU cache logic is in the critical path of the forward TBE kernel.
   Even when the LXU cache is not used, the kernel still checks whether a
   row should be fetched from the cache or HBM at runtime.  The branching
   logic should be harmless for the memory (subsystem) bound case.
   However, it could add significant overhead if TBE is conditional
   bound.  (We have observed that FP16 weight type is generally compute
   or conditional bound, while FP32 weight type is memory bound.)
   
   This diff adds a static conditional in the forward TBE kernel to
   enable/disable the LXU cache code path at compile time.  At runtime,
   the host selects the kernel with/without cache enabled based on
   whether the LXU cache is present.
   
   This diff also moves the conditional outside the D loop.  It should
   add a small benefit for the large D cases when cache is used.
   
   Reviewed By: jspark1105
   
   Differential Revision: D39353035
   
   fbshipit-source-id: bfd3d842971091e954e49c6c8fad034db1fcbc9b

   

   sryap committed Sep 10, 2022
   Configuration menu

   • View commit details
   • Copy full SHA for 9aff411
   • Browse repository at this point

   Copy the full SHA

   9aff411 View commit details

   Browse the repository at this point in the history