[AMDGPU] Improve FLAT scratch detection

[rampitec]

Summary

We were useing too broad check for isFLATScratch() which also
includes FLAT global.

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

517

What's wrong with another bit? Don't we already have it for global?

[rampitec]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

517

We do not have it for global, otherwise I would just check it is unset.

I just see that we may soon exhaust the bitfield width.

[rampitec]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

517

We are actually using 56 bits out of 64. If I would use a bit for every single feature in every singe GPU we would go over the roof a long time ago (we actually did, so I had to remove some of the bits and use some non-obvious ways to replace it).

[rampitec]

Check for isSegmentSpecificFLAT() first to keep it O(1) most of the times.

[foad]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

580

If getFlatScratchInst is just a table lookup, is there any need to do the isSegmentSpecificFLAT test first?

[rampitec]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

580

I have reetored it to have O(1) in most cases, table lookup will only run if it is already known segmented flat. It is purely optimization.

[foad]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

580

I see. I thought the table lookup was a direct O(1) lookup. I didn't realise it is a binary chop.

[rampitec]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

580

Right, it is O(log(N)). N is small but still.

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

580

I think using a bit here is fine. We're not out, and there are other bits we could prune out that are less important

[rampitec]

Changed to use TSFlags bit.

llvm/lib/Target/AMDGPU/FLATInstructions.td

