[mlir] Add LLVM loop codegen options to control software pipelining

Support specifying the II and disabling pipelining.

Reviewed By: ftynse

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

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

@@ -98,6 +98,15 @@ public:
   /// is provided the option is deleted.
   LoopOptionsAttrBuilder &setDisableUnroll(Optional<bool> value);
 
+  /// Set the `disable_pipeline` option to the provided value. If no value
+  /// is provided the option is deleted.
+  LoopOptionsAttrBuilder &setDisablePipeline(Optional<bool> value);
+
+  /// Set the `pipeline_initiation_interval` option to the provided value.
+  /// If no value is provided the option is deleted.
+  LoopOptionsAttrBuilder &
+  setPipelineInitiationInterval(Optional<uint64_t> count);
+
   /// Returns true if any option has been set.
   bool empty() { return options.empty(); }
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -50,11 +50,14 @@ def LLVM_FMFAttr : DialectAttr<
 def LOptDisableUnroll : I32EnumAttrCase<"disable_unroll", 1>;
 def LOptDisableLICM : I32EnumAttrCase<"disable_licm", 2>;
 def LOptInterleaveCount : I32EnumAttrCase<"interleave_count", 3>;
+def LOptDisablePipeline : I32EnumAttrCase<"disable_pipeline", 4>;
+def LOptPipelineInitiationInterval : I32EnumAttrCase<"pipeline_initiation_interval", 5>;
 
 def LoopOptionCase : I32EnumAttr<
     "LoopOptionCase",
     "LLVM loop option",
