[MLIR][LLVM] Add inlining support for loop annotations

[Dinistro]

This commit extends the LLVM dialect's inliner interface support updating loop annotation attributes. This is necessary because the loop annotations can contain debug locations, which are verified by LLVM's verifier. LLVM requires these locations to have the same scope as the function this attribute is contained in.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Christian Ulmann (Dinistro)

Changes

This commit extends the LLVM dialect's inliner interface support updating loop annotation attributes. This is necessary because the loop annotations can contain debug locations, which are verified by LLVM's verifier. LLVM requires these locations to have the same scope as the function this attribute is contained in.

---

Full diff: https://github.com/llvm/llvm-project/pull/94447.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp (+52)
• (added) mlir/test/Dialect/LLVMIR/inlining-loop-annotation.mlir (+56)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 4a6154ea6d300..010fde94dced6 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -513,6 +513,57 @@ static void handleAccessGroups(Operation *call,
           accessGroupOpInterface.getAccessGroupsOrNull(), accessGroups));
 }
 
+/// Updates locations inside loop annotations to reflect that they were inlined.
+static void
+handleLoopAnnotations(Operation *call,
+                      iterator_range<Region::iterator> inlinedBlocks) {
+  // Attempt to extract a DISubprogram from the callee.
+  auto func = call->getParentOfType<FunctionOpInterface>();
+  if (!func)
+    return;
+  LocationAttr funcLoc = func->getLoc();
+  auto fusedLoc = dyn_cast_if_present<FusedLoc>(funcLoc);
+  if (!fusedLoc)
+    return;
+  auto scope =
+      dyn_cast_if_present<LLVM::DISubprogramAttr>(fusedLoc.getMetadata());
+  if (!scope)
+    return;
+
+  // Helper to build a new fused location that reflects the inlining of the loop
+  // annotation.
+  auto updateLoc = [&](FusedLoc loc) -> FusedLoc {
+    if (!loc)
+      return {};
+    Location callSiteLoc = CallSiteLoc::get(loc, call->getLoc());
+    return FusedLoc::get(loc.getContext(), callSiteLoc, scope);
+  };
+
+  AttrTypeReplacer replacer;
+  replacer.addReplacement([&](LLVM::LoopAnnotationAttr loopAnnotation)
+                              -> std::pair<Attribute, WalkResult> {
+    FusedLoc newStartLoc = updateLoc(loopAnnotation.getStartLoc());
+    FusedLoc newEndLoc = updateLoc(loopAnnotation.getEndLoc());
+    if (!newStartLoc && !newEndLoc)
+      return {loopAnnotation, WalkResult::advance()};
+    auto newLoopAnnotation = LLVM::LoopAnnotationAttr::get(
+        loopAnnotation.getContext(), loopAnnotation.getDisableNonforced(),
+        loopAnnotation.getVectorize(), loopAnnotation.getInterleave(),
+        loopAnnotation.getUnroll(), loopAnnotation.getUnrollAndJam(),
+        loopAnnotation.getLicm(), loopAnnotation.getDistribute(),
+        loopAnnotation.getPipeline(), loopAnnotation.getPeeled(),
+        loopAnnotation.getUnswitch(), loopAnnotation.getMustProgress(),
+        loopAnnotation.getIsVectorized(), newStartLoc, newEndLoc,
+        loopAnnotation.getParallelAccesses());
+    // Needs to advance, as loop annotations can be nested.
+    return {newLoopAnnotation, WalkResult::advance()};
+  });
+
+  for (Block &block : inlinedBlocks)
+    for (Operation &op : block)
+      replacer.recursivelyReplaceElementsIn(&op);
+}
+
 /// 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.
@@ -787,6 +838,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
     handleInlinedAllocas(call, inlinedBlocks);
     handleAliasScopes(call, inlinedBlocks);
     handleAccessGroups(call, inlinedBlocks);
+    handleLoopAnnotations(call, inlinedBlocks);
   }
 
   // Keeping this (immutable) state on the interface allows us to look up
