Finish making address-space sensitive, and fix AMDGPU backend

Author: QuantumSegfault

llvm/include/llvm/CodeGen/BasicTTIImpl.h

@@ -28,6 +28,7 @@
 #include "llvm/Analysis/TargetTransformInfoImpl.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/ISDOpcodes.h"
+#include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/ValueTypes.h"
@@ -1244,9 +1245,14 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
         EVT LoadVT = EVT::getEVT(Src);
         unsigned LType =
           ((Opcode == Instruction::ZExt) ? ISD::ZEXTLOAD : ISD::SEXTLOAD);
-        if (DstLT.first == SrcLT.first &&
-            TLI->isLoadExtLegal(LType, ExtVT, LoadVT))
-          return 0;
+
+        if (I && isa<LoadInst>(I->getOperand(0))) {
+          auto *LI = cast<LoadInst>(I->getOperand(0));
+
+          if (DstLT.first == SrcLT.first &&
+              TLI->isLoadExtLegal(LType, ExtVT, LoadVT, LI->getPointerAddressSpace()))
+            return 0;
+        }
       }
       break;
     case Instruction::AddrSpaceCast:
@@ -1531,7 +1537,7 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
       if (Opcode == Instruction::Store)
         LA = getTLI()->getTruncStoreAction(LT.second, MemVT);
       else
-        LA = getTLI()->getLoadExtAction(ISD::EXTLOAD, LT.second, MemVT);
+        LA = getTLI()->getLoadExtAction(ISD::EXTLOAD, LT.second, MemVT, AddressSpace);
 
       if (LA != TargetLowering::Legal && LA != TargetLowering::Custom) {
         // This is a vector load/store for some illegal type that is scalarized.

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -1483,7 +1483,11 @@ class LLVM_ABI TargetLoweringBase {
     unsigned Shift = 4 * ExtType;
 
     if (!LoadExtActions.count(AddrSpace)) {
-      return Legal; // default
+      if (MemVT == MVT::i2 || MemVT == MVT::i4 || MemVT == MVT::v128i2 ||
+          MemVT == MVT::v64i4)
+        return Expand;
+
+      return Legal;
     }
     return (
         LegalizeAction)((LoadExtActions.at(AddrSpace)[ValI][MemI] >> Shift) &
@@ -2641,25 +2645,38 @@ class LLVM_ABI TargetLoweringBase {
   /// Indicate that the specified load with extension does not work with the
   /// specified type and indicate what to do about it.
   void setLoadExtAction(unsigned ExtType, MVT ValVT, MVT MemVT,
-                        LegalizeAction Action, unsigned AddrSpace) {
+                        LegalizeAction Action, unsigned AddrSpace = 0) {
     assert(ExtType < ISD::LAST_LOADEXT_TYPE && ValVT.isValid() &&
            MemVT.isValid() && "Table isn't big enough!");
     assert((unsigned)Action < 0x10 && "too many bits for bitfield array");
+
+    if (!LoadExtActions.count(AddrSpace)) {
+      LoadExtActions[AddrSpace]; // Initialize the map for the addrspace
+
+      for (MVT AVT : MVT::all_valuetypes()) {
+        for (MVT VT : {MVT::i2, MVT::i4, MVT::v128i2, MVT::v64i4}) {
+          setLoadExtAction(ISD::EXTLOAD, AVT, VT, Expand, AddrSpace);
+          setLoadExtAction(ISD::ZEXTLOAD, AVT, VT, Expand, AddrSpace);
+        }
+      }
+    }
+
     unsigned Shift = 4 * ExtType;
     LoadExtActions[AddrSpace][ValVT.SimpleTy][MemVT.SimpleTy] &=
         ~((uint16_t)0xF << Shift);
     LoadExtActions[AddrSpace][ValVT.SimpleTy][MemVT.SimpleTy] |=
         (uint16_t)Action << Shift;
   }
   void setLoadExtAction(ArrayRef<unsigned> ExtTypes, MVT ValVT, MVT MemVT,
-                        LegalizeAction Action) {
+                        LegalizeAction Action, unsigned AddrSpace = 0) {
     for (auto ExtType : ExtTypes)
-      setLoadExtAction(ExtType, ValVT, MemVT, Action);
+      setLoadExtAction(ExtType, ValVT, MemVT, Action, AddrSpace);
   }
   void setLoadExtAction(ArrayRef<unsigned> ExtTypes, MVT ValVT,
-                        ArrayRef<MVT> MemVTs, LegalizeAction Action) {
+                        ArrayRef<MVT> MemVTs, LegalizeAction Action,
+                        unsigned AddrSpace = 0) {
     for (auto MemVT : MemVTs)
-      setLoadExtAction(ExtTypes, ValVT, MemVT, Action);
+      setLoadExtAction(ExtTypes, ValVT, MemVT, Action, AddrSpace);
   }
 
   /// Let target indicate that an extending atomic load of the specified type
@@ -3135,7 +3152,7 @@ class LLVM_ABI TargetLoweringBase {
       LType = ISD::SEXTLOAD;
     }
 
-    return isLoadExtLegal(LType, VT, LoadVT);
+    return isLoadExtLegal(LType, VT, LoadVT, Load->getPointerAddressSpace());
   }
 
   /// Return true if any actual instruction that defines a value of type FromTy

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -7347,7 +7347,7 @@ bool CodeGenPrepare::optimizeLoadExt(LoadInst *Load) {
 
   // Reject cases that won't be matched as extloads.
   if (!LoadResultVT.bitsGT(TruncVT) || !TruncVT.isRound() ||
-      !TLI->isLoadExtLegal(ISD::ZEXTLOAD, LoadResultVT, TruncVT))
+      !TLI->isLoadExtLegal(ISD::ZEXTLOAD, LoadResultVT, TruncVT, Load->getPointerAddressSpace()))
     return false;
 
   IRBuilder<> Builder(Load->getNextNode());