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

[via llvm-commits]

================
@@ -1146,10 +1146,23 @@ bool AMDGPUDAGToDAGISel::isDSOffset2Legal(SDValue Base, unsigned Offset0,
   return CurDAG->SignBitIsZero(Base);
 }
 
-bool AMDGPUDAGToDAGISel::isFlatScratchBaseLegal(SDValue Base,
+bool AMDGPUDAGToDAGISel::isFlatScratchBaseLegal(SDValue Addr, SDValue Base,
                                                 uint64_t FlatVariant) const {
   if (FlatVariant != SIInstrFlags::FlatScratch)
     return true;
+
+  if (Addr.getOpcode() == ISD::ADD) {
+    // For `nuw` addition, we should not have negative base address.
+    if (Addr->getFlags().hasNoUnsignedWrap())
+      return true;
+
+    auto *RHS = dyn_cast<ConstantSDNode>(Addr.getOperand(1));
----------------
ruiling wrote:

I have made the suggested change locally, but it also means that we need an assert that `RHS` should be immediate as well. Because so far our hardware only allow negative offset as immediate offset. I feel that checking against negative immediate would be more clear here. But in case you think generalize to `SignBitIsZero` is preferred, I am not quite against it, but assertion that `RHS` is immediate should be here as well.

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