From 62d7d94c2ea4f04d9c6f3feaa525ac4971fef815 Mon Sep 17 00:00:00 2001 From: Christian Ulmann Date: Fri, 5 May 2023 07:51:46 +0000 Subject: [PATCH] [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 --- mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td | 4 ++- mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp | 38 ++++++++++++++++++---- mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h | 6 +++- .../Target/LLVMIR/LoopAnnotationTranslation.cpp | 25 +++++++++++++- mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h | 9 +++-- mlir/lib/Target/LLVMIR/ModuleImport.cpp | 2 +- mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 2 +- mlir/test/Dialect/LLVMIR/loop-metadata.mlir | 37 ++++++++++++++++++++- mlir/test/Target/LLVMIR/Import/metadata-loop.ll | 33 +++++++++++++++++++ mlir/test/Target/LLVMIR/loop-metadata.mlir | 29 +++++++++++++++++ 10 files changed, 171 insertions(+), 14 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td index b6b5668..05ee083 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 0003490..a9f080e 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 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 @@ struct LoopMetadataConversion { }; } // 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 @@ LoopMetadataConversion::convertParallelAccesses() { 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 @@ LoopAnnotationAttr LoopMetadataConversion::convert() { 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/LoopAnnotationImporter.h b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h index 5d69a63..7dcfc4c 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 @@ public: 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/LoopAnnotationTranslation.cpp b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp index 6abcf69..8e6906a 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(); + 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/LoopAnnotationTranslation.h b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h index e663f64..0e016d9 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 @@ public: /// 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 5c585e7..8c5fe35 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(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 index f3f5b24..f8854d7 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(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 index 8841a1f..631f271 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 #followup = #llvm.loop_annotation @@ -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 +// 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 index 859f737..dcd2004 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 1d668dd..f17cc00 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 +#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: -- 2.7.4