[AMDGPU][NFC] Refine matching SMRD offsets.

[kosarev]

Summary

Tell the matcher what we are looking for instead of matching everything
and then discarding the result if doesn't fit.

[foad]

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

1889

Can you add a comment and document the arguments?

1971

Can you add a comment and document the arguments?

[kosarev]

Added comments.

[arsenm]

Does GlobalISel need the equivalent change?

[kosarev]

In D128171#3601575, @arsenm wrote:

Does GlobalISel need the equivalent change?

Nope, it seems to have its own facilities to match such operands.

[kosarev]

Ping.

[foad]

I don't understand all the selection logic here, but your changes look reasonable, and I especially like the API change of not passing in a bool &Imm parameter.

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.h

bool SelectGlobalSAddr(SDNode *N, SDValue Addr, SDValue &SAddr,
                          SDValue &VOffset, SDValue &Offset) const;
   bool SelectScratchSAddr(SDNode *N, SDValue Addr, SDValue &SAddr,
                           SDValue &Offset) const;
   bool checkFlatScratchSVSSwizzleBug(SDValue VAddr, SDValue SAddr,
                                      uint64_t ImmOffset) const;
   bool SelectScratchSVAddr(SDNode *N, SDValue Addr, SDValue &VAddr,
                            SDValue &SAddr, SDValue &Offset) const;
 
-  bool SelectSMRDOffset(SDValue ByteOffsetNode, SDValue &Offset,
-                        bool &Imm) const;
+  bool SelectSMRDOffset(SDValue ByteOffsetNode, SDValue &Offset, bool Imm,
+                        bool Imm32Only) const;
   SDValue Expand32BitAddress(SDValue Addr) const;
-  bool SelectSMRD(SDValue Addr, SDValue &SBase, SDValue &Offset,
-                  bool &Imm) const;
+  bool SelectSMRD(SDValue Addr, SDValue &SBase, SDValue &Offset, bool Imm,
+                  bool Imm32Only = false) const;
   bool SelectSMRDImm(SDValue Addr, SDValue &SBase, SDValue &Offset) const;
   bool SelectSMRDImm32(SDValue Addr, SDValue &SBase, SDValue &Offset) const;
   bool SelectSMRDSgpr(SDValue Addr, SDValue &SBase, SDValue &Offset) const;
   bool SelectSMRDBufferImm(SDValue Addr, SDValue &Offset) const;
   bool SelectSMRDBufferImm32(SDValue Addr, SDValue &Offset) const;
   bool SelectMOVRELOffset(SDValue Index, SDValue &Base, SDValue &Offset) const;
 
   bool SelectVOP3Mods_NNaN(SDValue In, SDValue &Src, SDValue &SrcMods) const;

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

bool AMDGPUDAGToDAGISel::SelectScratchSVAddr(SDNode *N, SDValue Addr,
 
   if (checkFlatScratchSVSSwizzleBug(VAddr, SAddr, ImmOffset))
     return false;
   SAddr = SelectSAddrFI(CurDAG, SAddr);
   Offset = CurDAG->getTargetConstant(ImmOffset, SDLoc(), MVT::i16);
   return true;
 }
 
