[milr][llvm] Add remaining loop metadata support

[Dinistro]

Summary

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

[gysit]

LGTM module on test function name.

mlir/test/Target/LLVMIR/Import/metadata-loop.ll

246

nit: peeled -> unswitched?

[Dinistro]

fix test label name

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

def LoopPipelineAttr : LLVM_Attr<"LoopPipeline", "loop_pipeline"> {
   let parameters = (ins
     OptionalParameter<"BoolAttr">:$disable,
     OptionalParameter<"IntegerAttr">:$initiationinterval
   );
 
   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
     branches that are "latches" (loop backedges) and maps to the `!llvm.loop`
     metadatas: https://llvm.org/docs/LangRef.html#llvm-loop
     It stores annotations in attribute parameters and groups related options in
     nested attributes to provide structured access.
   }];
 
   let parameters = (ins
     OptionalParameter<"BoolAttr">:$disableNonforced,
     OptionalParameter<"LoopVectorizeAttr">:$vectorize,
     OptionalParameter<"LoopInterleaveAttr">:$interleave,
     OptionalParameter<"LoopUnrollAttr">:$unroll,
     OptionalParameter<"LoopUnrollAndJamAttr">:$unrollAndJam,
     OptionalParameter<"LoopLICMAttr">:$licm,
     OptionalParameter<"LoopDistributeAttr">:$distribute,
     OptionalParameter<"LoopPipelineAttr">:$pipeline,
+    OptionalParameter<"LoopPeeledAttr">:$peeled,
+    OptionalParameter<"LoopUnswitchAttr">:$unswitch,
     OptionalParameter<"BoolAttr">:$mustProgress,
     OptionalParameter<"BoolAttr">:$isVectorized,
     OptionalArrayRefParameter<"SymbolRefAttr">:$parallelAccesses
   );
 
   let assemblyFormat = "`<` struct(params) `>`";
 }
 

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
   AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
     return TypeSwitch<Attribute, AliasResult>(attr)
         .Case<DIVoidResultTypeAttr, DIBasicTypeAttr, DICompileUnitAttr,
               DICompositeTypeAttr, DIDerivedTypeAttr, DIFileAttr,
               DILexicalBlockAttr, DILexicalBlockFileAttr, DILocalVariableAttr,
               DISubprogramAttr, DISubroutineTypeAttr, LoopAnnotationAttr,
               LoopVectorizeAttr, LoopInterleaveAttr, LoopUnrollAttr,
               LoopUnrollAndJamAttr, LoopLICMAttr, LoopDistributeAttr,