-    [LOptDisableUnroll, LOptDisableLICM, LOptInterleaveCount
+    [LOptDisableUnroll, LOptDisableLICM, LOptInterleaveCount,
+     LOptDisablePipeline, LOptPipelineInitiationInterval
     ]> {
   let cppNamespace = "::mlir::LLVM";
 }

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

@@ -2431,6 +2431,20 @@ LoopOptionsAttrBuilder::setDisableUnroll(Optional<bool> value) {
   return setOption(LoopOptionCase::disable_unroll, value);
 }
 
+/// Set the `disable_pipeline` option to the provided value. If no value
+/// is provided the option is deleted.
+LoopOptionsAttrBuilder &
+LoopOptionsAttrBuilder::setDisablePipeline(Optional<bool> value) {
+  return setOption(LoopOptionCase::disable_pipeline, value);
+}
+
+/// Set the `pipeline_initiation_interval` option to the provided value.
+/// If no value is provided the option is deleted.
+LoopOptionsAttrBuilder &LoopOptionsAttrBuilder::setPipelineInitiationInterval(
+    Optional<uint64_t> count) {
+  return setOption(LoopOptionCase::pipeline_initiation_interval, count);
+}
+
 template <typename T>
 static Optional<T>
 getOption(ArrayRef<std::pair<LoopOptionCase, int64_t>> options,
@@ -2479,9 +2493,11 @@ void LoopOptionsAttr::print(DialectAsmPrinter &printer) const {
     switch (option.first) {
     case LoopOptionCase::disable_licm:
     case LoopOptionCase::disable_unroll:
+    case LoopOptionCase::disable_pipeline:
       printer << (option.second ? "true" : "false");
       break;
     case LoopOptionCase::interleave_count:
+    case LoopOptionCase::pipeline_initiation_interval:
       printer << option.second;
       break;
     }
@@ -2518,6 +2534,7 @@ Attribute LoopOptionsAttr::parse(MLIRContext *context, DialectAsmParser &parser,
     switch (*option) {
     case LoopOptionCase::disable_licm:
     case LoopOptionCase::disable_unroll:
+    case LoopOptionCase::disable_pipeline:
       if (succeeded(parser.parseOptionalKeyword("true")))
         value = 1;
       else if (succeeded(parser.parseOptionalKeyword("false")))
@@ -2529,6 +2546,7 @@ Attribute LoopOptionsAttr::parse(MLIRContext *context, DialectAsmParser &parser,
       }
       break;
     case LoopOptionCase::interleave_count:
+    case LoopOptionCase::pipeline_initiation_interval:
       if (failed(parser.parseInteger(value))) {
         parser.emitError(parser.getNameLoc(), "expected integer value");
         return {};

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

@@ -189,6 +189,15 @@ static llvm::MDNode *getLoopOptionMetadata(llvm::LLVMContext &ctx,
     cstValue = llvm::ConstantInt::get(
         llvm::IntegerType::get(ctx, /*NumBits=*/32), value);
     break;
+  case LoopOptionCase::disable_pipeline:
+    name = "llvm.loop.pipeline.disable";
+    cstValue = llvm::ConstantInt::getBool(ctx, value);
+    break;
+  case LoopOptionCase::pipeline_initiation_interval:
+    name = "llvm.loop.pipeline.initiationinterval";
+    cstValue = llvm::ConstantInt::get(
+        llvm::IntegerType::get(ctx, /*NumBits=*/32), value);
+    break;
   }
   return llvm::MDNode::get(ctx, {llvm::MDString::get(ctx, name),
                                  llvm::ConstantAsMetadata::get(cstValue)});

mlir/test/Dialect/LLVMIR/roundtrip.mlir

@@ -424,8 +424,8 @@ module {
   // CHECK-LABEL: @loopOptions
   llvm.func @loopOptions() {
     // CHECK: llvm.br
-    // CHECK-SAME: llvm.loop = {options = #llvm.loopopts<disable_unroll = true, disable_licm = true, interleave_count = 1>}, parallel_access = [@metadata::@group1]}
-    llvm.br ^bb1 {llvm.loop = {options = #llvm.loopopts<disable_unroll = true, disable_licm = true, interleave_count = 1>}, parallel_access = [@metadata::@group1]}
+    // CHECK-SAME: llvm.loop = {options = #llvm.loopopts<disable_unroll = true, disable_licm = true, interleave_count = 1, disable_pipeline = true, pipeline_initiation_interval = 1>}, parallel_access = [@metadata::@group1]}
+    llvm.br ^bb1 {llvm.loop = {options = #llvm.loopopts<disable_unroll = true, disable_licm = true, interleave_count = 1, disable_pipeline = true, pipeline_initiation_interval = 1>}, parallel_access = [@metadata::@group1]}
   ^bb1:
     llvm.return
   }

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -1480,13 +1480,13 @@ module {
     ^bb3(%1: i32):
       %2 = llvm.icmp "slt" %1, %arg1 : i32
       // CHECK: br i1 {{.*}} !llvm.loop ![[LOOP_NODE:[0-9]+]]
-      llvm.cond_br %2, ^bb4, ^bb5 {llvm.loop = {parallel_access = [@metadata::@group1, @metadata::@group2], options = #llvm.loopopts<disable_licm = true, disable_unroll = true, interleave_count = 1>}}
+      llvm.cond_br %2, ^bb4, ^bb5 {llvm.loop = {parallel_access = [@metadata::@group1, @metadata::@group2], options = #llvm.loopopts<disable_licm = true, disable_unroll = true, interleave_count = 1, disable_pipeline = true, pipeline_initiation_interval = 2>}}
     ^bb4:
       %3 = llvm.add %1, %arg2  : i32
       // CHECK: = load i32, i32* %{{.*}} !llvm.access.group ![[ACCESS_GROUPS_NODE:[0-9]+]]
       %5 = llvm.load %4 { access_groups = [@metadata::@group1, @metadata::@group2] } : !llvm.ptr<i32>
       // CHECK: br label {{.*}} !llvm.loop ![[LOOP_NODE]]
-      llvm.br ^bb3(%3 : i32) {llvm.loop = {parallel_access = [@metadata::@group1, @metadata::@group2], options = #llvm.loopopts<disable_unroll = true, disable_licm = true, interleave_count = 1>}}
+      llvm.br ^bb3(%3 : i32) {llvm.loop = {parallel_access = [@metadata::@group1, @metadata::@group2], options = #llvm.loopopts<disable_unroll = true, disable_licm = true, interleave_count = 1, disable_pipeline = true, pipeline_initiation_interval = 2>}}
     ^bb5:
       llvm.return
   }
@@ -1498,11 +1498,13 @@ module {
   }
 }
 
-// CHECK: ![[LOOP_NODE]] = distinct !{![[LOOP_NODE]], ![[PA_NODE:[0-9]+]], ![[UNROLL_DISABLE_NODE:[0-9]+]], ![[LICM_DISABLE_NODE:[0-9]+]], ![[INTERLEAVE_NODE:[0-9]+]]}
+// CHECK: ![[LOOP_NODE]] = distinct !{![[LOOP_NODE]], ![[PA_NODE:[0-9]+]], ![[UNROLL_DISABLE_NODE:[0-9]+]], ![[LICM_DISABLE_NODE:[0-9]+]], ![[INTERLEAVE_NODE:[0-9]+]], ![[PIPELINE_DISABLE_NODE:[0-9]+]], ![[II_NODE:[0-9]+]]}
 // CHECK: ![[PA_NODE]] = !{!"llvm.loop.parallel_accesses", ![[GROUP_NODE1:[0-9]+]], ![[GROUP_NODE2:[0-9]+]]}
 // CHECK: ![[GROUP_NODE1]] = distinct !{}
 // CHECK: ![[GROUP_NODE2]] = distinct !{}
 // CHECK: ![[UNROLL_DISABLE_NODE]] = !{!"llvm.loop.unroll.disable", i1 true}
 // CHECK: ![[LICM_DISABLE_NODE]] = !{!"llvm.licm.disable", i1 true}
 // CHECK: ![[INTERLEAVE_NODE]] = !{!"llvm.loop.interleave.count", i32 1}
+// CHECK: ![[PIPELINE_DISABLE_NODE]] = !{!"llvm.loop.pipeline.disable", i1 true}
+// CHECK: ![[II_NODE]] = !{!"llvm.loop.pipeline.initiationinterval", i32 2}
 // CHECK: ![[ACCESS_GROUPS_NODE]] = !{![[GROUP_NODE1]], ![[GROUP_NODE2]]}