[mlir][Transform] Fix implementation of the generic apply that is based on applyToOne.

[nicolasvasilache]

Summary

The result of applying an N-result producing transformation to M payload ops
is an M-wide result, each containing N result operations.
This requires a transposition of the results obtained by calling applyToOne.

This revision fixes the issue and adds more advanced tests that exercise the behavior.

[nicolasvasilache]

Finish impl and add tests.

[nicolasvasilache]

Drop spurious dump.

[ftynse]

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

799

static, otherwise we may run into multiple definitions

839–858

This error sounds like a misuse of the TransformEachTrait by the IR designer: the trait should not be attached to the transform ops that produce a different number of results for different inputs. I'd rather make it an assertion, or at least a definite failure, not a silenceable one. I understand the diagnostic interface is more convenient so we can emit a diagnostic and then llvm::report_fatal_error("trait misused") I suppose.

845

Nit: there's emitSilenceableError() declared in the interface.

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

61 ↗

(On Diff #439336)

Leftover debug.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

Address.

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
     which clean up the results. Therefore, it can only be applied to an op with
     the "isolated from above property". If finer granularity is required, it can
     be achieved by outlining the target part of the payload IR into, e.g., a
     function, performing the transformation, and inlining it back. This
     transformation only fails if the entire pattern rewriting failed, i.e., it
     does **not** fail when no ops were vectorized.
 
     Note that this transformation is invalidating the handles to any payload IR
-    operation that is contained inside the vectoriztaion target.
+    operation that is contained inside the vectorization target.
   }];
 
   let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<BoolAttr, "false">:$vectorize_padding);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<Operation *> applyToOne(
       ::mlir::Operation *target, TransformState &state);
   }];
 }
 
 #endif // LINALG_TRANSFORM_OPS

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

for (Operation *target : targets) {
     }
 
     auto result = transform(specificOp);
     if (failed(appendTransformResultToVector(result, results)))
       return DiagnosedSilenceableFailure::definiteFailure();
   }
   return DiagnosedSilenceableFailure::success();
 }
+
+/// Helper function to transform M ops with N results into N results of M ops.
+static inline SmallVector<SmallVector<Operation *, 1>>
+transposeResults(const SmallVector<SmallVector<Operation *>, 1> &m) {
+  SmallVector<SmallVector<Operation *, 1>> res;
+  if (m.empty())
+    return res;
+  int64_t rows = m.size(), cols = m[0].size();
+  for (int64_t j = 0; j < cols; ++j)
+    res.push_back(SmallVector<Operation *, 1>(rows, nullptr));
+  for (int64_t i = 0; i < rows; ++i) {
+    assert(static_cast<int64_t>(m[i].size()) == cols);
+    for (int64_t j = 0; j < cols; ++j) {
+      res[j][i] = m[i][j];
+    }
+  }
+  return res;
+}
 } // namespace detail
 } // namespace transform
 } // namespace mlir
 
 template <typename OpTy>
 mlir::DiagnosedSilenceableFailure
 mlir::transform::TransformEachOpTrait<OpTy>::apply(
     TransformResults &transformResults, TransformState &state) {
   using TransformOpType = typename llvm::function_traits<
       decltype(&OpTy::applyToOne)>::template arg_t<0>;
   ArrayRef<Operation *> targets =
       state.getPayloadOps(this->getOperation()->getOperand(0));
   SmallVector<SmallVector<Operation *>, 1> results;
   // In the multi-result case, collect the number of results each transform
   // produced.
   DiagnosedSilenceableFailure result = detail::applyTransformToEach(
       targets, results, [&](TransformOpType specificOp) {
         return static_cast<OpTy *>(this)->applyToOne(specificOp, state);
       });
   if (!result.succeeded())
     return result;
-  for (const SmallVector<Operation *> &oneTargetResults : results) {
-    if (OpTy::template hasTrait<OpTrait::ZeroResults>())
-      continue;
-    if (OpTy::template hasTrait<OpTrait::OneResult>()) {
-      transformResults.set(
-          this->getOperation()->getResult(0).template cast<OpResult>(),
-          oneTargetResults);
-      continue;
-    }
-    if (this->getOperation()->getNumResults() != oneTargetResults.size()) {
-      Diagnostic diag(this->getOperation()->getLoc(),
-                      DiagnosticSeverity::Error);
-      diag << "unexpected number of results (got " << oneTargetResults.size()
-           << " expected " << this->getOperation()->getNumResults() << ")";
-      return DiagnosedSilenceableFailure::silencableFailure(std::move(diag));
-    }
-    for (const auto &it :
-         llvm::zip(this->getOperation()->getResults(), oneTargetResults)) {
-      transformResults.set(std::get<0>(it).template cast<OpResult>(),
-                           std::get<1>(it));
-    }
-  }
+  if (results.empty())
+    return DiagnosedSilenceableFailure::success();
+
+  // Ensure all applications return the same number of results.
+  // Variadic cases are much trickier to handle in a generic fashion.
+  int64_t nRes = results[0].size();
+  if (llvm::any_of(results, [&](const auto &r) {
+        return static_cast<int64_t>(r.size()) != nRes;
+      })) {
+    return static_cast<OpTy *>(this)->emitSilenceableError()
+           << "expected all applications of " << OpTy::getOperationName()
+           << " to produce " << nRes
+           << " results.\n If you need variadic results, consider using a "
+              "generic `apply` instead of the specialized `applyToOne`";
+  }
+  // Ensure the number of results agrees with what the transform op expects.
+  if (this->getOperation()->getNumResults() != nRes) {
+    InFlightDiagnostic diag = static_cast<OpTy *>(this)->emitError()
+                              << "unexpected number of results (got " << nRes
+                              << " expected "
+                              << this->getOperation()->getNumResults() << ")";
+    return DiagnosedSilenceableFailure::definiteFailure();
+  }
+
+  // If no results, bail early.
+    return DiagnosedSilenceableFailure::success();
+
+  // Perform transposition of M applications producing N results each into N
+  // results for each of the M applications.
+  SmallVector<SmallVector<Operation *, 1>> transposedResults =
+      detail::transposeResults(results);
+  // Single result applies to M ops produces one single M-result.
+    assert(transposedResults.size() == 1 && "Expected single result");
+        transposedResults[0]);
+    return DiagnosedSilenceableFailure::success();
+  // M ops, N results each.
+       llvm::zip(this->getOperation()->getResults(), transposedResults)) {
   return DiagnosedSilenceableFailure::success();
 }
 
 template <typename OpTy>
 mlir::LogicalResult
 mlir::transform::TransformEachOpTrait<OpTy>::verifyTrait(Operation *op) {
   static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
                 "expected single-operand op");

