[MLIR][LLVM] Add inlining support for loop annotations (#94447)

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.

Author: Dinistro

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

@@ -511,6 +511,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.
@@ -784,6 +835,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

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)