AMDPGU/SI: Use AssertAnd node to mask high bit for scratch offsets

[• tstellarAMD]

Summary

We can safely assume that the high bit of scratch offsets will never
be set, because this would require at least 128 GB of GPU memory.

[arsenm]

lib/Target/AMDGPU/SIISelLowering.cpp

824–826 ↗

(On Diff #29791)

I think this should be wrapped under a subtarget feature (that defaults to on) and use a more conservative assumption if off. For a name idea -huge-memory-support or something like that

[arsenm]

lib/Target/AMDGPU/SIISelLowering.cpp

824–826 ↗

(On Diff #29791)

By default to on I mean this should default to happening rather than default the conservative assumption

[• tstellarAMD]

Updated patch: Use AssertZext instead of AssertAnd and add a subtarget attribute to
control this behavior.

[arsenm]

LGTM

llvm/trunk/lib/Target/AMDGPU/AMDGPU.td

def FeatureVGPRSpilling : SubtargetFeature<"vgpr-spilling",
         "true",
         "Enable spilling of VGPRs to scratch memory">;
 
 def FeatureSGPRInitBug : SubtargetFeature<"sgpr-init-bug",
         "SGPRInitBug",
         "true",
         "VI SGPR initilization bug requiring a fixed SGPR allocation size">;
 
+def FeatureEnableHugeScratchBuffer : SubtargetFeature<"huge-scratch-buffer",
+        "EnableHugeScratchBuffer",
+        "true",
+        "Enable scratch buffer sizes greater than 128 GB">;
+
 class SubtargetFeatureFetchLimit <string Value> :
                           SubtargetFeature <"fetch"#Value,
         "TexVTXClauseSize",
         Value,
         "Limit the maximum number of fetches in a clause to "#Value>;
 
 def FeatureFetchLimit8 : SubtargetFeatureFetchLimit <"8">;
 def FeatureFetchLimit16 : SubtargetFeatureFetchLimit <"16">;

llvm/trunk/lib/Target/AMDGPU/AMDGPUSubtarget.h

private:
   bool SGPRInitBug;
   bool IsGCN;
   bool GCN1Encoding;
   bool GCN3Encoding;
   bool CIInsts;
   bool FeatureDisable;
   int LDSBankCount;
   unsigned IsaVersion; 
+  bool EnableHugeScratchBuffer;
 
   AMDGPUFrameLowering FrameLowering;
   std::unique_ptr<AMDGPUTargetLowering> TLInfo;
   std::unique_ptr<AMDGPUInstrInfo> InstrInfo;
   InstrItineraryData InstrItins;
   Triple TargetTriple;
 
 public:
public:
   bool isTargetELF() const {
     return false;
   }
 
   StringRef getDeviceName() const {
     return DevName;
   }
 
+  bool enableHugeScratchBuffer() const {
+    return EnableHugeScratchBuffer;
+  }
+
   bool dumpCode() const {
     return DumpCode;
   }
   bool r600ALUEncoding() const {
     return R600ALUInst;
   }
   bool isAmdHsaOS() const {
     return TargetTriple.getOS() == Triple::AMDHSA;

llvm/trunk/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

: AMDGPUGenSubtargetInfo(TT, GPU, FS), DevName(GPU), Is64bit(false),
       TexVTXClauseSize(0), Gen(AMDGPUSubtarget::R600), FP64(false),
       FP64Denormals(false), FP32Denormals(false), FastFMAF32(false),
       CaymanISA(false), FlatAddressSpace(false), EnableIRStructurizer(true),
       EnablePromoteAlloca(false), EnableIfCvt(true), EnableLoadStoreOpt(false),
       EnableUnsafeDSOffsetFolding(false),
       WavefrontSize(0), CFALUBug(false), LocalMemorySize(0),
       EnableVGPRSpilling(false), SGPRInitBug(false), IsGCN(false),
       GCN1Encoding(false), GCN3Encoding(false), CIInsts(false), LDSBankCount(0),
-      IsaVersion(ISAVersion0_0_0),
+      IsaVersion(ISAVersion0_0_0), EnableHugeScratchBuffer(false),
       FrameLowering(TargetFrameLowering::StackGrowsUp,
                     64 * 16, // Maximum stack alignment (long16)
                     0),
       InstrItins(getInstrItineraryForCPU(GPU)), TargetTriple(TT) {
 
   initializeSubtargetDependencies(TT, GPU, FS);
 
   if (getGeneration() <= AMDGPUSubtarget::NORTHERN_ISLANDS) {

llvm/trunk/lib/Target/AMDGPU/SIISelLowering.cpp

for (SDNode::use_iterator I = Parent->use_begin(), E = Parent->use_end();
     if (I->getOpcode() == Opcode)
       return *I;
   }
   return nullptr;
 }
 
 SDValue SITargetLowering::LowerFrameIndex(SDValue Op, SelectionDAG &DAG) const {
 
+  SDLoc SL(Op);
   FrameIndexSDNode *FINode = cast<FrameIndexSDNode>(Op);
   unsigned FrameIndex = FINode->getIndex();
 
-  return DAG.getTargetFrameIndex(FrameIndex, MVT::i32);
+  // A FrameIndex node represents a 32-bit offset into scratch memory.  If
+  // the high bit of a frame index offset were to be set, this would mean
+  // that it represented an offset of ~2GB * 64 = ~128GB from the start of the
+  // scratch buffer, with 64 being the number of threads per wave.
+  //
+  // If we know the machine uses less than 128GB of scratch, then we can
+  // amrk the high bit of the FrameIndex node as known zero,
+  // which is important, because it means in most situations we can
+  // prove that values derived from FrameIndex nodes are non-negative.
+  // This enables us to take advantage of more addressing modes when
+  // accessing scratch buffers, since for scratch reads/writes, the register
+  // offset must always be positive.
+
+  SDValue TFI = DAG.getTargetFrameIndex(FrameIndex, MVT::i32);
+  if (Subtarget->enableHugeScratchBuffer())
+    return TFI;
+
+  return DAG.getNode(ISD::AssertZext, SL, MVT::i32, TFI,
+                    DAG.getValueType(EVT::getIntegerVT(*DAG.getContext(), 31)));
 }
 
 /// This transforms the control flow intrinsics to get the branch destination as
 /// last parameter, also switches branch target with BR if the need arise
 SDValue SITargetLowering::LowerBRCOND(SDValue BRCOND,
                                       SelectionDAG &DAG) const {
 
   SDLoc DL(BRCOND);
for (unsigned i = 0, Idx = AMDGPU::sub0; i < 4; ++i) {
     default: break;
     case AMDGPU::sub0: Idx = AMDGPU::sub1; break;
     case AMDGPU::sub1: Idx = AMDGPU::sub2; break;
     case AMDGPU::sub2: Idx = AMDGPU::sub3; break;
     }
   }
 }
 
+static bool isFrameIndexOp(SDValue Op) {
+  if (Op.getOpcode() == ISD::AssertZext)
+    Op = Op.getOperand(0);
+
+  return isa<FrameIndexSDNode>(Op);
+}
+
 /// \brief Legalize target independent instructions (e.g. INSERT_SUBREG)
 /// with frame index operands.
 /// LLVM assumes that inputs are to these instructions are registers.
 void SITargetLowering::legalizeTargetIndependentNode(SDNode *Node,
                                                      SelectionDAG &DAG) const {
 
   SmallVector<SDValue, 8> Ops;
   for (unsigned i = 0; i < Node->getNumOperands(); ++i) {
-    if (!isa<FrameIndexSDNode>(Node->getOperand(i))) {
+    if (!isFrameIndexOp(Node->getOperand(i))) {
       Ops.push_back(Node->getOperand(i));
       continue;
     }
 
     SDLoc DL(Node);
     Ops.push_back(SDValue(DAG.getMachineNode(AMDGPU::S_MOV_B32, DL,
                                      Node->getOperand(i).getValueType(),
                                      Node->getOperand(i)), 0));