[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

Author: Dinistro

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) `>`";

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

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 {

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.

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 {

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());
 }
 

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) &&

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
+}

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"}

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: