[MLIR][LLVM] Put byval static allocas in the entry block.

[gysit]

LGTM!

mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp

 }
 
 /// Copies the data from a byval pointer argument into newly alloca'ed memory
 /// and returns the value of the alloca.
 static Value handleByValArgumentInit(OpBuilder &builder, Location loc,
                                      Value argument, Type elementType,
                                      unsigned elementTypeSize,
                                      unsigned targetAlignment) {
+  Block *entryBlock = &(*argument.getParentRegion()->begin());
   // Allocate the new value on the stack.
+  Value allocaOp;
+  {
+    // Since this is a static alloca, we can put it directly in the entry block,
+    // so they can be absorbed into the prologue/epilogue at code generation.
+    OpBuilder::InsertionGuard insertionGuard(builder);
+    builder.setInsertionPointToStart(entryBlock);
     Value one = builder.create<LLVM::ConstantOp>(loc, builder.getI64Type(),
                                                  builder.getI64IntegerAttr(1));
-  Value allocaOp = builder.create<LLVM::AllocaOp>(
+    allocaOp = builder.create<LLVM::AllocaOp>(
         loc, argument.getType(), elementType, one, targetAlignment);
+  }
   // Copy the pointee to the newly allocated value.
   Value copySize = builder.create<LLVM::ConstantOp>(
       loc, builder.getI64Type(), builder.getI64IntegerAttr(elementTypeSize));
   Value isVolatile = builder.create<LLVM::ConstantOp>(
       loc, builder.getI1Type(), builder.getBoolAttr(false));
   builder.create<LLVM::MemcpyOp>(loc, allocaOp, argument, copySize, isVolatile);
   return allocaOp;
 }

mlir/test/Dialect/LLVMIR/inlining.mlir

 // -----
 
 llvm.func @with_byval_arg(%ptr : !llvm.ptr { llvm.byval = f64 }) {
   llvm.return
 }
 
 // CHECK-LABEL: llvm.func @test_byval
 // CHECK-SAME: %[[PTR:[a-zA-Z0-9_]+]]: !llvm.ptr
+llvm.func @test_byval(%ptr : !llvm.ptr) {
+  // Make sure the new static alloca goes to the entry block.
   // CHECK: %[[ALLOCA:.+]] = llvm.alloca %{{.+}} x f64
+  // CHECK: llvm.br ^[[BB1:[a-zA-Z0-9_]+]]
+  llvm.br ^bb1
+  // CHECK: ^[[BB1]]
+^bb1:
   // CHECK: "llvm.intr.memcpy"(%[[ALLOCA]], %[[PTR]]
-llvm.func @test_byval(%ptr : !llvm.ptr) {
   llvm.call @with_byval_arg(%ptr) : (!llvm.ptr) -> ()
+  llvm.br ^bb2
+^bb2:
   llvm.return
 }
 
 // -----
 
 llvm.func @with_byval_arg(%ptr : !llvm.ptr { llvm.byval = f64 }) attributes {memory = #llvm.memory_effects<other = readwrite, argMem = read, inaccessibleMem = readwrite>} {
   llvm.return
 }

 }
 
 // -----
 
 llvm.func @aligned_byval_arg(%ptr : !llvm.ptr { llvm.byval = i16, llvm.align = 16 }) attributes {memory = #llvm.memory_effects<other = read, argMem = read, inaccessibleMem = read>} {
   llvm.return
 }
 
-// CHECK-LABEL: llvm.func @test_byval_alloca
-llvm.func @test_byval_alloca() {
-  // Make sure only the unaligned alloca triggers a memcpy.
-  %size = llvm.mlir.constant(1 : i64) : i64
-  // CHECK: %[[ALLOCA:.+]] = llvm.alloca {{.+}}alignment = 1
-  // CHECK: "llvm.intr.memcpy"(%{{.+}}, %[[ALLOCA]]
+// CHECK-LABEL: llvm.func @test_byval_unaligned_alloca
+llvm.func @test_byval_unaligned_alloca() {
+  %size = llvm.mlir.constant(4 : i64) : i64
+  // CHECK-DAG: %[[SRC:.+]] = llvm.alloca {{.+}}alignment = 1 : i64
+  // CHECK-DAG: %[[DST:.+]] = llvm.alloca {{.+}}alignment = 16 : i64
+  // CHECK: "llvm.intr.memcpy"(%[[DST]], %[[SRC]]
   %unaligned = llvm.alloca %size x i16 { alignment = 1 } : (i64) -> !llvm.ptr
   llvm.call @aligned_byval_arg(%unaligned) : (!llvm.ptr) -> ()
+  llvm.return
+}
+
+// -----
+
+llvm.func @aligned_byval_arg(%ptr : !llvm.ptr { llvm.byval = i16, llvm.align = 16 }) attributes {memory = #llvm.memory_effects<other = read, argMem = read, inaccessibleMem = read>} {
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @test_byval_aligned_alloca
+llvm.func @test_byval_aligned_alloca() {
   // CHECK-NOT: memcpy
+  %size = llvm.mlir.constant(1 : i64) : i64
   %aligned = llvm.alloca %size x i16 { alignment = 16 } : (i64) -> !llvm.ptr
   llvm.call @aligned_byval_arg(%aligned) : (!llvm.ptr) -> ()
   llvm.return
 }
 
 // -----
 
 llvm.mlir.global private @unaligned_global(42 : i64) : i64
 llvm.mlir.global private @aligned_global(42 : i64) { alignment = 64 } : i64
 
 llvm.func @aligned_byval_arg(%ptr : !llvm.ptr { llvm.byval = i16, llvm.align = 16 }) attributes {memory = #llvm.memory_effects<other = read, argMem = read, inaccessibleMem = read>} {
   llvm.return
 }
 
 // CHECK-LABEL: llvm.func @test_byval_global
 llvm.func @test_byval_global() {
   // Make sure only the unaligned global triggers a memcpy.
-  // CHECK: %[[UNALIGNED:.+]] = llvm.mlir.addressof @unaligned_global
-  // CHECK: %[[ALLOCA:.+]] = llvm.alloca
+  // CHECK-DAG: %[[UNALIGNED:.+]] = llvm.mlir.addressof @unaligned_global
+  // CHECK-DAG: %[[ALLOCA:.+]] = llvm.alloca
   // CHECK: "llvm.intr.memcpy"(%[[ALLOCA]], %[[UNALIGNED]]
   // CHECK-NOT: llvm.alloca
   %unaligned = llvm.mlir.addressof @unaligned_global : !llvm.ptr
   llvm.call @aligned_byval_arg(%unaligned) : (!llvm.ptr) -> ()
   %aligned = llvm.mlir.addressof @aligned_global : !llvm.ptr
   llvm.call @aligned_byval_arg(%aligned) : (!llvm.ptr) -> ()
   llvm.return
 }