[mlir] improve and test TransformState::Extension

[ftynse]

Summary

Add the mechanism for TransformState extensions to update the mapping between
Transform IR values and Payload IR operations held by the state. The mechanism
is intentionally restrictive, similarly to how results of the transform op are
handled.

Introduce test ops that exercise a simple extension that maintains information
across the application of multiple transform ops.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

public:
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
   Operation *getTopLevel() const;
 
   /// Returns the list of ops that the given transform IR value corresponds to.
   /// This is helpful for transformations that apply to a particular handle.
   ArrayRef<Operation *> getPayloadOps(Value value) const;
 
+  /// Returns the Transform IR handle for the given Payload IR op if it exists
+  /// in the state, null otherwise.
+  Value getHandleForPayloadOp(Operation *op) const;
+
   /// Applies the transformation specified by the given transform op and updates
   /// the state accordingly.
   LogicalResult applyTransform(TransformOpInterface transform);
 
   /// Records the mapping between a block argument in the transform IR and a
   /// list of operations in the payload IR. The arguments must be defined in
   /// blocks of the currently processed transform IR region, typically after a
   /// region scope is defined.
#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
   protected:
     /// Constructs an extension of the given TransformState object.
     Extension(TransformState &state) : state(state) {}
 
     /// Provides read-only access to the parent TransformState object.
     const TransformState &getTransformState() const { return state; }
 
+    /// Replaces the given payload op with another op. If the replacement op is
+    /// null, removes the association of the payload op with its handle.
+    LogicalResult replacePayloadOp(Operation *op, Operation *replacement);
+
   private:
     /// Back-reference to the state that is being extended.
     TransformState &state;
   };
 
   /// Adds a new Extension of the type specified as template parameter,
   /// constructing it with the arguments provided. The extension is owned by the
   /// TransformState. It is expected that the state does not already have an
private:
 
   /// Forgets the payload IR ops associated with the given transform IR value.
   void removePayloadOps(Value value);
 
   /// Updates the payload IR ops associated with the given transform IR value.
   /// The callback function is called once per associated operation and is
   /// expected to return the modified operation or nullptr. In the latter case,
   /// the corresponding operation is no longer associated with the transform IR
