[mlir][Transform] Make applyToOne return a DiagnosedSilenceableFailure

This revision revisits the implementation of applyToOne and its handling
of recoverable errors as well as propagation of null handles.
The implementation is simplified to always require passing a vector<Operation*>
in which the results are returned, resulting in less template instantiation magic.

Reviewed By: ftynse

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

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

@@ -15,6 +15,7 @@
 
 namespace mlir {
 namespace linalg {
+class GenericOp;
 class LinalgOp;
 } // namespace linalg
 } // namespace mlir

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

@@ -22,18 +22,26 @@ def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
   let description = [{
     Decomposes named complex operations, such as higher-dimensional
     (depthwise) convolutions, into combinations of lower-dimensional equivalents
-    when possible. The operand handle must point to a list of such operations.
-    The returning handle points to the main produced computational operation,
-    such as the lower-dimensional convolution.
+    when possible.
+
+    Return modes:
+    =============
+    This operation ignores non-Linalg ops and drops them in the return.
+    If all the operations referred to by the `target` PDLOperation decompose
+    properly, the transform succeeds. Otherwise the transform silently fails.
+    The return handle points to only the subset of successfully produced 
+    computational operations, which can be empty.
   }];
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::LinalgOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -61,20 +69,27 @@ def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Transforms a named structued operation into the generic form with the
-    explicit attached region. The operand handle must point to a list of
-    structured operations, it is consumed by the transformation and is not
-    expected to be used afterwards. The resulting handle points to the list
-    of equivalent generic operations, in the same order as the original named
-    operations.
+    explicit attached region. 
+
+    Return modes:
+    =============
+    This operation ignores non-Linalg ops and drops them in the return.
+    If all the operations referred to by the `target` PDLOperation generalize
+    properly, the transform succeeds. Otherwise the transform silently fails.
+    The return handle points to only the subset of successfully produced 
+    equivalent generic operations, which can be empty or contain the original
+    ops if they were already in generic form.
   }];
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$transformed);
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::LinalgOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -84,6 +99,16 @@ def InterchangeOp : Op<Transform_Dialect, "structured.interchange",
   let description = [{
     Interchanges the iterators of the operations pointed to by the target handle
     using the iterator interchange attribute.
+
+    Return modes:
+    =============
+    This operation ignores non-linalg::Generic ops and drops them in the return.
+    This operation fails if the interchange attribute is invalid.
+    If all the operations referred to by the `target` PDLOperation interchange
+    properly, the transform succeeds. 
+    If any interchange fails, the transform definitely fails.
+    The return handle points to only the subset of successfully produced 
+    interchanged operations, which can be empty.
   }];
 
   let arguments =
@@ -95,8 +120,10 @@ def InterchangeOp : Op<Transform_Dialect, "structured.interchange",
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::GenericOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -106,6 +133,16 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
   let description = [{
     Pads the operations pointed to by the target handle using the options
     provides as operation attributes.
+
+    Return modes:
+    =============
+    This operation ignores non-Linalg ops and drops them in the return.
+    This operation may produce a definiteFailure if the padding fails for any
+    reason.
+    If all the operations referred to by the `target` PDLOperation pad
+    properly, the transform succeeds. Otherwise the transform silently fails.
+    The return handle points to only the subset of successfully produced 
+    padded operations, which can be empty.
   }];
 
   let arguments =
@@ -123,8 +160,10 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::LinalgOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -135,11 +174,23 @@ def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
     Indicates that ops of a specific kind in the given function should be
     scalarized (i.e. their dynamic dimensions tiled by 1).
 
-    This operation returns the tiled op but not the loops.
+    Return modes:
+    =============
+    This operation ignores non-Linalg ops and drops them in the return.
+    This operation produces `definiteFailure` if the scalarization fails for any
+    reason.
+    If all the operations referred to by the `target` PDLOperation scalarize
+    properly, the transform succeeds. Otherwise the transform silently fails.
 
+    The return handle points to only the subset of successfully produced 
+    tiled-by-1 operations, which can be empty.
+
+    This operation does not return handles to the tiled loop.
     We make this design choice because it is hard to know ahead of time the
     number of loops that will be produced (it depends on the number of dynamic
     dimensions after multiple transformations have been applied).
+    Loops can always be recovered by navigating from the tiled operations if
+    needed.
   }];
 
   let arguments = (ins PDL_Operation:$target);
@@ -148,8 +199,10 @@ def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::LinalgOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -206,7 +259,17 @@ def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
       - use_alloc: whether to use an alloc op to allocate the temporary 
         tensor (default: do not use alloc op)
 
-    This op returns 4 handles to:
+    Return modes:
+    =============
+    This operation ignores non-Linalg ops and drops them in the return.
+    This operation produces `definiteFailure` if the splitting fails for any
+    reason.
+
+    If all the operations referred to by the `target` PDLOperation split
+    properly, the transform succeeds. Otherwise the transform silently fails.
+    The 4 returned handles points to only the subset of successfully produced 
+    computational operations, which can all be empty.
+    This 4 returned handles point to:
       - the init op (or tensor_alloc op if use_alloc = true), 
       - the fill op used to initialize the neutral element, 
       - the split op and 
@@ -316,15 +379,18 @@ def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
                    DefaultValuedAttr<I64Attr, "{}">:$insert_split_dimension,
                    UnitAttr:$use_scaling_algorithm,
                    UnitAttr:$use_alloc);
