[llvm] [AMDGPU] Folding imm offset in more cases for scratch access (PR #70634)

[via llvm-commits]

================
@@ -1146,26 +1146,36 @@ bool AMDGPUDAGToDAGISel::isDSOffset2Legal(SDValue Base, unsigned Offset0,
   return CurDAG->SignBitIsZero(Base);
 }
 
-bool AMDGPUDAGToDAGISel::isFlatScratchBaseLegal(SDValue Addr, SDValue Base,
+// This is used to check whether the address of scratch_load/store in the
+// form of `base1 + (base2) + immediate offset` is legal with respect to the
+// hardware's requirement that the SGPR/VGPR address offset in the flat scratch
+// instruction should be unsigned.
+bool AMDGPUDAGToDAGISel::isFlatScratchBaseLegal(SDValue Addr, SDValue Base1,
+                                                SDValue Base2,
                                                 uint64_t FlatVariant) const {
   if (FlatVariant != SIInstrFlags::FlatScratch)
     return true;
 
-  if (Addr.getOpcode() == ISD::ADD) {
+  assert(Base1.getNode());
+  if (Addr.getOpcode() == ISD::ADD && !Base2.getNode()) {
+    assert(isa<ConstantSDNode>(Addr->getOperand(1)));
     // For `nuw` addition, we should not have negative base address.
     if (Addr->getFlags().hasNoUnsignedWrap())
       return true;
 
-    auto *RHS = dyn_cast<ConstantSDNode>(Addr.getOperand(1));
-    // If the immediate offset is negative, we should not have the base being
-    // negative as well.
-    if (RHS && RHS->getSExtValue() < 0)
+    // If the immediate offset is negative, the base address cannot also be
+    // negative.
+    if (CurDAG->SignBitIsZero(Addr.getOperand(1)))
----------------
ruiling wrote:

Should be `if (!CurDAG->SignBitIsZero(Addr.getOperand(1)))` here. The issue also applies to GlobalISel. Sorry for that.

https://github.com/llvm/llvm-project/pull/70634