[mlir][irdl] Add `irdl.attributes` operation for defining named attributes

[unterumarmung]

Summary

This commit introduces the irdl.attributes operation, which allows defining named attributes for the parent operation. Each attribute is defined with a name and a type constraint.

Example usage:

irdl.dialect @example {
  irdl.operation @attr_op {
    %0 = irdl.any
    %1 = irdl.is i64
    irdl.attributes {
      "attr1" = %0,
      "attr2" = %1
    }
  }
}

In this example the operation will expect an arbitrary attribute "attr1"
and an attribute "attr2" with value i64.

[math-fehr]

Thanks a lot for the patch!
I added some inline comments.

I think that one confusion is that constraints should be applied directly on the attributes, and should not constraint the attributes type.
For instance, irdl.attributes { attr1 = %x } should constraint attr1 to the %x constraint, and should not constraint attr1 type to %x.
The reasoning is that we might want to write constraints such as attr1 is AnyOf<"foo", "bar"> to constraint directly the attribute itself,
or something like attr1 is IntegerAttr<%param1, %param2> to constraint it to be an IntegerAttr with the given parameters.
Note that I'm working on making these constraints expressible, for now we can only state the AnyOf one!

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

74

I would get op->getAttrDictionary here, since it is already constructed, and is always sorted.

76–81

In MLIR, it is valid to have extra attributes on operations, so this part of the code should be removed.

85

If you use the dictionary, you can just write getNamed(attrDefName) here.

106–116

I would merge this part with the previous one, so in the end you only iterate to the list of attributes once.

112

This should just be attr.getValue() instead of TypeAttr::get(type).
This is because of the attribute is a TypeAttr, this will stay the same, and if the attribute is another Attribute, this version would fail.

mlir/test/Dialect/IRDL/testd.irdl.mlir

118–121

Here, you are actually constraining two attributes to be exactly i32 and i64.
However I think in your test you want to express that these are IntegerAttr with type i32 and i64.

The reason why we don't want to constraint only the type directly here, is that later on we would like to express things like
AnyOf("foo", "bar") to express that an attribute is either the "foo" string literal or the "bar" literal.

For expressing that we want to constraint the type of an IntegerAttr, we will need more feature in IRDL first :/

So here, I would just keep the same test, but test it with i32 and i64 instead of the IntegerAttr.

[math-fehr]

Oh also, could you write a verifier in the operation that attributeValues and attributeValueNames have the same length?

[unterumarmung]

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

106–116

I haven't merged this part with the previous one because the function consists of two logical parts: (1) Basic error checking and (2) Constraint verification. It seemed appropriate to keep them separate. Additionally, the constraint verification part appears to be more resource-intensive than the basic check. Therefore, I believe it would be more prudent to ensure that all necessary attributes are passed before performing the constraint checking.

Regarding the efficiency of the basic error checking part, I think we can potentially improve its time complexity. Currently, it has a time complexity of O(n*m) and a space complexity of O(1), where n and m represent the sizes of the actual attributes and declared ones, respectively. By storing both arrays in two sets and calculating their intersection using std::set_intersection, we could reduce the time complexity to O(m+n) and the space complexity to O(m+n). However, I understand that this approach may introduce some code complexity, so I refactored it into a separate function.

Nevertheless, I am more than willing to implement it according to your preferences. I just wanted to provide my reasoning and suggest a possible solution to enhance its performance.

112

If I do so, I get the following error

within split at mlir/test/Dialect/IRDL/testd.mlir:238 offset :5:3: error: unexpected error: expected 'i32' but got '42 : i32'
  "testd.attrs"() {attr1 = 42 : i32, attr2 = 42 : i32} : () -> ()
  ^

So, I guess to make it work I need to fix verifiers, right?
Could you give any pointer on where to look at or how to fix it the best way, please?

[Mogball]

mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td

290

please fit this into lines less than 80 characters wide

mlir/lib/Dialect/IRDL/IR/IRDL.cpp

79

please drop the 4. The practice is to not use explicit sizing unless there is a compelling reason to.

84–85

?

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

74

please spell out this auto

76

this one too. The rule is that auto should only be used if the type is spelled verbatim on the RHS like auto thing = dyn_cast<ThingOp>(op) or "too complicated" like DenseMap<int, int>::iterator

83

same here

86

llvm::find_if

