[mlir][transform] TransformStateExtension: Replace op/value handles separately

[nicolasvasilache]

Looks good but is this really NFC ?

I.e. replacePayloadOp used to implicitly also replacePayloadValues and now we fail hard.

Is this a standalone applicable first step or does it need to land at the same time as dependent changes to avoid breaking downstream?

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

25

nit: less -> fewer

[springerm]

In D147038#4227075, @nicolasvasilache wrote:

Looks good but is this really NFC ?

I.e. replacePayloadOp used to implicitly also replacePayloadValues and now we fail hard.

Is this a standalone applicable first step or does it need to land at the same time as dependent changes to avoid breaking downstream?

Removed the NFC. This requires downstream changes if replacePayloadOp is used. In the case of IREE, it should be integrated together with D147039, then the integration fix is simply deleting the IREE TrackingListener.

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

protected:
     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. Returns
     /// failure if the op is not associated with any handle.
+    ///
+    /// Note: This function does not update value handles. None of the original
+    /// op's results may be mapped to any value handle.
     LogicalResult replacePayloadOp(Operation *op, Operation *replacement);
 
+    /// Replaces the given payload value with another value. If the replacement
+    /// value is null, removes the association of the payload value with its
+    /// handle. Returns failure if the value is not associated with any handle.
+    LogicalResult replacePayloadValue(Value value, Value 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,
   /// as well as any association between value handles and the results of said
   /// payload IR op.
   void forgetMapping(Value opHandle, ValueRange origOpFlatResults);
 
   void forgetValueMapping(Value valueHandle,
                           ArrayRef<Operation *> payloadOperations);
 
-  /// 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. Value handles associated with the results of the operation are
-  /// also updated to be associated with the results of the new operation. For
-  /// this reason, the new operation must have the same number of results.
+  /// Replaces the given payload op with another op. If the replacement op is
+  /// null, removes the association of the payload op with its handle.
   ///
-  /// Returns failure if the payload does not satisfy the conditions associated
-  /// with the type of the handle value.
+  /// Note: This function does not update value handles. None of the original
+  /// op's results may be mapped to any value handle.
   LogicalResult replacePayloadOp(Operation *op, Operation *replacement);
 
+  /// Replaces the given payload value with another value. If the replacement
+  /// value is null, removes the association of the payload value with its
+  /// handle.
+  LogicalResult replacePayloadValue(Value value, Value replacement);
+
   /// If the operand is a handle consumed by the operation, i.e. has the "free"
   /// memory effect associated with it, identifies other handles that are
   /// pointing to payload IR operations nested in the operations pointed to by
   /// the consumed handle. Marks all such handles as invalidated to trigger
   /// errors if they are used. If `throughValue` is passed, record the fact that
   /// an op handle was invalidated because a value handle associated with
   /// results of the payload op or its block arguments was invalidated.
   void recordOpHandleInvalidation(OpOperand &consumingHandle,

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

transform::TransformState::replacePayloadOp(Operation *op,
   // care if there are on such handles.
   SmallVector<Value> opHandles;
   (void)getHandlesForPayloadOp(op, opHandles);
   for (Value handle : opHandles) {
     Mappings &mappings = getMapping(handle);
     dropMappingEntry(mappings.reverse, op, handle);
   }
 
-  // Drop the mapping between the op results and all value handles that point to
-  // them. Don't care if there are no such handles.
-  RaggedArray<Value> resultValueHandles;
+#ifndef NDEBUG
   for (Value opResult : op->getResults()) {
     SmallVector<Value> valueHandles;
     (void)getHandlesForPayloadValue(opResult, valueHandles);
-    for (Value handle : valueHandles) {
-      Mappings &localMappings = getMapping(handle);
-      dropMappingEntry(localMappings.reverseValues, opResult, handle);
-    }
-    resultValueHandles.push_back(std::move(valueHandles));
+    assert(valueHandles.empty() && "expected no mapping to old results");
   }
+#endif // NDEBUG
 
   // TODO: consider invalidating the handles to nested objects here.
 
   // If replacing with null, that is erasing the mapping, drop the mapping
   // between the handles and the IR objects and return.
   if (!replacement) {
     for (Value handle : opHandles) {
       Mappings &mappings = getMapping(handle);
       dropMappingEntry(mappings.direct, handle, op);
     }
-    for (Value opResult : op->getResults()) {
-      SmallVector<Value> valueHandles;
-      (void)getHandlesForPayloadValue(opResult, valueHandles);
-      for (Value handle : valueHandles) {
-        Mappings &localMappings = getMapping(handle);
-        dropMappingEntry(localMappings.values, handle, opResult);
-      }
-    }
     return success();
   }
 
   // Otherwise, replace the pointed-to object of all handles while preserving
   // their relative order. First, replace the mapped operation if present.
   for (Value handle : opHandles) {
     Mappings &mappings = getMapping(handle);
     auto it = mappings.direct.find(handle);
     if (it == mappings.direct.end())
       continue;
 
     SmallVector<Operation *, 2> &association = it->getSecond();
     // Note that an operation may be associated with the handle more than once.
     for (Operation *&mapped : association) {
       if (mapped == op)
         mapped = replacement;
     }
     mappings.reverse[replacement].push_back(handle);
   }
 
-  // Second, replace the mapped results of the operation.
-  for (auto [origResult, handleList] :
-       llvm::zip(op->getResults(), resultValueHandles)) {
-    // No handles to the value, skip even if there is no replacement.
-    if (handleList.empty())
-      continue;
-
-    unsigned resultNumber = origResult.getResultNumber();
-    if (resultNumber >= replacement->getNumResults()) {
-      return emitError(op->getLoc())
-             << "cannot replace an op with another op producing less results "
-                "while tracking handles";
-    }
+  return success();
 
-    Value replacementResult = replacement->getResult(resultNumber);
-    for (Value resultHandle : handleList) {
-      Mappings &mappings = getMapping(resultHandle);
-      auto it = mappings.values.find(resultHandle);
+LogicalResult
+transform::TransformState::replacePayloadValue(Value value, Value replacement) {
+  SmallVector<Value> valueHandles;
+  (void)getHandlesForPayloadValue(value, valueHandles);
+
+  for (Value handle : valueHandles) {
+    Mappings &mappings = getMapping(handle);
+    dropMappingEntry(mappings.reverseValues, value, handle);
+
+    // If replacing with null, that is erasing the mapping, drop the mapping
+    // between the handles and the IR objects
+    if (!replacement) {
+      dropMappingEntry(mappings.values, handle, value);
+    } else {
+      auto it = mappings.values.find(handle);
       if (it == mappings.values.end())
         continue;
 
       SmallVector<Value> &association = it->getSecond();
       for (Value &mapped : association) {
-        if (mapped == origResult)
-          mapped = replacementResult;
+        if (mapped == value)
+          mapped = replacement;
       }
-      mappings.reverseValues[replacementResult].push_back(resultHandle);
+      mappings.reverseValues[replacement].push_back(handle);
     }
   }
 
   return success();
 }
 
 void transform::TransformState::recordOpHandleInvalidationOne(
     OpOperand &consumingHandle, ArrayRef<Operation *> potentialAncestors,
transform::TransformState::Extension::replacePayloadOp(Operation *op,
   if (failed(state.getHandlesForPayloadOp(op, handles)))
     return failure();
 
   // TODO: we may need to invalidate handles to operations and values nested in
   // the operation being replaced.
   return state.replacePayloadOp(op, replacement);
 }
 
+LogicalResult
+transform::TransformState::Extension::replacePayloadValue(Value value,
+                                                          Value replacement) {
+  SmallVector<Value> handles;
+  if (failed(state.getHandlesForPayloadValue(value, handles)))
+    return failure();
+
+  return state.replacePayloadValue(value, replacement);
+}
+
 //===----------------------------------------------------------------------===//
 // TransformResults
 //===----------------------------------------------------------------------===//
 
 transform::TransformResults::TransformResults(unsigned numSegments) {
   operations.appendEmptyRows(numSegments);
   params.appendEmptyRows(numSegments);
   values.appendEmptyRows(numSegments);

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

module {
   }
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   test_add_test_extension "A"
-   // This is okay because we are replacing the top-level module opeation
+   // This is okay because we are replacing the top-level module operation
    // (0 results) with this operation that has _more_ (1) results.
   %dummy = test_remap_operand_to_self %arg0 : !pdl.operation
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):

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

 set(LLVM_TARGET_DEFINITIONS TestTransformDialectExtension.td)
 mlir_tablegen(TestTransformDialectExtension.h.inc -gen-op-decls)
 mlir_tablegen(TestTransformDialectExtension.cpp.inc -gen-op-defs)
 mlir_tablegen(TestTransformDialectExtensionTypes.h.inc -gen-typedef-decls -typedefs-dialect=transform)
 mlir_tablegen(TestTransformDialectExtensionTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=transform)
 add_public_tablegen_target(MLIRTestTransformDialectExtensionIncGen)
 
 add_mlir_library(MLIRTestTransformDialect
   TestTransformDialectExtension.cpp
   TestTransformDialectInterpreter.cpp
+  TestTransformStateExtension.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   DEPENDS
   MLIRTestTransformDialectExtensionIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRPDLDialect
   MLIRTransformDialect
   MLIRTransformDialectTransforms
 )

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

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);
-  }
+  LogicalResult updateMapping(Operation *previous, Operation *updated);
 
 private:
   StringAttr message;
 };
 } // namespace test
 } // namespace mlir
 
 #endif // MLIR_TEST_LIB_DIALECT_TRANSFORM_TESTTRANSFORMSTATEEXTENSION_H

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

+//===- TestTransformStateExtension.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestTransformStateExtension.h"
+
+using namespace mlir;
+
+LogicalResult
+test::TestTransformStateExtension::updateMapping(Operation *previous,
+                                                 Operation *updated) {
+  // Update value handles. The new ops should have at least as many results as
+  // the replacement op. Fewer results are acceptable, if those results are not
+  // mapped to any handle.
+  for (int r = updated->getNumResults(); r < previous->getNumResults(); ++r) {
+    SmallVector<Value> handles;
+    (void)getTransformState().getHandlesForPayloadValue(previous->getResult(r),
+                                                        handles);
+    if (!handles.empty())
+      return emitError(previous->getLoc())
+             << "cannot replace an op with another op producing less results "
+                "while tracking handles";
+  }
+
+  for (auto [oldValue, newValue] :
+       llvm::zip(previous->getResults(), updated->getResults()))
+    if (failed(replacePayloadValue(oldValue, newValue)))
+      return failure();
+
+  // Update op handle.
+  return replacePayloadOp(previous, updated);
+}