PNaCl: Add ExpandByVal pass for expanding out by-value struct args and results

This pass expands out the "byval" and "sret" argument attributes.

This will affect the calling conventions for PPAPI under PNaCl (for
passing PP_Var etc. by value), so the PNaCl PPAPI shims will need to
be updated in order to enable this pass by default.

BUG=https://code.google.com/p/nativeclient/issues/detail?id=3400
TEST=PNaCl toolchain trybots + GCC torture tests + LLVM test suite + Spec2k

Review URL: https://codereview.chromium.org/13973018

6 files changed, 318 insertions, 0 deletions

diff --git a/include/llvm/InitializePasses.h b/include/llvm/InitializePasses.h
index ff8bedc36f..9bcc8bf2a0 100644
--- a/include/llvm/InitializePasses.h
+++ b/include/llvm/InitializePasses.h
@@ -277,6 +277,7 @@ void initializeLoopVectorizePass(PassRegistry&);
 void initializeBBVectorizePass(PassRegistry&);
 void initializeMachineFunctionPrinterPassPass(PassRegistry&);
 // @LOCALMOD-BEGIN
+void initializeExpandByValPass(PassRegistry&);
 void initializeExpandConstantExprPass(PassRegistry&);
 void initializeExpandCtorsPass(PassRegistry&);
 void initializeExpandGetElementPtrPass(PassRegistry&);
diff --git a/include/llvm/Transforms/NaCl.h b/include/llvm/Transforms/NaCl.h
index 650779c058..ba765c1c21 100644
--- a/include/llvm/Transforms/NaCl.h
+++ b/include/llvm/Transforms/NaCl.h
@@ -19,6 +19,7 @@ class ModulePass;
 class Use;
 class Value;
 
+ModulePass *createExpandByValPass();
 FunctionPass *createExpandConstantExprPass();
 ModulePass *createExpandCtorsPass();
 BasicBlockPass *createExpandGetElementPtrPass();
diff --git a/lib/Transforms/NaCl/CMakeLists.txt b/lib/Transforms/NaCl/CMakeLists.txt
index fb3e30ece1..47d51f2940 100644
--- a/lib/Transforms/NaCl/CMakeLists.txt
+++ b/lib/Transforms/NaCl/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_llvm_library(LLVMNaClTransforms
+  ExpandByVal.cpp
   ExpandConstantExpr.cpp
   ExpandCtors.cpp
   ExpandGetElementPtr.cpp
diff --git a/lib/Transforms/NaCl/ExpandByVal.cpp b/lib/Transforms/NaCl/ExpandByVal.cpp
new file mode 100644
index 0000000000..cc2c73325e
--- /dev/null
+++ b/lib/Transforms/NaCl/ExpandByVal.cpp
@@ -0,0 +1,191 @@
+//===- ExpandByVal.cpp - Expand out use of "byval" and "sret" attributes---===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass expands out by-value passing of structs as arguments and
+// return values.  In LLVM IR terms, it expands out the "byval" and
+// "sret" function argument attributes.
+//
+// The semantics of the "byval" attribute are that the callee function
+// gets a private copy of the pointed-to argument that it is allowed
+// to modify.  In implementing this, we have a choice between making
+// the caller responsible for making the copy or making the callee
+// responsible for making the copy.  We choose the former, because
+// this matches how the normal native calling conventions work, and
+// because it often allows the caller to write struct contents
+// directly into the stack slot that it passes the callee, without an
+// additional copy.
+//
+// Note that this pass does not attempt to modify functions that pass
+// structs by value without using "byval" or "sret", such as:
+//
+//   define %struct.X @func()                           ; struct return
+//   define void @func(%struct.X %arg)                  ; struct arg
+//
+// The pass only handles functions such as:
+//
+//   define void @func(%struct.X* sret %result_buffer)  ; struct return
+//   define void @func(%struct.X* byval %ptr_to_arg)    ; struct arg
+//
+// This is because PNaCl Clang generates the latter and not the former.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/NaCl.h"
+
+using namespace llvm;
+
+namespace {
+  // This is a ModulePass so that it can strip attributes from
+  // declared functions as well as defined functions.
+  class ExpandByVal : public ModulePass {
+  public:
+    static char ID; // Pass identification, replacement for typeid
+    ExpandByVal() : ModulePass(ID) {
+      initializeExpandByValPass(*PassRegistry::getPassRegistry());
+    }
+
+    virtual bool runOnModule(Module &M);
+  };
+}
+
+char ExpandByVal::ID = 0;
+INITIALIZE_PASS(ExpandByVal, "expand-byval",
+                "Expand out by-value passing of structs",
+                false, false)
+
+// removeAttribute() currently does not work on Attribute::Alignment
+// (it fails with an assertion error), so we have to take a more
+// convoluted route to removing this attribute by recreating the
+// AttributeSet.
+AttributeSet RemoveAttrs(LLVMContext &Context, AttributeSet Attrs) {
+  SmallVector<AttributeSet, 8> AttrList;
+  for (unsigned Slot = 0; Slot < Attrs.getNumSlots(); ++Slot) {
+    unsigned Index = Attrs.getSlotIndex(Slot);
+    AttrBuilder AB;
+    for (AttributeSet::iterator Attr = Attrs.begin(Slot), E = Attrs.end(Slot);
+         Attr != E; ++Attr) {
+      if (!Attr->isAlignAttribute() &&
+          Attr->getKindAsEnum() != Attribute::ByVal &&
+          Attr->getKindAsEnum() != Attribute::StructRet) {
+        AB.addAttribute(*Attr);
+      }
+    }
+    AttrList.push_back(AttributeSet::get(Context, Index, AB));
+  }
+  return AttributeSet::get(Context, AttrList);
+}
+
+// ExpandCall() can take a CallInst or an InvokeInst.  It returns
+// whether the instruction was modified.
+template <class InstType>
+static bool ExpandCall(DataLayout *DL, InstType *Call) {
+  bool Modify = false;
+  AttributeSet Attrs = Call->getAttributes();
+  for (unsigned ArgIdx = 0; ArgIdx < Call->getNumArgOperands(); ++ArgIdx) {
+    unsigned AttrIdx = ArgIdx + 1;
+
+    if (Attrs.hasAttribute(AttrIdx, Attribute::StructRet))
+      Modify = true;
+
+    if (Attrs.hasAttribute(AttrIdx, Attribute::ByVal)) {
+      Modify = true;
+
+      Value *ArgPtr = Call->getArgOperand(ArgIdx);
+      Type *ArgType = ArgPtr->getType()->getPointerElementType();
+      ConstantInt *ArgSize = ConstantInt::get(
+          Call->getContext(), APInt(64, DL->getTypeStoreSize(ArgType)));
+      unsigned Alignment = Attrs.getParamAlignment(AttrIdx);
+      // In principle, using the alignment from the argument attribute
+      // should be enough.  However, Clang is not emitting this
+      // attribute for PNaCl.  LLVM alloca instructions do not use the
+      // ABI alignment of the type, so this must be specified
+      // explicitly.
+      // See https://code.google.com/p/nativeclient/issues/detail?id=3403
+      unsigned AllocAlignment =
+          std::max(Alignment, DL->getABITypeAlignment(ArgType));
+
+      // Make a copy of the byval argument.
+      Instruction *CopyBuf = new AllocaInst(ArgType, 0, AllocAlignment,
+                                            ArgPtr->getName() + ".byval_copy");
+      Function *Func = Call->getParent()->getParent();
+      Func->getEntryBlock().getInstList().push_front(CopyBuf);
+      IRBuilder<> Builder(Call);
+      Builder.CreateLifetimeStart(CopyBuf, ArgSize);
+      // Using the argument's alignment attribute for the memcpy
+      // should be OK because the LLVM Language Reference says that
+      // the alignment attribute specifies "the alignment of the stack
+      // slot to form and the known alignment of the pointer specified
+      // to the call site".
+      Instruction *MemCpy = Builder.CreateMemCpy(CopyBuf, ArgPtr, ArgSize,
+                                                 Alignment);
+      MemCpy->setDebugLoc(Call->getDebugLoc());
+
+      Call->setArgOperand(ArgIdx, CopyBuf);
+
+      // Mark the argument copy as unused using llvm.lifetime.end.
+      if (isa<CallInst>(Call)) {
+        BasicBlock::iterator It = BasicBlock::iterator(Call);
+        Builder.SetInsertPoint(++It);
+        Builder.CreateLifetimeEnd(CopyBuf, ArgSize);
+      } else if (InvokeInst *Invoke = dyn_cast<InvokeInst>(Call)) {
+        Builder.SetInsertPoint(Invoke->getNormalDest()->getFirstInsertionPt());
+        Builder.CreateLifetimeEnd(CopyBuf, ArgSize);
+        Builder.SetInsertPoint(Invoke->getUnwindDest()->getFirstInsertionPt());
+        Builder.CreateLifetimeEnd(CopyBuf, ArgSize);
+      }
+    }
+  }
+  if (Modify) {
+    Call->setAttributes(RemoveAttrs(Call->getContext(), Attrs));
+
+    if (CallInst *CI = dyn_cast<CallInst>(Call)) {
+      // This is no longer a tail call because the callee references
+      // memory alloca'd by the caller.
+      CI->setTailCall(false);
+    }
+  }
+  return Modify;
+}
+
+bool ExpandByVal::runOnModule(Module &M) {
+  bool Modified = false;
+  DataLayout DL(&M);
+
+  for (Module::iterator Func = M.begin(), E = M.end(); Func != E; ++Func) {
+    AttributeSet NewAttrs = RemoveAttrs(Func->getContext(),
+                                        Func->getAttributes());
+    Modified |= (NewAttrs != Func->getAttributes());
+    Func->setAttributes(NewAttrs);
+
+    for (Function::iterator BB = Func->begin(), E = Func->end();
+         BB != E; ++BB) {
+      for (BasicBlock::iterator Inst = BB->begin(), E = BB->end();
+           Inst != E; ++Inst) {
+        if (CallInst *Call = dyn_cast<CallInst>(Inst)) {
+          Modified |= ExpandCall(&DL, Call);
+        } else if (InvokeInst *Call = dyn_cast<InvokeInst>(Inst)) {
+          Modified |= ExpandCall(&DL, Call);
+        }
+      }
+    }
+  }
+
+  return Modified;
+}
+
+ModulePass *llvm::createExpandByValPass() {
+  return new ExpandByVal();
+}
diff --git a/test/Transforms/NaCl/expand-byval.ll b/test/Transforms/NaCl/expand-byval.ll
new file mode 100644
index 0000000000..12977a7d7d
--- /dev/null
+++ b/test/Transforms/NaCl/expand-byval.ll
@@ -0,0 +1,123 @@
+; RUN: opt -expand-byval %s -S | FileCheck %s
+
+target datalayout = "p:32:32:32"
+
+%MyStruct = type { i32, i8, i32 }
+%AlignedStruct = type { double, double }
+
+
+; Removal of "byval" attribute for passing structs arguments by value
+
+declare void @ext_func(%MyStruct*)
+
+define void @byval_receiver(%MyStruct* byval align 32 %ptr) {
+  call void @ext_func(%MyStruct* %ptr)
+  ret void
+}
+; Strip the "byval" and "align" attributes.
+; CHECK: define void @byval_receiver(%MyStruct* %ptr) {
+; CHECK-NEXT: call void @ext_func(%MyStruct* %ptr)
+
+
+declare void @ext_byval_func(%MyStruct* byval)
+; CHECK: declare void @ext_byval_func(%MyStruct*)
+
+define void @byval_caller(%MyStruct* %ptr) {
+  call void @ext_byval_func(%MyStruct* byval %ptr)
+  ret void
+}
+; CHECK: define void @byval_caller(%MyStruct* %ptr) {
+; CHECK-NEXT: %ptr.byval_copy = alloca %MyStruct, align 4
+; CHECK: call void @llvm.lifetime.start(i64 12, i8* %{{.*}})
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* %{{.*}}, i8* %{{.*}}, i64 12, i32 0, i1 false)
+; CHECK-NEXT: call void @ext_byval_func(%MyStruct* %ptr.byval_copy)
+
+
+define void @byval_tail_caller(%MyStruct* %ptr) {
+  tail call void @ext_byval_func(%MyStruct* byval %ptr)
+  ret void
+}
+; CHECK: define void @byval_tail_caller(%MyStruct* %ptr) {
+; CHECK: {{^}} call void @ext_byval_func(%MyStruct* %ptr.byval_copy)
+
+
+define void @byval_invoke(%MyStruct* %ptr) {
+  invoke void @ext_byval_func(%MyStruct* byval align 32 %ptr)
+      to label %cont unwind label %lpad
+cont:
+  ret void
+lpad:
+  %lp = landingpad { i8*, i32 } personality i8* null cleanup
+  ret void
+}
+; CHECK: define void @byval_invoke(%MyStruct* %ptr) {
+; CHECK: %ptr.byval_copy = alloca %MyStruct, align 32
+; CHECK: call void @llvm.lifetime.start(i64 12, i8* %{{.*}})
+; CHECK: invoke void @ext_byval_func(%MyStruct* %ptr.byval_copy)
+; CHECK: cont:
+; CHECK: call void @llvm.lifetime.end(i64 12, i8* %{{.*}})
+; CHECK: lpad:
+; CHECK: call void @llvm.lifetime.end(i64 12, i8* %{{.*}})
+
+
+; Check handling of alignment
+
+; Check that "align" is stripped for declarations too.
+declare void @ext_byval_func_align(%MyStruct* byval align 32)
+; CHECK: declare void @ext_byval_func_align(%MyStruct*)
+
+define void @byval_caller_align_via_attr(%MyStruct* %ptr) {
+  call void @ext_byval_func(%MyStruct* byval align 32 %ptr)
+  ret void
+}
+; CHECK: define void @byval_caller_align_via_attr(%MyStruct* %ptr) {
+; CHECK-NEXT: %ptr.byval_copy = alloca %MyStruct, align 32
+; The memcpy may assume that %ptr is 32-byte-aligned.
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* %2, i8* %3, i64 12, i32 32, i1 false)
+
+declare void @ext_byval_func_align_via_type(%AlignedStruct* byval)
+
+; %AlignedStruct contains a double so requires an alignment of 8 bytes.
+; Looking at the alignment of %AlignedStruct is a workaround for a bug
+; in pnacl-clang:
+; https://code.google.com/p/nativeclient/issues/detail?id=3403
+define void @byval_caller_align_via_type(%AlignedStruct* %ptr) {
+  call void @ext_byval_func_align_via_type(%AlignedStruct* byval %ptr)
+  ret void
+}
+; CHECK: define void @byval_caller_align_via_type(%AlignedStruct* %ptr) {
+; CHECK-NEXT: %ptr.byval_copy = alloca %AlignedStruct, align 8
+; Don't assume that %ptr is 8-byte-aligned when doing the memcpy.
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* %{{.*}}, i8* %{{.*}}, i64 16, i32 0, i1 false)
+
+
+; Removal of "sret" attribute for returning structs by value
+
+declare void @ext_sret_func(%MyStruct* sret align 32)
+; CHECK: declare void @ext_sret_func(%MyStruct*)
+
+define void @sret_func(%MyStruct* sret align 32 %buf) {
+  ret void
+}
+; CHECK: define void @sret_func(%MyStruct* %buf) {
+
+define void @sret_caller(%MyStruct* %buf) {
+  call void @ext_sret_func(%MyStruct* sret align 32 %buf)
+  ret void
+}
+; CHECK: define void @sret_caller(%MyStruct* %buf) {
+; CHECK-NEXT: call void @ext_sret_func(%MyStruct* %buf)
+
+
+; Check that other attributes are preserved
+
+define void @inreg_attr(%MyStruct* inreg %ptr) {
+  ret void
+}
+; CHECK: define void @inreg_attr(%MyStruct* inreg %ptr) {
+
+declare void @func_attrs() #0
+; CHECK: declare void @func_attrs() #0
+
+attributes #0 = { noreturn nounwind }
+; CHECK: attributes #0 = { noreturn nounwind }
diff --git a/tools/opt/opt.cpp b/tools/opt/opt.cpp
index ff6e5effeb..26a5e4a3c1 100644
--- a/tools/opt/opt.cpp
+++ b/tools/opt/opt.cpp
@@ -577,6 +577,7 @@ int main(int argc, char **argv) {
   initializeInstrumentation(Registry);
   initializeTarget(Registry);
   // @LOCALMOD-BEGIN
+  initializeExpandByValPass(Registry);
   initializeExpandConstantExprPass(Registry);
   initializeExpandCtorsPass(Registry);
   initializeExpandGetElementPtrPass(Registry);