[MLIR][LLVM] Copy byval attributes during inlining.

Support inlining of function calls with the byval attribute on function
arguments by copying the pointee into a newly alloca'ed pointer at the
callsite before inlining.

The alignment attribute is not yet taken into account.

Reviewed By: ftynse, gysit

Differential Revision: https://reviews.llvm.org/D146616

Author: definelicht

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

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "TypeDetail.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Dialect/LLVMIR/LLVMInterfaces.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
@@ -2854,6 +2855,39 @@ static void moveConstantAllocasToEntryBlock(
   }
 }
 
+static Value handleByValArgument(OpBuilder &builder, Operation *callable,
+                                 Value argument,
+                                 NamedAttribute byValAttribute) {
+  auto func = cast<LLVM::LLVMFuncOp>(callable);
+  LLVM::MemoryEffectsAttr memoryEffects = func.getMemoryAttr();
+  // If there is no memory effects attribute, assume that the function is
+  // not read-only.
+  bool isReadOnly = memoryEffects &&
+                    memoryEffects.getArgMem() != ModRefInfo::ModRef &&
+                    memoryEffects.getArgMem() != ModRefInfo::Mod;
+  if (isReadOnly)
+    return argument;
+  // Resolve the pointee type and its size.
+  auto ptrType = cast<LLVM::LLVMPointerType>(argument.getType());
+  Type elementType = cast<TypeAttr>(byValAttribute.getValue()).getValue();
+  unsigned int typeSize =
+      DataLayout(callable->getParentOfType<DataLayoutOpInterface>())
+          .getTypeSize(elementType);
+  // Allocate the new value on the stack.
+  Value one = builder.create<LLVM::ConstantOp>(
+      func.getLoc(), builder.getI64Type(), builder.getI64IntegerAttr(1));
+  Value allocaOp =
+      builder.create<LLVM::AllocaOp>(func.getLoc(), ptrType, elementType, one);
+  // Copy the pointee to the newly allocated value.
+  Value copySize = builder.create<LLVM::ConstantOp>(
+      func.getLoc(), builder.getI64Type(), builder.getI64IntegerAttr(typeSize));
+  Value isVolatile = builder.create<LLVM::ConstantOp>(
+      func.getLoc(), builder.getI1Type(), builder.getBoolAttr(false));
+  builder.create<LLVM::MemcpyOp>(func.getLoc(), allocaOp, argument, copySize,
+                                 isVolatile);
+  return allocaOp;
+}
+
 namespace {
 struct LLVMInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
@@ -2866,8 +2900,19 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
     auto funcOp = dyn_cast<LLVM::LLVMFuncOp>(callable);
     if (!callOp || !funcOp)
       return false;
-    // TODO: Handle argument and result attributes;
-    if (funcOp.getArgAttrs() || funcOp.getResAttrs())
+    if (auto attrs = funcOp.getArgAttrs()) {
+      for (Attribute attr : *attrs) {
+        auto attrDict = cast<DictionaryAttr>(attr);
+        for (NamedAttribute attr : attrDict) {
+          if (attr.getName() == LLVMDialect::getByValAttrName())
+            continue;
+          // TODO: Handle all argument attributes;
+          return false;
+        }
+      }
+    }
+    // TODO: Handle result attributes;
+    if (funcOp.getResAttrs())
       return false;
     // TODO: Handle exceptions.
     if (funcOp.getPersonality())
@@ -2942,6 +2987,14 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
       dst.replaceAllUsesWith(src);
   }
 
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (auto attr = argumentAttrs.getNamed(LLVMDialect::getByValAttrName()))
+      return handleByValArgument(builder, callable, argument, *attr);
+    return argument;
+  }
+
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const override {

mlir/test/Dialect/LLVMIR/inlining.mlir

@@ -187,20 +187,6 @@ llvm.func @caller() {
 
 // -----
 
-llvm.func @callee(%ptr : !llvm.ptr {llvm.byval = !llvm.ptr}) -> (!llvm.ptr) {
-  llvm.return %ptr : !llvm.ptr
-}
-
-// CHECK-LABEL: llvm.func @caller
-// CHECK-NEXT: llvm.call @callee
-// CHECK-NEXT: return
-llvm.func @caller(%ptr : !llvm.ptr) -> (!llvm.ptr) {
-  %0 = llvm.call @callee(%ptr) : (!llvm.ptr) -> (!llvm.ptr)
-  llvm.return %0 : !llvm.ptr
-}
-
-// -----
-
 llvm.func @static_alloca() -> f32 {
   %0 = llvm.mlir.constant(4 : i32) : i32
   %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
@@ -349,3 +335,47 @@ llvm.func @test_inline(%cond0 : i1, %cond1 : i1, %funcArg : f32) -> f32 {
 ^bb3(%blockArg: f32):
   llvm.return %blockArg : f32
 }
+
+// -----
+
+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
+// CHECK: %[[ALLOCA:.+]] = llvm.alloca %{{.+}} x f64
+// CHECK: "llvm.intr.memcpy"(%[[ALLOCA]], %[[PTR]]
+llvm.func @test_byval(%ptr : !llvm.ptr) {
+  llvm.call @with_byval_arg(%ptr) : (!llvm.ptr) -> ()
+  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
+}
+
+// CHECK-LABEL: llvm.func @test_byval_read_only
+// CHECK-NOT: llvm.call
+// CHECK-NEXT: llvm.return
+llvm.func @test_byval_read_only(%ptr : !llvm.ptr) {
+  llvm.call @with_byval_arg(%ptr) : (!llvm.ptr) -> ()
+  llvm.return
+}
+
+// -----
+
+llvm.func @with_byval_arg(%ptr : !llvm.ptr { llvm.byval = f64 }) attributes {memory = #llvm.memory_effects<other = readwrite, argMem = write, inaccessibleMem = readwrite>} {
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @test_byval_write_only
+// CHECK-SAME: %[[PTR:[a-zA-Z0-9_]+]]: !llvm.ptr
+// CHECK: %[[ALLOCA:.+]] = llvm.alloca %{{.+}} x f64
+// CHECK: "llvm.intr.memcpy"(%[[ALLOCA]], %[[PTR]]
+llvm.func @test_byval_write_only(%ptr : !llvm.ptr) {
+  llvm.call @with_byval_arg(%ptr) : (!llvm.ptr) -> ()
+  llvm.return
+}