+// Match an immediate (if Imm is true) or an SGPR (if Imm is false)
+// offset. If Imm32Only is true, match only 32-bit immediate offsets
+// available on CI.
 bool AMDGPUDAGToDAGISel::SelectSMRDOffset(SDValue ByteOffsetNode,
-                                          SDValue &Offset, bool &Imm) const {
+                                          SDValue &Offset, bool Imm,
+                                          bool Imm32Only) const {
   ConstantSDNode *C = dyn_cast<ConstantSDNode>(ByteOffsetNode);
   if (!C) {
+    if (Imm)
+      return false;
     if (ByteOffsetNode.getValueType().isScalarInteger() &&
         ByteOffsetNode.getValueType().getSizeInBits() == 32) {
       Offset = ByteOffsetNode;
-      Imm = false;
       return true;
     }
     if (ByteOffsetNode.getOpcode() == ISD::ZERO_EXTEND) {
       if (ByteOffsetNode.getOperand(0).getValueType().getSizeInBits() == 32) {
         Offset = ByteOffsetNode.getOperand(0);
-        Imm = false;
         return true;
       }
     }
     return false;
   }
 
   SDLoc SL(ByteOffsetNode);
   // GFX9 and GFX10 have signed byte immediate offsets.
   int64_t ByteOffset = C->getSExtValue();
   Optional<int64_t> EncodedOffset =
       AMDGPU::getSMRDEncodedOffset(*Subtarget, ByteOffset, false);
-  if (EncodedOffset) {
+  if (EncodedOffset && Imm && !Imm32Only) {
     Offset = CurDAG->getTargetConstant(*EncodedOffset, SL, MVT::i32);
-    Imm = true;
     return true;
   }
 
   // SGPR and literal offsets are unsigned.
   if (ByteOffset < 0)
     return false;
 
   EncodedOffset = AMDGPU::getSMRDEncodedLiteralOffset32(*Subtarget, ByteOffset);
-  if (EncodedOffset) {
+  if (EncodedOffset && Imm32Only) {
     Offset = CurDAG->getTargetConstant(*EncodedOffset, SL, MVT::i32);
     return true;
   }
 
   if (!isUInt<32>(ByteOffset) && !isInt<32>(ByteOffset))
     return false;
 
+  if (!Imm) {
     SDValue C32Bit = CurDAG->getTargetConstant(ByteOffset, SL, MVT::i32);
     Offset = SDValue(
         CurDAG->getMachineNode(AMDGPU::S_MOV_B32, SL, MVT::i32, C32Bit), 0);
-
     return true;
   }
 
+  return false;
+}
+
 SDValue AMDGPUDAGToDAGISel::Expand32BitAddress(SDValue Addr) const {
   if (Addr.getValueType() != MVT::i32)
     return Addr;
 
   // Zero-extend a 32-bit address.
   SDLoc SL(Addr);
 
   const MachineFunction &MF = CurDAG->getMachineFunction();
SDValue(CurDAG->getMachineNode(AMDGPU::S_MOV_B32, SL, MVT::i32, AddrHi),
             0),
     CurDAG->getTargetConstant(AMDGPU::sub1, SL, MVT::i32),
   };
 
   return SDValue(CurDAG->getMachineNode(AMDGPU::REG_SEQUENCE, SL, MVT::i64,
                                         Ops), 0);
 }
 
+// Match a base and an immediate (if Imm is true) or an SGPR
+// (if Imm is false) offset. If Imm32Only is true, match only 32-bit
+// immediate offsets available on CI.
 bool AMDGPUDAGToDAGISel::SelectSMRD(SDValue Addr, SDValue &SBase,
-                                     SDValue &Offset, bool &Imm) const {
+                                    SDValue &Offset, bool Imm,
+                                    bool Imm32Only) const {
   SDLoc SL(Addr);
 
   // A 32-bit (address + offset) should not cause unsigned 32-bit integer
   // wraparound, because s_load instructions perform the addition in 64 bits.
   if ((Addr.getValueType() != MVT::i32 ||
        Addr->getFlags().hasNoUnsignedWrap())) {
     SDValue N0, N1;
     // Extract the base and offset if possible.
     if (CurDAG->isBaseWithConstantOffset(Addr) ||
         Addr.getOpcode() == ISD::ADD) {
       N0 = Addr.getOperand(0);
       N1 = Addr.getOperand(1);
     } else if (getBaseWithOffsetUsingSplitOR(*CurDAG, Addr, N0, N1)) {
       assert(N0 && N1 && isa<ConstantSDNode>(N1));
     }
     if (N0 && N1) {
-      if (SelectSMRDOffset(N1, Offset, Imm)) {
+      if (SelectSMRDOffset(N1, Offset, Imm, Imm32Only)) {
         SBase = Expand32BitAddress(N0);
         return true;
       }
     }
+    return false;
   }
+  if (!Imm)
+    return false;
   SBase = Expand32BitAddress(Addr);
   Offset = CurDAG->getTargetConstant(0, SL, MVT::i32);
-  Imm = true;
   return true;
 }
 
 bool AMDGPUDAGToDAGISel::SelectSMRDImm(SDValue Addr, SDValue &SBase,
                                        SDValue &Offset) const {
-  bool Imm = false;
-  return SelectSMRD(Addr, SBase, Offset, Imm) && Imm;
+  return SelectSMRD(Addr, SBase, Offset, /* Imm */ true);
 }
 
 bool AMDGPUDAGToDAGISel::SelectSMRDImm32(SDValue Addr, SDValue &SBase,
                                          SDValue &Offset) const {
-
   assert(Subtarget->getGeneration() == AMDGPUSubtarget::SEA_ISLANDS);
-
-  bool Imm = false;
-  if (!SelectSMRD(Addr, SBase, Offset, Imm))
-    return false;
-
-  return !Imm && isa<ConstantSDNode>(Offset);
+  return SelectSMRD(Addr, SBase, Offset, /* Imm */ true, /* Imm32Only */ true);
 }
 
 bool AMDGPUDAGToDAGISel::SelectSMRDSgpr(SDValue Addr, SDValue &SBase,
                                         SDValue &Offset) const {
-  bool Imm = false;
-  return SelectSMRD(Addr, SBase, Offset, Imm) && !Imm &&
-         !isa<ConstantSDNode>(Offset);
+  return SelectSMRD(Addr, SBase, Offset, /* Imm */ false);
 }
 
 bool AMDGPUDAGToDAGISel::SelectSMRDBufferImm(SDValue Addr,
                                              SDValue &Offset) const {
   if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(Addr)) {
     // The immediate offset for S_BUFFER instructions is unsigned.
     if (auto Imm =
             AMDGPU::getSMRDEncodedOffset(*Subtarget, C->getZExtValue(), true)) {