[MLIR][LLVM] Realign allocas to avoid dynamic realignment in inliner.

[definelicht]

Summary

When the natural stack alignment is set and is larger than or equal to
the target alignment required by a read-only byval argument defined by
an alloca, avoid the copy by just realigning the alloca to the target
alignment.

This also includes a bugfix for passing a null DataLayoutInterface to
the DataLayout constructor when no parent op defines
DataLayoutInterface, and will now pass ModuleOp instead in this case.

[definelicht]

Remove assert, since this seems to often not be satisfied.

[definelicht]

Use existing closest function instead of reinventing it.

[definelicht]

Add fix to accessing scope when it's null in DataLayout.

[gysit]

LGTM

mlir/test/Dialect/LLVMIR/inlining.mlir

528

nit: maybe rename declared_func to use_ptr to make sure it is obvious we are not testing that declarations are not inlined here. Using a load may also be a good option to avoid DCE.

[definelicht]

Implement review comment and extend documentation a bit.

[definelicht]

Optimistically assume that natural stack alignment is not exceeded if it is not specified. Reorganize the code to avoid redundant casts.

[gysit]

Still LGTM!

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

 //===- LLVMInlining.cpp - LLVM inlining interface and logic -----*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Logic for inlining LLVM functions and the definition of the
 // LLVMInliningInterface.
 //
 //===----------------------------------------------------------------------===//
 
 #include "LLVMInlining.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "llvm-inliner"
 
 using namespace mlir;
 
 /// Check whether the given alloca is an input to a lifetime intrinsic,
static Value handleByValArgumentInit(OpBuilder &builder, Location loc,
   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;
 }
 
+/// Returns the data layout defined by the closest defining scope of `op` that
+/// exposes DataLayoutOpInterface (or the ModuleOp if no such scope exists).
+static DataLayout getDataLayoutOf(Operation *op) {
+  if (auto interface = op->getParentOfType<DataLayoutOpInterface>())
+    return DataLayout(interface);
+  return DataLayout(op->getParentOfType<ModuleOp>());
+}
+
 /// Handles a function argument marked with the byval attribute by introducing a
 /// memcpy if necessary, either due to the pointee being writeable in the
 /// callee, and/or due to an alignment mismatch. `requestedAlignment` specifies
 /// the alignment set in the "align" argument attribute (or 1 if no align
 /// attribute was set).
 static Value handleByValArgument(OpBuilder &builder, Operation *callable,
                                  Value argument, Type elementType,
                                  unsigned requestedAlignment) {
   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() != LLVM::ModRefInfo::ModRef &&
                     memoryEffects.getArgMem() != LLVM::ModRefInfo::Mod;
   // Check if there's an alignment mismatch requiring us to copy.
-  DataLayout dataLayout(callable->getParentOfType<DataLayoutOpInterface>());
+  DataLayout dataLayout = getDataLayoutOf(callable);
   unsigned minimumAlignment = dataLayout.getTypeABIAlignment(elementType);
-  if (isReadOnly && (requestedAlignment <= minimumAlignment ||
-                     getAlignmentOf(argument) >= requestedAlignment))
+  if (isReadOnly) {
+    if (requestedAlignment <= minimumAlignment)
+      return argument;
+    if (getAlignmentOf(argument) >= requestedAlignment)
+      return argument;
+    // If this is a stack-allocated value and the requested alignment does not
+    // exceed the natural stack alignment, realign the alloca to avoid a copy.
+    // If the requested alignment exceeds the natural stack alignment, this will
+    // trigger a dynamic stack realignment, so we don't do this unless it's
+    // really necessary. If no stack alignment is specified, conservatively
+    // don't try realigning the alloca.
+    unsigned stackAlignment = dataLayout.getStackAlignment() >> 3;
+    if (requestedAlignment <= stackAlignment) {
+      if (auto alloca = argument.getDefiningOp<LLVM::AllocaOp>()) {
+        alloca.setAlignment(requestedAlignment);
         return argument;
+      }
+    }
+  }
   unsigned targetAlignment = std::max(requestedAlignment, minimumAlignment);
   return handleByValArgumentInit(builder, func.getLoc(), argument, elementType,
                                  dataLayout.getTypeSize(elementType),
                                  targetAlignment);
 }
 
 /// Returns true if the given argument or result attribute is supported by the
 /// inliner, false otherwise.

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

void checkMissingLayout(DataLayoutSpecInterface originalLayout, OpTy op) {
   }
 }
 
 mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
 
 mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
     : originalLayout(getCombinedDataLayout(op)), scope(op),
       allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
+  assert(op);
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
 #endif
 }
 
 mlir::DataLayout::DataLayout(ModuleOp op)
     : originalLayout(getCombinedDataLayout(op)), scope(op),
       allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
+  assert(op);
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
 #endif
 }
 
 mlir::DataLayout mlir::DataLayout::closest(Operation *op) {
   // Search the closest parent either being a module operation or implementing

mlir/test/Dialect/LLVMIR/inlining.mlir

llvm.func @test_byval_global() {
   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
 }
 
 // -----
 
+module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.stack_alignment", 32 : i32>>} {
+
+llvm.func @aligned_byval_arg(%ptr : !llvm.ptr { llvm.byval = i16, llvm.align = 4 }) attributes {memory = #llvm.memory_effects<other = read, argMem = read, inaccessibleMem = read>} {
+  llvm.return
+}
+
+llvm.func @declared_func(%ptr : !llvm.ptr)
+
+// CHECK-LABEL: llvm.func @test_byval_realign_alloca
+llvm.func @test_byval_realign_alloca() {
+  %size = llvm.mlir.constant(1 : i64) : i64
+  // Because the natural stack alignment is larger than the target alignment,
+  // the alloca will just be realigned instead of doing a copy.
+  // CHECK-NOT: llvm.alloca {{.+}}alignment = 2
+  // CHECK: llvm.alloca {{.+}}alignment = 4
+  // CHECK-NOT: memcpy
+  %ptr = llvm.alloca %size x i16 { alignment = 2 } : (i64) -> !llvm.ptr
+  llvm.call @aligned_byval_arg(%ptr) : (!llvm.ptr) -> ()
+  llvm.call @declared_func(%ptr) : (!llvm.ptr) -> ()
+  llvm.return
+}
+
+}
+
+// -----
+
 llvm.func @ignored_attrs(%ptr : !llvm.ptr { llvm.inreg, llvm.nocapture, llvm.nofree, llvm.preallocated = i32, llvm.returned, llvm.alignstack = 32 : i64, llvm.writeonly, llvm.noundef, llvm.nonnull }, %x : i32 { llvm.zeroext }) -> (!llvm.ptr { llvm.noundef, llvm.inreg, llvm.nonnull }) {
   llvm.return %ptr : !llvm.ptr
 }
 
 // CHECK-LABEL: @test_ignored_attrs
 // CHECK-NOT: llvm.call
 // CHECK-NEXT: llvm.return
 llvm.func @test_ignored_attrs(%ptr : !llvm.ptr, %x : i32) {