[llvm] [WebAssembly] Fold unsigned offset (PR #145829)

[via llvm-commits]

https://github.com/badumbatish updated https://github.com/llvm/llvm-project/pull/145829

>From e1cd7019bde37cb3d6ad3540f4357717a0077c0b Mon Sep 17 00:00:00 2001
From: badumbatish <jjasmine at igalia.com>
Date: Wed, 25 Jun 2025 14:29:56 -0700
Subject: [PATCH 1/2] Precommit test for folding a load of unsigned offset

---
 .../CodeGen/WebAssembly/fold-load-unsigned.ll | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 llvm/test/CodeGen/WebAssembly/fold-load-unsigned.ll

diff --git a/llvm/test/CodeGen/WebAssembly/fold-load-unsigned.ll b/llvm/test/CodeGen/WebAssembly/fold-load-unsigned.ll
new file mode 100644
index 0000000000000..18aadb6c0de80
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/fold-load-unsigned.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -O3 -wasm-keep-registers | FileCheck %s
+target triple = "wasm64"
+
+ at data = hidden local_unnamed_addr global [12 x i8] c"Hello world\00", align 1
+
+define hidden signext i8 @fold_load_unsigned(i32 noundef %idx) local_unnamed_addr {
+; CHECK-LABEL: fold_load_unsigned:
+; CHECK:         .functype fold_load_unsigned (i32) -> (i32)
+; CHECK-NEXT:  # %bb.0: # %entry
+; CHECK-NEXT:    local.get $push4=, 0
+; CHECK-NEXT:    i64.extend_i32_s $push0=, $pop4
+; CHECK-NEXT:    i64.const $push1=, data
+; CHECK-NEXT:    i64.add $push2=, $pop0, $pop1
+; CHECK-NEXT:    i32.load $push3=, 0($pop2):p2align=0
+; CHECK-NEXT:    # fallthrough-return
+entry:
+  %arrayidx = getelementptr inbounds nuw [12 x i8], ptr @data, i32 0, i32 %idx
+  %0 = load i8, ptr %arrayidx, align 1
+  ret i8 %0
+}

>From 9db18a17c7c7841ea320383966b0b92a0fa2de10 Mon Sep 17 00:00:00 2001
From: badumbatish <jjasmine at igalia.com>
Date: Wed, 25 Jun 2025 18:58:55 -0700
Subject: [PATCH 2/2] Attempt to fold no unsigned wrap offset load

---
 .../WebAssembly/WebAssemblyISelLowering.cpp   | 47 +++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
index ec77154d17caa..f099df8b815de 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
@@ -19,11 +19,14 @@
 #include "WebAssemblyTargetMachine.h"
 #include "WebAssemblyUtilities.h"
 #include "llvm/CodeGen/CallingConvLower.h"
+#include "llvm/CodeGen/ISDOpcodes.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/SDPatternMatch.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/IR/DiagnosticInfo.h"
@@ -32,10 +35,13 @@
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsWebAssembly.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetOptions.h"
+#include <iostream>
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-lower"
@@ -178,6 +184,9 @@ WebAssemblyTargetLowering::WebAssemblyTargetLowering(
       for (auto T : {MVT::i32, MVT::i64})
         setOperationAction(Op, T, Custom);
 
+  // Folding adds with global addresses into load offset
+  setTargetDAGCombine(ISD::LOAD);
+
   // SIMD-specific configuration
   if (Subtarget->hasSIMD128()) {
 
@@ -3212,6 +3221,42 @@ static SDValue performTruncateCombine(SDNode *N,
   return truncateVectorWithNARROW(OutVT, In, DL, DAG);
 }
 
+static SDValue performLoadCombine(SDNode *N, SelectionDAG &DAG) {
+  using namespace SDPatternMatch;
+
+  SDLoc DL(N);
+  SDValue Add = N->getOperand(1);
+  SDValue Offset = Add.getOperand(0), Global = Add.getOperand(1);
+
+  // INFO: we're folding this
+  //     t4: i64 = sign_extend t2
+  //     t5: i64 = add nuw t4, GlobalAddress:i64<ptr @data> 0
+  //     t8: i8,ch = load<(load (s8) from %ir.arrayidx)> t0, t5, undef:i64
+  // into this:
+  //     t4: i64 = sign_extend t2
+  //     t8: i8,ch = load<(load (s8) from %ir.arrayidx)> t0,
+  //     GlobalAddress:i64<ptr @data>, t4
+
+  if (Add->getFlags().hasNoUnsignedWrap()) {
+    if (isa<GlobalAddressSDNode>(Global)) {
+      auto *LD = dyn_cast<LoadSDNode>(N);
+      // SDValue SelectionDAG::getLoad(ISD::MemIndexedMode AM, ISD::LoadExtType
+      // ExtType,
+      //                               EVT VT, const SDLoc &dl, SDValue Chain,
+      //                               SDValue Ptr, SDValue Offset, EVT MemVT,
+      //                               MachineMemOperand *MMO) {
+
+      // INFO: Use the above signature for now, the previous getLoad signatuer
+      // automatically creates undef
+      return DAG.getLoad(ISD::MemIndexedMode::PRE_INC,
+                         ISD::LoadExtType::NON_EXTLOAD, LD->getValueType(0), DL,
+                         LD->getChain(), Global, Offset, LD->getValueType(0),
+                         LD->getMemOperand());
+    }
+  }
+  return SDValue();
+}
+
 static SDValue performBitcastCombine(SDNode *N,
                                      TargetLowering::DAGCombinerInfo &DCI) {
   auto &DAG = DCI.DAG;
@@ -3408,6 +3453,8 @@ WebAssemblyTargetLowering::PerformDAGCombine(SDNode *N,
   switch (N->getOpcode()) {
   default:
     return SDValue();
+  case ISD::LOAD:
+    return performLoadCombine(N, DCI.DAG);
   case ISD::BITCAST:
     return performBitcastCombine(N, DCI);
   case ISD::SETCC: