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

[definelicht]

Summary

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.

[ftynse]

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

2961–2963

Please expand auto unless the type is obvious from context (e.g., there's a cast on the RHS) or impossible to spell (e.g., lambdas).

2968

Nit: consider early return here.

[gysit]

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

2967

Should this also be a check for write-only accesses?

memoryEffects &&
memoryEffects.getArgMem() != ModRefInfo::ModRef &&
memoryEffects.getArgMem() != ModRefInfo::Mod

Or does ModRef also match pure write accesses?

[definelicht]

Addressed comments by moving byval handling into a separate function so we can do early returns without breaking the control flow of other potential function argument attribute handling.

[definelicht]

Add missing check for write-only.

[gysit]

Thanks LGTM!

[definelicht]

Get data layout from callable, omit call argument, and remove auto.

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

struct LLVMInlinerInterface : public DialectInlinerInterface {
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
     if (!wouldBeCloned)
       return false;
     auto callOp = dyn_cast<LLVM::CallOp>(call);
     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())
       return false;
     if (funcOp.getPassthrough()) {
       // TODO: Used attributes should not be passthrough.
       DenseSet<StringAttr> disallowed(
           {StringAttr::get(funcOp->getContext(), "noduplicate"),
void handleTerminator(Operation *op,
 
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &[dst, src] :
          llvm::zip(valuesToRepl, returnOp.getOperands()))
       dst.replaceAllUsesWith(src);
   }
 
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    Value value = argument;
+    // Handle the byval function argument attribute.
+    if (auto attr = argumentAttrs.getNamed(LLVMDialect::getByValAttrName())) {
+      auto func = cast<LLVM::LLVMFuncOp>(callable);
+      auto 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;
+      if (!isReadOnly) {
+        // Resolve the pointee type and its size.
+        auto ptrType = cast<LLVM::LLVMPointerType>(argument.getType());
+        auto elementType = cast<TypeAttr>(attr->getValue()).getValue();
+        unsigned int typeSize =
+            DataLayout(call->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);
+        value = allocaOp;
+      }
+    }
+    return value;
+  }
+
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const override {
     // Alloca operations with a constant size that were in the entry block of
     // the callee should be moved to the entry block of the caller, as this will
     // fold into prologue/epilogue code during code generation.
     // This is not implemented as a standalone pattern because we need to know
     // which newly inlined block was previously the entry block of the callee.

mlir/test/Dialect/LLVMIR/inlining.mlir

llvm.func @caller() {
   llvm.call @callee_presplitcoroutine() : () -> ()
   llvm.call @callee_returns_twice() : () -> ()
   llvm.call @callee_strictfp() : () -> ()
   llvm.return
 }
 
 // -----
 
-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
   %2 = llvm.load %1 : !llvm.ptr -> f32
   llvm.return %2 : f32
 }
 
 llvm.func @dynamic_alloca(%size : i32) -> f32 {
^bb1:
   // CHECK-NOT: llvm.intr.lifetime.end
   %0 = llvm.call @alloca_with_lifetime(%cond1) : (i1) -> f32
   llvm.br ^bb3(%0: f32)
 ^bb2:
   llvm.br ^bb3(%funcArg: 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
+// 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
+}