[Mlir-commits] [mlir] 62d7d94 - [mlir][LLVM] Support locations in loop annotation

Christian Ulmann llvmlistbot at llvm.org
Fri May 5 00:54:51 PDT 2023


Author: Christian Ulmann
Date: 2023-05-05T07:54:22Z
New Revision: 62d7d94c2ea4f04d9c6f3feaa525ac4971fef815

URL: https://github.com/llvm/llvm-project/commit/62d7d94c2ea4f04d9c6f3feaa525ac4971fef815
DIFF: https://github.com/llvm/llvm-project/commit/62d7d94c2ea4f04d9c6f3feaa525ac4971fef815.diff

LOG: [mlir][LLVM] Support locations in loop annotation

This commit introduces support for locations as part of the loop
annotation attribute. These locations indicate the start and the end of
the loop.

Reviewed By: gysit

Differential Revision: https://reviews.llvm.org/D149858

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
    mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp
    mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h
    mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp
    mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h
    mlir/lib/Target/LLVMIR/ModuleImport.cpp
    mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
    mlir/test/Dialect/LLVMIR/loop-metadata.mlir
    mlir/test/Target/LLVMIR/Import/metadata-loop.ll
    mlir/test/Target/LLVMIR/loop-metadata.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index b6b56685df863..05ee0832360dc 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -206,7 +206,9 @@ def LoopAnnotationAttr : LLVM_Attr<"LoopAnnotation", "loop_annotation"> {
     OptionalParameter<"LoopUnswitchAttr">:$unswitch,
     OptionalParameter<"BoolAttr">:$mustProgress,
     OptionalParameter<"BoolAttr">:$isVectorized,
-    OptionalArrayRefParameter<"SymbolRefAttr">:$parallelAccesses
+    OptionalArrayRefParameter<"SymbolRefAttr">:$parallelAccesses,
+    OptionalParameter<"FusedLoc">:$startLoc,
+    OptionalParameter<"FusedLoc">:$endLoc
   );
 
   let assemblyFormat = "`<` struct(params) `>`";

