[mlir][ods] Add materialize derived attribute method

[jpienaar]

Summary

Generate method to generate a DictionaryAttr with attribute values of derived
attribute. If a conversion back from the derived attribute C++ type to
Attribute is not defined, then attempting to materialize such an op's derived
attributes would result in runtime failure.

This allows to treat derived attributes and attributes of an op in more uniform manner where needed. The derived attributes are not added to the operation but returned as new attribute instead.

[rriddle]

mlir/include/mlir/Interfaces/DerivedAttributeOpInterface.td

35

Can you indent this like the others?

mlir/test/lib/Dialect/Test/TestPatterns.cpp

146

Seems easier to just do:

emitRemark() << d.first << " = " << d.second;

[jpienaar]

Addressing review suggestions and add missing cmake dep

[jpienaar]

Updated, thanks

[jpienaar]

rebase to head to avoid unrelated failure at previous head

[rriddle]

mlir/include/mlir/IR/OpBase.td

1434

Can we add a substitution for a builder instance? The builder makes constructing many of the standard attributes much easier.

mlir/include/mlir/Interfaces/DerivedAttributeOpInterface.td

37

Weird indent here.

mlir/test/lib/Dialect/Test/CMakeLists.txt

33 ↗

(On Diff #258742)

Can we have these sorted?

mlir/test/lib/Dialect/Test/TestPatterns.cpp

146

Is the strref necessary here?

[jpienaar]

Add builder as replacement pattern

mlir/include/mlir/IR/OpBase.td

 // Derive attribute kinds
 
 // DerivedAttr are attributes whose value is computed from properties
 // of the operation. They do not require additional storage and are
 // materialized as needed.
 // Note: All derived attributes should be materializable as an Attribute. E.g.,
 // do not use DerivedAttr for things that could not have been stored as
 // Attribute.
-class DerivedAttr<code ret, code b> : Attr<CPred<"true">, "derived attribute"> {
+//
+class DerivedAttr<code ret, code b, code convert = ""> :
+    Attr<CPred<"true">, "derived attribute"> {
   let returnType = ret;
   code body = b;
+
+  // Specify how to convert from the derived attribute to an attibute.
+  //
+  // ## Special placeholders
+  //
+  // Special placeholders can be used to refer to entities during conversion:
+  //
+  // * `$_ctx` will be replaced by the MLIRContext* instance.
+  // * `$_self` will be replaced with the derived attribute (value produces
+  //    `returnType`).
+  let convertFromStorage = convert;
 }
 
 // Derived attribute that returns a mlir::Type.
-class DerivedTypeAttr<code body> : DerivedAttr<"Type", body>;
+class DerivedTypeAttr<code body> : DerivedAttr<"Type", body> {
+  let convertFromStorage = "TypeAttr::get($_self)";
+}
 
 //===----------------------------------------------------------------------===//
 // Constant attribute kinds
 
 // Represents a constant attribute of specific Attr type. A constant
 // attribute can be specified only of attributes that have a constant
 // builder call defined. The constant value is specified as a string.
 //

mlir/include/mlir/Interfaces/DerivedAttributeOpInterface.td

def DerivedAttributeOpInterface : OpInterface<"DerivedAttributeOpInterface"> {
 
   let methods = [
     StaticInterfaceMethod<
       /*desc=*/"Returns whether name corresponds to a derived attribute.",
       /*retTy=*/"bool",
       /*methodName=*/"isDerivedAttribute",
       /*args=*/(ins "StringRef":$name)
     >,
+    InterfaceMethod<
+      /*desc=*/[{Materializes the derived attributes. Returns null attribute
+                 where unable to materialize a derived attribute as attribute.}],
+      /*retTy=*/"DictionaryAttr",
+      /*methodName=*/"materializeDerivedAttributes"
+    >,
   ];
 }
 
 #endif // MLIR_DERIVEDATTRIBUTEOPINTERFACE

mlir/test/lib/Dialect/Test/TestDialect.h

 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/DerivedAttributeOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffects.h"
 
 #include "TestOpEnums.h.inc"
 
 namespace mlir {
 
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_TESTDIALECT_H

mlir/test/lib/Dialect/Test/TestOps.td

 
 def RankedIntElementsAttrOp : TEST_Op<"ranked_int_elements_attr"> {
   let arguments = (ins
       RankedI32ElementsAttr<[2]>:$vector_i32_attr,
       RankedI64ElementsAttr<[4, 8]>:$matrix_i64_attr
   );
 }
 
+def DerivedTypeAttrOp : TEST_Op<"derived_type_attr", []> {
+  let results = (outs AnyTensor:$output);
+  DerivedTypeAttr element_dtype =
+    DerivedTypeAttr<"return getElementTypeOrSelf(output().getType());">;
+  DerivedAttr size = DerivedAttr<"int",
+    "return output().getType().cast<ShapedType>().getSizeInBits();",
+    "IntegerAttr::get(IntegerType::get(32, $_ctx), $_self)">;
+}
+
 //===----------------------------------------------------------------------===//
 // Test Attribute Constraints
 //===----------------------------------------------------------------------===//
 
 def SymbolRefOp : TEST_Op<"symbol_ref_attr"> {
   let arguments = (ins
     Confined<FlatSymbolRefAttr, [ReferToOp<"FuncOp">]>:$symbol
   );

mlir/test/lib/Dialect/Test/TestPatterns.cpp

 //===- TestPatterns.cpp - Test dialect pattern driver ---------------------===//
 //
 // Part of the LLVM Project, 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
 #include "mlir/Conversion/StandardToStandard/StandardToStandard.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/Support/FormatVariadic.h"
 using namespace mlir;
 
 // Native function for testing NativeCodeCall
 static Value chooseOperand(Value input1, Value input2, BoolAttr choice) {
   return choice.getValue() ? input1 : input2;
 }
 
 static void createOpI(PatternRewriter &rewriter, Value input) {
if (getFunction().getName() == "testReifyFunctions") {
       // Generate test patterns for each, but skip terminator.
       for (auto *op : ops)
         reifyReturnShape(op);
     }
   }
 };
 } // end anonymous namespace
 
+namespace {
+struct TestDerivedAttributeDriver
+    : public PassWrapper<TestDerivedAttributeDriver, FunctionPass> {
+  void runOnFunction() override;
+};
+} // end anonymous namespace
+
+void TestDerivedAttributeDriver::runOnFunction() {
+  getFunction().walk([](DerivedAttributeOpInterface dOp) {
+    auto dAttr = dOp.materializeDerivedAttributes();
+    if (!dAttr)
+      return;
+    for (auto d : dAttr) {
+      dOp.emitRemark(llvm::formatv("{0} = {1}", d.first.strref(), d.second));
+    }
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // Legalization Driver.
 //===----------------------------------------------------------------------===//
 
 namespace {
 //===----------------------------------------------------------------------===//
 // Region-Block Rewrite Testing
 

 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerPatternsTestPass() {
   mlir::PassRegistration<TestReturnTypeDriver>("test-return-type",
                                                "Run return type functions");
 
+  mlir::PassRegistration<TestDerivedAttributeDriver>(
+      "test-derived-attr", "Run test derived attributes");
+
   mlir::PassRegistration<TestPatternDriver>("test-patterns",
                                             "Run test dialect patterns");
 
   mlir::PassRegistration<TestLegalizePatternDriver>(
       "test-legalize-patterns", "Run test dialect legalization patterns", [] {
         return std::make_unique<TestLegalizePatternDriver>(
             legalizerConversionMode);
       });
 
   PassRegistration<TestRemappedValue>(
       "test-remapped-value",
       "Test public remapped value mechanism in ConversionPatternRewriter");
 }
 } // namespace mlir

mlir/test/mlir-tblgen/op-attribute.td

 
 // DECL: class DerivedTypeAttrOp : public Op
 // DECL-SAME: DerivedAttributeOpInterface::Trait
 // DECL: static bool isDerivedAttribute
 // DEF: bool DerivedTypeAttrOp::isDerivedAttribute(StringRef name) {
 // DEF:   if (name == "element_dtype") return true;
 // DEF:   return false;
 // DEF: }
+// DEF: DerivedTypeAttrOp::materializeDerivedAttributes
 
 // Test that only default valued attributes at the end of the arguments
 // list get default values in the builder signature
 // ---
 
 def EOp : NS_Op<"e_op", []> {
   let arguments = (ins
     I32Attr:$i32_attr,

mlir/test/mlir-tblgen/op-derived-attribute.mlir

+// RUN: mlir-opt -test-derived-attr -verify-diagnostics %s | FileCheck %s --dump-input-on-failure
+
+// CHECK-LABEL: verifyDerivedAttributes
+func @verifyDerivedAttributes() {
+  // expected-remark @+2 {{element_dtype = f32}}
+  // expected-remark @+1 {{size = 320}}
+  %0 = "test.derived_type_attr"() : () -> tensor<10xf32>
+  // expected-remark @+2 {{element_dtype = i79}}
+  // expected-remark @+1 {{size = 948}}
+  %1 = "test.derived_type_attr"() : () -> tensor<12xi79>
+
+  return
+}

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

for (auto &namedAttr : op.getAttributes()) {
     if (attr.isDerivedAttr()) {
       emitDerivedAttr(name, attr);
     } else {
       emitAttrWithStorageType(name, attr);
       emitAttrWithReturnType(name, attr);
     }
   }
 
-  // Generate helper method to query whether a named attribute is a derived
-  // attribute. This enables, for example, avoiding adding an attribute that
-  // overlaps with a derived attribute.
-  auto derivedAttr = make_filter_range(op.getAttributes(),
+  auto derivedAttrs = make_filter_range(op.getAttributes(),
                                         [](const NamedAttribute &namedAttr) {
                                           return namedAttr.attr.isDerivedAttr();
                                         });
-  if (!derivedAttr.empty()) {
+  if (!derivedAttrs.empty()) {
     opClass.addTrait("DerivedAttributeOpInterface::Trait");
+    // Generate helper method to query whether a named attribute is a derived
+    // attribute. This enables, for example, avoiding adding an attribute that
+    // overlaps with a derived attribute.
+    {
       auto &method = opClass.newMethod("bool", "isDerivedAttribute",
                                        "StringRef name", OpMethod::MP_Static);
       auto &body = method.body();
-    for (auto namedAttr : derivedAttr)
-      body << "    if (name == \"" << namedAttr.name << "\") return true;\n";
+      for (auto namedAttr : derivedAttrs)
+        body << "  if (name == \"" << namedAttr.name << "\") return true;\n";
       body << " return false;";
     }
+    // Generate method to materialize derived attributes as a DictionaryAttr.
+    {
+      OpMethod &method =
+          opClass.newMethod("DictionaryAttr", "materializeDerivedAttributes");
+      auto &body = method.body();
+
+      auto nonMaterializable =
+          make_filter_range(derivedAttrs, [](const NamedAttribute &namedAttr) {
+            return namedAttr.attr.getConvertFromStorageCall().empty();
+          });
+      if (!nonMaterializable.empty()) {
+        std::string attrs;
+        llvm::raw_string_ostream os(attrs);
+        interleaveComma(nonMaterializable, os,
+                        [&](const NamedAttribute &attr) { os << attr.name; });
+        PrintWarning(
+            op.getLoc(),
+            formatv(
+                "op has non-materialzable derived attributes '{0}', skipping",
+                os.str()));
+        body << formatv("  emitOpError(\"op has non-materializable derived "
+                        "attributes '{0}'\");\n",
+                        attrs);
+        body << "  return nullptr;";
+        return;
+      }
+
+      body << "  MLIRContext* ctx = getContext();\n";
+      body << "  return DictionaryAttr::get({\n";
+      interleave(
+          derivedAttrs, body,
+          [&](const NamedAttribute &namedAttr) {
+            auto tmpl = namedAttr.attr.getConvertFromStorageCall();
+            body << "    {Identifier::get(\"" << namedAttr.name << "\", ctx),\n"
+                 << tgfmt(tmpl, &fctx.withSelf(namedAttr.name + "()")
+                                     .addSubst("_ctx", "ctx"))
+                 << "}";
+          },
+          ",\n");
+      body << "\n    }, ctx);";
+    }
+  }
 }
 
 void OpEmitter::genAttrSetters() {
   // Generate raw named setter type. This is a wrapper class that allows setting
   // to the attributes via setters instead of having to use the string interface
   // for better compile time verification.
   auto emitAttrWithStorageType = [&](StringRef name, Attribute attr) {
     auto &method = opClass.newMethod("void", (name + "Attr").str(),
for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
     body << "  " << builderOpState << ".addOperands(" << argName << ");\n";
   }
 
   // If the operation has the operand segment size attribute, add it here.
   if (op.getTrait("OpTrait::AttrSizedOperandSegments")) {
     body << "  " << builderOpState
          << ".addAttribute(\"operand_segment_sizes\", "
             "odsBuilder->getI32VectorAttr({";
-    llvm::interleaveComma(
-        llvm::seq<int>(0, op.getNumOperands()), body, [&](int i) {
-          if (op.getOperand(i).isOptional())
-            body << "(" << getArgumentName(op, i) << " ? 1 : 0)";
-          else if (op.getOperand(i).isVariadic())
-            body << "static_cast<int32_t>(" << getArgumentName(op, i)
-                 << ".size())";
-          else
-            body << "1";
-        });
+    interleaveComma(llvm::seq<int>(0, op.getNumOperands()), body, [&](int i) {
+        body << "static_cast<int32_t>(" << getArgumentName(op, i) << ".size())";
     body << "}));\n";
   }
 
   // Push all attributes to the result.
   for (const auto &namedAttr : op.getAttributes()) {
     auto &attr = namedAttr.attr;
     if (!attr.isDerivedAttr()) {
       bool emitNotNullCheck = attr.isOptional();
if (!opTrait || !opTrait->shouldDeclareMethods())
       continue;
     auto interface = opTrait->getOpInterface();
     for (auto method : interface.getMethods()) {
       // Don't declare if the method has a body or a default implementation.
       if (method.getBody() || method.getDefaultImplementation())
         continue;
       std::string args;
       llvm::raw_string_ostream os(args);
-      llvm::interleaveComma(method.getArguments(), os,
-                            [&](const OpInterfaceMethod::Argument &arg) {
-                              os << arg.type << " " << arg.name;
-                            });
+      interleaveComma(method.getArguments(), os,
       opClass.newMethod(method.getReturnType(), method.getName(), os.str(),
                         method.isStatic() ? OpMethod::MP_Static
                                           : OpMethod::MP_None,
                         /*declOnly=*/true);
     }
   }
 }
 
for (auto *def : defs) {
     }
   }
 }
 
 // Emits a comma-separated list of the ops.
 static void emitOpList(const std::vector<Record *> &defs, raw_ostream &os) {
   IfDefScope scope("GET_OP_LIST", os);
 
-  llvm::interleave(
+  interleave(
       // TODO: We are constructing the Operator wrapper instance just for
       // getting it's qualified class name here. Reduce the overhead by having a
       // lightweight version of Operator class just for that purpose.
       defs, [&os](Record *def) { os << Operator(def).getQualCppClassName(); },
       [&os]() { os << ",\n"; });
 }
 
 static bool emitOpDecls(const RecordKeeper &recordKeeper, raw_ostream &os) {