Add ArgumentCopyElision MI pass

[rnk]

Summary

LLVM frequently copies arguments from the stack when generating
unoptimized code or when an argument's address is taken. This pass
identifies stores of loads of argument stack slots in the entry block,
deletes the store, and replaces all uses of the frame index with the
original load's frame index.

It is currently disabled by default under the flag
-enable-arg-copy-elim.

In a 32-bit release build of Chrome, this saves 20K from both
chrome_child.dll and chrome.dll. Each DLL is about 50MB, so this is less
than 0.05% for each, so it's not what I'd hoped, but it's simple enough
to keep. In the future, this pass could be extended to handle aggregates
that have been split into multiple LLVM arguments in clang.

The potential wins in a debug build are much greater because all
arguments will live in memory, but I think we should do this whole thing
in fast isel instead. First, x86 fast isel needs to know how to lower
arguments in memory. When we do that, it should be easy to add the
IR-equivalent pattern matching to identify allocas that are immediately
initialized by arguments passed in memory and use the appropriate
argument frame index from the beginning of ISel.

Fixes PR26328

[inglorion]

Besides the effect on binary size, do we know what other effects this has? For example, what is the effect on build time and run-time performance?

[rnk]

I haven't measured either yet.

Build time effect should be below the noise because it is O(entry block instructions) if there are no address taken parameters, which seems to be the common case given that this doesn't have much impact on overall size. If there are address-taken parameters, it iterates all instructions once to replace frame indices. I don't think it's common practice to measure for compile time regressions when adding one more linear pass every function.