88

why not pointer compare the StringAttrs? Instead of the StringRefs

110

[math-fehr]

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

106–116

I see!

So one thing to note is that it is fine if the cost is a bit bigger when the verifier fail. After all, this happens at most once, since it will fail the entire compilation pipeline. So merging both will make it faster for the cases that matter, and only a bit slower on the others.

Then, if we want to actually optimize more the performance, we can actually rely on the fact that DictionaryAttr is sorted, so we can just iterate on both the dictionary and the constraints once, if we sort the constraints beforehand. That way we get O(m+n) time and still O(1) size.

112

Yes, you would need to change the tests as well.
The idea is that in your current code, you are constraining the attribute to be a typed attribute of type i32.
However, in our verifiers, we want to directly constraint the attribute itself, not its type.

We don't have examples here yet of that, since we only constraint types for now, but the idea is that you just call verifier.verify({emitError}, attr, constraint) instead of extracting a Type subelement and then passing it to a TypeAttr.

[unterumarmung]

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

112

So, if IRDL doesn't support constraining anything but values, how should we test that it can constraint values also?
Or should I just remove the test with the constraint verification for now?
Or maybe there is a possible fix in the ConstraintVerifier that just searches for the type if it sees a type constraint?

[unterumarmung]

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

88

I'm not very familiar with MLIR in general. Comparing 2 unrelated (at least, for the MLIR engine) StringAttrs by pointers never occured to me.
Is it guaranteed somehow that attributes with the same value are created only once and then reused?

112

anything but values

I meant "anything but types"

[math-fehr]

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

88

Yes, whenever a StringAttr is created, it is uniqued by the MLIRContext (a StringAttr is actually just a pointer to the actual value).
So if you have two StringAttr that encodes the same value, you can just compare them by equality, and this will be done by a pointer equality.

112

So the current way in the ConstraintVerifier is currently intended, so in your case you should just test it with i32 and i64 values (so directly "attr1" = i32).

We need to add more support to constraint directly the type of an attribute, but I would need to think about it some more to know what is the best way (I believe for now the best way would be to add external named constraints).

[unterumarmung]

Updated code according to review

[unterumarmung]

@math-fehr @Mogball Thank you for the review! I've reworked the code accroding to your suggestions and it got quite nicer. Please take a look.

---

In D152618#4412153, @math-fehr wrote:

Oh also, could you write a verifier in the operation that attributeValues and attributeValueNames have the same length?

I'm not sure if it's needed because this rule is enforced by the operation parser, isn't it?
But on the other hand I may be missing something. If you think that the verifier is needed, could you please provide a test case that tests this check? I couldn't come up with any.

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

83

The type here llvm::detail::DenseMapPair<mlir::StringAttr, unsigned long>
I'm not sure that it should be spelled out since it's too verbose

Also, while refactoring according the review, I got rid of it

85

That's markedly better, thanks!

86

Thanks, I was looking for llvm STL algorithms but some how missed STLExtras.h header

106–116

If you mean that we can zip them and iterate checking if the pair we got contains the same-named attributes, I'm not sure it's possible because, as you mentioned, MLIR allows extra attributes to be passed, so they would interfere.
If you meant that we can use "two pointers" technique - that is possible, yes. But I'm not sure whether this approach does not reduce the code quality.

I merged the two parts as you suggested and took advantage of the fact that attributes on actual operation are sorted, so the whole algorithm now is O(log(n)*m)
Is this approach acceptable?

[unterumarmung]

Fix formatting

[math-fehr]

I'm not sure if it's needed because this rule is enforced by the operation parser, isn't it?
But on the other hand I may be missing something. If you think that the verifier is needed, could you please provide a test case that tests this check? I couldn't come up with any.

It's enforced by the operation parser, but only if you use the custom format.
If you use the generic format, or if you create the operation in C++, this is not enforced anymore.
For a test, you can just write it using the generic format (something like "irdl.attributes"() { attributeValueNames = ["name"] }).

Other than this, this looks good to me!

mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td

287–288

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

106–116

Yes, I meant a "two pointers" algorithm.
But I'm fine with a O(log(n) * m) complexity, as anyway we should have an even lower dialect if we want better performances!

[unterumarmung]

Rebased and fixed review comments

[unterumarmung]

In D152618#4431013, @math-fehr wrote:

If you use the generic format, or if you create the operation in C++, this is not enforced anymore.

Yes, you are right! Totally forgot about the generic format.

All review comments are fixed, please take a look.

[math-fehr]

You should add the check in the verifier rather than in the loading function.
When we create any IRDL program (even if we don't load it), it should trigger an error.
Can you also add a test for that?

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

176–182

This should be added in the operation verifier rather than here.
Also, when we load an IRDL program, we expect the program to verify, so we won't need it there as well.

[Mogball]

mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td

291

292

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

90–92

drop trivial braces

173–174

please spell out these autos

[Mogball]

LGTM but please wait for @math-fehr's approval

[unterumarmung]

Address review comments

[unterumarmung]

@math-fehr You are right that I should've implemented the verifier method instead of checking the operation validnesses at loading.
I've done so and added a new test file for this purpose because I couldn't find the appropriate one

@Mogball thank you for the approve, I've addressed your comments

[math-fehr]

Perfect! Thanks a lot!

mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td

 }
 
 //===----------------------------------------------------------------------===//
 // IRDL Operation definition
 //===----------------------------------------------------------------------===//
 
 def IRDL_OperationOp : IRDL_Op<"operation",
     [HasParent<"DialectOp">, NoTerminator, NoRegionArguments,
-    AtMostOneChildOf<"OperandsOp, ResultsOp">, Symbol]> {
+    AtMostOneChildOf<"OperandsOp, ResultsOp, AttributesOp">, Symbol]> {
   let summary = "Define a new operation";
   let description = [{
     `irdl.operation` defines a new operation belonging to the `irdl.dialect`
     parent.
 
     Operations can define constraints on their operands and results with the
     `irdl.results` and `irdl.operands` operations. If these operations are not
     present in the region, the results or operands are expected to be empty.
let description = [{
     The operation will expect one operand of the `cmath.complex` type, and two
     results that have the underlying type of the `cmath.complex`.
   }];
 
   let arguments = (ins Variadic<IRDL_AttributeType>:$args);
   let assemblyFormat = " `(` $args `)` attr-dict ";
 }
 
+def IRDL_AttributesOp : IRDL_Op<"attributes", [HasParent<"OperationOp">]> {
+  let summary = "Define the attributes of an operation";
+
+  let description = [{
+    `irdl.attributes` defines the attributes of the `irdl.operation` parent 
+    operation definition.
+
+    In the following example, `irdl.attributes` defines the attributes of the
+    `attr_op` operation:
+
+    ```mlir
+    irdl.dialect @example {
+
+      irdl.operation @attr_op {
+        %0 = irdl.any
+        %1 = irdl.is i64
+        irdl.attibutes {
+          "attr1" = %0,
+          "attr2" = %1
+        }
+      }
+    }
+    ```
+
+    The operation will expect attribute "attr1" of any type
+    and attribute "attr2" of type `i64` type.
+  }];
+
+  let arguments = (ins Variadic<IRDL_AttributeType>:$attributeValues, StrArrayAttr:$attributeValueNames);
+  let assemblyFormat = "custom<AttributesOp>($attributeValues, $attributeValueNames) attr-dict ";
+}
+
 //===----------------------------------------------------------------------===//
 // IRDL Constraint operations
 //===----------------------------------------------------------------------===//
 
 class IRDL_ConstraintOp<string mnemonic, list<Trait> traits = []>
     : IRDL_Op<mnemonic, [VerifyConstraintInterface,
         DeclareOpInterfaceMethods<VerifyConstraintInterface>] # traits> {
 }

mlir/lib/Dialect/IRDL/IR/IRDL.cpp

 }
 
 LogicalResult DialectOp::verify() {
   if (!Dialect::isValidNamespace(getName()))
     return emitOpError("invalid dialect name");
   return success();
 }
 
+static ParseResult
+parseAttributesOp(OpAsmParser &p,
+                  SmallVectorImpl<OpAsmParser::UnresolvedOperand> &attrOperands,
+                  ArrayAttr &attrNamesAttr) {
+  Builder &builder = p.getBuilder();
+  SmallVector<Attribute> attrNames;
+  if (succeeded(p.parseOptionalLBrace())) {
+    auto parseOperands = [&]() {
+      if (p.parseAttribute(attrNames.emplace_back()) || p.parseEqual() ||
+          p.parseOperand(attrOperands.emplace_back()))
+        return failure();
+      return success();
+    };
+    if (p.parseCommaSeparatedList(parseOperands) || p.parseRBrace())
+      return failure();
+  }
+  attrNamesAttr = builder.getArrayAttr(attrNames);
+  return success();
+}
+
+static void printAttributesOp(OpAsmPrinter &p, AttributesOp op,
+                              OperandRange attrArgs, ArrayAttr attrNames) {
+  if (attrNames.empty())
+    return;
+  p << "{";
+  interleaveComma(llvm::seq<int>(0, attrNames.size()), p,
+                  [&](int i) { p << attrNames[i] << " = " << attrArgs[i]; });
+  p << '}';
+}
+
 #include "mlir/Dialect/IRDL/IR/IRDLInterfaces.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/IRDL/IR/IRDLTypesGen.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/IRDL/IR/IRDLOps.cpp.inc"

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

 //===- IRDLLoading.cpp - IRDL dialect loading --------------------- C++ -*-===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Manages the loading of MLIR objects from IRDL operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/IRDL/IRDLLoading.h"
 #include "mlir/Dialect/IRDL/IR/IRDL.h"
 #include "mlir/Dialect/IRDL/IR/IRDLInterfaces.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/SMLoc.h"
 
 using namespace mlir;
 using namespace mlir::irdl;
 
irdlAttrOrTypeVerifier(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
 /// Verify that the given operation satisfies the given constraints.
 /// This encodes the logic of the verification method for operations defined
 /// with IRDL.
 static LogicalResult
 irdlOpVerifier(Operation *op, ArrayRef<std::unique_ptr<Constraint>> constraints,
-               ArrayRef<size_t> operandConstrs,
-               ArrayRef<size_t> resultConstrs) {
+               ArrayRef<size_t> operandConstrs, ArrayRef<size_t> resultConstrs,
+               const DenseMap<StringAttr, size_t> &attributeConstrs) {
   /// Check that we have the right number of operands.
   unsigned numOperands = op->getNumOperands();
   size_t numExpectedOperands = operandConstrs.size();
   if (numOperands != numExpectedOperands)
     return op->emitOpError() << numExpectedOperands
                              << " operands expected, but got " << numOperands;
 
   /// Check that we have the right number of results.
   unsigned numResults = op->getNumResults();
   size_t numExpectedResults = resultConstrs.size();
   if (numResults != numExpectedResults)
     return op->emitOpError()
            << numExpectedResults << " results expected, but got " << numResults;
 
-  auto emitError = [op]() { return op->emitError(); };
+  auto emitError = [op] { return op->emitError(); };
 
   ConstraintVerifier verifier(constraints);
 
+  /// Сheck that we have all needed attributes passed and satisfy the constraints.
+  DictionaryAttr actualAttrs = op->getAttrDictionary();
+
+  for (auto [name, constraint] : attributeConstrs) {
+    /// First, check if the attribute actually passed.
+    std::optional<NamedAttribute> actual = actualAttrs.getNamed(name);
+    if (!actual.has_value()) {
+      return op->emitOpError()
+             << "attribute " << name << " is expected but not provided";
+    }
+
+    /// Then, check if the attribute value satisfies the constraint.
+    if (failed(verifier.verify({emitError}, actual->getValue(), constraint))) {
+      return failure();
+    }
+  }
+
   /// Check that all operands satisfy the constraints.
   for (auto [i, operandType] : enumerate(op->getOperandTypes()))
     if (failed(verifier.verify({emitError}, TypeAttr::get(operandType),
                                operandConstrs[i])))
       return failure();
 
   /// Check that all results satisfy the constraints.
   for (auto [i, resultType] : enumerate(op->getResultTypes()))
     if (failed(verifier.verify({emitError}, TypeAttr::get(resultType),
                                resultConstrs[i])))
       return failure();
 
   return success();
 }
 
 /// Define and load an operation represented by a `irdl.operation`
 /// operation.
 static WalkResult loadOperation(
     OperationOp op, ExtensibleDialect *dialect,
     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> &types,
     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> &attrs) {
   // Resolve SSA values to verifier constraint slots
   SmallVector<Value> constrToValue;
   for (Operation &op : op->getRegion(0).getOps()) {
     if (isa<VerifyConstraintInterface>(op)) {
       if (op.getNumResults() != 1)
         return op.emitError()
                << "IRDL constraint operations must have exactly one result";
       constrToValue.push_back(op.getResult(0));
     }
   }
for (Value result : resultsOp->getArgs()) {
         if (constr == result) {
           resultConstraints.push_back(i);
           break;
         }
       }
     }
   }
 
