[llvm] [RISCV] Add a pass to remove ADDI by reassociating to fold into load/store address. (PR #127151)

[Luke Lau via llvm-commits]

================
@@ -0,0 +1,282 @@
+//===- RISCVFoldMemOffset.cpp - Fold ADDI into memory offsets ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===---------------------------------------------------------------------===//
+//
+// Look for ADDIs that can be removed by folding their immediate into later
+// load/store addresses. There may be other arithmetic instructions between the
+// addi and load/store that we need to reassociate through. If the final result
+// of the arithmetic is only used by load/store addresses, we can fold the
+// offset into the all the load/store as long as it doesn't create an offset
+// that is too large.
+//
+//===---------------------------------------------------------------------===//
+
+#include "RISCV.h"
+#include "RISCVSubtarget.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include <queue>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "riscv-fold-mem-offset"
+#define RISCV_FOLD_MEM_OFFSET_NAME "RISC-V Fold Memory Offset"
+
+namespace {
+
+class RISCVFoldMemOffset : public MachineFunctionPass {
+public:
+  static char ID;
+
+  RISCVFoldMemOffset() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  bool foldOffset(Register OrigReg, int64_t InitialOffset,
+                  const MachineRegisterInfo &MRI,
+                  DenseMap<MachineInstr *, int64_t> &FoldableInstrs);
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  StringRef getPassName() const override { return RISCV_FOLD_MEM_OFFSET_NAME; }
+};
+
+// Wrapper class around a std::optional to allow accumulation.
+class FoldableOffset {
+  std::optional<int64_t> Offset;
+
+public:
+  bool hasValue() const { return Offset.has_value(); }
+  int64_t getValue() const { return *Offset; }
+
+  FoldableOffset &operator=(int64_t RHS) {
+    Offset = RHS;
+    return *this;
+  }
+
+  FoldableOffset &operator+=(int64_t RHS) {
+    if (!Offset)
+      Offset = 0;
+    Offset = (uint64_t)*Offset + (uint64_t)RHS;
+    return *this;
+  }
+
+  int64_t operator*() { return *Offset; }
+};
+
+} // end anonymous namespace
+
+char RISCVFoldMemOffset::ID = 0;
+INITIALIZE_PASS(RISCVFoldMemOffset, DEBUG_TYPE, RISCV_FOLD_MEM_OFFSET_NAME,
+                false, false)
+
+FunctionPass *llvm::createRISCVFoldMemOffsetPass() {
+  return new RISCVFoldMemOffset();
+}
+
+// Walk forward from the ADDI looking for arithmetic instructions we can
+// analyze or memory instructions that use it as part of their address
+// calculation. For each arithmetic instruction we lookup how the offset
+// contributes to the value in that register use that information to
+// calculate the contribution to the output of this instruction.
+// Only addition and left shift are supported.
+// FIXME: Add multiplication by constant. The constant will be in a register.
+bool RISCVFoldMemOffset::foldOffset(
+    Register OrigReg, int64_t InitialOffset, const MachineRegisterInfo &MRI,
+    DenseMap<MachineInstr *, int64_t> &FoldableInstrs) {
+  // Map to hold how much the offset contributes to the value of this register.
+  DenseMap<Register, int64_t> RegToOffsetMap;
+
+  // Insert root offset into the map.
+  RegToOffsetMap[OrigReg] = InitialOffset;
+
+  std::queue<Register> Worklist;
+  Worklist.push(OrigReg);
+
+  while (!Worklist.empty()) {
+    Register Reg = Worklist.front();
+    Worklist.pop();
+
+    if (!Reg.isVirtual())
+      return false;
+
+    for (auto &User : MRI.use_nodbg_instructions(Reg)) {
+      FoldableOffset Offset;
+
+      switch (User.getOpcode()) {
+      default:
+        return false;
+      case RISCV::ADD:
----------------
lukel97 wrote:

Do we ever end up with another ADDI? I presume it's usually folded away

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