class FLAT_Global_Store_AddTid_Pseudo <string opName, RegisterClass vdataClass,
   let has_saddr = 1;
   let enabled_saddr = 1;
   let maybeAtomic = 1;
 }
 
 class FlatScratchInst <string sv_op, string mode> {
   string SVOp = sv_op;
   string Mode = mode;
+  string Op = sv_op#!if(!eq(mode, "SS"), "_SADDR",
+                    !if(!eq(mode, "SV"), "",
+                        "_"#mode));
+  bit IsFlatScratch = 1;
 }
 
 class FLAT_Scratch_Load_Pseudo <string opName, RegisterClass regClass,
   bit HasTiedOutput = 0,
   bit EnableSaddr = 0,
   bit EnableVaddr = !not(EnableSaddr)>
   : FLAT_Pseudo<
   opName,

llvm/lib/Target/AMDGPU/SIInstrInfo.h

 
 class APInt;
 class MachineDominatorTree;
 class MachineRegisterInfo;
 class RegScavenger;
 class GCNSubtarget;
 class TargetRegisterClass;
 
+namespace AMDGPU {
+
+  LLVM_READONLY
+  int getVOPe64(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getVOPe32(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getSDWAOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getDPPOp32(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getBasicFromSDWAOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getCommuteRev(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getCommuteOrig(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getAddr64Inst(uint16_t Opcode);
+
+  /// Check if \p Opcode is an Addr64 opcode.
+  ///
+  /// \returns \p Opcode if it is an Addr64 opcode, otherwise -1.
+  LLVM_READONLY
+  int getIfAddr64Inst(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getMUBUFNoLdsInst(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getAtomicRetOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getAtomicNoRetOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getSOPKOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getGlobalSaddrOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getVCMPXNoSDstOp(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getFlatScratchInstSTfromSS(uint16_t Opcode);
+
+  LLVM_READONLY
+  int getFlatScratchInst(uint16_t Opcode);
+
+  const uint64_t RSRC_DATA_FORMAT = 0xf00000000000LL;
+  const uint64_t RSRC_ELEMENT_SIZE_SHIFT = (32 + 19);
+  const uint64_t RSRC_INDEX_STRIDE_SHIFT = (32 + 21);
+  const uint64_t RSRC_TID_ENABLE = UINT64_C(1) << (32 + 23);
+
+} // end namespace AMDGPU
+
 class SIInstrInfo final : public AMDGPUGenInstrInfo {
 private:
   const SIRegisterInfo RI;
   const GCNSubtarget &ST;
   TargetSchedModel SchedModel;
 
   // The inverse predicate should have the negative value.
   enum BranchPredicate {
static bool isSegmentSpecificFLAT(const MachineInstr &MI) {
     return (Flags & SIInstrFlags::FLAT) && !(Flags & SIInstrFlags::LGKM_CNT);
   }
 
   bool isSegmentSpecificFLAT(uint16_t Opcode) const {
     auto Flags = get(Opcode).TSFlags;
     return (Flags & SIInstrFlags::FLAT) && !(Flags & SIInstrFlags::LGKM_CNT);
   }
 
-  // FIXME: Make this more precise
   static bool isFLATScratch(const MachineInstr &MI) {
-    return isSegmentSpecificFLAT(MI);
+    return isSegmentSpecificFLAT(MI) &&
+           AMDGPU::getFlatScratchInst(MI.getOpcode()) >= 0;
   }
 
   bool isFLATScratch(uint16_t Opcode) const {
-    return isSegmentSpecificFLAT(Opcode);
+    return isSegmentSpecificFLAT(Opcode) &&
+           AMDGPU::getFlatScratchInst(Opcode) >= 0;
   }
 
   // Any FLAT encoded instruction, including global_* and scratch_*.
   bool isFLAT(uint16_t Opcode) const {
     return get(Opcode).TSFlags & SIInstrFlags::FLAT;
   }
 
   static bool isEXP(const MachineInstr &MI) {

 
 /// \brief Return false if EXEC is not changed between the def of \p VReg at \p
 /// DefMI and all its uses. Should be run on SSA. Currently does not attempt to
 /// track between blocks.
 bool execMayBeModifiedBeforeAnyUse(const MachineRegisterInfo &MRI,
                                    Register VReg,
                                    const MachineInstr &DefMI);
 
-namespace AMDGPU {
-
-  LLVM_READONLY
-  int getVOPe64(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getVOPe32(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getSDWAOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getDPPOp32(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getBasicFromSDWAOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getCommuteRev(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getCommuteOrig(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getAddr64Inst(uint16_t Opcode);
-
-  /// Check if \p Opcode is an Addr64 opcode.
-  ///
-  /// \returns \p Opcode if it is an Addr64 opcode, otherwise -1.
-  LLVM_READONLY
-  int getIfAddr64Inst(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getMUBUFNoLdsInst(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getAtomicRetOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getAtomicNoRetOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getSOPKOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getGlobalSaddrOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getVCMPXNoSDstOp(uint16_t Opcode);
-
-  LLVM_READONLY
-  int getFlatScratchInstSTfromSS(uint16_t Opcode);
-
-  const uint64_t RSRC_DATA_FORMAT = 0xf00000000000LL;
-  const uint64_t RSRC_ELEMENT_SIZE_SHIFT = (32 + 19);
-  const uint64_t RSRC_INDEX_STRIDE_SHIFT = (32 + 21);
-  const uint64_t RSRC_TID_ENABLE = UINT64_C(1) << (32 + 23);
-
-} // end namespace AMDGPU
-
 namespace SI {
 namespace KernelInputOffsets {
 
 /// Offsets in bytes from the start of the input buffer
 enum Offsets {
   NGROUPS_X = 0,
   NGROUPS_Y = 4,
   NGROUPS_Z = 8,

llvm/lib/Target/AMDGPU/SIInstrInfo.td

 def getFlatScratchInstSTfromSS : InstrMapping {
   let FilterClass = "FlatScratchInst";
   let RowFields = ["SVOp"];
   let ColFields = ["Mode"];
   let KeyCol = ["SS"];
   let ValueCols = [["ST"]];
 }
 
+// Maps a flat scratch instruction to itself.
+// Used to check if an instruction is a flat scratch.
+def getFlatScratchInst : InstrMapping {
+  let FilterClass = "FlatScratchInst";
+  let RowFields = ["Op"];
+  let ColFields = ["IsFlatScratch"];
+  let KeyCol = ["1"];
+  let ValueCols = [["1"]];
+}
 
 include "SIInstructions.td"
 
 include "DSInstructions.td"
 include "MIMGInstructions.td"