[Mlir-commits] [mlir] 6168e82 - [MLIR][LLVM] Add inlining support for loop annotations (#94447)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 5 05:08:05 PDT 2024
Author: Christian Ulmann
Date: 2024-06-05T14:08:00+02:00
New Revision: 6168e82c1e17c58ebc9c0b6c8f2273fd9a610977
URL: https://github.com/llvm/llvm-project/commit/6168e82c1e17c58ebc9c0b6c8f2273fd9a610977
DIFF: https://github.com/llvm/llvm-project/commit/6168e82c1e17c58ebc9c0b6c8f2273fd9a610977.diff
LOG: [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.
Added:
mlir/test/Dialect/LLVMIR/inlining-loop-annotation.mlir
Modified:
mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 5552dc5e244b8..cf3369d053fae 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/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
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)
More information about the Mlir-commits
mailing list