[llvm-branch-commits] [llvm] [AMDGPU] Support one immediate folding for global load (PR #178608)

[via llvm-branch-commits]

================
@@ -2037,13 +2037,43 @@ bool AMDGPUDAGToDAGISel::SelectGlobalSAddr(SDNode *N, SDValue Addr,
     LHS = Addr.getOperand(0);
 
     if (!LHS->isDivergent()) {
-      // add (i64 sgpr), (*_extend (i32 vgpr))
       RHS = Addr.getOperand(1);
-      ScaleOffset = SelectScaleOffset(N, RHS, Subtarget->hasSignedGVSOffset());
+
       if (SDValue ExtRHS = matchExtFromI32orI32(
               RHS, Subtarget->hasSignedGVSOffset(), CurDAG)) {
+        // add (i64 sgpr), (*_extend (scale (i32 vgpr)))
         SAddr = LHS;
         VOffset = ExtRHS;
+        if (NeedIOffset && !ImmOffset &&
+            CurDAG->isBaseWithConstantOffset(ExtRHS)) {
+          // add (i64 sgpr), (zero_extend (add (scale (i32 vgpr)), (i32 imm)))
+          int64_t COffset =
+              cast<ConstantSDNode>(ExtRHS.getOperand(1))->getSExtValue();
+
+          if (TII->isLegalFLATOffset(COffset, AMDGPUAS::GLOBAL_ADDRESS,
+                                     SIInstrFlags::FlatGlobal)) {
+            // If the MSB of the first operand of the addition is known to be
+            // zero, which is followed by zext, we are sure overflow would not
+            // happen during addition.
+            if (RHS.getOpcode() == ISD::ZERO_EXTEND &&
+                CurDAG->SignBitIsZero(ExtRHS.getOperand(0))) {
----------------
ruiling wrote:

I agree there are different ways to prove it is safe to to do the addition folding. Here I choose to check the first operand of the addition is non-negative. With this, we would know no-overflow would happen for the addition. `zext nneg` just means the result of `voffset + imm` is non-negative. It is not strong enough to prove the overflow not happen. It maybe still useful combined with other checks. But I do not figure out how to do that properly. So I would like we leave that as future improvement.

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