Rework LLVM Dialect LoopOptions attribute

[mehdi_amini]

Summary

Instead of storing an array of LoopOpt attributes, which were just
wrapping std::pair<enum, int> anyway, we can have an attribute storing
a sorted ArrayRef<std::pair<enum, int>> as a single unit. This improves
here the textual format and the general API. Note that we're limiting
the options to fit into an int64_t by design, but this isn't a new
constraint.

Building the LoopOptions attribute is likely worth a specific builder
for efficient reason, that'll be the subject of a future patch.

[ftynse]

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

29

Could we have a better comment?

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

2424

This assumption doesn't seem documented anywhere.

More generally, I am not convinced the additional complexity and use awkwardness that comes with requiring the user to maintain the list of options sorted is worth the savings. We may have maybe a dozen attributes tops stored consecutively in memory, linear scan should be just fine here.

mlir/test/Dialect/LLVMIR/invalid.mlir

793–794

I find this double error message counterintuitive. It's even worse because the second message claims the attribute type is unknown instead of saying it failed to parse.

[mehdi_amini]

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

29

Ah right, that was a bad copy/paste :)

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

2424

As I mentioned in the description, I don't expect anyone to use this API directly, we should offer a builder API with convenient access. Sketch would be:

LoopOptionsAttrBuilder b; // can be constructed from a LoopOptionsAttr
b.setDisableLICM(true) 
  .setInterleaveCount(12)
  .setDisableUnroll(false);
LoopOptionsAttr opts = b.create();

So you can even write this:

br_op.setLoopOpts(LoopOptionsAttBuilder(br_op.loopOpts()).setDisableLICM(true).setInterleaveCount(12).setDisableUnroll(false))

mlir/test/Dialect/LLVMIR/invalid.mlir

793–794

Yes I noticed, and actually sent a fix here yesterday: https://reviews.llvm.org/D98162

[rriddle]

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

2402

The predicate of lower_bound should return if the value is less(ordered before), not equal. Use less_first here?

2453

I know this isn't new code, but we don't really use failed for these in ifs(because of chaining).

2491

while (succeeded(parser.parseOptionalComma()))?

2519–2522

2519–2522

If the attribute has a mnemonic it should be included in the generatedAttributePrinter. Is that not happening?

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

227

I wonder if we should just add a templated get to DictionaryAttr.

[mehdi_amini]

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

2402

Oh right! I don't think I can use less_first because only the first argument is a pair.

2424

@rriddle an opinion here?

(right now I add an API doc and renamed the parameter to sortedOptions)

2491

Thanks, much better :)

[mehdi_amini]

Address comments

[rriddle]

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

29

Can you move this into the description field instead?

mlir/include/mlir/IR/BuiltinAttributes.td

206 ↗