-              LoopPipelineAttr>([&](auto attr) {
+              LoopPipelineAttr, LoopPeeledAttr, LoopUnswitchAttr>(
+            [&](auto attr) {
               os << decltype(attr)::getMnemonic();
               return AliasResult::OverridableAlias;
             })
         .Default([](Attribute) { return AliasResult::NoAlias; });
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // DialectInlinerInterface
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp

struct LoopMetadataConversion {
   /// Conversion functions for sub-attributes.
   FailureOr<LoopVectorizeAttr> convertVectorizeAttr();
   FailureOr<LoopInterleaveAttr> convertInterleaveAttr();
   FailureOr<LoopUnrollAttr> convertUnrollAttr();
   FailureOr<LoopUnrollAndJamAttr> convertUnrollAndJamAttr();
   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;
   const llvm::MDNode *node;
   Location loc;
   LoopAnnotationImporter &loopAnnotationImporter;
   MLIRContext *ctx;
 };

 
 FailureOr<LoopPipelineAttr> LoopMetadataConversion::convertPipelineAttr() {
   FailureOr<BoolAttr> disable = lookupBoolNode("llvm.loop.pipeline.disable");
   FailureOr<IntegerAttr> initiationinterval =
       lookupIntNode("llvm.loop.pipeline.initiationinterval");
   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 =
       lookupMDNodes("llvm.loop.parallel_accesses");
   if (failed(nodes))
     return failure();
   SmallVector<SymbolRefAttr> refs;
   for (llvm::MDNode *node : *nodes) {
FailureOr<BoolAttr> disableNonForced =
       lookupUnitNode("llvm.loop.disable_nonforced");
   FailureOr<LoopVectorizeAttr> vecAttr = convertVectorizeAttr();
   FailureOr<LoopInterleaveAttr> interleaveAttr = convertInterleaveAttr();
   FailureOr<LoopUnrollAttr> unrollAttr = convertUnrollAttr();
   FailureOr<LoopUnrollAndJamAttr> unrollAndJamAttr = convertUnrollAndJamAttr();
   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");
   FailureOr<SmallVector<SymbolRefAttr>> parallelAccesses =
       convertParallelAccesses();
 
   // Drop the metadata if there are parts that cannot be imported.
   if (!propertyMap.empty()) {
     for (auto name : propertyMap.keys())
       emitWarning(loc) << "unknown loop annotation " << name;
     return {};
   }
 
   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
 LoopAnnotationImporter::translateLoopAnnotation(const llvm::MDNode *node,
                                                 Location loc) {
   if (!node)
     return {};
 

mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp

struct LoopAnnotationConversion {
   /// Conversion functions for each for each loop annotation sub-attribute.
   void convertLoopOptions(LoopVectorizeAttr options);
   void convertLoopOptions(LoopInterleaveAttr options);
   void convertLoopOptions(LoopUnrollAttr options);
   void convertLoopOptions(LoopUnrollAndJamAttr options);
   void convertLoopOptions(LoopLICMAttr options);
   void convertLoopOptions(LoopDistributeAttr options);
   void convertLoopOptions(LoopPipelineAttr options);
+  void convertLoopOptions(LoopPeeledAttr options);
+  void convertLoopOptions(LoopUnswitchAttr options);
 
   LoopAnnotationAttr attr;
   ModuleTranslation &moduleTranslation;
   Operation *op;
   LoopAnnotationTranslation &loopAnnotationTranslation;
   llvm::LLVMContext &ctx;
   llvm::SmallVector<llvm::Metadata *> metadataNodes;
 };

 }
 
 void LoopAnnotationConversion::convertLoopOptions(LoopPipelineAttr options) {
   convertBoolNode("llvm.loop.pipeline.disable", options.getDisable());
   convertI32Node("llvm.loop.pipeline.initiationinterval",
                  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.
   auto dummy = llvm::MDNode::getTemporary(ctx, std::nullopt);
   metadataNodes.push_back(dummy.get());
 
   addUnitNode("llvm.loop.disable_nonforced", attr.getDisableNonforced());
   addUnitNode("llvm.loop.mustprogress", attr.getMustProgress());
llvm::MDNode *LoopAnnotationConversion::convert() {
   if (auto options = attr.getUnrollAndJam())
     convertLoopOptions(options);
   if (auto options = attr.getLicm())
     convertLoopOptions(options);
   if (auto options = attr.getDistribute())
     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()) {
     SmallVector<llvm::Metadata *> parallelAccess;
     parallelAccess.push_back(
         llvm::MDString::get(ctx, "llvm.loop.parallel_accesses"));
     for (SymbolRefAttr accessGroupRef : parallelAccessGroups)
       parallelAccess.push_back(

mlir/test/Dialect/LLVMIR/loop-metadata.mlir

 #distribute = #llvm.loop_distribute<
   disable = true, followupCoincident = #followup, followupSequential = #followup,
   followupFallback = #followup, followupAll = #followup
 >
 
 // 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
 // CHECK-DAG: unroll = #[[UNROLL]]
 // CHECK-DAG: unrollAndJam = #[[UNROLL_AND_JAM]]
 // 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,
         mustProgress = true,
         vectorize = #vectorize,
         interleave = #interleave,
         unroll = #unroll,
         unrollAndJam = #unrollAndJam,
         licm = #licm,
         distribute = #distribute,
         pipeline = #pipeline,
+        peeled = #peeled,
+        unswitch = #unswitch,
         isVectorized = false,
         parallelAccesses = @metadata::@group1, @metadata::@group2>
 
 // CHECK: llvm.func @loop_annotation
 llvm.func @loop_annotation() {
   // CHECK: llvm.br ^bb1 {llvm.loop = #[[LOOP_ANNOT]]
   llvm.br ^bb1 {llvm.loop = #loopMD}
 ^bb1:
   llvm.return
 }
 
 llvm.metadata @metadata {
   llvm.access_group @group1
   llvm.access_group @group2
 }

mlir/test/Target/LLVMIR/Import/metadata-loop.ll

 }
 
 !1 = distinct !{!1, !2, !3}
 !2 = !{!"llvm.loop.pipeline.disable", i1 0}
 !3 = !{!"llvm.loop.pipeline.initiationinterval", i32 2}
 
 ; // -----
 
+; 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 {
 ; CHECK:   llvm.access_group @[[GROUP0]]
 
 ; CHECK-LABEL: @parallel_accesses
 define void @parallel_accesses(ptr %arg) {
 entry:

mlir/test/Target/LLVMIR/loop-metadata.mlir

 }
 
 // CHECK: ![[LOOP_NODE]] = distinct !{![[LOOP_NODE]], !{{[0-9]+}}, !{{[0-9]+}}}
 // CHECK-DAG: ![[VEC_NODE0:[0-9]+]] = !{!"llvm.loop.pipeline.disable", i1 false}
 // CHECK-DAG: ![[VEC_NODE0:[0-9]+]] = !{!"llvm.loop.pipeline.initiationinterval", i32 1}
 
 // -----
 
+// 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
     %4 = llvm.alloca %arg1 x i32 : (i32) -> (!llvm.ptr<i32>)
     llvm.br ^bb3(%0 : i32)
   ^bb3(%1: i32):
     %2 = llvm.icmp "slt" %1, %arg1 : i32
     // CHECK: br i1 {{.*}} !llvm.loop ![[LOOP_NODE:[0-9]+]]