mlir/test/Dialect/Transform/selective-targeting.mlir

+// RUN:  mlir-opt %s -test-transform-dialect-interpreter --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @matmul_tensors_1(
+func.func @matmul_tensors_1(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
+  %arg2: tensor<128x128xf32> {linalg.inplaceable = true})
+    -> tensor<128x128xf32> {
+  // This operation is marked for tiling only.
+  // CHECK-COUNT-3: scf.for
+  // CHECK-COUNT-3: tensor.extract_slice
+  // CHECK: linalg.matmul
+  // CHECK-SAME: -> tensor<4x4xf32>
+  %0 = linalg.matmul { test.attrA }
+                      ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  func.return %0 : tensor<128x128xf32>
+}
+
+func.func @matmul_tensors_2(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
+  %arg2: tensor<128x128xf32> {linalg.inplaceable = true})
+    -> tensor<128x128xf32> {
+  // This operation is marked f
+  // This operation is marked for tiling and vectorization.
+  // CHECK-COUNT-3: scf.for
+  // CHECK-COUNT-3: vector.transfer_read
+  // CHECK:       vector.contract
+  // CHECK-NOT:   linalg.matmul
+  // CHECK:       vector.transfer_write
+  %0 = linalg.matmul { test.attrA, test.attrC }
+                      ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  func.return %0 : tensor<128x128xf32>
+}
+
+func.func @matmul_tensors_3(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
+  %arg2: tensor<128x128xf32> {linalg.inplaceable = true})
+    -> tensor<128x128xf32> {
+  // This operation is marked for vectorization only.
+  // CHECK-NOT: scf.for
+  // CHECK-COUNT-3: vector.transfer_read
+  // CHECK: vector.contract
+  // CHECK-SAME: into vector<128x128xf32>
+  // CHECK: vector.transfer_write
+  %0 = linalg.matmul { test.attrC }
+                      ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  func.return %0 : tensor<128x128xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  // Match matmul operations inside @matmul_tensors with test.attrA set.
+  pdl.pattern @pdl_target_attrA : benefit(1) {
+    %args = operands
+    %results = types
+    %attr = attribute
+    %0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrA" = %attr}-> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  // Match matmul operations inside @matmul_tensors with test.attrC set.
+  pdl.pattern @pdl_target_attrC : benefit(1) {
+    %args = operands
+    %results = types
+    %attr = attribute
+    %0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrC" = %attr}-> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target_attrA in %arg1
+    transform.structured.tile %0 {sizes = [4, 4, 4]}
+    %1 = pdl_match @pdl_target_attrC in %arg1
+    %2 = transform.get_closest_isolated_parent %1
+    transform.structured.vectorize %2
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @vectorize_one
+func.func @vectorize_one(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
+  %arg2: tensor<128x128xf32> {linalg.inplaceable = true})
+    -> tensor<128x128xf32> {
+  // CHECK: vector.contract
+  %0 = linalg.matmul {test.attrA}
+                     ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  func.return %0 : tensor<128x128xf32>
+}
+
+func.func @vectorize_none(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
+  %arg2: tensor<128x128xf32> {linalg.inplaceable = true})
+    -> tensor<128x128xf32> {
+  // CHECK: linalg.matmul
+  %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  func.return %0 : tensor<128x128xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %attr = attribute
+    %0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrA" = %attr}-> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    %1 = get_closest_isolated_parent %0
+    transform.structured.vectorize %1
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @vectorize_all
+func.func @vectorize_all(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>,
+  %arg3: tensor<128x128xf32> {linalg.inplaceable = true})
+    -> tensor<128x128xf32> {
+  // CHECK: vector.contract
+  %0 = linalg.matmul {test.attrA}
+                     ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  // CHECK: vector.contract
+  %1 = linalg.matmul ins(%arg0, %0: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg3: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  return %1 : tensor<128x128xf32>
+}
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  transform.structured.vectorize %arg0
+}

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

-// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file --verify-diagnostics
+// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
 
 // expected-remark @below {{applying transformation}}
 transform.test_transform_op
 
 // -----
 
 %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
 // expected-remark @below {{succeeded}}

 }
 // -----
 
 transform.sequence {
 ^bb0(%arg0: !pdl.operation):
   // expected-error @below {{unexpected number of results (got 0 expected 3)}}
   transform.test_wrong_number_of_results %arg0
 }
+
+// -----
+
+func.func @foo() {
+  "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 {
+  ^bb0(%arg1: !pdl.operation):
+    %0 = pdl_match @some in %arg1
+    // expected-error @below {{expected all applications of transform.test_wrong_number_of_multi_results to produce 1 results}}
+    transform.test_wrong_number_of_multi_results %0
+  }
+}
+
+// -----
+
+func.func @foo() {
+  "op" () : () -> ()
+  "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 {
+  ^bb0(%arg1: !pdl.operation):
+    %0 = pdl_match @some in %arg1
+    // Transform matches 3 ops and produces 2 results.
+    %1:2 = transform.test_correct_number_of_multi_results %0
+  }
+}

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

 }
 
 FailureOr<SmallVector<Operation *>>
 mlir::test::TestWrongNumberOfResultsOp::applyToOne(
     Operation *, transform::TransformState &state) {
   return SmallVector<Operation *>{};
 }
 