-  /// value.
-  void updatePayloadOps(Value value,
-                        function_ref<Operation *(Operation *)> callback);
+  /// value. May fail if the operation produced by the update callback is
+  /// already associated with a different Transform IR handle value.
+  LogicalResult
+  updatePayloadOps(Value value,
 
+  /// Attempts to record the mapping between the given Payload IR operation and
+  /// the given Transform IR handle. Fails and reports an error if the operation
+  /// is already tracked by another handle.
+  static LogicalResult tryEmplaceReverseMapping(Mappings &map, Operation *op,
+                                                Value handle);
+
   /// The mappings between transform IR values and payload IR ops, aggregated by
   /// the region in which the transform IR values are defined.
   llvm::SmallDenseMap<Region *, Mappings> mappings;
 
   /// Extensions attached to the TransformState, identified by the TypeID of
   /// their type. Only one extension of any given type is allowed.
   DenseMap<TypeID, std::unique_ptr<Extension>> extensions;
 

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

 ArrayRef<Operation *>
 transform::TransformState::getPayloadOps(Value value) const {
   const TransformOpMapping &operationMapping = getMapping(value).direct;
   auto iter = operationMapping.find(value);
   assert(iter != operationMapping.end() && "unknown handle");
   return iter->getSecond();
 }
 
+Value transform::TransformState::getHandleForPayloadOp(Operation *op) const {
+  for (const Mappings &mapping : llvm::make_second_range(mappings)) {
+    if (Value handle = mapping.reverse.lookup(op))
+      return handle;
+  }
+  return Value();
+}
+
+LogicalResult transform::TransformState::tryEmplaceReverseMapping(
+    Mappings &map, Operation *operation, Value handle) {
+  auto insertionResult = map.reverse.insert({operation, handle});
+  if (!insertionResult.second) {
+    InFlightDiagnostic diag = operation->emitError()
+                              << "operation tracked by two handles";
+    diag.attachNote(handle.getLoc()) << "handle";
+    diag.attachNote(insertionResult.first->second.getLoc()) << "handle";
+    return diag;
+  }
+  return success();
+}
+
 LogicalResult
 transform::TransformState::setPayloadOps(Value value,
                                          ArrayRef<Operation *> targets) {
   assert(value != kTopLevelValue &&
          "attempting to reset the transformation root");
 
   if (value.use_empty())
     return success();
 
   // Setting new payload for the value without cleaning it first is a misuse of
   // the API, assert here.
   SmallVector<Operation *> storedTargets(targets.begin(), targets.end());
   Mappings &mappings = getMapping(value);
   bool inserted =
       mappings.direct.insert({value, std::move(storedTargets)}).second;
   assert(inserted && "value is already associated with another list");
   (void)inserted;
 
   // Having multiple handles to the same operation is an error in the transform
   // expressed using the dialect and may be constructed by valid API calls from
   // valid IR. Emit an error here.
   for (Operation *op : targets) {
-    auto insertionResult = mappings.reverse.insert({op, value});
-    if (!insertionResult.second) {
-      InFlightDiagnostic diag = op->emitError()
-                                << "operation tracked by two handles";
-      diag.attachNote(value.getLoc()) << "handle";
-      diag.attachNote(insertionResult.first->second.getLoc()) << "handle";
-      return diag;
-    }
+    if (failed(tryEmplaceReverseMapping(mappings, op, value)))
+      return failure();
   }
 
   return success();
 }
 
 void transform::TransformState::removePayloadOps(Value value) {
   Mappings &mappings = getMapping(value);
   for (Operation *op : mappings.direct[value])
     mappings.reverse.erase(op);
   mappings.direct.erase(value);
 }
 
-void transform::TransformState::updatePayloadOps(
+LogicalResult transform::TransformState::updatePayloadOps(
     Value value, function_ref<Operation *(Operation *)> callback) {
-  auto it = getMapping(value).direct.find(value);
-  assert(it != getMapping(value).direct.end() && "unknown handle");
+  Mappings &mappings = getMapping(value);
+  auto it = mappings.direct.find(value);
+  assert(it != mappings.direct.end() && "unknown handle");
   SmallVector<Operation *> &association = it->getSecond();
   SmallVector<Operation *> updated;
   updated.reserve(association.size());
 
-  for (Operation *op : association)
-    if (Operation *updatedOp = callback(op))
+  for (Operation *op : association) {
+    mappings.reverse.erase(op);
+    if (Operation *updatedOp = callback(op)) {
       updated.push_back(updatedOp);
+      if (failed(tryEmplaceReverseMapping(mappings, updatedOp, value)))
+        return failure();
+    }
+  }
 
   std::swap(association, updated);
+  return success();
 }
 
 LogicalResult
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   transform::TransformResults results(transform->getNumResults());
   if (failed(transform.apply(results, *this)))
     return failure();
 
assert(en.value().getDefiningOp() == transform.getOperation() &&
            "current transform op");
     if (failed(setPayloadOps(en.value(), results.get(en.index()))))
       return failure();
   }
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// TransformState::Extension
+//===----------------------------------------------------------------------===//
+
 transform::TransformState::Extension::~Extension() = default;
 
+LogicalResult
+transform::TransformState::Extension::replacePayloadOp(Operation *op,
+                                                       Operation *replacement) {
+  return state.updatePayloadOps(state.getHandleForPayloadOp(op),
+                                [&](Operation *current) {
+                                  return current == op ? replacement : current;
+                                });
+}
+
 //===----------------------------------------------------------------------===//
 // TransformResults
 //===----------------------------------------------------------------------===//
 
 transform::TransformResults::TransformResults(unsigned numSegments) {
   segments.resize(numSegments,
                   ArrayRef<Operation *>(nullptr, static_cast<size_t>(0)));
 }

mlir/test/Dialect/Transform/transform-state-extension.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter -verify-diagnostics -split-input-file
+
+// expected-note @below {{associated payload op}}
+module {
+  transform.sequence {
+  ^bb0(%arg0: !pdl.operation):
+    // expected-remark @below {{extension absent}}
+    test_check_if_test_extension_present %arg0
+    test_add_test_extension "A"
+    // expected-remark @below {{extension present, A}}
+    test_check_if_test_extension_present %arg0
+    test_remove_test_extension
+    // expected-remark @below {{extension absent}}
+    test_check_if_test_extension_present %arg0
+  }
+}
+
+// -----
+
+// expected-note @below {{associated payload op}}
+module {
+  transform.sequence {
+  ^bb0(%arg0: !pdl.operation):
+    test_add_test_extension "A"
+    test_remove_test_extension
+    test_add_test_extension "B"
+    // expected-remark @below {{extension present, B}}
+    test_check_if_test_extension_present %arg0
+  }
+}
+
+// -----
+
+// expected-note @below {{associated payload op}}
+module {
+  transform.sequence {
+  ^bb0(%arg0: !pdl.operation):
+    test_add_test_extension "A"
+    // expected-remark @below {{extension present, A}}
+    test_check_if_test_extension_present %arg0
+    // expected-note @below {{associated payload op}}
+    test_remap_operand_to_self %arg0
+    // expected-remark @below {{extension present, A}}
+    test_check_if_test_extension_present %arg0
+  }
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

 //===- TestTransformDialectExtension.cpp ----------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines an extension of the MLIR Transform dialect for testing
 // purposes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestTransformDialectExtension.h"
+#include "TestTransformStateExtension.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
-#include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple transform op defined outside of the dialect. Just emits a remark when
 /// applied. This op is defined in C++ to test that C++ definitions also work
 /// for op injection into the Transform dialect.
LogicalResult mlir::test::TestPrintRemarkAtOperandOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   ArrayRef<Operation *> payload = state.getPayloadOps(getOperand());
   for (Operation *op : payload)
     op->emitRemark() << getMessage();
 
   return success();
 }
 
+LogicalResult
+mlir::test::TestAddTestExtensionOp::apply(transform::TransformResults &results,
+                                          transform::TransformState &state) {
+  state.addExtension<TestTransformStateExtension>(getMessageAttr());
+  return success();
+}
+
+LogicalResult mlir::test::TestCheckIfTestExtensionPresentOp::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  auto *extension = state.getExtension<TestTransformStateExtension>();
+  if (!extension) {
+    emitRemark() << "extension absent";
+    return success();
+  }
+
+  InFlightDiagnostic diag = emitRemark()
+                            << "extension present, " << extension->getMessage();
+  for (Operation *payload : state.getPayloadOps(getOperand())) {
+    diag.attachNote(payload->getLoc()) << "associated payload op";
+    assert(state.getHandleForPayloadOp(payload) == getOperand() &&
+           "inconsistent mapping between transform IR handles and payload IR "
+           "operations");
+  }
+
+  return success();
+}
+
+LogicalResult mlir::test::TestRemapOperandPayloadToSelfOp::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  auto *extension = state.getExtension<TestTransformStateExtension>();
+  if (!extension)
+    return emitError() << "TestTransformStateExtension missing";
+
+  return extension->updateMapping(state.getPayloadOps(getOperand()).front(),
+                                  getOperation());
+}
+
+LogicalResult mlir::test::TestRemoveTestExtensionOp::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  state.removeExtension<TestTransformStateExtension>();
+  return success();
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