-  let results = (outs PDL_Operation:$fill_op,
+  let results = (outs PDL_Operation:$init_or_alloc_op,
+                      PDL_Operation:$fill_op,
                       PDL_Operation:$split_linalg_op,
                       PDL_Operation:$combining_linalg_op);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::LinalgOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -372,6 +438,13 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
 
     Note that this transformation is invalidating the handles to any payload IR
     operation that is contained inside the vectorization target.
+
+    Return modes:
+    =============
+    This operation produces `definiteFailure` if vectorization fails for any
+    reason.
+    The operation always returns the handle to the target op that is expected 
+    to be isolated from above.
   }];
 
   let arguments = (ins PDL_Operation:$target,
@@ -381,8 +454,10 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<Operation *> applyToOne(
-      ::mlir::Operation *target, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

@@ -239,6 +239,8 @@ tileAndFuseLinalgOps(OpBuilder &builder, ArrayRef<LinalgOp> ops,
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`op.rank` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
+///
+/// Return failure if the permutation is not valid.
 FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
                                           GenericOp genericOp,
                                           ArrayRef<unsigned> interchangeVector);

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td

@@ -50,7 +50,7 @@ def LoopOutlineOp : Op<Transform_Dialect, "loop.outline",
      outlined into a separate function. The provided name is used as a _base_
      for forming actual function names following SymbolTable auto-renaming
      scheme to avoid duplicate symbols. Expects that all ops in the Payload IR
-     have a SymbolTable ancestor (typically true because of the top-level 
+     have a SymbolTable ancestor (typically true because of the top-level
      module). Returns the handle to the list of outlined functions in the same
      order as the operand handle.
   }];
@@ -68,28 +68,40 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
   let summary = "Peels the last iteration of the loop";
   let description = [{
      Updates the given loop so that its step evenly divides its range and puts
-     the remaining iteration into a separate loop or a conditional. Note that
-     even though the Payload IR modification may be performed in-place, this
-     operation consumes the operand handle and produces a new one. Applies to
-     each loop associated with the operand handle individually. The results
-     follow the same order as the operand.
-
-     Note: If it can be proven statically that the step already evenly divides
-     the range, this op is a no-op. In the absence of sufficient static
-     information, this op may peel a loop, even if the step always divides the
-     range evenly at runtime.
+     the remaining iteration into a separate loop or a conditional.
+
+     In the absence of sufficient static information, this op may peel a loop,
+     even if the step always divides the range evenly at runtime.
+
+     Return modes:
+     =============
+     This operation ignores non-scf::ForOp ops and drops them in the return.
+
+     This operation always succeeds and returns the scf::ForOp with the
+     postcondition: "the loop trip count is divisible by the step".
+     This operation may return the same unmodified loop handle when peeling did
+     not modify the IR (i.e. the loop trip count was already divisible).
+
+     Note that even though the Payload IR modification may be performed
+     in-place, this operation consumes the operand handle and produces a new
+     one.
+
+     TODO: Return both the peeled loop and the remainder loop.
   }];
 
   let arguments =
       (ins PDL_Operation:$target,
            DefaultValuedAttr<BoolAttr, "false">:$fail_if_already_divisible);
+  // TODO: Return both the peeled loop and the remainder loop. 
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(
-      ::mlir::scf::ForOp loop, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::scf::ForOp target, 
+        ::llvm::SmallVector<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -102,10 +114,21 @@ def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
     each of them. That is, performs some amount of reads from memory before the
     loop rather than inside the loop, the same amount of writes into memory
     after the loop, and updates each iteration to read the data for a following
-    iteration rather than the current one. The amount is specified by the
-    attributes. The values read and about to be stored are transferred as loop
-    iteration arguments. Currently supports memref and vector transfer
-    operations as memory reads/writes.
+    iteration rather than the current one. 
+
+    The amount is specified by the attributes. 
+
+    The values read and about to be stored are transferred as loop iteration
+    arguments. Currently supports memref and vector transfer operations as 
+    memory reads/writes.
+
+    Return modes:
+    =============
+    This operation ignores non-scf::For ops and drops them in the return.
+    If all the operations referred to by the `target` PDLOperation pipeline
+    properly, the transform succeeds. Otherwise the transform silently fails.
+    The return handle points to only the subset of successfully produced 
+    pipelined loops, which can be empty.
   }];
 
   let arguments = (ins PDL_Operation:$target,
@@ -116,8 +139,10 @@ def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::FailureOr<::mlir::scf::ForOp> applyToOne(
-      ::mlir::scf::ForOp loop, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::scf::ForOp target, 
+        ::llvm::SmallVector<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }
 
@@ -126,11 +151,18 @@ def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Unrolls the given loop with the given unroll factor";
   let description = [{
-     Unrolls each loop associated with the given handle to have up to the given
-     number of loop body copies per iteration. If the unroll factor is larger
-     than the loop trip count, the latter is used as the unroll factor instead.
-     Does not produce a new handle as the operation may result in the loop being
-     removed after a full unrolling.
+    Unrolls each loop associated with the given handle to have up to the given
+    number of loop body copies per iteration. If the unroll factor is larger
+    than the loop trip count, the latter is used as the unroll factor instead.
+
+    Return modes:
+    ==============
+    This operation ignores non-scf::For ops and drops them in the return.
+    If all the operations referred to by the `target` PDLOperation unroll
+    properly, the transform succeeds. Otherwise the transform silently fails.
+
+    Does not return handles as the operation may result in the loop being
+    removed after a full unrolling.
   }];
 
   let arguments = (ins PDL_Operation:$target,
@@ -139,8 +171,10 @@ def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
-    ::mlir::LogicalResult applyToOne(
-      ::mlir::scf::ForOp loop, TransformState &state);
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::scf::ForOp target, 
+        ::llvm::SmallVector<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
   }];
 }