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

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

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)