The runtime effects will definitely be positive, since we're eliding copies, but given how low the size impact was, I don't expect it will be measurable either :(. Maybe we should chat about different ways we could measure it.

[inglorion]

Right. The reason I asked is that it seems to me that enabling the pass will be neutral or better for code size and run time, while not costing much or possibly even being beneficial at compile time (we do a bit of extra analysis on a few instructions, but then we lower the number of instructions that the rest of the build operates on). If that is indeed the case, I would be in favor of enabling the new pass by default.

[rnk]

• Remove leftover fastisel test
• Add ARM test
• Clarify comments

[rnk]

Any thoughts on this? I think the basic algorithm is sound, but I'd appreciate it if someone took a quick glance.

[MatzeB]

Can you explain why this needs an own pass. My first intuition would have been that this should be simple to solve at the SelectionDAG level.

[inglorion]

The implementation looks good to me. I think the ARM test is missing a few cases.

test/CodeGen/ARM/argument-copy-elision.ll

5

I expect this should have the retaddr_taken and size_mismatch tests, too. If not, I think you can remove this declaration.

[rnk]

In D28388#645541, @MatzeB wrote:

Can you explain why this needs an own pass. My first intuition would have been that this should be simple to solve at the SelectionDAG level.

Initially I wanted it to be independent of ISel because we have 3 (!) ISel implementations. In particular, I was hoping to improve our -O0 code so that it's less bulky. However, I realized that x86 fastisel doesn't even handle arguments in memory, so I gave up on that. So, I guess it would be better to do this as part of SDAG.

Ideally, we'd do this as part of FunctionLoweringInfo::set so that the StaticAlloca map is populated with the fixed argument frame indices from the beginning of ISel, but we need to get the target to lower formal arguments first, which happens later. I could try lowering formal arguments up front, but that seemed more invasive than doing a standalone pass.

[rnk]

I rewrote this as an IR analysis and uploaded that as D29668

include/llvm/CodeGen/MachineFrameInfo.h

bool isImmutableObjectIndex(int ObjectIdx) const {
     // Tail calling functions can clobber their function arguments.
     if (HasTailCall)
       return false;
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     return Objects[ObjectIdx+NumFixedObjects].isImmutable;
   }
 
+  /// setIsImmutableObjectIndex - Marks the immutability of an object.
+  void setIsImmutableObjectIndex(int ObjectIdx, bool Immutable) {
+    assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
+           "Invalid Object Idx!");
+    Objects[ObjectIdx+NumFixedObjects].isImmutable = Immutable;
+  }
+
   /// Returns true if the specified index corresponds to a spill slot.
   bool isSpillSlotObjectIndex(int ObjectIdx) const {
     assert(unsigned(ObjectIdx+NumFixedObjects) < Objects.size() &&
            "Invalid Object Idx!");
     return Objects[ObjectIdx+NumFixedObjects].isSpillSlot;
   }
 
   bool isStatepointSpillSlotObjectIndex(int ObjectIdx) const {

include/llvm/CodeGen/Passes.h

/// MachineDominanaceFrontier - This pass is a machine dominators analysis pass.
   /// slots relative to one another and allocates base registers to access them
   /// when it is estimated by the target to be out of range of normal frame
   /// pointer or stack pointer index addressing.
   extern char &LocalStackSlotAllocationID;
 
   /// ExpandISelPseudos - This pass expands pseudo-instructions.
   extern char &ExpandISelPseudosID;
 
+  /// ArgumentCopyElision - This pass elides copies from argument stack memory
+  /// to local allocas.
+  FunctionPass *createArgumentCopyElisionPass();
+
   /// createExecutionDependencyFixPass - This pass fixes execution time
   /// problems with dependent instructions, such as switching execution
   /// domains to match.
   ///
   /// The pass will examine instructions using and defining registers in RC.
   ///
   FunctionPass *createExecutionDependencyFixPass(const TargetRegisterClass *RC);
 

include/llvm/InitializePasses.h

 void initializeADCELegacyPassPass(PassRegistry&);
 void initializeAddDiscriminatorsLegacyPassPass(PassRegistry&);
 void initializeAddressSanitizerModulePass(PassRegistry&);
 void initializeAddressSanitizerPass(PassRegistry&);
 void initializeAliasSetPrinterPass(PassRegistry&);
 void initializeAlignmentFromAssumptionsPass(PassRegistry&);
 void initializeAlwaysInlinerLegacyPassPass(PassRegistry&);
 void initializeArgPromotionPass(PassRegistry&);
+void initializeArgumentCopyElisionPass(PassRegistry&);
 void initializeAssumptionCacheTrackerPass(PassRegistry &);
 void initializeAtomicExpandPass(PassRegistry&);
 void initializeBBVectorizePass(PassRegistry&);
 void initializeBDCELegacyPassPass(PassRegistry &);
 void initializeBarrierNoopPass(PassRegistry&);
 void initializeBasicAAWrapperPassPass(PassRegistry&);
 void initializeBlockExtractorPassPass(PassRegistry&);
 void initializeBlockFrequencyInfoWrapperPassPass(PassRegistry&);

lib/CodeGen/ArgumentCopyElision.cpp

+//===-- llvm/CodeGen/ArgumentCopyElision.cpp --------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Eliminate copies from argument stack memory to local variable stack slots.
+// This can run at -O0 to simplify the generated machine code.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetInstrInfo.h"
+#include "llvm/Target/TargetLowering.h"
+#include "llvm/Target/TargetSubtargetInfo.h"
+using namespace llvm;
+
+#define DEBUG_TYPE "arg-copy-elision"
+
+static cl::opt<bool>
+    EnableArgCopyElim("enable-arg-copy-elim", cl::Hidden, cl::init(false),
+                      cl::desc("enable stack argument copy elimination"));
+
+namespace {
+class ArgumentCopyElision : public MachineFunctionPass {
+  const TargetInstrInfo *TII;
+  MachineFrameInfo *MFI;
+  const MachineRegisterInfo *MRI;
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+  ArgumentCopyElision() : MachineFunctionPass(ID) {}
+
+private:
+  bool isStoreOfArgumentLoad(MachineInstr &MI, unsigned &SrcReg, int &ArgIndex,
+                             int &CopyIndex);
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+} // end anonymous namespace
+
+char ArgumentCopyElision::ID = 0;
+INITIALIZE_PASS(ArgumentCopyElision, "arg-copy-elision",
+                "Elide copies of arguments passed on the stack", false, false)
+
+bool ArgumentCopyElision::isStoreOfArgumentLoad(MachineInstr &MI,
+                                                unsigned &SrcReg, int &ArgIndex,
+                                                int &CopyIndex) {
+  SrcReg = TII->isStoreToStackSlot(MI, CopyIndex);
+  if (SrcReg == 0)
+    return false;
+
+  // Check that this store initializes the entire stack object.
+  int64_t StoreSize = (*MI.memoperands_begin())->getSize();
+  if (StoreSize != MFI->getObjectSize(CopyIndex))
+    return false;
+
+  // Don't elide copies from one fixed stack slot to another.
+  if (MFI->isFixedObjectIndex(CopyIndex))
+    return false;
+
+  // Check if the source register is an argument load.
+  assert(MRI->hasOneDef(SrcReg) && "should be in SSA");
+  MachineInstr &SrcLoad = *MRI->def_instr_begin(SrcReg);
+
+  // FIXME: Ask the target if this is a register with a home location, like on
+  // Win64.
+
+  // Check if this is a load from a fixed stack slot.
+  if (!TII->isLoadFromStackSlot(SrcLoad, ArgIndex) ||
+      !MFI->isFixedObjectIndex(ArgIndex))
+    return false;
+
+  // Look for immutable stack objects, which indicate a non-byval argument slot.
+  if (!MFI->isImmutableObjectIndex(ArgIndex)) {
+    DEBUG(dbgs() << "  rejecting elision candidate due to mutable argument "
+                 << MachineOperand::CreateFI(ArgIndex) << '\n';);
+    return false;
+  }
+
+  return true;
+}
+
+/// State of a stack object during dataflow over the entry block. Each stack
+/// object starts out unknown. If we see a use of the frame index in a store or
+/// memcpy that fully initializes the stack object, transition to SOS_Argument.
+/// Any other use of the frame index means we stop matching that stack object.
+
+struct ElisionCandidate {
+  enum State {
+    Unknown = INT_MAX,
+    Clobbered = INT_MAX - 1,
+  };
+
+  ElisionCandidate() : ArgIndexOrState(Unknown) {}
+  ElisionCandidate(int ArgIndex) : ArgIndexOrState(ArgIndex) {}
+
+  static ElisionCandidate clobbered() { return ElisionCandidate(Clobbered); }
+
+  /// Either a stack object index, or two invalid sentinel values indicating
+  /// that this stack object was not initialized from an argument in memory.
+  int ArgIndexOrState = Unknown;
+
+  bool isUnknown() const { return ArgIndexOrState == Unknown; };
+  bool isClobbered() const { return ArgIndexOrState == Clobbered; };
+  bool isArgument() const { return ArgIndexOrState < Clobbered; };
+};
+
+bool ArgumentCopyElision::runOnMachineFunction(MachineFunction &MF) {
+  if (!EnableArgCopyElim)
+    return false;
+
+  const TargetSubtargetInfo &ST = MF.getSubtarget();
+  TII = ST.getInstrInfo();
+  MRI = &MF.getRegInfo();
+  MFI = &MF.getFrameInfo();
+
+  DEBUG(dbgs() << "\nBegin ArgumentCopyElision\n");
+
+  // Build a map from stack object index to clobber state with a vector. Don't
+  // include fixed stack objects, since we never elide from one to another.
+  SmallVector<ElisionCandidate, 8> ElisionCandidates;
+  ElisionCandidates.resize(MFI->getObjectIndexEnd());
+  SmallVector<int, 4> CapturesSinceLastStore;
+
+  // Iterate through the entry block, looking for copies to elide.
+  bool FoundElisions = false;
+  SmallVector<MachineInstr *, 4> DeadInstrs;
+  for (MachineInstr &MI : MF.front()) {
+    int ArgIndex, CopyIndex;
+    unsigned SrcReg;
+    if (isStoreOfArgumentLoad(MI, SrcReg, ArgIndex, CopyIndex)) {
+      if (ElisionCandidates[CopyIndex].isUnknown()) {
+        // Mark these frame indicies for copy elision.
+        DEBUG(dbgs() << "  eliding argument copy from "
+                     << MachineOperand::CreateFI(ArgIndex) << " to "
+                     << MachineOperand::CreateFI(CopyIndex) << '\n');
+        ElisionCandidates[CopyIndex] = ElisionCandidate(ArgIndex);
+        assert(ElisionCandidates[CopyIndex].isArgument());
+        FoundElisions = true;
+
+        // Forget that the copied object was captured. This instruction stores
+        // to it.
+        CapturesSinceLastStore.erase(std::remove(CapturesSinceLastStore.begin(),
+                                                 CapturesSinceLastStore.end(),
+                                                 CopyIndex),
+                                     CapturesSinceLastStore.end());
+
+        // Mark the copy destination as dead, as we are about to replace all
+        // uses of it.
+        MFI->RemoveStackObject(CopyIndex);
+
+        // Mark the argument object as mutable. By eliding the copy, we allow
+        // the user to modify the argument memory.
+        MFI->setIsImmutableObjectIndex(ArgIndex, false);
+
+        // Mark the store for deletion.
+        DeadInstrs.push_back(&MI);
+
+        // Mark the load for deletion if this was the only user.
+        if (MRI->hasOneUse(SrcReg))
+          DeadInstrs.push_back(&*MRI->def_instr_begin(SrcReg));
+      }
+    } else {
+      // Conservatively mark any frame index operands as captured if they
+      // haven't already been initialized with an argument.
+      for (const MachineOperand &MO : MI.operands()) {
+        if (!MO.isFI())
+          continue;
+        int FI = MO.getIndex();
+        if (!MFI->isFixedObjectIndex(FI) && ElisionCandidates[FI].isUnknown())
+          CapturesSinceLastStore.push_back(FI);
+      }
+
+      // If this instruction stores, mark all captures as clobbered.
+      if (MI.mayStore() && !CapturesSinceLastStore.empty()) {
+        DEBUG(dbgs() << "  storing instruction clobbering all captured "
+                        "argument frame indices:\n"
+                     << "    " << MI);
+        DEBUG(dbgs() << "      frame indices:");
+        for (int FI : CapturesSinceLastStore) {
+          DEBUG(dbgs() << ' ' << MachineOperand::CreateFI(FI));
+          ElisionCandidates[FI] = ElisionCandidate::clobbered();
+        }
+        DEBUG(dbgs() << '\n');
+        CapturesSinceLastStore.clear();
+      }
+    }
+  }
+
+  if (!FoundElisions) {
+    DEBUG(dbgs() << "No argument copy elision candidates\n\n");
+    return false;
+  }
+
+  for (MachineInstr *MI : DeadInstrs) {
+    DEBUG(dbgs() << "  ArgumentCopyElision: DELETING: " << *MI);
+    MI->eraseFromParent();
+  }
+
+  // Replace all uses of CopyIndex with ArgIndex. Assert that there are no other
+  // uses of ArgIndex, which should be the case, since all instruction selection
+  // schemes appear to load the argument into a virtual register and then split
+  // the live range during register allocation by rematerializing from the
+  // immutable stack slot. By changing the stack slot to mutable, we will force
+  // the register allocator to spill if it needs to.
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : MBB) {
+      for (MachineOperand &MO : MI.operands()) {
+        // Find frame index operands that we want to copy elide and replace
+        // them.
+        if (!MO.isFI())
+          continue;
+        int FI = MO.getIndex();
+        if (MFI->isFixedObjectIndex(FI) || !ElisionCandidates[FI].isArgument())
+          continue;
+        DEBUG(dbgs() << "  replacing " << MO << " in " << MI);
+        MO.setIndex(ElisionCandidates[FI].ArgIndexOrState);
+      }
+    }
+  }
+
+  DEBUG(dbgs() << "End ArgumentCopyElision\n");
+
+  return true;
+}
+
+FunctionPass *llvm::createArgumentCopyElisionPass() {
+  return new ArgumentCopyElision();
+}

