[mlir] reallow null results in TransformEachOpTrait

[ftynse]

Summary

Previous changes in 98acd7468307b6099e7deae206a749af324ff95f were overly
eager to disallow null payload everywhere. The semantics of
TransformEachOpTrait allows individual applications to return null
payloads as means of filtering out the operations to which they are not
applicable without emitting even a silenceable failure. This is a
questionable choice, but one apparently relied upon. Null payloads are
not supposed to leak outside of the trait.

[qcolombet]

LGTM

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

mlir::transform::TransformEachOpTrait<OpTy>::apply(
   // propagate gracefully.
   // In this case, we fill all results with an empty vector.
   if (targets.empty()) {
     SmallVector<Operation *> emptyPayload;
     SmallVector<Attribute> emptyParams;
     for (OpResult r : this->getOperation()->getResults()) {
       if (r.getType().isa<TransformParamTypeInterface>())
         transformResults.setParams(r, emptyParams);
+      else if (r.getType().isa<TransformValueHandleTypeInterface>())
+        transformResults.setValues(r, ValueRange());
       else
         transformResults.set(r, emptyPayload);
     }
     return DiagnosedSilenceableFailure::success();
   }
 
   // Step 2. Call applyToOne on each target and record newly produced ops in its
   // corresponding results entry.

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

diag.attachNote(transformOpLoc)
         << "if you need variadic results, consider a generic `apply` "
         << "instead of the specialized `applyToOne`.";
     return failure();
   }
 
   // Check that the right kind of value was produced.
   for (const auto &[ptr, res] :
        llvm::zip(partialResult, transformOp->getResults())) {
-    if (ptr.isNull()) {
-      return emitDiag() << "null result #" << res.getResultNumber()
-                        << " produced";
+    if (ptr.isNull())
+      continue;
+    if (res.getType().template isa<TransformHandleTypeInterface>() &&
+        !ptr.is<Operation *>()) {
+      return emitDiag() << "application of " << transformOpName
+                        << " expected to produce an Operation * for result #"
+                        << res.getResultNumber();
     }
-    if (ptr.is<Operation *>() &&
-        !res.getType().template isa<TransformHandleTypeInterface>()) {
+    if (res.getType().template isa<TransformParamTypeInterface>() &&
+        !ptr.is<Attribute>()) {
       return emitDiag() << "application of " << transformOpName
                         << " expected to produce an Attribute for result #"
                         << res.getResultNumber();
     }
-    if (ptr.is<Attribute>() &&
-        !res.getType().template isa<TransformParamTypeInterface>()) {
+    if (res.getType().template isa<TransformValueHandleTypeInterface>() &&
+        !ptr.is<Value>()) {
       return emitDiag() << "application of " << transformOpName
-                        << " expected to produce an Operation * for result #"
+                        << " expected to produce a Value for result #"
                         << res.getResultNumber();
     }
   }
   return success();
 }
 
+template <typename T>
+static SmallVector<T> castVector(ArrayRef<transform::MappedValue> range) {
+  return llvm::to_vector(llvm::map_range(
+      range, [](transform::MappedValue value) { return value.get<T>(); }));
+}
+
 void transform::detail::setApplyToOneResults(
     Operation *transformOp, TransformResults &transformResults,
     ArrayRef<ApplyToEachResultList> results) {
+  SmallVector<SmallVector<MappedValue>> transposed;
+  transposed.resize(transformOp->getNumResults());
+  for (const ApplyToEachResultList &partialResults : results) {
+    if (llvm::any_of(partialResults,
+                     [](MappedValue value) { return value.isNull(); }))
+      continue;
+    assert(transformOp->getNumResults() == partialResults.size() &&
+           "expected as many partial results as op as results");
+    for (auto &[i, value] : llvm::enumerate(partialResults))
+      transposed[i].push_back(value);
+  }
+
   for (OpResult r : transformOp->getResults()) {
+    unsigned position = r.getResultNumber();
     if (r.getType().isa<TransformParamTypeInterface>()) {
-      auto params = llvm::to_vector(
-          llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
-            return oneResult[r.getResultNumber()].get<Attribute>();
-          }));
-      transformResults.setParams(r, params);
+      transformResults.setParams(r,
+                                 castVector<Attribute>(transposed[position]));
     } else if (r.getType().isa<TransformValueHandleTypeInterface>()) {
-      auto values = llvm::to_vector(
-          llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
-            return oneResult[r.getResultNumber()].get<Value>();
-          }));
-      transformResults.setValues(r, values);
+      transformResults.setValues(r, castVector<Value>(transposed[position]));
     } else {
-      auto payloads = llvm::to_vector(
-          llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
-            return oneResult[r.getResultNumber()].get<Operation *>();
-          }));
-      transformResults.set(r, payloads);
+      transformResults.set(r, castVector<Operation *>(transposed[position]));
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Utilities for PossibleTopLevelTransformOpTrait.
 //===----------------------------------------------------------------------===//
 

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

^bb0(%arg1: !pdl.operation):
     %0 = pdl_match @some in %arg1 : (!pdl.operation) -> !pdl.operation
     // Transform fails to match any but still produces 2 results.
     %1:2 = transform.test_correct_number_of_multi_results %0
   }
 }
 
 // -----
 
+// This should not fail.
+
 func.func @foo() {
-  // expected-note @below {{when applied to this op}}
   "op" () : () -> ()
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   pdl.pattern @some : benefit(1) {
     %0 = pdl.operands
     %1 = pdl.types
     %2 = pdl.operation "op"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
     pdl.rewrite %2 with "transform.dialect"
   }
 
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = pdl_match @some in %arg1 : (!pdl.operation) -> !pdl.operation
-    // expected-error @below {{null result #0 produced}}
     transform.test_mixed_null_and_non_null_results %0
   }
 }
 
 // -----
 
 // Expecting to match all operations by merging the handles that matched addi
 // and subi separately.
^bb0(%arg0: !transform.any_op):
     transform.test_produce_transform_param_or_forward_operand %arg0
       { first_result_is_param }
       : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
   }
 }
 
 // -----
 
-// expected-note @below {{when applied to this op}}
+// Should not fail.
+
 module {
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
-    // expected-error @below {{null result #0 produced}}
     transform.test_produce_transform_param_or_forward_operand %arg0
       { first_result_is_null }
       : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
   }
 }
 
 // -----
 
 // expected-note @below {{when applied to this op}}
 module {
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
     // expected-error @below {{expected to produce an Attribute for result #1}}
     transform.test_produce_transform_param_or_forward_operand %arg0
       { second_result_is_handle }
       : (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
   }
 }
 
 // -----
 
+// expected-note @below {{when applied to this op}}
+module {
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-error @below {{expected to produce a Value for result #0}}
+    transform.test_produce_transform_param_or_forward_operand %arg0
+      { second_result_is_handle }
+      : (!transform.any_op) -> (!transform.any_value, !transform.param<i64>)
+  }
+}
+
+// -----
+
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   // expected-error @below {{attempting to assign a null payload op to this transform value}}
   %0 = transform.test_produce_null_payload : !transform.any_op
 }
 
 // -----
 

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

 def TestProduceTransformParamOrForwardOperandOp
   : Op<Transform_Dialect, "test_produce_transform_param_or_forward_operand",
        [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
         TransformEachOpTrait, TransformOpInterface]> {
   let arguments = (ins TransformHandleTypeInterface:$in,
                        UnitAttr:$first_result_is_param,
                        UnitAttr:$first_result_is_null,
                        UnitAttr:$second_result_is_handle);
-  let results = (outs TransformHandleTypeInterface:$out,
+  let results = (outs AnyType:$out,
                       TransformParamTypeInterface:$param);
   let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
   let cppNamespace = "::mlir::test";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,