diff --git a/mlir/test/Dialect/LLVMIR/inlining-loop-annotation.mlir b/mlir/test/Dialect/LLVMIR/inlining-loop-annotation.mlir
new file mode 100644
index 0000000000000..e218c151400c7
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/inlining-loop-annotation.mlir
@@ -0,0 +1,56 @@
+// RUN: mlir-opt %s -inline -split-input-file | FileCheck %s
+
+#di_file = #llvm.di_file<"file.mlir" in "/">
+
+// CHECK: #[[START_ORIGINAL:.*]] = loc({{.*}}:42
+#loc1 = loc("test.mlir":42:4)
+// CHECK: #[[END_ORIGINAL:.*]] = loc({{.*}}:52
+#loc2 = loc("test.mlir":52:4)
+#loc3 = loc("test.mlir":62:4)
+// CHECK: #[[CALL_ORIGINAL:.*]] = loc({{.*}}:72
+#loc4 = loc("test.mlir":72:4)
+
+#di_compile_unit = #llvm.di_compile_unit<id = distinct[0]<>, sourceLanguage = DW_LANG_C, file = #di_file, isOptimized = false, emissionKind = None>
+// CHECK: #[[CALLEE_DI:.*]] = #llvm.di_subprogram<{{.*}}, name = "callee"
+#di_subprogram_callee = #llvm.di_subprogram<compileUnit = #di_compile_unit, scope = #di_file, name = "callee", file = #di_file, subprogramFlags = Definition>
+
+// CHECK: #[[CALLER_DI:.*]] = #llvm.di_subprogram<{{.*}}, name = "caller"
+#di_subprogram_caller = #llvm.di_subprogram<compileUnit = #di_compile_unit, scope = #di_file, name = "caller", file = #di_file, subprogramFlags = Definition>
+
+// CHECK: #[[START_FUSED_ORIGINAL:.*]] = loc(fused<#[[CALLEE_DI]]>[#[[START_ORIGINAL]]
+#start_loc_fused = loc(fused<#di_subprogram_callee>[#loc1])
+// CHECK: #[[END_FUSED_ORIGINAL:.*]] = loc(fused<#[[CALLEE_DI]]>[#[[END_ORIGINAL]]
+#end_loc_fused= loc(fused<#di_subprogram_callee>[#loc2])
+#caller_loc= loc(fused<#di_subprogram_caller>[#loc3])
+// CHECK: #[[CALL_FUSED:.*]] = loc(fused<#[[CALLER_DI]]>[#[[CALL_ORIGINAL]]
+#call_loc= loc(fused<#di_subprogram_caller>[#loc4])
+
+#loopMD = #llvm.loop_annotation<
+        startLoc = #start_loc_fused,
+        endLoc = #end_loc_fused>
+
+// CHECK: #[[START_CALLSITE_LOC:.*]] = loc(callsite(#[[START_FUSED_ORIGINAL]] at #[[CALL_FUSED]]
+// CHECK: #[[END_CALLSITE_LOC:.*]] = loc(callsite(#[[END_FUSED_ORIGINAL]] at #[[CALL_FUSED]]
+// CHECK: #[[START_FUSED_LOC:.*]] = loc(fused<#[[CALLER_DI]]>[#[[START_CALLSITE_LOC]]
+// CHECK: #[[END_FUSED_LOC:.*]] = loc(fused<#[[CALLER_DI]]>[
+// CHECK: #[[LOOP_ANNOT:.*]] = #llvm.loop_annotation<
+// CHECK-SAME: startLoc = #[[START_FUSED_LOC]], endLoc = #[[END_FUSED_LOC]]>
+
+llvm.func @cond() -> i1
+
+llvm.func @callee() {
+  llvm.br ^head
+^head:
+  %c = llvm.call @cond() : () -> i1
+  llvm.cond_br %c, ^head, ^exit {loop_annotation = #loopMD}
+^exit:
+  llvm.return
+}
+
+// CHECK: @loop_annotation
+llvm.func @loop_annotation() {
+  // CHECK: llvm.cond_br
+  // CHECK-SAME: {loop_annotation = #[[LOOP_ANNOT]]
+  llvm.call @callee() : () -> () loc(#call_loc)
+  llvm.return
+} loc(#caller_loc)

[gysit]

LGTM