diff  --git a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp
index 000349091afce..a9f080ef0044f 100644
--- a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp
+++ b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp
@@ -23,7 +23,8 @@ struct LoopMetadataConversion {
   /// Converts this structs loop metadata node into a LoopAnnotationAttr.
   LoopAnnotationAttr convert();
 
-  LogicalResult initPropertyMap();
+  /// Initializes the shared state for the conversion member functions.
+  LogicalResult initConversionState();
 
   /// Helper function to get and erase a property.
   const llvm::MDNode *lookupAndEraseProperty(StringRef name);
@@ -53,7 +54,10 @@ struct LoopMetadataConversion {
   FailureOr<LoopPeeledAttr> convertPeeledAttr();
   FailureOr<LoopUnswitchAttr> convertUnswitchAttr();
   FailureOr<SmallVector<SymbolRefAttr>> convertParallelAccesses();
+  FusedLoc convertStartLoc();
+  FailureOr<FusedLoc> convertEndLoc();
 
+  llvm::SmallVector<llvm::DILocation *, 2> locations;
   llvm::StringMap<const llvm::MDNode *> propertyMap;
   const llvm::MDNode *node;
   Location loc;
@@ -62,16 +66,17 @@ struct LoopMetadataConversion {
 };
 } // namespace
 
-LogicalResult LoopMetadataConversion::initPropertyMap() {
+LogicalResult LoopMetadataConversion::initConversionState() {
   // Check if it's a valid node.
   if (node->getNumOperands() == 0 ||
       dyn_cast<llvm::MDNode>(node->getOperand(0)) != node)
     return emitWarning(loc) << "invalid loop node";
 
   for (const llvm::MDOperand &operand : llvm::drop_begin(node->operands())) {
-    // Skip over DILocations.
-    if (isa<llvm::DILocation>(operand))
+    if (auto *diLoc = dyn_cast<llvm::DILocation>(operand)) {
+      locations.push_back(diLoc);
       continue;
+    }
 
     auto *property = dyn_cast<llvm::MDNode>(operand);
     if (!property)
@@ -405,8 +410,25 @@ LoopMetadataConversion::convertParallelAccesses() {
   return refs;
 }
 
+FusedLoc LoopMetadataConversion::convertStartLoc() {
+  if (locations.empty())
+    return {};
+  return dyn_cast<FusedLoc>(
+      loopAnnotationImporter.moduleImport.translateLoc(locations[0]));
+}
+
+FailureOr<FusedLoc> LoopMetadataConversion::convertEndLoc() {
+  if (locations.size() < 2)
+    return FusedLoc();
+  if (locations.size() > 2)
+    return emitError(loc)
+           << "expected loop metadata to have at most two DILocations";
+  return dyn_cast<FusedLoc>(
+      loopAnnotationImporter.moduleImport.translateLoc(locations[1]));
+}
+
 LoopAnnotationAttr LoopMetadataConversion::convert() {
-  if (failed(initPropertyMap()))
+  if (failed(initConversionState()))
     return {};
 
   FailureOr<BoolAttr> disableNonForced =
@@ -433,10 +455,14 @@ LoopAnnotationAttr LoopMetadataConversion::convert() {
     return {};
   }
 
+  FailureOr<FusedLoc> startLoc = convertStartLoc();
+  FailureOr<FusedLoc> endLoc = convertEndLoc();
+
   return createIfNonNull<LoopAnnotationAttr>(
       ctx, disableNonForced, vecAttr, interleaveAttr, unrollAttr,
       unrollAndJamAttr, licmAttr, distributeAttr, pipelineAttr, peeledAttr,
-      unswitchAttr, mustProgress, isVectorized, parallelAccesses);
+      unswitchAttr, mustProgress, isVectorized, parallelAccesses, startLoc,
+      endLoc);
 }
 
 LoopAnnotationAttr

diff  --git a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h
index 5d69a63a21502..7dcfc4c750d98 100644
--- a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h
+++ b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h
@@ -26,7 +26,8 @@ namespace detail {
 /// AccessGroupMetadataOps.
 class LoopAnnotationImporter {
 public:
-  explicit LoopAnnotationImporter(OpBuilder &builder) : builder(builder) {}
+  LoopAnnotationImporter(ModuleImport &moduleImport, OpBuilder &builder)
+      : moduleImport(moduleImport), builder(builder) {}
   LoopAnnotationAttr translateLoopAnnotation(const llvm::MDNode *node,
                                              Location loc);
 
@@ -44,6 +45,9 @@ class LoopAnnotationImporter {
   FailureOr<SmallVector<SymbolRefAttr>>
   lookupAccessGroupAttrs(const llvm::MDNode *node) const;
 
+  /// The ModuleImport owning this instance.
+  ModuleImport &moduleImport;
+
 private:
   /// Returns the LLVM metadata corresponding to a llvm loop metadata attribute.
   LoopAnnotationAttr lookupLoopMetadata(const llvm::MDNode *node) const {

diff  --git a/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp
index 6abcf6964abed..8e6906a215325 100644
--- a/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "LoopAnnotationTranslation.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
@@ -32,6 +33,7 @@ struct LoopAnnotationConversion {
   void convertBoolNode(StringRef name, BoolAttr attr, bool negated = false);
   void convertI32Node(StringRef name, IntegerAttr attr);
   void convertFollowupNode(StringRef name, LoopAnnotationAttr attr);
+  void convertLocation(FusedLoc attr);
 
   /// Conversion functions for each for each loop annotation sub-attribute.
   void convertLoopOptions(LoopVectorizeAttr options);
@@ -186,12 +188,33 @@ void LoopAnnotationConversion::convertLoopOptions(LoopUnswitchAttr options) {
               options.getPartialDisable());
 }
 
-llvm::MDNode *LoopAnnotationConversion::convert() {
+void LoopAnnotationConversion::convertLocation(FusedLoc location) {
+  auto localScopeAttr =
+      location.getMetadata().dyn_cast_or_null<DILocalScopeAttr>();
+  if (!localScopeAttr)
+    return;
+  auto *localScope = dyn_cast<llvm::DILocalScope>(
+      loopAnnotationTranslation.moduleTranslation.translateDebugInfo(
+          localScopeAttr));
+  if (!localScope)
+    return;
+  const llvm::Metadata *loc =
+      loopAnnotationTranslation.moduleTranslation.translateLoc(location,
+                                                               localScope);
+  metadataNodes.push_back(const_cast<llvm::Metadata *>(loc));
+}
 
+llvm::MDNode *LoopAnnotationConversion::convert() {
   // Reserve operand 0 for loop id self reference.
   auto dummy = llvm::MDNode::getTemporary(ctx, std::nullopt);
   metadataNodes.push_back(dummy.get());
 
+  if (FusedLoc startLoc = attr.getStartLoc())
+    convertLocation(startLoc);
+
+  if (FusedLoc endLoc = attr.getEndLoc())
+    convertLocation(endLoc);
+
   addUnitNode("llvm.loop.disable_nonforced", attr.getDisableNonforced());
   addUnitNode("llvm.loop.mustprogress", attr.getMustProgress());
   // "isvectorized" is encoded as an i32 value.

diff  --git a/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h
index e663f64b061fc..0e016d9cbeef0 100644
--- a/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h
+++ b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h
@@ -25,8 +25,10 @@ namespace detail {
 /// into a corresponding llvm::MDNodes.
 class LoopAnnotationTranslation {
 public:
-  LoopAnnotationTranslation(Operation *mlirModule, llvm::Module &llvmModule)
-      : mlirModule(mlirModule), llvmModule(llvmModule) {}
+  LoopAnnotationTranslation(ModuleTranslation &moduleTranslation,
+                            Operation *mlirModule, llvm::Module &llvmModule)
+      : moduleTranslation(moduleTranslation), mlirModule(mlirModule),
+        llvmModule(llvmModule) {}
 
   llvm::MDNode *translateLoopAnnotation(LoopAnnotationAttr attr, Operation *op);
 
@@ -43,6 +45,9 @@ class LoopAnnotationTranslation {
   /// referenced by the AccessGroupOpInterface or null if there are none.
   llvm::MDNode *getAccessGroups(AccessGroupOpInterface op) const;
 
+  /// The ModuleTranslation owning this instance.
+  ModuleTranslation &moduleTranslation;
+
 private:
   /// Returns the LLVM metadata corresponding to a llvm loop metadata attribute.
   llvm::MDNode *lookupLoopMetadata(Attribute options) const {

diff  --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 5c585e7104b67..8c5fe35a6f0dc 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -153,7 +153,7 @@ ModuleImport::ModuleImport(ModuleOp mlirModule,
       typeTranslator(*mlirModule->getContext()),
       debugImporter(std::make_unique<DebugImporter>(mlirModule)),
       loopAnnotationImporter(
-          std::make_unique<LoopAnnotationImporter>(builder)) {
+          std::make_unique<LoopAnnotationImporter>(*this, builder)) {
   builder.setInsertionPointToStart(mlirModule.getBody());
 }
 

diff  --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index f3f5b24876a20..f8854d75f972b 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -444,7 +444,7 @@ ModuleTranslation::ModuleTranslation(Operation *module,
       debugTranslation(
           std::make_unique<DebugTranslation>(module, *this->llvmModule)),
       loopAnnotationTranslation(std::make_unique<LoopAnnotationTranslation>(
-          module, *this->llvmModule)),
+          *this, module, *this->llvmModule)),
       typeTranslator(this->llvmModule->getContext()),
       iface(module->getContext()) {
   assert(satisfiesLLVMModule(mlirModule) &&

diff  --git a/mlir/test/Dialect/LLVMIR/loop-metadata.mlir b/mlir/test/Dialect/LLVMIR/loop-metadata.mlir
index 8841a1fad9186..631f27139ba9f 100644
--- a/mlir/test/Dialect/LLVMIR/loop-metadata.mlir
+++ b/mlir/test/Dialect/LLVMIR/loop-metadata.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// RUN: mlir-opt %s --split-input-file | mlir-opt --split-input-file | FileCheck %s
 
 // CHECK-DAG: #[[FOLLOWUP:.*]] = #llvm.loop_annotation<disableNonforced = true>
 #followup = #llvm.loop_annotation<disableNonforced = true>
@@ -80,3 +80,38 @@ llvm.metadata @metadata {
   llvm.access_group @group1
   llvm.access_group @group2
 }
+
+// -----
+
+#di_file = #llvm.di_file<"metadata-loop.ll" in "/">
+
+// CHECK: #[[START_LOC:.*]] = loc("loop-metadata.mlir":42:4)
+#loc1 = loc("loop-metadata.mlir":42:4)
+// CHECK: #[[END_LOC:.*]] = loc("loop-metadata.mlir":52:4)
+#loc2 = loc("loop-metadata.mlir":52:4)
+
+#di_compile_unit = #llvm.di_compile_unit<sourceLanguage = DW_LANG_C, file = #di_file, isOptimized = false, emissionKind = None>
+// CHECK: #[[SUBPROGRAM:.*]] = #llvm.di_subprogram<
+#di_subprogram = #llvm.di_subprogram<compileUnit = #di_compile_unit, scope = #di_file, name = "loop_locs", file = #di_file, subprogramFlags = Definition>
+
+// CHECK: #[[START_LOC_FUSED:.*]] = loc(fused<#[[SUBPROGRAM]]>[#[[START_LOC]]]
+#start_loc_fused = loc(fused<#di_subprogram>[#loc1])
+// CHECK: #[[END_LOC_FUSED:.*]] = loc(fused<#[[SUBPROGRAM]]>[#[[END_LOC]]]
+#end_loc_fused= loc(fused<#di_subprogram>[#loc2])
+
+// CHECK: #[[LOOP_ANNOT:.*]] = #llvm.loop_annotation<
+// CHECK-DAG: disableNonforced = false
+// CHECK-DAG: startLoc = #[[START_LOC_FUSED]]
+// CHECK-DAG: endLoc = #[[END_LOC_FUSED]]
+#loopMD = #llvm.loop_annotation<disableNonforced = false,
+        mustProgress = true,
+        startLoc = #start_loc_fused,
+        endLoc = #end_loc_fused>
+
+// CHECK: llvm.func @loop_annotation_with_locs
+llvm.func @loop_annotation_with_locs() {
+  // CHECK: llvm.br ^bb1 {loop_annotation = #[[LOOP_ANNOT]]
+  llvm.br ^bb1 {loop_annotation = #loopMD}
+^bb1:
+  llvm.return
+}

diff  --git a/mlir/test/Target/LLVMIR/Import/metadata-loop.ll b/mlir/test/Target/LLVMIR/Import/metadata-loop.ll
index 859f7377a3eb1..dcd2004040abb 100644
--- a/mlir/test/Target/LLVMIR/Import/metadata-loop.ll
+++ b/mlir/test/Target/LLVMIR/Import/metadata-loop.ll
@@ -343,3 +343,36 @@ end:
 !1 = distinct !{!1, !2}
 !2 = !{!"llvm.loop.parallel_accesses", !0, !3}
 !3 = distinct !{}
+
+; // -----
+
+; CHECK: #[[start_loc:.*]] = loc("metadata-loop.ll":1:2)
+; CHECK: #[[end_loc:.*]] = loc("metadata-loop.ll":2:2)
+; CHECK: #[[SUBPROGRAM:.*]] = #llvm.di_subprogram<
+; CHECK: #[[start_loc_fused:.*]] = loc(fused<#[[SUBPROGRAM]]>[#[[start_loc]]])
+; CHECK: #[[end_loc_fused:.*]] = loc(fused<#[[SUBPROGRAM]]>[#[[end_loc]]])
+; CHECK: #[[$ANNOT_ATTR:.*]] = #llvm.loop_annotation<
+; CHECK-SAME: mustProgress = true
+; CHECK-SAME: startLoc = #[[start_loc_fused]]
+; CHECK-SAME: endLoc = #[[end_loc_fused]]
+
+; CHECK-LABEL: @loop_locs
+define void @loop_locs(i64 %n, ptr %A) {
+entry:
+; CHECK: llvm.br ^{{.*}} {loop_annotation = #[[$ANNOT_ATTR]]}
+  br label %end, !llvm.loop !6
+end:
+  ret void
+}
+
+!llvm.dbg.cu = !{!1}
+!llvm.module.flags = !{!0}
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = distinct !DICompileUnit(language: DW_LANG_C, file: !2)
+!2 = !DIFile(filename: "metadata-loop.ll", directory: "/")
+!3 = distinct !DISubprogram(name: "loop_locs", scope: !2, file: !2, spFlags: DISPFlagDefinition, unit: !1)
+!4 = !DILocation(line: 1, column: 2, scope: !3)
+!5 = !DILocation(line: 2, column: 2, scope: !3)
+
+!6 = distinct !{!6, !4, !5, !7}
+!7 = !{!"llvm.loop.mustprogress"}

diff  --git a/mlir/test/Target/LLVMIR/loop-metadata.mlir b/mlir/test/Target/LLVMIR/loop-metadata.mlir
index 1d668dd56a991..f17cc003aa0ae 100644
--- a/mlir/test/Target/LLVMIR/loop-metadata.mlir
+++ b/mlir/test/Target/LLVMIR/loop-metadata.mlir
@@ -288,3 +288,32 @@ llvm.metadata @metadata {
 // CHECK-DAG: ![[PIPELINE_DISABLE_NODE:[0-9]+]] = !{!"llvm.loop.pipeline.disable", i1 true}
 // CHECK-DAG: ![[II_NODE:[0-9]+]] = !{!"llvm.loop.pipeline.initiationinterval", i32 2}
 // CHECK-DAG: ![[ACCESS_GROUPS_NODE:[0-9]+]] = !{![[GROUP_NODE1]], ![[GROUP_NODE2]]}
+
+// -----
+
+#di_file = #llvm.di_file<"metadata-loop.ll" in "/">
+
+#loc1 = loc("loop-metadata.mlir":42:4)
+#loc2 = loc("loop-metadata.mlir":52:4)
+
+#di_compile_unit = #llvm.di_compile_unit<sourceLanguage = DW_LANG_C, file = #di_file, isOptimized = false, emissionKind = None>
+#di_subprogram = #llvm.di_subprogram<compileUnit = #di_compile_unit, scope = #di_file, name = "loop_locs", file = #di_file, subprogramFlags = Definition>
+
+#start_loc_fused = loc(fused<#di_subprogram>[#loc1])
+#end_loc_fused= loc(fused<#di_subprogram>[#loc2])
+
+#loopMD = #llvm.loop_annotation<disableNonforced = false,
+        startLoc = #start_loc_fused,
+        endLoc = #end_loc_fused>
+
+// CHECK-LABEL: @loop_annotation_with_locs
+llvm.func @loop_annotation_with_locs() {
+// CHECK: br {{.*}} !llvm.loop ![[LOOP_NODE:[0-9]+]]
+  llvm.br ^bb1 {loop_annotation = #loopMD}
+^bb1:
+  llvm.return
+}
+
+// CHECK: ![[LOOP_NODE]] = distinct !{![[LOOP_NODE]], ![[START_LOC:.*]], ![[END_LOC:.*]]}
+// CHECK: ![[START_LOC]] = !DILocation(line: 42, column: 4, scope:
+// CHECK: ![[END_LOC]] = !DILocation(line: 52, column: 4, scope:


        


More information about the Mlir-commits mailing list