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

[arpith-jacob]

Summary

Support specifying the II and disabling pipelining.

[arpith-jacob]

Thank you for the quick review. Please commit when you get the chance.

[billyroberts11]

In D98420#2619500, @arpith-jacob wrote:

Thank you for the quick review. Please commit when you get the chance.

Digital Design & Custom Development Company
The company Northell.design is one of the leading software development companies in Ukraine, providing custom software development, workplace automation, and IT infrastructure development. The company's team of experts has successfully implemented more than 5,000 projects. This makes them a leading system integrator that has completed large-scale projects in IT infrastructure development, unified communications system implementation, information security, and mobile solutions development.

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

public:
   /// Set the `interleave_count` option to the provided value. If no value
   /// is provided the option is deleted.
   LoopOptionsAttrBuilder &setInterleaveCount(Optional<uint64_t> count);
 
   /// Set the `disable_unroll` option to the provided value. If no value
   /// 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(); }
 
 private:
   template <typename T>
   LoopOptionsAttrBuilder &setOption(LoopOptionCase tag, Optional<T> value);
 
   friend class LoopOptionsAttr;
   SmallVector<LoopOptionsAttr::OptionValuePair> options;
 };
 
 } // end namespace LLVM
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_

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

def LLVM_FMFAttr : DialectAttr<
   let convertFromStorage = "$_self.getFlags()";
   let constBuilderCall =
           "::mlir::LLVM::FMFAttr::get($_builder.getContext(), $0)";
 }
 
 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";
 }
 
 class LLVM_Builder<string builder> {
   string llvmBuilder = builder;
 }
 

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

 
 /// Set the `disable_unroll` option to the provided value. If no value
 /// is provided the option is deleted.
 LoopOptionsAttrBuilder &
 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,
           LoopOptionCase option) {
   auto it =
       lower_bound(options, option, [](auto optionPair, LoopOptionCase option) {
         return optionPair.first < option;
       });

 
 void LoopOptionsAttr::print(DialectAsmPrinter &printer) const {
   printer << getMnemonic() << "<";
   llvm::interleaveComma(getOptions(), printer, [&](auto option) {
     printer << stringifyEnum(option.first) << " = ";
     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;
     }
   });
   printer << ">";
 }
 
 Attribute LoopOptionsAttr::parse(MLIRContext *context, DialectAsmParser &parser,
do {
     }
     if (failed(parser.parseEqual()))
       return {};
 
     int64_t value;
     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")))
         value = 0;
       else {
         parser.emitError(parser.getNameLoc(),
                          "expected boolean value 'true' or 'false'");
         return {};
       }
       break;
     case LoopOptionCase::interleave_count:
+    case LoopOptionCase::pipeline_initiation_interval:
       if (failed(parser.parseInteger(value))) {
         parser.emitError(parser.getNameLoc(), "expected integer value");
         return {};
       }
       break;
     }
     options.push_back(std::make_pair(*option, value));
   } while (succeeded(parser.parseOptionalComma()));

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

case LoopOptionCase::disable_unroll:
     name = "llvm.loop.unroll.disable";
     cstValue = llvm::ConstantInt::getBool(ctx, value);
     break;
   case LoopOptionCase::interleave_count:
     name = "llvm.loop.interleave.count";
     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)});
 }
 
 static void setLoopMetadata(Operation &opInst, llvm::Instruction &llvmInst,
                             llvm::IRBuilderBase &builder,
                             LLVM::ModuleTranslation &moduleTranslation) {

mlir/test/Dialect/LLVMIR/roundtrip.mlir

// CHECK: {{.*}} = llvm.fneg %arg0 : f32
   %10 = llvm.fneg %arg0 {fastmathFlags = #llvm.fastmath<>} : f32
   return
 }
 
 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
   }
   // CHECK: llvm.metadata @metadata attributes {test_attribute} {
   llvm.metadata @metadata attributes {test_attribute} {
     // CHECK: llvm.access_group @group1
     llvm.access_group @group1
     llvm.return
   }
 }

mlir/test/Target/LLVMIR/llvmir.mlir

 module {
   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]+]]
-      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
   }
 
   llvm.metadata @metadata {
     llvm.access_group @group1
     llvm.access_group @group2
     llvm.return
   }
 }
 
-// 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]]}