def TestPrintRemarkAtOperandOp
   let arguments = (ins 
     Arg<PDL_Operation, "",
         [TransformMappingRead, PayloadIRRead]>:$operand,
     StrAttr:$message);
   let assemblyFormat = "$operand `,` $message attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
+def TestAddTestExtensionOp
+  : Op<Transform_Dialect, "test_add_test_extension",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        NoSideEffect]> {
+  let arguments = (ins StrAttr:$message);
+  let assemblyFormat = "$message attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
+def TestCheckIfTestExtensionPresentOp
+  : Op<Transform_Dialect, "test_check_if_test_extension_present",
+       [DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins
+    Arg<PDL_Operation, "", [TransformMappingRead, PayloadIRRead]>:$operand);
+  let assemblyFormat = "$operand attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
+def TestRemapOperandPayloadToSelfOp
+  : Op<Transform_Dialect, "test_remap_operand_to_self",
+       [DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins
+    Arg<PDL_Operation, "",
+        [TransformMappingRead, TransformMappingWrite, PayloadIRRead]>:$operand);
+  let assemblyFormat = "$operand attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
+def TestRemoveTestExtensionOp
+  : Op<Transform_Dialect, "test_remove_test_extension",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        NoSideEffect]> {
+  let assemblyFormat = "attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD

mlir/test/lib/Dialect/Transform/TestTransformStateExtension.h

+//===- TestTransformStateExtension.h - Test Utility -------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines an TransformState extension for the purpose of testing the
+// relevant APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H
+#define MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace test {
+class TestTransformStateExtension
+    : public transform::TransformState::Extension {
+public:
+  TestTransformStateExtension(transform::TransformState &state,
+                              StringAttr message)
+      : Extension(state), message(message) {}
+
+  StringRef getMessage() const { return message.getValue(); }
+
+  LogicalResult updateMapping(Operation *previous, Operation *updated) {
+    return replacePayloadOp(previous, updated);
+  }
+
+private:
+  StringAttr message;
+};
+} // namespace test
+} // namespace mlir
+
+#endif // MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H