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

[Matt Arsenault via llvm-commits]

================
@@ -1146,13 +1146,57 @@ bool AMDGPUDAGToDAGISel::isDSOffset2Legal(SDValue Base, unsigned Offset0,
   return CurDAG->SignBitIsZero(Base);
 }
 
-bool AMDGPUDAGToDAGISel::isFlatScratchBaseLegal(SDValue Base,
-                                                uint64_t FlatVariant) const {
+// Whether we can infer the operands are non-negative if the result is
+// non-negative.
+static bool canInferNonNegativeOperands(SDValue Addr) {
+  return (Addr.getOpcode() == ISD::ADD &&
+          Addr->getFlags().hasNoUnsignedWrap()) ||
+         Addr->getOpcode() == ISD::OR;
+}
+
+// Check that the address value of flat scratch load/store being put into
+// SGPR/VGPR is legal (i.e. unsigned per hardware requirement). When \p
+// CheckBothOperands is set, we will check both operands in the instruction,
+// otherwise only check the first source operand.
+bool AMDGPUDAGToDAGISel::isFlatScratchBaseLegal(SDValue Addr,
+                                                uint64_t FlatVariant,
+                                                bool CheckBothOperands) 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 CurDAG->SignBitIsZero(Base);
+
+  if (canInferNonNegativeOperands(Addr))
+    return true;
+
+  auto LHS = Addr.getOperand(0);
+  auto RHS = Addr.getOperand(1);
+  if (CheckBothOperands)
+    return CurDAG->SignBitIsZero(RHS) && CurDAG->SignBitIsZero(LHS);
+
+  // If the immediate offset is negative, the base address cannot also be
+  // negative. Although this is not strictly true in theory (in the case of
+  // wrap around), but given the range of legal immediate offset, this is true
+  // to get a final valid scratch address.
----------------
arsenm wrote:

I think the comment needs to be clearer about this, that such values are filtered out because we happen to know isLegalFLATOffset will also be checked at some point 

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