[EVM] Support for spilling constants to memory

Author: PavelKopyl

llvm/lib/Target/EVM/CMakeLists.txt

@@ -44,6 +44,7 @@ add_llvm_target(EVMCodeGen
   EVMSHA3ConstFolding.cpp
   EVMSingleUseExpression.cpp
   EVMSplitCriticalEdges.cpp
+  EVMConstantSpiller.cpp
   EVMStackModel.cpp
   EVMStackShuffler.cpp
   EVMStackSolver.cpp

llvm/lib/Target/EVM/EVMConstantSpiller.cpp

@@ -0,0 +1,136 @@
+//===----- EVMConstantSpiller.cpp - Spill constants to memory --*- C++ -*--===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+//
+//===----------------------------------------------------------------------===//
+
+#include "EVMConstantSpiller.h"
+#include "EVMInstrInfo.h"
+#include "EVMSubtarget.h"
+#include "MCTargetDesc/EVMMCTargetDesc.h"
+#include "TargetInfo/EVMTargetInfo.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Module.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "evm-spill-constants"
+
+constexpr uint64_t SpillSlotSize = 32;
+
+static cl::opt<unsigned> ConstantSpillThreshold(
+    "evm-constant-spill-threshold", cl::Hidden, cl::init(30),
+    cl::desc("Minimum number of uses of a constant across the module required "
+             "before spilling it to memory is considered profitable"));
+
+static MachineInstr *emitPush(MachineBasicBlock &MBB,
+                              MachineBasicBlock::iterator InsertBefore,
+                              const EVMInstrInfo *TII, const APInt &Imm,
+                              LLVMContext &Ctx, const DebugLoc &DL) {
+  unsigned Opc = EVM::getPUSHOpcode(Imm);
+  auto MI = BuildMI(MBB, InsertBefore, DL, TII->get(EVM::getStackOpcode(Opc)));
+  if (Opc != EVM::PUSH0)
+    MI.addCImm(ConstantInt::get(Ctx, Imm));
+  return MI;
+}
+
+uint64_t EVMConstantSpiller::getSpillSize() const {
+  return ConstantUseCount.size() * SpillSlotSize;
+}
+
+void EVMConstantSpiller::emitConstantSpills(uint64_t SpillOffset,
+                                            MachineFunction &EntryMF) {
+  LLVMContext &Ctx = EntryMF.getFunction().getContext();
+  const EVMInstrInfo *TII = EntryMF.getSubtarget<EVMSubtarget>().getInstrInfo();
+
+  DenseMap<APInt, uint64_t> ConstantToSpillOffset;
+  for (const auto &KV : ConstantUseCount) {
+    ConstantToSpillOffset[KV.first] = SpillOffset;
+    SpillOffset += SpillSlotSize;
+  }
+
+  // Emit constant stores in prologue of the '__entry' function.
+  MachineBasicBlock &SpillMBB = EntryMF.front();
+  auto InsertBefore = SpillMBB.begin();
+  for (const auto &[Imm, Offset] : ConstantToSpillOffset) {
+    LLVM_DEBUG({
+      dbgs() << "Spilling constant: " << Imm
+             << ", number of uses: " << ConstantUseCount.at(Imm)
+             << ", at offset: " << Offset << '\n';
+    });
+
+    BuildMI(SpillMBB, InsertBefore, DebugLoc(), TII->get(EVM::MSTORE_S));
+    emitPush(SpillMBB, InsertBefore, TII, APInt(256, Offset), Ctx, DebugLoc());
+    emitPush(SpillMBB, InsertBefore, TII, Imm, Ctx, DebugLoc());
+  }
+
+  // Reload spilled constants.
+  for (MachineInstr *MI : ReloadCandidates) {
+    const APInt Imm = MI->getOperand(0).getCImm()->getValue().zext(256);
+    uint64_t Offset = ConstantToSpillOffset.at(Imm);
+
+    MachineBasicBlock *MBB = MI->getParent();
+    emitPush(*MBB, MI, TII, APInt(256, Offset), Ctx, MI->getDebugLoc());
+    auto Load = BuildMI(*MBB, MI, MI->getDebugLoc(), TII->get(EVM::MLOAD_S));
+    Load->setAsmPrinterFlag(MachineInstr::ReloadReuse);
+    MI->eraseFromParent();
+  }
+}
+
+static bool shouldSkip(const MachineInstr &MI) {
+  if (!EVMInstrInfo::isPush(&MI) || (MI.getOpcode() == EVM::PUSH0_S))
+    return true;
+
+  const APInt Imm = MI.getOperand(0).getCImm()->getValue();
+  return Imm.getActiveBits() < 8 * 8;
+}
+
+void EVMConstantSpiller::analyzeModule(Module &M, MachineModuleInfo &MMI) {
+  for (Function &F : M) {
+    MachineFunction *MF = MMI.getMachineFunction(F);
+    if (!MF)
+      continue;
+
+    for (MachineBasicBlock &MBB : *MF) {
+      for (MachineInstr &MI : MBB) {
+        if (shouldSkip(MI))
+          continue;
+
+        const APInt Imm = MI.getOperand(0).getCImm()->getValue().zext(256);
+        if (Imm.isAllOnes())
+          continue;
+
+        ConstantUseCount[Imm]++;
+        ReloadCandidates.push_back(&MI);
+      }
+    }
+  }
+}
+
+void EVMConstantSpiller::filterCandidates() {
+  SmallVector<APInt> ImmToRemove;
+  for (const auto &[Imm, NumUses] : ConstantUseCount)
+    if (NumUses < ConstantSpillThreshold)
+      ImmToRemove.push_back(Imm);
+
+  for (const APInt &Imm : ImmToRemove)
+    ConstantUseCount.erase(Imm);
+
+  erase_if(ReloadCandidates, [this](const MachineInstr *MI) {
+    const APInt &Imm = MI->getOperand(0).getCImm()->getValue().zext(256);
+    return !ConstantUseCount.contains(Imm);
+  });
+}
+
+EVMConstantSpiller::EVMConstantSpiller(Module &M, MachineModuleInfo &MMI) {
+  analyzeModule(M, MMI);
+  filterCandidates();
+}

llvm/lib/Target/EVM/EVMConstantSpiller.h

@@ -0,0 +1,51 @@
+//===----- EVMConstantSpiller.h - Spill constants to memory ----*- C++ -*--===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file identifies frequently used large constants across the module.
+// If a constant’s usage exceeds the threshold and its bit-width is sufficient,
+// it is spilled at the beginning of the __entry function and reloaded at use
+// sites.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_EVM_EVMCONSTANTSPILLER_H
+#define LLVM_LIB_TARGET_EVM_EVMCONSTANTSPILLER_H
+
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/IR/Module.h"
+
+namespace llvm {
+
+class MachineInstr;
+
+class EVMConstantSpiller {
+public:
+  EVMConstantSpiller(Module &M, MachineModuleInfo &MMI);
+
+  /// Return the total size needed for the spill area.
+  uint64_t getSpillSize() const;
+
+  /// Insert constant spills into the first basic block of the __entry
+  /// function and insert reloads at their use sites.
+  void emitConstantSpills(uint64_t SpillOffset, MachineFunction &EntryMF);
+
+private:
+  /// Maps each APInt constant to the number of times it appears across all
+  /// functions in the module
+  SmallDenseMap<APInt, unsigned> ConstantUseCount;
+  /// PUSH instructions that need to be reloaded
+  SmallVector<MachineInstr *> ReloadCandidates;
+
+  /// Filters out constants whose usage count is below the threshold.
+  void filterCandidates();
+
+  void analyzeModule(Module &M, MachineModuleInfo &MMI);
+};
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_EVM_EVMCONSTANTSPILLER_H

llvm/lib/Target/EVM/EVMConstantUnfolding.cpp

@@ -34,7 +34,6 @@
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/Module.h"
-#include "llvm/IR/PassInstrumentation.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 #include <functional>

llvm/lib/Target/EVM/EVMFinalizeStackFrames.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "EVM.h"
+#include "EVMConstantSpiller.h"
 #include "MCTargetDesc/EVMMCTargetDesc.h"
 #include "TargetInfo/EVMTargetInfo.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
@@ -147,12 +148,26 @@ bool EVMFinalizeStackFrames::runOnModule(Module &M) {
   MachineModuleInfo &MMI = getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
   SmallVector<std::pair<MachineFunction *, uint64_t>, 8> ToReplaceFI;
 
-  // Calculate the stack size for each function.
-  for (Function &F : M) {
-    MachineFunction *MF = MMI.getMachineFunction(F);
-    if (!MF)
-      continue;
+  SmallVector<MachineFunction *> MFs;
+  for_each(M.getFunctionList(), [&MFs, &MMI](Function &F) {
+    if (MachineFunction *MF = MMI.getMachineFunction(F))
+      MFs.push_back(MF);
+  });
+
+  if (MFs.size() &&
+      MFs.front()->getFunction().hasFnAttribute("evm-entry-function")) {
+    EVMConstantSpiller ConstSpiller(M, MMI);
+    if (ConstSpiller.getSpillSize()) {
+      MachineFunction *MF = MFs.front();
+      uint64_t StackSize = calculateFrameObjectOffsets(*MF);
+      MF->getFrameInfo().CreateSpillStackObject(ConstSpiller.getSpillSize(),
+                                                Align(32));
+      ConstSpiller.emitConstantSpills(StackSize, *MF);
+    }
+  }
 
+  // Calculate the stack size for each function.
+  for (MachineFunction *MF : MFs) {
     uint64_t StackSize = calculateFrameObjectOffsets(*MF);
     if (StackSize == 0)
       continue;

llvm/test/CodeGen/EVM/constant-spilling.ll

@@ -0,0 +1,79 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -O3 --evm-stack-region-size=32 --evm-stack-region-offset=128 --evm-constant-spill-threshold=2 < %s | FileCheck %s
+target datalayout = "E-p:256:256-i256:256:256-S256-a:256:256"
+target triple = "evm-unknown-unknown"
+
+declare void @llvm.evm.return(ptr addrspace(1), i256) noreturn
+
+define void @entry() #1 {
+; CHECK-LABEL: entry:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    MSTORE
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    PUSH4 0x4E487B71
+; CHECK-NEXT:    PUSH1 0xE0
+; CHECK-NEXT:    SHL
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    MLOAD ; Reload Reuse
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    RETURN
+entry:
+  tail call void @llvm.evm.return(ptr addrspace(1) null, i256 35408467139433450592217433187231851964531694900788300625387963629091585785856)
+  unreachable
+}
+
+define void @test_spill() #2 {
+; CHECK-LABEL: test_spill:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    MLOAD ; Reload Reuse
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    RETURN
+entry:
+  tail call void @llvm.evm.return(ptr addrspace(1) null, i256 35408467139433450592217433187231851964531694900788300625387963629091585785856)
+  unreachable
+}
+
+; Check that -1 constant is not spilled.
+define void @test_not_spiller() #2 {
+; CHECK-LABEL: test_not_spiller:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    NOT
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    RETURN
+entry:
+  tail call void @llvm.evm.return(ptr addrspace(1) null, i256 -1)
+  unreachable
+}
+
+; Check that small constants are not spilled.
+define void @test_not_spilled2() #2 {
+; CHECK-LABEL: test_not_spilled2:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH5 0x100000000
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    RETURN
+entry:
+  tail call void @llvm.evm.return(ptr addrspace(1) null, i256 4294967296)
+  unreachable
+}
+
+; Check that small constants are not spilled.
+define void @test_not_spilled3() #2 {
+; CHECK-LABEL: test_not_spilled3:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH5 0x100000000
+; CHECK-NEXT:    PUSH0
+; CHECK-NEXT:    RETURN
+entry:
+  tail call void @llvm.evm.return(ptr addrspace(1) null, i256 4294967296)
+  unreachable
+}
+
+attributes #1 = { nofree noreturn nounwind "evm-entry-function" }
+attributes #2 = { nofree nounwind noreturn }