[Mlir-commits] [mlir] [MLIR][LLVM] Add inlining support for loop annotations (PR #94447)
Christian Ulmann
llvmlistbot at llvm.org
Wed Jun 5 02:21:22 PDT 2024
https://github.com/Dinistro created https://github.com/llvm/llvm-project/pull/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.
>From 8a13d5fa14b0448f1ca636570818f8a86e03e68d Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Wed, 5 Jun 2024 09:16:51 +0000
Subject: [PATCH] [MLIR][LLVM] Add inlining support for loop annotations
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.
---
mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp | 52 +++++++++++++++++
.../LLVMIR/inlining-loop-annotation.mlir | 56 +++++++++++++++++++
2 files changed, 108 insertions(+)
create mode 100644 mlir/test/Dialect/LLVMIR/inlining-loop-annotation.mlir
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)
More information about the Mlir-commits
mailing list