[llvm] [CodeGen] Make LoadExtActions address-space aware (PR #162407)

[Matt Arsenault via llvm-commits]

================
@@ -1472,27 +1472,40 @@ class LLVM_ABI TargetLoweringBase {
   /// Return how this load with extension should be treated: either it is legal,
   /// needs to be promoted to a larger size, needs to be expanded to some other
   /// code sequence, or the target has a custom expander for it.
-  LegalizeAction getLoadExtAction(unsigned ExtType, EVT ValVT,
-                                  EVT MemVT) const {
-    if (ValVT.isExtended() || MemVT.isExtended()) return Expand;
-    unsigned ValI = (unsigned) ValVT.getSimpleVT().SimpleTy;
-    unsigned MemI = (unsigned) MemVT.getSimpleVT().SimpleTy;
+  LegalizeAction getLoadExtAction(unsigned ExtType, EVT ValVT, EVT MemVT,
+                                  unsigned AddrSpace) const {
+    if (ValVT.isExtended() || MemVT.isExtended())
+      return Expand;
+    unsigned ValI = (unsigned)ValVT.getSimpleVT().SimpleTy;
+    unsigned MemI = (unsigned)MemVT.getSimpleVT().SimpleTy;
     assert(ExtType < ISD::LAST_LOADEXT_TYPE && ValI < MVT::VALUETYPE_SIZE &&
            MemI < MVT::VALUETYPE_SIZE && "Table isn't big enough!");
     unsigned Shift = 4 * ExtType;
+
+    uint64_t OverrideKey = ((uint64_t)(ValI & 0xFF) << 40) |
+                           ((uint64_t)(MemI & 0xFF) << 32) |
+                           (uint64_t)AddrSpace;
+
+    if (LoadExtActionOverrides.count(OverrideKey)) {
+      return (
+          LegalizeAction)((LoadExtActionOverrides.at(OverrideKey) >> Shift) &
+                          0xf);
----------------
arsenm wrote:

Avoid double map lookup 

https://github.com/llvm/llvm-project/pull/162407