+FailureOr<SmallVector<Operation *>>
+mlir::test::TestWrongNumberOfMultiResultsOp::applyToOne(
+    Operation *op, transform::TransformState &state) {
+  static int count = 0;
+  if (count++ > 0)
+    return SmallVector<Operation *>{};
+  OperationState opState(op->getLoc(), "foo");
+  return SmallVector<Operation *>{OpBuilder(op).create(opState)};
+}
+
+FailureOr<SmallVector<Operation *>>
+mlir::test::TestCorrectNumberOfMultiResultsOp::applyToOne(
+    Operation *op, transform::TransformState &state) {
+  OperationState opState(op->getLoc(), "foo");
+  return SmallVector<Operation *>{OpBuilder(op).create(opState),
+                                  OpBuilder(op).create(opState)};
+}
+
 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 TestWrongNumberOfResultsOp
   let assemblyFormat = "$target attr-dict";
   let cppNamespace = "::mlir::test";
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
         ::mlir::Operation *target, transform::TransformState &state);
   }];
 }
 
+def TestWrongNumberOfMultiResultsOp
+  : Op<Transform_Dialect, "test_wrong_number_of_multi_results",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface, 
+     TransformEachOpTrait, TransformOpInterface]> {
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::test";
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
+        ::mlir::Operation *target, transform::TransformState &state);
+  }];
+}
+
+def TestCorrectNumberOfMultiResultsOp
+  : Op<Transform_Dialect, "test_correct_number_of_multi_results",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface, 
+     TransformEachOpTrait, TransformOpInterface]> {
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result1,
+                      PDL_Operation:$result2);
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::test";
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
+        ::mlir::Operation *target, transform::TransformState &state);
+  }];
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD