[mlir] relax value handle updates when operation is replaced

[ftynse]

Summary

The initial implementaiton of value handle update when the payload
operation defining the values associated with value handles was being
replaced required the replacement operation to have the same number of
results. This is not strictly necessary. The replacement operation may
have more results, or less results provided that there are no handles to
the results that have no equivalent in the replacement op.

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

for (Value opResult : op->getResults()) {
         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.
-  if (op->getNumResults() != replacement->getNumResults()) {
-    return emitError(op->getLoc())
-           << "cannot replace an op with another op producing a different "
-              "number of results while tracking handles";
-  }
-
-  // Replace the mapped operation if present.
+  // 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);
   }
 
-  // Replace the mapped results of the operation.
-  for (auto [origResult, replacementResult, handleList] : llvm::zip(
-           op->getResults(), replacement->getResults(), resultValueHandles)) {
+  // 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";
+    }
+
+    Value replacementResult = replacement->getResult(resultNumber);
     for (Value resultHandle : handleList) {
       Mappings &mappings = getMapping(resultHandle);
       auto it = mappings.values.find(resultHandle);
       if (it == mappings.values.end())
         continue;
 
       SmallVector<Value> &association = it->getSecond();
       for (Value &mapped : association) {

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

^bb0(%arg0: !pdl.operation):
     test_remap_operand_to_self %arg0
     // expected-remark @below {{extension present, A}}
     test_check_if_test_extension_present %arg0
   }
 }
 
 // -----
 
-// expected-error @below {{cannot replace an op with another op producing a different number of results while tracking handles}}
-module {
-  transform.sequence failures(propagate) {
-  ^bb0(%arg0: !pdl.operation):
-    test_add_test_extension "A"
-    %dummy = test_remap_operand_to_self %arg0 : !transform.any_op
-  }
+   // This is okay because we are replacing the top-level module opeation
+   // (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):
+  test_add_test_extension "A"
+  %dummy = test_remap_operand_to_self %arg0 : !pdl.operation
+  // This is still okay. Even though we are replacing the previous
+  // operation with (1 result) with this operation that has less (0) results,
+  // there is no handle to the result, hence no issue with value handle update.
+  test_remap_operand_to_self %dummy
 }
 
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  test_add_test_extension "A"
+  // expected-error @below {{cannot replace an op with another op producing less results while tracking handles}}
+  %dummy = test_remap_operand_to_self %arg0 : !pdl.operation
+  %valuehandle = transform.get_result %dummy[0] : (!pdl.operation) -> !transform.any_value
+  test_remap_operand_to_self %dummy
+}
 
 // -----
 
 module {
   transform.sequence failures(suppress) {
   ^bb0(%arg0: !pdl.operation):
     // expected-error @below {{TestTransformStateExtension missing}}
     test_remap_operand_to_self %arg0
   }
 }

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

DiagnosedSilenceableFailure mlir::test::TestRemapOperandPayloadToSelfOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   auto *extension = state.getExtension<TestTransformStateExtension>();
   if (!extension)
     return emitDefiniteFailure("TestTransformStateExtension missing");
 
   if (failed(extension->updateMapping(state.getPayloadOps(getOperand()).front(),
                                       getOperation())))
     return DiagnosedSilenceableFailure::definiteFailure();
+  if (getNumResults() > 0)
+    results.set(getResult(0).cast<OpResult>(), getOperation());
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestRemapOperandPayloadToSelfOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getOperand(), effects);
   transform::producesHandle(getOut(), effects);
   transform::onlyReadsPayload(effects);