lib/CodeGen/CMakeLists.txt

 add_llvm_library(LLVMCodeGen
   AggressiveAntiDepBreaker.cpp
   AllocationOrder.cpp
   Analysis.cpp
+  ArgumentCopyElision.cpp
   AtomicExpandPass.cpp
   BasicTargetTransformInfo.cpp
   BranchFolding.cpp
   BranchRelaxation.cpp
   BuiltinGCs.cpp
   CalcSpillWeights.cpp
   CallingConvLower.cpp
   CodeGen.cpp

lib/CodeGen/TargetPassConfig.cpp

void TargetPassConfig::addMachinePasses() {
   AddingMachinePasses = false;
 }
 
 /// Add passes that optimize machine instructions in SSA form.
 void TargetPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
   addPass(&EarlyTailDuplicateID);
 
+  // Eliminate copies of address-taken arguments now that we've lowered formal
+  // arguments.
+  addPass(createArgumentCopyElisionPass());
+
   // Optimize PHIs before DCE: removing dead PHI cycles may make more
   // instructions dead.
   addPass(&OptimizePHIsID, false);
 
   // This pass merges large allocas. StackSlotColoring is a different pass
   // which merges spill slots.
   addPass(&StackColoringID, false);
 

test/CodeGen/ARM/argument-copy-elision.ll

+; RUN: llc -enable-arg-copy-elim -O2 < %s -mtriple=arm-linux-gnueabi | FileCheck %s
+
+; Check that argument copy elision works on ARM.
+
+declare i8* @llvm.returnaddress(i32)
+
+declare void @capture(i32*)
+
+define void @basic_addrtaken(i32, i32, i32, i32, i32 %x) {
+entry:
+  %x.addr = alloca i32, align 4
+  store i32 %x, i32* %x.addr, align 4
+  call void @capture(i32* %x.addr)
+  ret void
+}
+
+; CHECK-LABEL: basic_addrtaken:
+; CHECK-NOT: ldr
+; CHECK-NOT: str
+; CHECK: add r0, sp, #8
+; CHECK: bl capture
+; CHECK: mov pc, lr

test/CodeGen/X86/argument-copy-elision.ll

+; RUN: llc -enable-arg-copy-elim -O2 < %s -mtriple=i686-windows | FileCheck %s
+
+declare i8* @llvm.returnaddress(i32)
+
+declare void @capture(i32*)
+
+define void @basic_addrtaken(i32 %x) {
+entry:
+  %x.addr = alloca i32, align 4
+  store i32 %x, i32* %x.addr, align 4
+  call void @capture(i32* %x.addr)
+  ret void
+}
+
+; The lea should be a *positive* offset from ESP, indicating that the address of
+; original argument is passed to capture.
+
+; CHECK-LABEL: _basic_addrtaken:
+; CHECK-NOT: movl{{.*}}(%esp){{.*}}
+; CHECK: leal {{[0-9]+}}(%esp), %[[reg:[^ ]*]]
+; CHECK: pushl %[[reg]]
+; CHECK: calll _capture
+; CHECK: retl
+
+; Make sure we don't elide copies from the return address.
+
+define void @retaddr_taken() {
+entry:
+  %x.addr = alloca i32, align 4
+  %ra = call i8* @llvm.returnaddress(i32 0)
+  %ra.i32 = ptrtoint i8* %ra to i32
+  store i32 %ra.i32, i32* %x.addr, align 4
+  call void @capture(i32* %x.addr)
+  ret void
+}
+
+; CHECK-LABEL: _retaddr_taken:
+; CHECK: movl{{.*}}(%esp){{.*}}
+; CHECK: pushl %{{.*}}
+; CHECK: calll _capture
+; CHECK: retl
+
+; Make sure we look at the size of the stores.
+; If the stores were reversed we could try to be more clever here.
+
+define void @size_mismatch(i32 %x, i32 %y) {
+entry:
+  %aggr = alloca {i32, i32}, align 4
+  %x.addr = getelementptr {i32,i32}, {i32,i32}* %aggr, i32 0, i32 0
+  store i32 %x, i32* %x.addr, align 4
+  %y.addr = getelementptr {i32,i32}, {i32,i32}* %aggr, i32 0, i32 1
+  store i32 %y, i32* %y.addr, align 4
+  call void @capture(i32* %x.addr)
+  ret void
+}
+
+; CHECK-LABEL: _size_mismatch:
+; CHECK: subl $8, %esp
+; CHECK: movl{{.*}}(%esp){{.*}}
+; CHECK: pushl %{{.*}}
+; CHECK: calll _capture
+; CHECK: retl
+