[mlir][LLVM] Handle access groups during inlining

[zero9178]

Summary

Handling access groups is luckily rather trivial: Any access groups from the call instruction are simply appended to any memory operations.
This is similar to one of the steps when handling alias scopes.
This patch nevertheless implements it as a separate function purely for readability purposes as it uses a different interface than alias scopes.

[gysit]

LGTM modulo minor comment.

mlir/test/Dialect/LLVMIR/inlining.mlir

588

Can you maybe add another memory operation in the caller that is also marked with the #caller interface. That way we test the same access group is used for both operations.

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

 
 /// Handles all interactions with alias scopes during inlining.
 static void handleAliasScopes(Operation *call,
                               iterator_range<Region::iterator> inlinedBlocks) {
   deepCloneAliasScopes(inlinedBlocks);
   appendCallOpAliasScopes(call, inlinedBlocks);
 }
 
+/// Appends any access groups of the call operation to any inlined memory
+/// operation.
+static void handleAccessGroups(Operation *call,
+                               iterator_range<Region::iterator> inlinedBlocks) {
+  auto callAccessGroupInterface = dyn_cast<LLVM::AccessGroupOpInterface>(call);
+  if (!callAccessGroupInterface)
+    return;
+
+  auto accessGroups = callAccessGroupInterface.getAccessGroupsOrNull();
+  if (!accessGroups)
+    return;
+
+  // Simply append the call op's access groups to any operation implementing
+  // AccessGroupOpInterface.
+  for (Block &block : inlinedBlocks)
+    for (auto accessGroupOpInterface :
+         block.getOps<LLVM::AccessGroupOpInterface>())
+      accessGroupOpInterface.setAccessGroups(concatArrayAttr(
+          accessGroupOpInterface.getAccessGroupsOrNull(), accessGroups));
+}
+
 /// If `requestedAlignment` is higher than the alignment specified on `alloca`,
 /// realigns `alloca` if this does not exceed the natural stack alignment.
 /// Returns the post-alignment of `alloca`, whether it was realigned or not.
 static unsigned tryToEnforceAllocaAlignment(LLVM::AllocaOp alloca,
                                             unsigned requestedAlignment,
                                             DataLayout const &dataLayout) {
   unsigned allocaAlignment = alloca.getAlignment().value_or(1);
   if (requestedAlignment <= allocaAlignment)
struct LLVMInlinerInterface : public DialectInlinerInterface {
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
   /// Conservative allowlist of operations supported so far.
   bool isLegalToInline(Operation *op, Region *, bool, IRMapping &) const final {
     if (isPure(op))
       return true;
-    // Some attributes on memory operations require handling during
-    // inlining. Since this is not yet implemented, refuse to inline memory
-    // operations that have any of these attributes.
-    if (auto iface = dyn_cast<LLVM::AccessGroupOpInterface>(op)) {
-      if (iface.getAccessGroupsOrNull()) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "Cannot inline: unhandled access group metadata\n");
-        return false;
-      }
-    }
     // clang-format off
     if (isa<LLVM::AllocaOp,
             LLVM::AssumeOp,
             LLVM::AtomicRMWOp,
             LLVM::AtomicCmpXchgOp,
             LLVM::CallOp,
             LLVM::DbgDeclareOp,
             LLVM::DbgValueOp,
Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
     return argument;
   }
 
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const override {
     handleInlinedAllocas(call, inlinedBlocks);
     handleAliasScopes(call, inlinedBlocks);
+    handleAccessGroups(call, inlinedBlocks);
   }
 
   // Keeping this (immutable) state on the interface allows us to look up
   // StringAttrs instead of looking up strings, since StringAttrs are bound to
   // the current context and thus cannot be initialized as static fields.
   const DenseSet<StringAttr> disallowedFunctionAttrs;
 };
 
 } // end anonymous namespace
 
 void LLVM::detail::addLLVMInlinerInterface(LLVM::LLVMDialect *dialect) {
   dialect->addInterfaces<LLVMInlinerInterface>();
 }

mlir/test/Dialect/LLVMIR/inlining.mlir

 // CHECK: llvm.unreachable
 // CHECK: llvm.intr.stackrestore %[[STACK]]
 func.func @test_inline(%ptr : !llvm.ptr) -> i32 {
   %0 = call @inner_func_inlinable(%ptr) : (!llvm.ptr) -> i32
   return %0 : i32
 }
 
 // -----