(On Diff #329130)

nit: Drop the extra newline.

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

2424

I don't really see a reason *not* to have the list be sorted. It seems like a simple invariant to impose on the user, especially if the current "raw" API is not something most users should have to interface with. We also have to provide the values in a deterministic way that allows for attributes with the same elements to be uniqued, gives deterministic iteration to users, etc.

@mehdi_amini Can you just add the Builder in this revision or is this something you were going to add in a followup?

[mehdi_amini]

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

2424

I can add the Builder now, I didn't because there are no users right now and so the "testing" aspect isn't straightforward: any suggestion on this?

[mehdi_amini]

Add LoopOptionsAttrBuilder

[ftynse]

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

2424

I am fine just having a sketch of the builder in a comment with TODO, so however hits the need sees what was planned.

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

def FastmathFlagsAttr : LLVM_Attr<"FMF"> {
 
   // List of type parameters.
   let parameters = (
     ins
     "FastmathFlags":$flags
   );
 }
 
+// A more complex parameterized attribute.
+def LoopOptionsAttr : LLVM_Attr<"LoopOptions"> {
+  let mnemonic = "loopopts";
+
+  // List of type parameters.
+  let parameters = (
+    ins
+    ArrayRefParameter<"std::pair<LoopOptionCase, int64_t>", "">:$options
+  );
+
+  let extraClassDeclaration = [{
+    Optional<bool> disableUnroll();
+    Optional<bool> disableLICM();
+    Optional<int64_t> interleaveCount();
+  }];
+
+  let builders = [
+    AttrBuilder<(ins "ArrayRef<std::pair<LoopOptionCase, int64_t>>":$options)>
+  ];
+  let skipDefaultBuilders = 1;
+}
+
 #endif // LLVMIR_ATTRDEFS

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

 
 namespace mlir {
 namespace LLVM {
 class LLVMDialect;
 
 namespace detail {
 struct LLVMTypeStorage;
 struct LLVMDialectImpl;
-struct LoopOptionAttrStorage;
 } // namespace detail
-
-/// An attribute that specifies LLVM loop codegen options.
-class LoopOptionAttr
-    : public Attribute::AttrBase<LoopOptionAttr, Attribute,
-                                 detail::LoopOptionAttrStorage> {
-public:
-  using Base::Base;
-
-  /// Specifies the llvm.loop.unroll.disable metadata.
-  static LoopOptionAttr getDisableUnroll(MLIRContext *context,
-                                         bool disable = true);
-
-  /// Specifies the llvm.licm.disable metadata.
-  static LoopOptionAttr getDisableLICM(MLIRContext *context,
-                                       bool disable = true);
-
-  /// Specifies the llvm.loop.interleave.count metadata.
-  static LoopOptionAttr getInterleaveCount(MLIRContext *context, int32_t count);
-
-  /// Returns the loop option, e.g. parallel_access.
-  LoopOptionCase getCase() const;
-
-  /// Returns if the loop option is activated. Only valid for boolean options.
-  bool getBool() const;
-
-  /// Returns the integer value associated with a loop option. Only valid for
-  /// integer options.
-  int32_t getInt() const;
-
-  void print(DialectAsmPrinter &p) const;
-  static Attribute parse(DialectAsmParser &parser);
-};
-
 } // namespace LLVM
 } // namespace mlir
 
 ///// Ops /////
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOps.h.inc"
 
 #include "mlir/Dialect/LLVMIR/LLVMOpsDialect.h.inc"

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

 static constexpr const char kVolatileAttrName[] = "volatile_";
 static constexpr const char kNonTemporalAttrName[] = "nontemporal";
 
 #include "mlir/Dialect/LLVMIR/LLVMOpsEnums.cpp.inc"
 #include "mlir/Dialect/LLVMIR/LLVMOpsInterfaces.cpp.inc"
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOpsAttrDefs.cpp.inc"
 
-namespace mlir {
-namespace LLVM {
-namespace detail {
-
-struct LoopOptionAttrStorage : public AttributeStorage {
-  using KeyTy = std::pair<uint64_t, int32_t>;
-
-  explicit LoopOptionAttrStorage(uint64_t option, int32_t value)
-      : option(option), value(value) {}
-
-  bool operator==(const KeyTy &key) const {
-    return key == KeyTy(option, value);
-  }
-
-  static LoopOptionAttrStorage *
-  construct(mlir::AttributeStorageAllocator &allocator, const KeyTy &key) {
-    return new (allocator.allocate<LoopOptionAttrStorage>())
-        LoopOptionAttrStorage(key.first, key.second);
-  }
-
-  uint64_t option;
-  int32_t value;
-};
-} // namespace detail
-} // namespace LLVM
-} // namespace mlir
-
 static auto processFMFAttr(ArrayRef<NamedAttribute> attrs) {
   SmallVector<NamedAttribute, 8> filteredAttrs(
       llvm::make_filter_range(attrs, [&](NamedAttribute attr) {
         if (attr.first == "fastmathFlags") {
           auto defAttr = FMFAttr::get(attr.second.getContext(), {});
           return defAttr != attr.second;
         }
         return true;
static LogicalResult verify(FenceOp &op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 void LLVMDialect::initialize() {
-  addAttributes<FMFAttr, LoopOptionAttr>();
+  addAttributes<FMFAttr, LoopOptionsAttr>();
 
   // clang-format off
   addTypes<LLVMVoidType,
            LLVMPPCFP128Type,
            LLVMX86MMXType,
            LLVMTokenType,
            LLVMLabelType,
            LLVMMetadataType,
if (parallelAccessGroup.hasValue()) {
         if (!accessGroupOp)
           return op->emitOpError()
                  << "expected '" << attr << "' to reference an access_group op";
       }
     }
 
     Optional<NamedAttribute> loopOptions =
         loopAttr.getNamed(LLVMDialect::getLoopOptionsAttrName());
-    if (loopOptions.hasValue()) {
-      auto options = loopOptions->second.dyn_cast<ArrayAttr>();
-      if (!options)
-        return op->emitOpError()
-               << "expected '" << LLVMDialect::getLoopOptionsAttrName()
-               << "' to be an array attribute";
-      if (!llvm::all_of(options, [](Attribute option) {
-            return option.isa<LoopOptionAttr>();
-          }))
-        return op->emitOpError() << "invalid loop options list " << options;
-    }
+    if (loopOptions.hasValue() && !loopOptions->second.isa<LoopOptionsAttr>())
+             << "' to be a `loopopts` attribute";
   }
 
   // If the data layout attribute is present, it must use the LLVM data layout
   // syntax. Try parsing it and report errors in case of failure. Users of this
   // attribute may assume it is well-formed and can pass it to the (asserting)
   // llvm::DataLayout constructor.
   if (attr.first.strref() != LLVM::LLVMDialect::getDataLayoutAttrName())
     return success();
if (failed(parser.parseOptionalGreater())) {
 
     if (failed(parser.parseGreater()))
       return {};
   }
 
   return FMFAttr::get(parser.getBuilder().getContext(), flags);
 }
 
-LoopOptionAttr LoopOptionAttr::getDisableUnroll(MLIRContext *context,
-                                                bool disable) {
-  auto option = LoopOptionCase::disable_unroll;
-  return Base::get(context, static_cast<uint64_t>(option),
-                   static_cast<int32_t>(disable));
-}
-
-LoopOptionAttr LoopOptionAttr::getDisableLICM(MLIRContext *context,
-                                              bool disable) {
-  auto option = LoopOptionCase::disable_licm;
-  return Base::get(context, static_cast<uint64_t>(option),
-                   static_cast<int32_t>(disable));
+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;
+      });
+  if (it == options.end())
+    return {};
+  return static_cast<T>(it->second);
 }
 
-LoopOptionAttr LoopOptionAttr::getInterleaveCount(MLIRContext *context,
-                                                  int32_t count) {
-  auto option = LoopOptionCase::interleave_count;
-  return Base::get(context, static_cast<uint64_t>(option),
-                   static_cast<int32_t>(count));
+Optional<bool> LoopOptionsAttr::disableUnroll() {
+  return getOption<bool>(getOptions(), LoopOptionCase::disable_unroll);
 }
 
-LoopOptionCase LoopOptionAttr::getCase() const {
-  return static_cast<LoopOptionCase>(getImpl()->option);
+Optional<bool> LoopOptionsAttr::disableLICM() {
+  return getOption<bool>(getOptions(), LoopOptionCase::disable_licm);
 }
 
-bool LoopOptionAttr::getBool() const {
-  LoopOptionCase option = getCase();
-  (void)option;
-  assert(option == LoopOptionCase::disable_licm ||
-         option == LoopOptionCase::disable_unroll &&
-             "expected a boolean loop option");
-  return static_cast<bool>(getImpl()->value);
+Optional<int64_t> LoopOptionsAttr::interleaveCount() {
+  return getOption<int64_t>(getOptions(), LoopOptionCase::interleave_count);
 }
 
-int32_t LoopOptionAttr::getInt() const {
-  LoopOptionCase option = getCase();
-  (void)option;
-  assert(option == LoopOptionCase::interleave_count &&
-         "expected an integer loop option");
-  return getImpl()->value;
+LoopOptionsAttr
+LoopOptionsAttr::get(MLIRContext *context,
+                     ArrayRef<std::pair<LoopOptionCase, int64_t>> options) {
+  assert(llvm::is_sorted(options, llvm::less_first()));
+  return Base::get(context, options);
 }
 
-void LoopOptionAttr::print(DialectAsmPrinter &printer) const {
-  printer << "loopopt<" << stringifyEnum(getCase()) << " = ";
-  switch (getCase()) {
+void LoopOptionsAttr::print(DialectAsmPrinter &printer) const {
+  printer << getMnemonic() << "<";
+  llvm::interleaveComma(getOptions(), printer, [&](auto option) {
+    switch (option.first) {
     case LoopOptionCase::disable_licm:
     case LoopOptionCase::disable_unroll:
-    printer << (getBool() ? "true" : "false");
+      printer << (option.second ? "true" : "false");
       break;
     case LoopOptionCase::interleave_count:
-    printer << getInt();
+      printer << option.second;
       break;
     }
+  });
   printer << ">";
 }
 
-Attribute LoopOptionAttr::parse(DialectAsmParser &parser) {
+Attribute LoopOptionsAttr::parse(MLIRContext *context, DialectAsmParser &parser,
+                                 Type type) {
   if (failed(parser.parseLess()))
     return {};
 
+  SmallVector<std::pair<LoopOptionCase, int64_t>> options;
+  llvm::SmallDenseSet<LoopOptionCase> seenOptions;
+  do {
     StringRef optionName;
     if (failed(parser.parseKeyword(&optionName)))
       return {};
 
     auto option = symbolizeLoopOptionCase(optionName);
     if (!option) {
       parser.emitError(parser.getNameLoc(), "unknown loop option: ")
           << optionName;
       return {};
     }
-
+    if (!seenOptions.insert(*option).second) {
+      parser.emitError(parser.getNameLoc(), "loop option present twice");
+      return {};
+    }
     if (failed(parser.parseEqual()))
       return {};
 
-  int32_t value;
+    int64_t value;
     switch (*option) {
     case LoopOptionCase::disable_licm:
     case LoopOptionCase::disable_unroll:
       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:
       if (failed(parser.parseInteger(value))) {
         parser.emitError(parser.getNameLoc(), "expected integer value");
         return {};
       }
       break;
     }
-
+    options.push_back(std::make_pair(*option, value));
+  } while (!parser.parseOptionalComma());
   if (failed(parser.parseGreater()))
     return {};
 
-  return Base::get(parser.getBuilder().getContext(),
-                   static_cast<uint64_t>(*option), value);
+  llvm::sort(options, llvm::less_first());
+  return get(parser.getBuilder().getContext(), options);
 }
 
 Attribute LLVMDialect::parseAttribute(DialectAsmParser &parser,
                                       Type type) const {
   if (type) {
     parser.emitError(parser.getNameLoc(), "unexpected type");
     return {};
   }
   StringRef attrKind;
   if (parser.parseKeyword(&attrKind))
     return {};
   if (auto attr =
           generatedAttributeParser(getContext(), parser, attrKind, type))
     return attr;
 
-  if (attrKind == "loopopt")
-    return LoopOptionAttr::parse(parser);
-
   parser.emitError(parser.getNameLoc(), "Unknown attribute type: ") << attrKind;
   return {};
 }
 
 void LLVMDialect::printAttribute(Attribute attr, DialectAsmPrinter &os) const {
   if (succeeded(generatedAttributePrinter(attr, os)))
     return;
-  if (auto lopt = attr.dyn_cast<LoopOptionAttr>())
+  if (auto lopt = attr.dyn_cast<LoopOptionsAttr>())
     lopt.print(os);
   else
     llvm_unreachable("Unknown attribute type");
 }

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

for (auto it : handlers)
     if (bitEnumContains(fmf, it.first))
       (ret.*(it.second))(true);
   return ret;
 }
 
 /// Returns an LLVM metadata node corresponding to a loop option. This metadata
 /// is attached to an llvm.loop node.
 static llvm::MDNode *getLoopOptionMetadata(llvm::LLVMContext &ctx,
-                                           LoopOptionAttr option) {
+                                           LoopOptionCase option,
+                                           int64_t value) {
   StringRef name;
-  llvm::Constant *value = nullptr;
-  switch (option.getCase()) {
+  llvm::Constant *cstValue = nullptr;
+  switch (option) {
   case LoopOptionCase::disable_licm:
     name = "llvm.licm.disable";
-    value = llvm::ConstantInt::getBool(ctx, option.getBool());
+    cstValue = llvm::ConstantInt::getBool(ctx, value);
     break;
   case LoopOptionCase::disable_unroll:
     name = "llvm.loop.unroll.disable";
-    value = llvm::ConstantInt::getBool(ctx, option.getBool());
+    cstValue = llvm::ConstantInt::getBool(ctx, value);
     break;
   case LoopOptionCase::interleave_count:
     name = "llvm.loop.interleave.count";
-    value = llvm::ConstantInt::get(llvm::IntegerType::get(ctx, /*NumBits=*/32),
-                                   option.getInt());
+    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(value)});
+                                 llvm::ConstantAsMetadata::get(cstValue)});
 }
 
 static void setLoopMetadata(Operation &opInst, llvm::Instruction &llvmInst,
                             llvm::IRBuilderBase &builder,
                             LLVM::ModuleTranslation &moduleTranslation) {
   if (Attribute attr = opInst.getAttr(LLVMDialect::getLoopAttrName())) {
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::MDNode *loopMD = moduleTranslation.lookupLoopOptionsMetadata(attr);
if (!loopMD) {
         for (SymbolRefAttr accessGroupRef :
              parallelAccessGroup->second.cast<ArrayAttr>()
                  .getAsRange<SymbolRefAttr>())
           parallelAccess.push_back(
               moduleTranslation.getAccessGroup(opInst, accessGroupRef));
         loopOptions.push_back(llvm::MDNode::get(ctx, parallelAccess));
       }
 
-      auto loopOptionsAttr =
-          loopAttr.getNamed(LLVMDialect::getLoopOptionsAttrName());
-      if (loopOptionsAttr.hasValue()) {
-        for (LoopOptionAttr loopOption :
-             loopOptionsAttr->second.cast<ArrayAttr>()
-                 .getAsRange<LoopOptionAttr>())
-          loopOptions.push_back(getLoopOptionMetadata(ctx, loopOption));
+      if (LoopOptionsAttr loopOptionsAttr =
+              loopAttr.get(LLVMDialect::getLoopOptionsAttrName())
+                  .dyn_cast_or_null<LoopOptionsAttr>()) {
+        for (auto option : loopOptionsAttr.getOptions())
+          loopOptions.push_back(
+              getLoopOptionMetadata(ctx, option.first, option.second));
       }
 
       // Create loop options and set the first operand to itself.
       loopMD = llvm::MDNode::get(ctx, loopOptions);
       loopMD->replaceOperandWith(0, loopMD);
 
       // Store a map from this Attribute to the LLVM metadata in case we
       // encounter it again.

mlir/test/Dialect/LLVMIR/invalid.mlir

llvm.metadata @metadata {
     llvm.return
   }
 }
 
 // -----
 
 module {
   llvm.func @loopOptions() {
-      // expected-error@below {{expected 'options' to be an array attribute}}
+      // expected-error@below {{expected 'options' to be a `loopopts` attribute}}
       llvm.br ^bb4 {llvm.loop = {options = "name"}}
     ^bb4:
       llvm.return
   }
 }
 
 // -----
 
 module {
   llvm.func @loopOptions() {
-      // expected-error@below {{invalid loop options list}}
-      llvm.br ^bb4 {llvm.loop = {options = ["name"]}}
+      // expected-error@+2 {{unknown loop option: name}}
+      // expected-error@below {{Unknown attribute type: loopopts}}
+      llvm.br ^bb4 {llvm.loop = {options = #llvm.loopopts<name>}}
+    ^bb4:
+      llvm.return
+  }
+}
+
+// -----
+
+module {
+  llvm.func @loopOptions() {
+      // expected-error@+2 {{loop option present twice}}
+      // expected-error@below {{Unknown attribute type: loopopts}}
+      llvm.br ^bb4 {llvm.loop = {options = #llvm.loopopts<disable_licm = true, disable_licm = true>}}
     ^bb4:
       llvm.return
   }
 }
 
 // -----
 
 module {

mlir/test/Target/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.loopopt<disable_unroll = true>, #llvm.loopopt<disable_licm = true>, #llvm.loopopt<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>}}
     ^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.loopopt<disable_unroll = true>, #llvm.loopopt<disable_licm = true>, #llvm.loopopt<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>}}
     ^bb5:
       llvm.return
   }
 
   llvm.metadata @metadata {
     llvm.access_group @group1
     llvm.access_group @group2
     llvm.return