Add a new Dialect interface to control implicit attribute propagation.

[mehdi_amini]

Summary

This interface allows for dialect attributes to be implicitly propagated through
rewrites, by augmenting the OpBuilder with a DictionnaryAttr member. It is
implicitly initialized to the current op attributes during pattern rewriting or
when the builder is created from an operation.

[bondhugula]

mlir/include/mlir/IR/AttributePropagationInterface.h

23

Typo: implicitly

31

the attr -> `attr`

39

Typo: Dictionary

mlir/include/mlir/IR/Builders.h

202

Needs a comment.

481

Typo: spell fix.

mlir/lib/IR/AttributePropagationInterface.cpp

6

Missing file doc comment - here and for the header file.

18

Can hoist this while getting it from op:

MLIRContext *ctx = op->getContext()

20

Comment somewhere for this block.

mlir/lib/IR/Builders.cpp

400

Needs a comment.

mlir/test/Transforms/test-canonicalize.mlir

55

Nit: Drop extra new line.

[bondhugula]

mlir/lib/IR/Builders.cpp

397

This doc comment will need an update. It's not just creating an operation from the fields in state.

[bondhugula]

mlir/include/mlir/IR/AttributePropagationInterface.h

32

-> propagate ... to op.

mlir/include/mlir/IR/Builders.h

199

This doc comment needs an update as well.

499

Rephrase for clarity. 'considered for propagation' based on what? Mention about the interface.

[rriddle]

mlir/include/mlir/IR/AttributePropagationInterface.h

11

Some of these are unnecessary.

42

nit: propagateDialectAttributes?

mlir/include/mlir/IR/Builders.h

202

Can you just call setPropagatedAttributes here?

mlir/lib/IR/AttributePropagationInterface.cpp

22

I'm afraid that this is going to be really really slow in-practice. Is this an indicator that it's about time we make "Dialect-Prefixed Attributes" more "first-class" in the API? Right now we use Identifier for attribute names but we could beef that up to be something more stateful. We already try to query state on these dialect attributes, and already feel the pain of the clunky string slicing+dialect lookup+interface lookup.

[stephenneuendorffer]

This seems to me like a 'tip of the iceberg' kind of solution. What happens if the attribute needs to be transformed? What happens if there is not a 1:1 correspondance between input operations and output operations in a transformation? etc. etc. I'm not sure what the right solution is, but this seems to be a narrowly scoped infrastructure solution to a tensorflow problem. I'd be much happier if this solved other obvious problems.

mlir/include/mlir/IR/Builders.h

485

"dictionary'

[bondhugula]

mlir/include/mlir/IR/Builders.h

202

All created ops would get the attributes of an op that was used for the insert point: This looks arbitrary outside of all pattern rewriter use cases! (where there is no replacement but say just IR injection happening) It looks like you'll need a new OpBuilder ctor or explicitly using setPropagatedAttributes on the builder when needed?

500

camelBack.

mlir/lib/IR/AttributePropagationInterface.cpp

17–28

