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

[Matt Arsenault via llvm-commits]

================
@@ -4486,14 +4490,66 @@ bool AMDGPUInstructionSelector::isDSOffset2Legal(Register Base, int64_t Offset0,
   return KB->signBitIsZero(Base);
 }
 
+// Check that the address value of flat scratch load/store being put into
+// SGPR/VGPR is legal with respect to hardware's requirement that address in
+// SGPR/VGPR should be unsigned. When \p CheckTwoInstrs is set, we will check
+// against the instruction that defines \p Addr as well as the instruction that
+// defines the base address. When \p CheckTwoOperands is set, we will check both
+// operands (In case of two instructions, they are the operands from the
+// instruction that defines the base address).
 bool AMDGPUInstructionSelector::isFlatScratchBaseLegal(
-    Register Base, uint64_t FlatVariant) const {
+    Register Addr, bool CheckTwoInstrs, bool CheckTwoOperands,
+    uint64_t FlatVariant) const {
   if (FlatVariant != SIInstrFlags::FlatScratch)
     return true;
 
-  // When value in 32-bit Base can be negative calculate scratch offset using
-  // 32-bit add instruction, otherwise use Base(unsigned) + offset.
-  return KB->signBitIsZero(Base);
+  // Whether we can infer the operands are non-negative if the result is
+  // non-negative.
+  auto HasOnlyNonNegativeOperands = [](MachineInstr *Addr) -> bool {
+    return Addr->getOpcode() == TargetOpcode::G_OR ||
+           (Addr->getOpcode() == TargetOpcode::G_PTR_ADD &&
+            Addr->getFlag(MachineInstr::NoUWrap));
+  };
+
+  auto AddrDef = getDefSrcRegIgnoringCopies(Addr, *MRI);
+  auto *AddrMI = AddrDef->MI;
+  if (CheckTwoInstrs) {
+    Register Base = AddrMI->getOperand(1).getReg();
+    auto BaseDef = getDefSrcRegIgnoringCopies(Base, *MRI);
+    auto RHSOffset = getIConstantVRegValWithLookThrough(
+        AddrMI->getOperand(2).getReg(), *MRI);
+    assert(RHSOffset);
+
+    if (HasOnlyNonNegativeOperands(BaseDef->MI) &&
+        (HasOnlyNonNegativeOperands(AddrMI) ||
+         RHSOffset->Value.getSExtValue() < 0))
+      return true;
+
+    Register LHS = BaseDef->MI->getOperand(1).getReg();
+    Register RHS = BaseDef->MI->getOperand(2).getReg();
+    return KB->signBitIsZero(LHS) && KB->signBitIsZero(RHS);
----------------
arsenm wrote:

Probably should check RHS first, as that's canonically simpler

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