[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

525

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

[rampitec]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

525

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

525

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

525

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

[arsenm]

llvm/lib/Target/AMDGPU/SIInstrInfo.h

525

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/AsmParser/AMDGPUAsmParser.cpp

Error(getFlatOffsetLoc(Operands),
           "flat offset modifier is not supported on this GPU");
     return false;
   }
 
   // Address offset is 12-bit signed for GFX10, 13-bit for GFX9.
   // For FLAT segment the offset must be positive;
   // MSB is ignored and forced to zero.
   unsigned OffsetSize = isGFX9() ? 13 : 12;
-  if (TSFlags & SIInstrFlags::IsNonFlatSeg) {
+  if (TSFlags & (SIInstrFlags::IsFlatGlobal | SIInstrFlags::IsFlatScratch)) {
     if (!isIntN(OffsetSize, Op.getImm())) {
       Error(getFlatOffsetLoc(Operands),
             isGFX9() ? "expected a 13-bit signed offset" :
                        "expected a 12-bit signed offset");
       return false;
     }
   } else {
     if (!isUIntN(OffsetSize - 1, Op.getImm())) {

llvm/lib/Target/AMDGPU/FLATInstructions.td

class FLAT_Pseudo<string opName, dag outs, dag ins,
   let Uses = !if(is_flat_global, [EXEC], [EXEC, FLAT_SCR]);
 
   // Internally, FLAT instruction are executed as both an LDS and a
   // Buffer instruction; so, they increment both VM_CNT and LGKM_CNT
   // and are not considered done until both have been decremented.
   let VM_CNT = 1;
   let LGKM_CNT = !not(!or(is_flat_global, is_flat_scratch));
 
-  let IsNonFlatSeg = !or(is_flat_global, is_flat_scratch);
+  let IsFlatGlobal = is_flat_global;
+
+  let IsFlatScratch = is_flat_scratch;
 }
 
 class FLAT_Real <bits<7> op, FLAT_Pseudo ps> :
   InstSI <ps.OutOperandList, ps.InOperandList, ps.Mnemonic # ps.AsmOperands, []>,
   Enc64 {
 
   let isPseudo = 0;
   let isCodeGenOnly = 0;

llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUInstPrinter.cpp

 void AMDGPUInstPrinter::printFlatOffset(const MCInst *MI, unsigned OpNo,
                                         const MCSubtargetInfo &STI,
                                         raw_ostream &O) {
   uint16_t Imm = MI->getOperand(OpNo).getImm();
   if (Imm != 0) {
     O << " offset:";
 
     const MCInstrDesc &Desc = MII.get(MI->getOpcode());
-    bool IsFlatSeg = !(Desc.TSFlags & SIInstrFlags::IsNonFlatSeg);
+    bool IsFlatSeg = !(Desc.TSFlags &
+        (SIInstrFlags::IsFlatGlobal | SIInstrFlags::IsFlatScratch));
 
     if (IsFlatSeg) { // Unsigned offset
       printU16ImmDecOperand(MI, OpNo, O);
     } else {         // Signed offset
       if (AMDGPU::isGFX10(STI)) {
         O << formatDec(SignExtend32<12>(MI->getOperand(OpNo).getImm()));
       } else {
         O << formatDec(SignExtend32<13>(MI->getOperand(OpNo).getImm()));

llvm/lib/Target/AMDGPU/SIDefines.h

// TODO: Should this be spilt into VOP3 a and b?
   ClampHi = UINT64_C(1) << 48,
 
   // Is a packed VOP3P instruction.
   IsPacked = UINT64_C(1) << 49,
 
   // Is a D16 buffer instruction.
   D16Buf = UINT64_C(1) << 50,
 
-  // FLAT instruction accesses FLAT_GLBL or FLAT_SCRATCH segment.
-  IsNonFlatSeg = UINT64_C(1) << 51,
+  // FLAT instruction accesses FLAT_GLBL segment.
+  IsFlatGlobal = UINT64_C(1) << 51,
 
   // Uses floating point double precision rounding mode
   FPDPRounding = UINT64_C(1) << 52,
 
   // Instruction is FP atomic.
   FPAtomic = UINT64_C(1) << 53,
 
   // Is a MFMA instruction.
   IsMAI = UINT64_C(1) << 54,
 
   // Is a DOT instruction.
-  IsDOT = UINT64_C(1) << 55
+  IsDOT = UINT64_C(1) << 55,
+
+  // FLAT instruction accesses FLAT_SCRATCH segment.
+  IsFlatScratch = UINT64_C(1) << 56
 };
 
 // v_cmp_class_* etc. use a 10-bit mask for what operation is checked.
 // The result is true if any of these tests are true.
 enum ClassFlags : unsigned {
   S_NAN = 1 << 0,        // Signaling NaN
   Q_NAN = 1 << 1,        // Quiet NaN
   N_INFINITY = 1 << 2,   // Negative infinity

llvm/lib/Target/AMDGPU/SIInstrFormats.td

class InstSI <dag outs, dag ins, string asm = "",
   field bit ClampHi = 0;
 
   // This bit indicates that this is a packed VOP3P instruction
   field bit IsPacked = 0;
 
   // This bit indicates that this is a D16 buffer instruction.
   field bit D16Buf = 0;
 
-  // This field indicates that FLAT instruction accesses FLAT_GLBL or
-  // FLAT_SCRATCH segment. Must be 0 for non-FLAT instructions.
-  field bit IsNonFlatSeg = 0;
+  // This field indicates that FLAT instruction accesses FLAT_GLBL segment.
+  // Must be 0 for non-FLAT instructions.
+  field bit IsFlatGlobal = 0;
 
   // Reads the mode register, usually for FP environment.
   field bit ReadsModeReg = 0;
 
   // This bit indicates that this uses the floating point double precision
   // rounding mode flags
   field bit FPDPRounding = 0;
 
   // Instruction is FP atomic.
   field bit FPAtomic = 0;
 
   // This bit indicates that this is one of MFMA instructions.
   field bit IsMAI = 0;
 
   // This bit indicates that this is one of DOT instructions.
   field bit IsDOT = 0;
 
+  // This field indicates that FLAT instruction accesses FLAT_SCRATCH segment.
+  // Must be 0 for non-FLAT instructions.
+  field bit IsFlatScratch = 0;
+
   // These need to be kept in sync with the enum in SIInstrFlags.
   let TSFlags{0} = SALU;
   let TSFlags{1} = VALU;
 
   let TSFlags{2} = SOP1;
   let TSFlags{3} = SOP2;
   let TSFlags{4} = SOPC;
   let TSFlags{5} = SOPK;
class InstSI <dag outs, dag ins, string asm = "",
   let TSFlags{46} = IntClamp;
   let TSFlags{47} = ClampLo;
   let TSFlags{48} = ClampHi;
 
   let TSFlags{49} = IsPacked;
 
   let TSFlags{50} = D16Buf;
 
-  let TSFlags{51} = IsNonFlatSeg;
+  let TSFlags{51} = IsFlatGlobal;
 
   let TSFlags{52} = FPDPRounding;
 
   let TSFlags{53} = FPAtomic;
 
   let TSFlags{54} = IsMAI;
 
   let TSFlags{55} = IsDOT;
 
+  let TSFlags{56} = IsFlatScratch;
+
   let SchedRW = [Write32Bit];
 
   field bits<1> DisableSIDecoder = 0;
   field bits<1> DisableVIDecoder = 0;
   field bits<1> DisableDecoder = 0;
 
   let isAsmParserOnly = !ne(DisableDecoder{0}, {0});
   let AsmVariantName = AMDGPUAsmVariants.Default;

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
   static bool isFLAT(const MachineInstr &MI) {
     return MI.getDesc().TSFlags & SIInstrFlags::FLAT;
   }
 
   // Is a FLAT encoded instruction which accesses a specific segment,
   // i.e. global_* or scratch_*.
   static bool isSegmentSpecificFLAT(const MachineInstr &MI) {
     auto Flags = MI.getDesc().TSFlags;
-    return (Flags & SIInstrFlags::FLAT) && !(Flags & SIInstrFlags::LGKM_CNT);
+    return Flags & (SIInstrFlags::IsFlatGlobal | SIInstrFlags::IsFlatScratch);
   }
 
   bool isSegmentSpecificFLAT(uint16_t Opcode) const {
     auto Flags = get(Opcode).TSFlags;
-    return (Flags & SIInstrFlags::FLAT) && !(Flags & SIInstrFlags::LGKM_CNT);
+    return Flags & (SIInstrFlags::IsFlatGlobal | SIInstrFlags::IsFlatScratch);
+  }
+
+  static bool isFLATGlobal(const MachineInstr &MI) {
+    return MI.getDesc().TSFlags & SIInstrFlags::IsFlatGlobal;
+  }
+
+  bool isFLATGlobal(uint16_t Opcode) const {
+    return get(Opcode).TSFlags & SIInstrFlags::IsFlatGlobal;
   }
 
-  // FIXME: Make this more precise
   static bool isFLATScratch(const MachineInstr &MI) {
-    return isSegmentSpecificFLAT(MI);
+    return MI.getDesc().TSFlags & SIInstrFlags::IsFlatScratch;
   }
 
   bool isFLATScratch(uint16_t Opcode) const {
-    return isSegmentSpecificFLAT(Opcode);
+    return get(Opcode).TSFlags & SIInstrFlags::IsFlatScratch;
   }
 
   // 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) {