This overhead on the creation of every op is a concern. Even if there are no dialect name prefixed attributes, you are cycling through all of the "insert point op's" attributes (op's intrinsic attributes as well as any other attributes) in all cases where the OpBuilder was created this way. It is also weird that there would be a workaround to avoid this overhead by creating the OpBuilder via its "insert point" ctor (propagated_attrs would be empty in that case) or some of the other ctors?

mlir/test/Transforms/test-canonicalize.mlir

59

Missing : after CHECK-NOT

62

Testing missing an attribute without any dialect prefix.

[mehdi_amini]

In D95394#2521877, @stephenneuendorffer wrote:

This seems to me like a 'tip of the iceberg' kind of solution. What happens if the attribute needs to be transformed? What happens if there is not a 1:1 correspondance between input operations and output operations in a transformation?

There are two categories of attributes: dependent and dialect.

It seems to me that your comments are mostly applicable to "dependent attributes", i.e. attributes that are associated to an operation. This is different from the dialect attributes which intended to added as annotation on existing operations: these come with their own verifier for example, while dependent attributes are supposed to be verified by the operation verifier.

(I think it is best discussed in the RFC thread though, and leave implementation details comment to the revision here)

[bondhugula]

In D95394#2521877, @stephenneuendorffer wrote:

This seems to me like a 'tip of the iceberg' kind of solution. What happens if the attribute needs to be transformed? What happens if there is not a 1:1 correspondance between input operations and output operations in a transformation? etc. etc. I'm not sure what the right solution is, but this seems to be a narrowly scoped infrastructure solution to a tensorflow problem. I'd be much happier if this solved other obvious problems.

I don't see this as a solution to a narrow or a tensorflow problem. In fact, this is a pretty generic solution to propagate attributes that are currently being lost -- and the situations this approach addresses is that of the attributes logically associated to a specific dialect. This would also encourage folks to prefix free attribute names with dialect name prefixes for preservation -- but that may not always be meaningful.

[bondhugula]

@mehdi_amini Anything further on this?

[mehdi_amini]

Not just now: we've been looking into alternatives, it takes time to pan out because it could be much more intrusive. I need to get back to this...

mlir/include/mlir/IR/AttributePropagationInterface.h

+//===- AttributePropagationInterface.h - Propagation Interface --*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_IR_ATTRIBUTEPROPAGATIONINTERFACE_H_
+#define MLIR_IR_ATTRIBUTEPROPAGATIONINTERFACE_H_
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/Identifier.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace mlir {
+class Attribute;
+class OpFoldResult;
+class Region;
+
+/// The AttributePropagationInterface allows dialects to provide a hook that is
+/// invoked after operation creation to implicity add/propagate dialect
+/// attributes (i.e. a NamedAttribute for which the key is prefixed with the
+/// dialect name).
+class AttributePropagationInterface
+    : public DialectInterface::Base<AttributePropagationInterface> {
+public:
+  AttributePropagationInterface(Dialect *dialect) : Base(dialect) {}
+
+  /// Invoked to allow the dialect to inspect the attr and the op and decide to
+  /// propagate or not this attribute on the operation. It is invoked by the
+  /// builder on op creation during rewrite.
+  virtual void propagateAttribute(NamedAttribute attr, Operation *op) const {
+    return;
+  }
+};
+
+/// Propagate the attributes from the provided Dictionnary attribute on an
+/// operation by calling back into the dialect interface for each
+/// NamedAttribute.
+void propagateAttributes(DictionaryAttr propagatedAttr, Operation *op);
+
+} // end namespace mlir
+
+#endif // MLIR_IR_ATTRIBUTEPROPAGATIONINTERFACE_H_

mlir/include/mlir/IR/Builders.h

explicit OpBuilder(Region *region, Listener *listener = nullptr)
       : OpBuilder(region->getContext(), listener) {
     if (!region->empty())
       setInsertionPoint(&region->front(), region->front().begin());
   }
   explicit OpBuilder(Region &region, Listener *listener = nullptr)
       : OpBuilder(&region, listener) {}
 
   /// Create a builder and set insertion point to the given operation, which
   /// will cause subsequent insertions to go right before it.
   explicit OpBuilder(Operation *op, Listener *listener = nullptr)
       : OpBuilder(op->getContext(), listener) {
+    propagated_attrs = op->getAttrDictionary();
     setInsertionPoint(op);
   }
 
   OpBuilder(Block *block, Block::iterator insertPoint,
             Listener *listener = nullptr)
       : OpBuilder(block->getParent()->getContext(), listener) {
     setInsertionPoint(block, insertPoint);
   }
public:
   }
   Operation *cloneWithoutRegions(Operation &op) {
     return insert(op.cloneWithoutRegions());
   }
   template <typename OpT> OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
+  /// Return the current dictionnary of attributes considered for propagation
+  /// on operations created in this builder.
+  DictionaryAttr getPropagatedAttributes() { return propagated_attrs; }
+
+  /// Set the dictionnary of attributes that will be considered for propagation
+  /// on operations created in this builder.
+  void setPropagatedAttributes(DictionaryAttr new_propagated_attrs) {
+    propagated_attrs = new_propagated_attrs;
+  }
+
 private:
   /// The current block this builder is inserting into.
   Block *block = nullptr;
   /// The insertion point within the block that this builder is inserting
   /// before.
   Block::iterator insertPoint;
   /// The optional listener for events of this builder.
   Listener *listener;
+  /// Attributes considered for propagation on every created operation.
+  DictionaryAttr propagated_attrs;
 };
 
 } // namespace mlir
 
 #endif

mlir/lib/IR/AttributePropagationInterface.cpp

+//===- AttributePropagationInterface.cpp - Attribute Propagation --*- C++ -*-=//
+//
+// 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 "mlir/IR/AttributePropagationInterface.h"
+#include "mlir/IR/Dialect.h"
+
+using namespace mlir;
+
+void mlir::propagateAttributes(DictionaryAttr propagatedAttr, Operation *op) {
+  if (!propagatedAttr)
+    return;
+  for (NamedAttribute attr : propagatedAttr) {
+    MLIRContext *ctx = attr.second.getContext();
+    auto dialectNamePair = attr.first.strref().split('.');
+    if (dialectNamePair.second.empty())
+      continue;
+    Dialect *dialect = ctx->getLoadedDialect(dialectNamePair.first);
+    if (!dialect)
+      continue;
+    if (auto *iface =
+            dialect->getRegisteredInterface<AttributePropagationInterface>())
+      iface->propagateAttribute(attr, op);
+  }
+}

mlir/lib/IR/Builders.cpp

 //===- Builders.cpp - Helpers for constructing MLIR Classes ---------------===//
 //
 // 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 "mlir/IR/Builders.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/AttributePropagationInterface.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/Support/raw_ostream.h"
 

 /// Add new block with 'argTypes' arguments and set the insertion point to the
 /// end of it.  The block is placed before 'insertBefore'.
 Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes) {
   assert(insertBefore && "expected valid insertion block");
   return createBlock(insertBefore->getParent(), Region::iterator(insertBefore),
                      argTypes);
 }
 
 /// Create an operation given the fields represented as an OperationState.
 Operation *OpBuilder::createOperation(const OperationState &state) {
-  return insert(Operation::create(state));
+  Operation *newOp = Operation::create(state);
+  propagateAttributes(getPropagatedAttributes(), newOp);
+  return insert(newOp);
 }
 
 /// Attempts to fold the given operation and places new results within
 /// 'results'. Returns success if the operation was folded, failure otherwise.
 /// Note: This function does not erase the operation on a successful fold.
 LogicalResult OpBuilder::tryFold(Operation *op,
                                  SmallVectorImpl<Value> &results) {
   results.reserve(op->getNumResults());

mlir/lib/Rewrite/PatternApplicator.cpp

while (true) {
     // Check that the pattern can be applied.
     if (canApply && !canApply(*bestPattern))
       continue;
 
     // Try to match and rewrite this pattern. The patterns are sorted by
     // benefit, so if we match we can immediately rewrite. For PDL patterns, the
     // match has already been performed, we just need to rewrite.
     rewriter.setInsertionPoint(op);
+    rewriter.setPropagatedAttributes(op->getAttrDictionary());
     LogicalResult result = success();
     if (pdlMatch) {
       bytecode->rewrite(rewriter, *pdlMatch, *mutableByteCodeState);
     } else {
       result = static_cast<const RewritePattern *>(bestPattern)
                    ->matchAndRewrite(op, rewriter);
     }
     if (succeeded(result) && (!onSuccess || succeeded(onSuccess(*bestPattern))))
       return success();
 
     // Perform any necessary cleanups.
     if (onFailure)
       onFailure(*bestPattern);
   }
   return failure();
 }

mlir/test/Transforms/test-canonicalize.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(canonicalize)' | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline='func(canonicalize)' | FileCheck %s
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_pattern
 func @remove_op_with_inner_ops_pattern() {
   // CHECK-NEXT: return
   "test.op_with_region_pattern"() ({
-    "foo.op_with_region_terminator"() : () -> ()
+    "test.op_with_region_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_fold_no_side_effect
 func @remove_op_with_inner_ops_fold_no_side_effect() {
   // CHECK-NEXT: return
   "test.op_with_region_fold_no_side_effect"() ({
-    "foo.op_with_region_terminator"() : () -> ()
+    "test.op_with_region_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_fold
 // CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32)
 func @remove_op_with_inner_ops_fold(%arg0 : i32) -> (i32) {
   // CHECK-NEXT: return %[[ARG_0]]
   %0 = "test.op_with_region_fold"(%arg0) ({
-    "foo.op_with_region_terminator"() : () -> ()
+    "test.op_with_region_terminator"() : () -> ()
   }) : (i32) -> (i32)
   return %0 : i32
 }
 
 // CHECK-LABEL: func @remove_op_with_variadic_results_and_folder
 // CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
 func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
   // CHECK-NEXT: return %[[ARG_0]], %[[ARG_1]]
func @test_commutative_multi(%arg0: i32, %arg1: i32) -> (i32, i32) {
   // CHECK-NEXT: %[[O0:.*]] = "test.op_commutative"(%[[ARG_0]], %[[ARG_1]], %[[C42]], %[[C43]]) : (i32, i32, i32, i32) -> i32
   %y = "test.op_commutative"(%c42_i32, %arg0, %arg1, %c43_i32) : (i32, i32, i32, i32) -> i32
 
   // CHECK-NEXT: %[[O1:.*]] = "test.op_commutative"(%[[ARG_0]], %[[ARG_1]], %[[C42]], %[[C43]]) : (i32, i32, i32, i32) -> i32
   %z = "test.op_commutative"(%arg0, %c42_i32, %c43_i32, %arg1): (i32, i32, i32, i32) -> i32
   // CHECK-NEXT: return %[[O0]], %[[O1]]
   return %y, %z: i32, i32
 }
+
+
+// CHECK-LABEL: func @propagate_attribute
+func private @propagate_attribute_callee()
+func @propagate_attribute() {
+  // CHECK-NOT test.foo = "bar"
+  // CHECK: call @propagate_attribute_callee()
+  // CHECK-SAME test.propagate = 42 : i32
+  "test.fold_to_call_op"() {test.foo = "bar", test.propagate = 42 : i32, callee=@propagate_attribute_callee} : () -> ()
+  return
+}

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

 //===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//
 //
 // 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 "TestTypes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/AttributePropagationInterface.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSwitch.h"
 
struct TestDialectFoldInterface : public DialectFoldInterface {
   /// operation that is *not* isolated from above, should be used when
   /// materializing constants.
   bool shouldMaterializeInto(Region *region) const final {
     // If this is a one region operation, then insert into it.
     return isa<OneRegionOp>(region->getParentOp());
   }
 };
 
+struct TestDialectFoldAttributePropagationInterface
+    : public AttributePropagationInterface {
+  using AttributePropagationInterface::AttributePropagationInterface;
+  void propagateAttribute(NamedAttribute attr, Operation *op) const final {
+    if (attr.first == "test.propagate")
+      op->setAttr(attr.first, attr.second);
+  }
+};
+
 /// This class defines the interface for handling inlining with standard
 /// operations.
 struct TestInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//

 // TestDialect
 //===----------------------------------------------------------------------===//
 
 void TestDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
-  addInterfaces<TestOpAsmInterface, TestDialectFoldInterface,
+  addInterfaces<TestDialectFoldAttributePropagationInterface,
+                TestOpAsmInterface, TestDialectFoldInterface,
                 TestInlinerInterface>();
   addTypes<TestType, TestRecursiveType,
 #define GET_TYPEDEF_LIST
 #include "TestTypeDefs.cpp.inc"
            >();
   allowUnknownOperations();
 }