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

[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 in theory we could get positive value by adding two
+  // negative values (like in the case: 0x80000010 + 0x80000020 = 0x30), but
+  // isLegalFLATOffset() will help filter out such values later.
----------------
ruiling wrote:

But it maybe not as fragile as you think. Do you see a situation this might be broken? Or do you think adding a check here would be helpful? like checking the `Imm > -0x40000000` which should be enough to be helpful, then base address should be non-negative. Otherwise, the result address would be either negative or too large that would fall out of the region a thread can access.

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