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 @@ -168,6 +168,32 @@ let assemblyFormat = "`<` struct(params) `>`"; } +def LoopPeeledAttr : LLVM_Attr<"LoopPeeled", "loop_peeled"> { + let description = [{ + This attribute defines pipelining specific loop annotations that map to + the "!llvm.loop.peeled" metadata. + }]; + + let parameters = (ins + OptionalParameter<"IntegerAttr">:$count + ); + + let assemblyFormat = "`<` struct(params) `>`"; +} + +def LoopUnswitchAttr : LLVM_Attr<"LoopUnswitch", "loop_unswitch"> { + let description = [{ + This attribute defines pipelining specific loop annotations that map to + the "!llvm.loop.unswitch" metadata. + }]; + + let parameters = (ins + OptionalParameter<"BoolAttr">:$partialDisable + ); + + let assemblyFormat = "`<` struct(params) `>`"; +} + def LoopAnnotationAttr : LLVM_Attr<"LoopAnnotation", "loop_annotation"> { let description = [{ This attributes encapsulates "loop metadata". It is meant to decorate @@ -186,6 +212,8 @@ OptionalParameter<"LoopLICMAttr">:$licm, OptionalParameter<"LoopDistributeAttr">:$distribute, OptionalParameter<"LoopPipelineAttr">:$pipeline, + OptionalParameter<"LoopPeeledAttr">:$peeled, + OptionalParameter<"LoopUnswitchAttr">:$unswitch, OptionalParameter<"BoolAttr">:$mustProgress, OptionalParameter<"BoolAttr">:$isVectorized, OptionalArrayRefParameter<"SymbolRefAttr">:$parallelAccesses diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp @@ -2666,10 +2666,11 @@ DISubprogramAttr, DISubroutineTypeAttr, LoopAnnotationAttr, LoopVectorizeAttr, LoopInterleaveAttr, LoopUnrollAttr, LoopUnrollAndJamAttr, LoopLICMAttr, LoopDistributeAttr, - LoopPipelineAttr>([&](auto attr) { - os << decltype(attr)::getMnemonic(); - return AliasResult::OverridableAlias; - }) + LoopPipelineAttr, LoopPeeledAttr, LoopUnswitchAttr>( + [&](auto attr) { + os << decltype(attr)::getMnemonic(); + return AliasResult::OverridableAlias; + }) .Default([](Attribute) { return AliasResult::NoAlias; }); } }; 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 @@ -50,6 +50,8 @@ FailureOr convertLICMAttr(); FailureOr convertDistributeAttr(); FailureOr convertPipelineAttr(); + FailureOr convertPeeledAttr(); + FailureOr convertUnswitchAttr(); FailureOr> convertParallelAccesses(); llvm::StringMap propertyMap; @@ -373,6 +375,17 @@ return createIfNonNull(ctx, disable, initiationinterval); } +FailureOr LoopMetadataConversion::convertPeeledAttr() { + FailureOr count = lookupIntNode("llvm.loop.peeled.count"); + return createIfNonNull(ctx, count); +} + +FailureOr LoopMetadataConversion::convertUnswitchAttr() { + FailureOr partialDisable = + lookupUnitNode("llvm.loop.unswitch.partial.disable"); + return createIfNonNull(ctx, partialDisable); +} + FailureOr> LoopMetadataConversion::convertParallelAccesses() { FailureOr> nodes = @@ -403,6 +416,8 @@ FailureOr licmAttr = convertLICMAttr(); FailureOr distributeAttr = convertDistributeAttr(); FailureOr pipelineAttr = convertPipelineAttr(); + FailureOr peeledAttr = convertPeeledAttr(); + FailureOr unswitchAttr = convertUnswitchAttr(); FailureOr mustProgress = lookupUnitNode("llvm.loop.mustprogress"); FailureOr isVectorized = lookupIntNodeAsBoolAttr("llvm.loop.isvectorized"); @@ -418,8 +433,8 @@ return createIfNonNull( ctx, disableNonForced, vecAttr, interleaveAttr, unrollAttr, - unrollAndJamAttr, licmAttr, distributeAttr, pipelineAttr, mustProgress, - isVectorized, parallelAccesses); + unrollAndJamAttr, licmAttr, distributeAttr, pipelineAttr, peeledAttr, + unswitchAttr, mustProgress, isVectorized, parallelAccesses); } LoopAnnotationAttr 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 @@ -42,6 +42,8 @@ void convertLoopOptions(LoopLICMAttr options); void convertLoopOptions(LoopDistributeAttr options); void convertLoopOptions(LoopPipelineAttr options); + void convertLoopOptions(LoopPeeledAttr options); + void convertLoopOptions(LoopUnswitchAttr options); LoopAnnotationAttr attr; ModuleTranslation &moduleTranslation; @@ -176,6 +178,15 @@ options.getInitiationinterval()); } +void LoopAnnotationConversion::convertLoopOptions(LoopPeeledAttr options) { + convertI32Node("llvm.loop.peeled.count", options.getCount()); +} + +void LoopAnnotationConversion::convertLoopOptions(LoopUnswitchAttr options) { + addUnitNode("llvm.loop.unswitch.partial.disable", + options.getPartialDisable()); +} + llvm::MDNode *LoopAnnotationConversion::convert() { // Reserve operand 0 for loop id self reference. @@ -202,6 +213,10 @@ convertLoopOptions(options); if (auto options = attr.getPipeline()) convertLoopOptions(options); + if (auto options = attr.getPeeled()) + convertLoopOptions(options); + if (auto options = attr.getUnswitch()) + convertLoopOptions(options); ArrayRef parallelAccessGroups = attr.getParallelAccesses(); if (!parallelAccessGroups.empty()) { 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 @@ -36,6 +36,12 @@ // CHECK-DAG: #[[PIPELINE:.*]] = #llvm.loop_pipeline #pipeline = #llvm.loop_pipeline +// CHECK-DAG: #[[PEELED:.*]] = #llvm.loop_peeled +#peeled = #llvm.loop_peeled + +// CHECK-DAG: #[[UNSWITCH:.*]] = #llvm.loop_unswitch +#unswitch = #llvm.loop_unswitch + // CHECK: #[[LOOP_ANNOT:.*]] = #llvm.loop_annotation< // CHECK-DAG: disableNonforced = false // CHECK-DAG: mustProgress = true @@ -44,6 +50,8 @@ // CHECK-DAG: licm = #[[LICM]] // CHECK-DAG: distribute = #[[DISTRIBUTE]] // CHECK-DAG: pipeline = #[[PIPELINE]] +// CHECK-DAG: peeled = #[[PEELED]] +// CHECK-DAG: unswitch = #[[UNSWITCH]] // CHECK-DAG: isVectorized = false // CHECK-DAG: parallelAccesses = @metadata::@group1, @metadata::@group2> #loopMD = #llvm.loop_annotation 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 @@ -221,6 +221,40 @@ ; // ----- +; CHECK-DAG: #[[PEELED_ATTR:.*]] = #llvm.loop_peeled +; CHECK-DAG: #[[$ANNOT_ATTR:.*]] = #llvm.loop_annotation + +; CHECK-LABEL: @peeled +define void @peeled(i64 %n, ptr %A) { +entry: +; CHECK: llvm.br ^{{.*}} {llvm.loop = #[[$ANNOT_ATTR]]} + br label %end, !llvm.loop !1 +end: + ret void +} + +!1 = distinct !{!1, !2} +!2 = !{!"llvm.loop.peeled.count", i32 5} + +; // ----- + +; CHECK-DAG: #[[UNSWITCH_ATTR:.*]] = #llvm.loop_unswitch +; CHECK-DAG: #[[$ANNOT_ATTR:.*]] = #llvm.loop_annotation + +; CHECK-LABEL: @unswitched +define void @unswitched(i64 %n, ptr %A) { +entry: +; CHECK: llvm.br ^{{.*}} {llvm.loop = #[[$ANNOT_ATTR]]} + br label %end, !llvm.loop !1 +end: + ret void +} + +!1 = distinct !{!1, !2} +!2 = !{!"llvm.loop.unswitch.partial.disable"} + +; // ----- + ; CHECK: #[[$ANNOT_ATTR:.*]] = #llvm.loop_annotation ; CHECK: llvm.metadata @__llvm_global_metadata { 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 @@ -207,6 +207,32 @@ // ----- +// CHECK-LABEL: @peeledOptions +llvm.func @peeledOptions() { + // CHECK: br {{.*}} !llvm.loop ![[LOOP_NODE:[0-9]+]] + llvm.br ^bb1 {llvm.loop = #llvm.loop_annotation>} +^bb1: + llvm.return +} + +// CHECK: ![[LOOP_NODE]] = distinct !{![[LOOP_NODE]], !{{[0-9]+}}} +// CHECK-DAG: ![[VEC_NODE0:[0-9]+]] = !{!"llvm.loop.peeled.count", i32 3} + +// ----- + +// CHECK-LABEL: @unswitchOptions +llvm.func @unswitchOptions() { + // CHECK: br {{.*}} !llvm.loop ![[LOOP_NODE:[0-9]+]] + llvm.br ^bb1 {llvm.loop = #llvm.loop_annotation>} +^bb1: + llvm.return +} + +// CHECK: ![[LOOP_NODE]] = distinct !{![[LOOP_NODE]], !{{[0-9]+}}} +// CHECK-DAG: ![[VEC_NODE0:[0-9]+]] = !{!"llvm.loop.unswitch.partial.disable"} + +// ----- + // CHECK-LABEL: @loopOptions llvm.func @loopOptions(%arg1 : i32, %arg2 : i32) { %0 = llvm.mlir.constant(0 : i32) : i32