+  // Gather which constraint slots correspond to attributes constraints
+  DenseMap<StringAttr, size_t> attributesContraints;
+  auto attributesOp = op.getOp<AttributesOp>();
+  if (attributesOp.has_value()) {
+    const auto values = attributesOp->getAttributeValues();
+    const auto names = attributesOp->getAttributeValueNames();
+
+    for (const auto &[name, value] : llvm::zip(names, values)) {
+      for (auto [i, constr] : enumerate(constrToValue)) {
+        if (constr == value) {
+          attributesContraints[name.cast<StringAttr>()] = i;
+          break;
+        }
+      }
+    }
+  }
+
   // IRDL does not support defining custom parsers or printers.
   auto parser = [](OpAsmParser &parser, OperationState &result) {
     return failure();
   };
   auto printer = [](Operation *op, OpAsmPrinter &printer, StringRef) {
     printer.printGenericOp(op);
   };
 
   auto verifier =
       [constraints{std::move(constraints)},
        operandConstraints{std::move(operandConstraints)},
-       resultConstraints{std::move(resultConstraints)}](Operation *op) {
+       resultConstraints{std::move(resultConstraints)},
+       attributesContraints{std::move(attributesContraints)}](Operation *op) {
         return irdlOpVerifier(op, constraints, operandConstraints,
-                              resultConstraints);
+                              resultConstraints, attributesContraints);
       };
 
   // IRDL does not support defining regions.
   auto regionVerifier = [](Operation *op) { return success(); };
 
   auto opDef = DynamicOpDefinition::get(
       op.getName(), dialect, std::move(verifier), std::move(regionVerifier),
       std::move(parser), std::move(printer));

mlir/test/Dialect/IRDL/testd.irdl.mlir

irdl.dialect @testd {
   // CHECK:   irdl.results(%[[v2]], %[[v2]])
   // CHECK: }
   irdl.operation @constraint_vars {
     %0 = irdl.is i32
     %1 = irdl.is i64
     %2 = irdl.any_of(%0, %1)
     irdl.results(%2, %2)
   }
+
+
+  // CHECK: irdl.operation @attrs {
+  // CHECK:   %[[v0:[^ ]*]] = irdl.is i32
+  // CHECK:   %[[v1:[^ ]*]] = irdl.is i64
+  // CHECK:   irdl.attributes {"attr1" = %[[v0]], "attr2" = %[[v1]]}
+  // CHECK: }
+  irdl.operation @attrs {
+    %0 = irdl.is i32
+    %1 = irdl.is i64
+
+    irdl.attributes {
+      "attr1" = %0,
+      "attr2" = %1
+    }
+  }
 }

mlir/test/Dialect/IRDL/testd.mlir

 
 // -----
 
 func.func @failedConstraintVars() {
   // expected-error@+1 {{expected 'i64' but got 'i32'}}
   "testd.constraint_vars"() : () -> (i64, i32)
   return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Constraint attributes
+//===----------------------------------------------------------------------===//
+
+func.func @succeededAttrs() {
+  // CHECK: "testd.attrs"() {attr1 = i32, attr2 = i64} : () -> ()
+  "testd.attrs"() {attr1 = i32, attr2 = i64} : () -> ()
+  return
+}
+
+// -----
+
+func.func @failedAttrsMissingAttr() {
+  // expected-error@+1 {{attribute "attr2" is expected but not provided}}
+  "testd.attrs"() {attr1 = i32} : () -> ()
+  return
+}
+
+// -----
+
+func.func @failedAttrsConstraint() {
+  // expected-error@+1 {{expected 'i32' but got 'i64'}}
+  "testd.attrs"() {attr1 = i64, attr2 = i64} : () -> ()
+  return
+}
+
+// -----
+
+func.func @failedAttrsConstraint2() {
+  // expected-error@+1 {{expected 'i64' but got 'i32'}}
+  "testd.attrs"() {attr1 = i32, attr2 = i32} : () -> ()
+  return
+}