-
-#group = #llvm.access_group<id = distinct[0]<>>
-
-llvm.func @inlinee(%ptr : !llvm.ptr) -> i32 {
-  %0 = llvm.load %ptr { access_groups = [#group] } : !llvm.ptr -> i32
-  llvm.return %0 : i32
-}
-
-// CHECK-LABEL: func @test_not_inline
-llvm.func @test_not_inline(%ptr : !llvm.ptr) -> i32 {
-  // CHECK-NEXT: llvm.call @inlinee
-  %0 = llvm.call @inlinee(%ptr) : (!llvm.ptr) -> (i32)
-  llvm.return %0 : i32
-}
-
-// -----
-
-#group = #llvm.access_group<id = distinct[0]<>>
-
-func.func private @with_mem_attr(%ptr : !llvm.ptr) {
-  %0 = llvm.mlir.constant(42 : i32) : i32
-  // Do not inline load/store operations that carry attributes requiring
-  // handling while inlining, until this is supported by the inliner.
-  llvm.store %0, %ptr { access_groups = [#group] }: i32, !llvm.ptr
-  return
-}
-
-// CHECK-LABEL: func.func @test_not_inline
-// CHECK-NEXT: call @with_mem_attr
-// CHECK-NEXT: return
-func.func @test_not_inline(%ptr : !llvm.ptr) {
-  call @with_mem_attr(%ptr) : (!llvm.ptr) -> ()
-  return
-}
-
-// -----
 // Check that llvm.return is correctly handled
 
 func.func @func(%arg0 : i32) -> i32  {
   llvm.return %arg0 : i32
 }
 // CHECK-LABEL: @llvm_ret
 // CHECK-NOT: call
 // CHECK:  return %arg0

 }
 
 // CHECK-LABEL: @test_disallow_arg_attr
 // CHECK-NEXT: llvm.call
 llvm.func @test_disallow_arg_attr(%ptr : !llvm.ptr) {
   llvm.call @disallowed_arg_attr(%ptr) : (!llvm.ptr) -> ()
   llvm.return
 }
+
+// -----
+
+#callee = #llvm.access_group<id = distinct[0]<>>
+#caller = #llvm.access_group<id = distinct[1]<>>
+
+llvm.func @inlinee(%ptr : !llvm.ptr) -> i32 {
+  %0 = llvm.load %ptr { access_groups = [#callee] } : !llvm.ptr -> i32
+  llvm.return %0 : i32
+}
+
+// CHECK-DAG: #[[$CALLEE:.*]] = #llvm.access_group<id = {{.*}}>
+// CHECK-DAG: #[[$CALLER:.*]] = #llvm.access_group<id = {{.*}}>
+
+// CHECK-LABEL: func @caller
+// CHECK: llvm.load
+// CHECK-SAME: access_groups = [#[[$CALLEE]], #[[$CALLER]]]
+llvm.func @caller(%ptr : !llvm.ptr) -> i32 {
+  %0 = llvm.call @inlinee(%ptr) { access_groups = [#caller] } : (!llvm.ptr) -> (i32)
+  llvm.return %0 : i32
+}
+
+// -----
+
+#caller = #llvm.access_group<id = distinct[1]<>>
+
+llvm.func @inlinee(%ptr : !llvm.ptr) -> i32 {
+  %0 = llvm.load %ptr : !llvm.ptr -> i32
+  llvm.return %0 : i32
+}
+
+// CHECK-DAG: #[[$CALLER:.*]] = #llvm.access_group<id = {{.*}}>
+
+// CHECK-LABEL: func @caller
+// CHECK: llvm.load
+// CHECK-SAME: access_groups = [#[[$CALLER]]]
+// CHECK: llvm.store
+// CHECK-SAME: access_groups = [#[[$CALLER]]]
+llvm.func @caller(%ptr : !llvm.ptr) -> i32 {
+  %c5 = llvm.mlir.constant(5 : i32) : i32
+  %0 = llvm.call @inlinee(%ptr) { access_groups = [#caller] } : (!llvm.ptr) -> (i32)
+  llvm.store %c5, %ptr { access_groups = [#caller] } : i32, !llvm.ptr
+  llvm.return %0 : i32
+}