[mlir] reallow null results in TransformEachOpTrait

Previous changes in 98acd74 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.

Reviewed By: qcolombet

Differential Revision: https://reviews.llvm.org/D143904

Author: ftynse

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

@@ -1015,6 +1015,8 @@ mlir::transform::TransformEachOpTrait<OpTy>::apply(
     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);
     }

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

@@ -921,48 +921,60 @@ transform::detail::checkApplyToOne(Operation *transformOp,
   // 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]));
     }
   }
 }

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

@@ -495,8 +495,9 @@ transform.with_pdl_patterns {
 
 // -----
 
+// This should not fail.
+
 func.func @foo() {
-  // expected-note @below {{when applied to this op}}
   "op" () : () -> ()
   return
 }
@@ -513,7 +514,6 @@ transform.with_pdl_patterns {
   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
   }
 }
@@ -1053,11 +1053,11 @@ module {
 
 // -----
 
-// 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>)
@@ -1079,6 +1079,19 @@ module {
 
 // -----
 
+// 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}}

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

@@ -371,7 +371,7 @@ def TestProduceTransformParamOrForwardOperandOp
                        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";