[AMDGPU] Add module flag to control TBUFFER combining

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

@@ -62,6 +62,7 @@
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "SIDefines.h"
+#include "SIMachineFunctionInfo.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/InitializePasses.h"
@@ -2507,6 +2508,14 @@ SILoadStoreOptimizer::collectMergeableInsts(
         LLVM_DEBUG(dbgs() << "Skip tbuffer with unknown format: " << MI);
         continue;
       }
+
+      const MachineFunction *MF = MI.getParent()->getParent();
+      const auto *MFI = MF->getInfo<SIMachineFunctionInfo>();
+      if (!MFI->isRelaxedTBufferOOBMode()) {
+        LLVM_DEBUG(
+            dbgs() << "Skip tbuffer combine: relaxed mode not enabled\n");
+        continue;
+      }
     }
 
     CombineInfo CI;

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

@@ -197,6 +197,12 @@ SIMachineFunctionInfo::SIMachineFunctionInfo(const Function &F,
   }
 
   ClusterDims = AMDGPU::ClusterDimsAttr::get(F);
+
+  // Enable relaxed TBUFFER OOB mode if amdgpu.oob.mode has bit 0x2 set.
+  if (const auto *CI = mdconst::extract_or_null<ConstantInt>(
+          F.getParent()->getModuleFlag("amdgpu.oob.mode"));
+      CI && (CI->getZExtValue() & 0x2))
+    setRelaxedTBufferOOBMode(true);
 }
 
 MachineFunctionInfo *SIMachineFunctionInfo::clone(
@@ -744,6 +750,7 @@ yaml::SIMachineFunctionInfo::SIMachineFunctionInfo(
       MaxMemoryClusterDWords(MFI.getMaxMemoryClusterDWords()),
       Mode(MFI.getMode()), HasInitWholeWave(MFI.hasInitWholeWave()),
       IsWholeWaveFunction(MFI.isWholeWaveFunction()),
+      RelaxedTBufferOOBMode(MFI.isRelaxedTBufferOOBMode()),
       DynamicVGPRBlockSize(MFI.getDynamicVGPRBlockSize()),
       ScratchReservedForDynamicVGPRs(MFI.getScratchReservedForDynamicVGPRs()) {
   for (Register Reg : MFI.getSGPRSpillPhysVGPRs())
@@ -793,6 +800,7 @@ bool SIMachineFunctionInfo::initializeBaseYamlFields(
   BytesInStackArgArea = YamlMFI.BytesInStackArgArea;
   ReturnsVoid = YamlMFI.ReturnsVoid;
   IsWholeWaveFunction = YamlMFI.IsWholeWaveFunction;
+  RelaxedTBufferOOBMode = YamlMFI.RelaxedTBufferOOBMode;
 
   if (YamlMFI.ScavengeFI) {
     auto FIOrErr = YamlMFI.ScavengeFI->getFI(MF.getFrameInfo());

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

@@ -301,6 +301,7 @@ struct SIMachineFunctionInfo final : public yaml::MachineFunctionInfo {
 
   bool HasInitWholeWave = false;
   bool IsWholeWaveFunction = false;
+  bool RelaxedTBufferOOBMode = false;
 
   unsigned DynamicVGPRBlockSize = 0;
   unsigned ScratchReservedForDynamicVGPRs = 0;
@@ -362,6 +363,8 @@ template <> struct MappingTraits<SIMachineFunctionInfo> {
     YamlIO.mapOptional("scratchReservedForDynamicVGPRs",
                        MFI.ScratchReservedForDynamicVGPRs, 0);
     YamlIO.mapOptional("isWholeWaveFunction", MFI.IsWholeWaveFunction, false);
+    YamlIO.mapOptional("RelaxedTBufferOOBMode", MFI.RelaxedTBufferOOBMode,
+                       false);
   }
 };
 
@@ -528,6 +531,9 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
   // scheduler stage.
   unsigned MaxMemoryClusterDWords = DefaultMemoryClusterDWordsLimit;
 
+  // Enable relaxed TBUFFER out-of-bounds mode. Default is false.
+  bool RelaxedTBufferOOBMode = false;
+
   MCPhysReg getNextUserSGPR() const;
 
   MCPhysReg getNextSystemSGPR() const;
@@ -1212,6 +1218,11 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
   unsigned getMaxNumWorkGroupsZ() const { return MaxNumWorkGroups[2]; }
 
   AMDGPU::ClusterDimsAttr getClusterDims() const { return ClusterDims; }
+
+  bool isRelaxedTBufferOOBMode() const { return RelaxedTBufferOOBMode; }
+  void setRelaxedTBufferOOBMode(bool Enabled) {
+    RelaxedTBufferOOBMode = Enabled;
+  }
 };
 
 } // end namespace llvm