[mlir][transform] Add op for adding attributes to payload IR

[qedawkins]

Summary

The ability to add attributes to payload IR is useful functionality
independent of any dialect. This is added here through transform.annotate
by enabling attributes tied to a TransformParamTypeInterface (which
internally refers to an Attribute) to be added to a target operation by
name.

The AnnotateOp does not produce a new handle as no existing handles
should be affected by adding an attribute. Existing attributes on
the payload with the same name will be overwritten.

[nicolasvasilache]

nice

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

def AlternativesOp : TransformDialectOp<"alternatives",
   let regions = (region VariadicRegion<SizedRegion<1>>:$alternatives);
 
   let assemblyFormat =
     "($scope^ `:` type($scope))? (`->` type($results)^)? "
     "attr-dict-with-keyword regions";
   let hasVerifier = 1;
 }
 
+def AnnotateOp : TransformDialectOp<"annotate",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Annotates the target operation with an attribute by name";
+  let description = [{
+    Adds an attribute with the given `name` to the `target` operation. An optional
+    `param` handle can be provided to give the attribute a specific value, else a
+    UnitAttr is added. Attributes can either be added with a single attribute in the
+    `param` payload broadcasted to all target operations, or the attributes will be
+    mapped 1:1 based on the order within the handles.
+
+    Fails silently if the length of the parameter payload does not match the length of
+    the parameter handles. Does not consume the provided handles.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       StrAttr:$name,
+                       Optional<TransformParamTypeInterface>:$param);
+  let results = (outs);
+
+  let assemblyFormat =
+    "$target $name attr-dict (`=` $param^)?"
+    "`:` type($target) (`,` type($param)^)?";
+}
+
 def CastOp : TransformDialectOp<"cast",
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<CastOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let arguments = (ins TransformHandleTypeInterface:$input);
   let results = (outs TransformHandleTypeInterface:$output);
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

if (terminator->getOperands().getTypes() != getResults().getTypes()) {
       return diag;
     }
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// AnnotateOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::AnnotateOp::apply(transform::TransformResults &results,
+                             transform::TransformState &state) {
+  SmallVector<Operation *> targets =
+      llvm::to_vector(state.getPayloadOps(getTarget()));
+
+  Attribute attr = UnitAttr::get(getContext());
+  if (auto paramH = getParam()) {
+    ArrayRef<Attribute> params = state.getParams(paramH);
+    if (params.size() != 1) {
+      if (targets.size() != params.size()) {
+        return emitSilenceableError()
+               << "parameter and target have different payload lengths ("
+               << params.size() << " vs " << targets.size() << ")";
+      }
+      for (auto &&[target, attr] : llvm::zip_equal(targets, params))
+        target->setAttr(getName(), attr);
+      return DiagnosedSilenceableFailure::success();
+    }
+    attr = params[0];
+  }
+  for (auto target : targets)
+    target->setAttr(getName(), attr);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::AnnotateOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getTarget(), effects);
+  onlyReadsHandle(getParam(), effects);
+  modifiesPayload(effects);
+}
+
+//===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
                               transform::TransformState &state) {
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();

mlir/test/Dialect/Transform/test-interpreter.mlir

 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["test.replace_me"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-remark @below {{2 iterations}}
   transform.test_tracked_rewrite %0 : (!transform.any_op) -> ()
   // One replacement op (test.drop_mapping) is dropped from the mapping.
   // expected-remark @below {{2}}
   test_print_number_of_associated_payload_ir_ops %0 : !transform.any_op
 }
+
+
+// -----
+
+// CHECK-LABEL: func @test_annotation()
+//  CHECK-NEXT:   "test.annotate_me"()
+//  CHECK-SAME:                        broadcast_attr = 2 : i64
+//  CHECK-SAME:                        new_attr = 1 : i32
+//  CHECK-SAME:                        unit_attr
+//  CHECK-NEXT:   "test.annotate_me"()
+//  CHECK-SAME:                        broadcast_attr = 2 : i64
+//  CHECK-SAME:                        existing_attr = "test"
+//  CHECK-SAME:                        new_attr = 1 : i32
+//  CHECK-SAME:                        unit_attr
+//  CHECK-NEXT:   "test.annotate_me"()
+//  CHECK-SAME:                        broadcast_attr = 2 : i64
+//  CHECK-SAME:                        new_attr = 1 : i32
+//  CHECK-SAME:                        unit_attr
+func.func @test_annotation() {
+  %0 = "test.annotate_me"() : () -> (i1)
+  %1 = "test.annotate_me"() {existing_attr = "test"} : () -> (i1)
+  %2 = "test.annotate_me"() {new_attr = 0} : () -> (i1)
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %0 = transform.structured.match ops{["test.annotate_me"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.test_produce_param_with_number_of_test_ops %0 : !transform.any_op
+  transform.annotate %0 "new_attr" = %1 : !transform.any_op, !transform.test_dialect_param
+
+  %2 = transform.param.constant 2 -> !transform.param<i64>
+  transform.annotate %0 "broadcast_attr" = %2 : !transform.any_op, !transform.param<i64>
+  transform.annotate %0 "unit_attr" : !transform.any_op
+}