[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

2995–2997

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).

3002

Nit: consider early return here.

[gysit]

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

3001

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

 //===- LLVMDialect.cpp - LLVM IR Ops and Dialect registration -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the types and operation details for the LLVM IR dialect in
 // MLIR, and the LLVM IR dialect.  It also registers the dialect.
 //
 //===----------------------------------------------------------------------===//
 #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"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
         continue;
       builder.create<LLVM::LifetimeEndOp>(
           allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
           allocaOp.getResult());
     }
   }
 }
 
+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;
 
   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 {
+    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 {
     // 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.
     moveConstantAllocasToEntryBlock(inlinedBlocks);
   }
 };
 } // end anonymous namespace
 

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_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
+}