diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td @@ -206,7 +206,9 @@ 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.h b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h --- a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h +++ b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h @@ -26,7 +26,8 @@ /// 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 @@ FailureOr> 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/LoopAnnotationImporter.cpp b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp --- a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp +++ b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp @@ -23,7 +23,8 @@ /// 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 @@ FailureOr convertPeeledAttr(); FailureOr convertUnswitchAttr(); FailureOr> convertParallelAccesses(); + FusedLoc convertStartLoc(); + FailureOr convertEndLoc(); + llvm::SmallVector locations; llvm::StringMap propertyMap; const llvm::MDNode *node; Location loc; @@ -62,16 +66,17 @@ }; } // namespace -LogicalResult LoopMetadataConversion::initPropertyMap() { +LogicalResult LoopMetadataConversion::initConversionState() { // Check if it's a valid node. if (node->getNumOperands() == 0 || dyn_cast(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(operand)) + if (auto *diLoc = dyn_cast(operand)) { + locations.push_back(diLoc); continue; + } auto *property = dyn_cast(operand); if (!property) @@ -405,8 +410,25 @@ return refs; } +FusedLoc LoopMetadataConversion::convertStartLoc() { + if (locations.empty()) + return {}; + return dyn_cast( + loopAnnotationImporter.moduleImport.translateLoc(locations[0])); +} + +FailureOr 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( + loopAnnotationImporter.moduleImport.translateLoc(locations[1])); +} + LoopAnnotationAttr LoopMetadataConversion::convert() { - if (failed(initPropertyMap())) + if (failed(initConversionState())) return {}; FailureOr disableNonForced = @@ -433,10 +455,14 @@ return {}; } + FailureOr startLoc = convertStartLoc(); + FailureOr endLoc = convertEndLoc(); + return createIfNonNull( 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/LoopAnnotationTranslation.h b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h --- a/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h +++ b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h @@ -25,8 +25,10 @@ /// 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 @@ /// 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/LoopAnnotationTranslation.cpp b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp --- 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 @@ 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 @@ options.getPartialDisable()); } -llvm::MDNode *LoopAnnotationConversion::convert() { +void LoopAnnotationConversion::convertLocation(FusedLoc location) { + auto localScopeAttr = + location.getMetadata().dyn_cast_or_null(); + if (!localScopeAttr) + return; + auto *localScope = dyn_cast( + loopAnnotationTranslation.moduleTranslation.translateDebugInfo( + localScopeAttr)); + if (!localScope) + return; + const llvm::Metadata *loc = + loopAnnotationTranslation.moduleTranslation.translateLoc(location, + localScope); + metadataNodes.push_back(const_cast(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/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -153,7 +153,7 @@ typeTranslator(*mlirModule->getContext()), debugImporter(std::make_unique(mlirModule)), loopAnnotationImporter( - std::make_unique(builder)) { + std::make_unique(*this, builder)) { builder.setInsertionPointToStart(mlirModule.getBody()); } diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -444,7 +444,7 @@ debugTranslation( std::make_unique(module, *this->llvmModule)), loopAnnotationTranslation(std::make_unique( - 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 --- 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 #followup = #llvm.loop_annotation @@ -80,3 +80,38 @@ 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 +// CHECK: #[[SUBPROGRAM:.*]] = #llvm.di_subprogram< +#di_subprogram = #llvm.di_subprogram + +// 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 + +// 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 --- a/mlir/test/Target/LLVMIR/Import/metadata-loop.ll +++ b/mlir/test/Target/LLVMIR/Import/metadata-loop.ll @@ -343,3 +343,36 @@ !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 --- a/mlir/test/Target/LLVMIR/loop-metadata.mlir +++ b/mlir/test/Target/LLVMIR/loop-metadata.mlir @@ -288,3 +288,32 @@ // 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 +#di_subprogram = #llvm.di_subprogram + +#start_loc_fused = loc(fused<#di_subprogram>[#loc1]) +#end_loc_fused= loc(fused<#di_subprogram>[#loc2]) + +#loopMD = #llvm.loop_annotation + +// 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: