[milr][llvm] Add remaining loop metadata support

This commit adds support for the last two loop metadata nodes produced
anywhere in the llvm-project.

Reviewed By: gysit

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

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 06783b161bf50543b7b43fbe81886505c64490f9..a6d8c703baaf8c67209689589aaf192d9fb17ef6 100644 (file)
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -168,6 +168,32 @@ def LoopPipelineAttr : LLVM_Attr<"LoopPipeline", "loop_pipeline"> {
   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 @@ def LoopAnnotationAttr : LLVM_Attr<"LoopAnnotation", "loop_annotation"> {
     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
index c19d268d745f6ed451e1cef9cdbc87ee0d819359..8d654df7d2179d632b1fa8dd98807fdd9a75180a 100644 (file)
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2666,10 +2666,11 @@ struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
               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
index 53e6f9c50e892de80e7d29f4a2361cb4bedba66e..2986b64538c2597fd4ce426b9a6fe452cb6707f5 100644 (file)
--- a/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp
+++ b/mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp
@@ -50,6 +50,8 @@ struct LoopMetadataConversion {
   FailureOr<LoopLICMAttr> convertLICMAttr();
   FailureOr<LoopDistributeAttr> convertDistributeAttr();
   FailureOr<LoopPipelineAttr> convertPipelineAttr();
+  FailureOr<LoopPeeledAttr> convertPeeledAttr();
+  FailureOr<LoopUnswitchAttr> convertUnswitchAttr();
   FailureOr<SmallVector<SymbolRefAttr>> convertParallelAccesses();
 
   llvm::StringMap<const llvm::MDNode *> propertyMap;
@@ -373,6 +375,17 @@ FailureOr<LoopPipelineAttr> LoopMetadataConversion::convertPipelineAttr() {
   return createIfNonNull<LoopPipelineAttr>(ctx, disable, initiationinterval);
 }
 
+FailureOr<LoopPeeledAttr> LoopMetadataConversion::convertPeeledAttr() {
+  FailureOr<IntegerAttr> count = lookupIntNode("llvm.loop.peeled.count");
+  return createIfNonNull<LoopPeeledAttr>(ctx, count);
+}
+
+FailureOr<LoopUnswitchAttr> LoopMetadataConversion::convertUnswitchAttr() {
+  FailureOr<BoolAttr> partialDisable =
+      lookupUnitNode("llvm.loop.unswitch.partial.disable");
+  return createIfNonNull<LoopUnswitchAttr>(ctx, partialDisable);
+}
+
 FailureOr<SmallVector<SymbolRefAttr>>
 LoopMetadataConversion::convertParallelAccesses() {
   FailureOr<SmallVector<llvm::MDNode *>> nodes =
@@ -403,6 +416,8 @@ LoopAnnotationAttr LoopMetadataConversion::convert() {
   FailureOr<LoopLICMAttr> licmAttr = convertLICMAttr();
   FailureOr<LoopDistributeAttr> distributeAttr = convertDistributeAttr();
   FailureOr<LoopPipelineAttr> pipelineAttr = convertPipelineAttr();
+  FailureOr<LoopPeeledAttr> peeledAttr = convertPeeledAttr();
+  FailureOr<LoopUnswitchAttr> unswitchAttr = convertUnswitchAttr();
   FailureOr<BoolAttr> mustProgress = lookupUnitNode("llvm.loop.mustprogress");
   FailureOr<BoolAttr> isVectorized =
       lookupIntNodeAsBoolAttr("llvm.loop.isvectorized");
@@ -418,8 +433,8 @@ LoopAnnotationAttr LoopMetadataConversion::convert() {
 
   return createIfNonNull<LoopAnnotationAttr>(
       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
index b02433b33edfe43c37c0f590c5573d46f2e5e198..5f27f97f2fd4596c9d2c0ac8baab1c9bff0f21d1 100644 (file)
--- a/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp
@@ -42,6 +42,8 @@ struct LoopAnnotationConversion {
   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 @@ void LoopAnnotationConversion::convertLoopOptions(LoopPipelineAttr options) {
                  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 @@ llvm::MDNode *LoopAnnotationConversion::convert() {
     convertLoopOptions(options);
   if (auto options = attr.getPipeline())
     convertLoopOptions(options);
+  if (auto options = attr.getPeeled())
+    convertLoopOptions(options);
+  if (auto options = attr.getUnswitch())
+    convertLoopOptions(options);
 
   ArrayRef<SymbolRefAttr> parallelAccessGroups = attr.getParallelAccesses();
   if (!parallelAccessGroups.empty()) {

diff --git a/mlir/test/Dialect/LLVMIR/loop-metadata.mlir b/mlir/test/Dialect/LLVMIR/loop-metadata.mlir
index 0d434323f8709e637b12cf73cbabab4094fbc115..8841a1fad9186c1adcfe0af5ac0669a69b94c450 100644 (file)
--- 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<disable = true, initiationinterval = 1 : i32>
 #pipeline = #llvm.loop_pipeline<disable = true, initiationinterval = 1 : i32>
 
+// CHECK-DAG: #[[PEELED:.*]] = #llvm.loop_peeled<count = 8 : i32>
+#peeled = #llvm.loop_peeled<count = 8 : i32>
+
+// CHECK-DAG: #[[UNSWITCH:.*]] = #llvm.loop_unswitch<partialDisable = true>
+#unswitch = #llvm.loop_unswitch<partialDisable = true>
+
 // 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<disableNonforced = false,
@@ -55,6 +63,8 @@
         licm = #licm,
         distribute = #distribute,
         pipeline = #pipeline,
+        peeled = #peeled,
+        unswitch = #unswitch,
         isVectorized = false,
         parallelAccesses = @metadata::@group1, @metadata::@group2>
 

diff --git a/mlir/test/Target/LLVMIR/Import/metadata-loop.ll b/mlir/test/Target/LLVMIR/Import/metadata-loop.ll
index 9aecb13c0095349c29cf70f8ca32875b859b7d0a..315852663dd09bee0755fac49c74d6652da4c300 100644 (file)
--- a/mlir/test/Target/LLVMIR/Import/metadata-loop.ll
+++ b/mlir/test/Target/LLVMIR/Import/metadata-loop.ll
@@ -221,6 +221,40 @@ end:
 
 ; // -----
 
+; CHECK-DAG: #[[PEELED_ATTR:.*]] = #llvm.loop_peeled<count = 5 : i32>
+; CHECK-DAG: #[[$ANNOT_ATTR:.*]] = #llvm.loop_annotation<peeled = #[[PEELED_ATTR]]>
+
+; 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<partialDisable = true>
+; CHECK-DAG: #[[$ANNOT_ATTR:.*]] = #llvm.loop_annotation<unswitch = #[[UNSWITCH_ATTR]]>
+
+; 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<parallelAccesses = @__llvm_global_metadata::@[[GROUP0:.*]]>
 
 ; CHECK: llvm.metadata @__llvm_global_metadata {

diff --git a/mlir/test/Target/LLVMIR/loop-metadata.mlir b/mlir/test/Target/LLVMIR/loop-metadata.mlir
index 9bed3ae3916298a5c1afcd0e17d82502d4a469ca..e7d06775bf9f9e1988719604a557b2e0bfede404 100644 (file)
--- a/mlir/test/Target/LLVMIR/loop-metadata.mlir
+++ b/mlir/test/Target/LLVMIR/loop-metadata.mlir
@@ -207,6 +207,32 @@ llvm.func @pipelineOptions() {
 
 // -----
 
+// CHECK-LABEL: @peeledOptions
+llvm.func @peeledOptions() {
+  // CHECK: br {{.*}} !llvm.loop ![[LOOP_NODE:[0-9]+]]
+  llvm.br ^bb1 {llvm.loop = #llvm.loop_annotation<peeled = <count = 3 : i32>>}
+^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<unswitch = <partialDisable = true>>}
+^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