[mlir][Transform] Add a transform.split_handles operation and fix general silenceable bugs.

[nicolasvasilache]

Summary

The transform.split_handles op is useful for ensuring a statically known number of operations are
tracked by the source handle and to extract them into individual handles
that can be further manipulated in isolation.

In the process of making the op robust wrt to silenceable errors and the suppress mode, issues were
uncovered and fixed.

The main issue was that silenceable errors were short-circuited too early and the payloads were not
set. This resulted in suppressed silenceable errors not propagating correctly.
Fixing the issue triggered a few test failures: silenceable error returns now must properly set the results state.

[springerm]

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

229–231

Why not a silenceable error?

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

488

drop

490

Documentation says that this operation always succeeds.

492–493

nit: Could use llvm::enumerate

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

775

Check number of payload ops:

// expected-remark @below {{1}}
transform.test_print_number_of_associated_payload_ir_ops %h#0

etc.

[ftynse]

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

490

It feels strange if the op fails to do what its name suggests it would (split the handle) and still succeeds, so I think the documentation is wrong. Also, this should be a silenceable failure, we are not irreversibly modifying the IR. There seems to be a confusion as to what a silenceable failure means: it's not just a compiler warning that the interpreter will step over and continue, it is still a failure and the interpretation will be interrupted. It can just recover and start in a different place.

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

779

Also check the user-visible error message.

[nicolasvasilache]

Address.

[nicolasvasilache]

Investigate deeper into suppres mode.

[nicolasvasilache]

Bugfixes.

[nicolasvasilache]

Fix commit msg.

[ftynse]

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

295–296

Leftover debug?

[nicolasvasilache]

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

295–296

no, it's legit

But maybe I can fold it as part of emitRemark() << state.getPayloadOps(getHandle()).size();

will try

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

def MergeHandlesOp : TransformDialectOp<"merge_handles",
 
   let arguments = (ins Variadic<PDL_Operation>:$handles,
                        UnitAttr:$deduplicate);
   let results = (outs PDL_Operation:$result);
   let assemblyFormat = "($deduplicate^)? $handles attr-dict";
   let hasFolder = 1;
 }
 
+def SplitHandlesOp : TransformDialectOp<"split_handles",
+    [FunctionalStyleTransformOpTrait,
+     DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Splits handles from a union of payload ops to a list";
+  let description = [{
+    Creates `num_result_handles` transform IR handles extracted from the 
+    `handle` operand. The resulting Payload IR operation handles are listed
+    in the same order as the operations appear in the source `handle`.
+    This is useful for ensuring a statically known number of operations are
+    tracked by the source `handle` and to extract them into individual handles
+    that can be further manipulated in isolation.
+
+    This operation succeeds and returns `num_result_handles` if the statically
+    specified `num_result_handles` corresponds to the dynamic number of 
+    operations contained in the source `handle`. Otherwise it silently fails.
+  }];
+
+  let arguments = (ins PDL_Operation:$handle,
+                       I64Attr:$num_result_handles);
+  let results = (outs Variadic<PDL_Operation>:$results);
+  let assemblyFormat = [{
+    $handle `in` `[` $num_result_handles `]` 
+    custom<StaticNumPDLResults>(type($results), ref($num_result_handles))
+    attr-dict
+  }];
+}
+
 def PDLMatchOp : TransformDialectOp<"pdl_match",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Finds ops that match the named PDL pattern";
   let description = [{
     Find Payload IR ops nested within the Payload IR op associated with the
     operand that match the PDL pattern identified by its name. The pattern is
     expected to be defined in the closest surrounding `WithPDLPatternsOp`.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 transform::MatchOp::apply(transform::TransformResults &results,
                           transform::TransformState &state) {
   llvm::StringSet<> strs;
   if (getOps().has_value())
     strs.insert(getOps()->getAsValueRange<StringAttr>().begin(),
                 getOps()->getAsValueRange<StringAttr>().end());
 
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
-  if (payloadOps.size() != 1)
+  if (payloadOps.size() != 1) {
+    results.set(getResult().cast<OpResult>(), {});
     return DiagnosedSilenceableFailure(
         this->emitOpError("requires exactly one target handle"));
+  }
 
   SmallVector<Operation *> res;
   auto matchFun = [&](Operation *op) {
     if (getOps().has_value() && !strs.contains(op->getName().getStringRef()))
       return;
 
     // Interfaces cannot be matched by name, just by ID.
     // So we specifically encode the interfaces we care about for this op.
splitPoints = llvm::to_vector(llvm::map_range(
               !op->getResult(0).getType().isIndex()) {
             diag = emitSilenceableError()
                    << "expected dynamic split point handle to point to a "
                       "single-result index-typed op";
             diag.attachNote(op->getLoc()) << "dynamic split point";
           }
           return OpFoldResult(op->getResult(0));
         }));
-    if (!diag.succeeded())
+    if (diag.isSilenceableFailure()) {
+      results.set(getFirst().cast<OpResult>(), {});
+      results.set(getSecond().cast<OpResult>(), {});
       return diag;
+    }
 
     if (splitPoints.size() != payload.size()) {
       emitError() << "expected the dynamic split point handle to point to as "
                      "many operations ("
                   << splitPoints.size() << ") as the target handle ("
                   << payload.size() << ")";
       return DiagnosedSilenceableFailure::definiteFailure();
     }
   } else {
     splitPoints.resize(payload.size(),
                        rewriter.getIndexAttr(getStaticSplitPoint()));
   }
 
   // Split each target operation.
   SmallVector<Operation *> first, second;
   for (const auto &pair : llvm::zip(payload, splitPoints)) {
     Operation *target = std::get<0>(pair);
     auto linalgOp = dyn_cast<LinalgOp>(target);
     if (!linalgOp) {
       auto diag = emitSilenceableError() << "only applies to structured ops";
       diag.attachNote(target->getLoc()) << "target op";
+      results.set(getFirst().cast<OpResult>(), {});
+      results.set(getSecond().cast<OpResult>(), {});
       return diag;
     }
 
     if (getDimension() >= linalgOp.getNumLoops()) {
       auto diag = emitSilenceableError() << "dimension " << getDimension()
                                          << " does not exist in target op";
       diag.attachNote(target->getLoc()) << "target op";
+      results.set(getFirst().cast<OpResult>(), {});
+      results.set(getSecond().cast<OpResult>(), {});
       return diag;
     }
 
     rewriter.setInsertionPoint(linalgOp);
     std::tie(first.emplace_back(), second.emplace_back()) = linalg::splitOp(
         rewriter, cast<TilingInterface>(linalgOp.getOperation()),
         getDimension(), std::get<1>(pair));
   }

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

for (Operation *target : state.getPayloadOps(getTarget())) {
     for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
       loop = current->getParentOfType<scf::ForOp>();
       if (!loop) {
         DiagnosedSilenceableFailure diag = emitSilenceableError()
                                            << "could not find an '"
                                            << scf::ForOp::getOperationName()
                                            << "' parent";
         diag.attachNote(target->getLoc()) << "target op";
+        results.set(getResult().cast<OpResult>(), {});
         return diag;
       }
       current = loop;
     }
     parents.insert(loop);
   }
   results.set(getResult().cast<OpResult>(), parents.getArrayRef());
   return DiagnosedSilenceableFailure::success();
for (Operation *target : state.getPayloadOps(getTarget())) {
     Location location = target->getLoc();
     Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(target);
     SimpleRewriter rewriter(getContext());
     scf::ExecuteRegionOp exec = wrapInExecuteRegion(rewriter, target);
     if (!exec) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
                                          << "failed to outline";
       diag.attachNote(target->getLoc()) << "target op";
+      results.set(getTransformed().cast<OpResult>(), {});
       return diag;
     }
     func::CallOp call;
     FailureOr<func::FuncOp> outlined = outlineSingleBlockRegion(
         rewriter, location, exec.getRegion(), getFuncName(), &call);
 
     if (failed(outlined)) {
       (void)reportUnknownTransformError(target);

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

for (OpOperand &operand : transform->getOpOperands()) {
           diag.attachNote(op->getLoc()) << "repeated target op";
           return diag;
         }
       }
     }
   }
 
   transform::TransformResults results(transform->getNumResults());
+  // Compute the result but do not short-circuit the silenceable failure case as
+  // we still want the handles to propagate properly so the "suppress" mode can
+  // proceed on a best effort basis.
   DiagnosedSilenceableFailure result(transform.apply(results, *this));
-  if (!result.succeeded())
+  if (result.isDefiniteFailure())
     return result;
 
   // Remove the mapping for the operand if it is consumed by the operation. This
   // allows us to catch use-after-free with assertions later on.
   auto memEffectInterface =
       cast<MemoryEffectOpInterface>(transform.getOperation());
   SmallVector<MemoryEffects::EffectInstance, 2> effects;
   for (OpOperand &target : transform->getOpOperands()) {
for (OpResult result : transform->getResults()) {
     setPayloadOps(result, results.get(result.getResultNumber()));
   }
 
   printOnFailureRAII.release();
   DEBUG_WITH_TYPE(DEBUG_PRINT_AFTER_ALL, {
     DBGS() << "Top-level payload:\n";
     getTopLevel()->print(llvm::dbgs());
   });
-  return DiagnosedSilenceableFailure::success();
+  return result;
 }
 
 //===----------------------------------------------------------------------===//
 // TransformState::Extension
 //===----------------------------------------------------------------------===//
 
 transform::TransformState::Extension::~Extension() = default;
 

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

 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
 
 using namespace mlir;
 
+/// Custom parser for ReplicateOp.
 static ParseResult parsePDLOpTypedResults(
     OpAsmParser &parser, SmallVectorImpl<Type> &types,
     const SmallVectorImpl<OpAsmParser::UnresolvedOperand> &handles) {
   types.resize(handles.size(), pdl::OperationType::get(parser.getContext()));
   return success();
 }
 
+/// Custom printer for ReplicateOp.
 static void printPDLOpTypedResults(OpAsmPrinter &, Operation *, TypeRange,
                                    ValueRange) {}
 
+/// Custom parser for SplitHandlesOp.
+static ParseResult parseStaticNumPDLResults(OpAsmParser &parser,
+                                            SmallVectorImpl<Type> &types,
+                                            IntegerAttr numHandlesAttr) {
+  types.resize(numHandlesAttr.getInt(),
+               pdl::OperationType::get(parser.getContext()));
+  return success();
+}
+
+/// Custom printer for SplitHandlesOp.
+static void printStaticNumPDLResults(OpAsmPrinter &, Operation *, TypeRange,
+                                     IntegerAttr) {}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // PatternApplicatorExtension
 //===----------------------------------------------------------------------===//
 
 namespace {
if (getDeduplicate() || getHandles().size() != 1)
     return {};
 
   // If deduplication is not required and there is only one operand, it can be
   // used directly instead of merging.
   return getHandles().front();
 }
 
 //===----------------------------------------------------------------------===//
+// SplitHandlesOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::SplitHandlesOp::apply(transform::TransformResults &results,
+                                 transform::TransformState &state) {
+  int64_t numResultHandles =
+      getHandle() ? state.getPayloadOps(getHandle()).size() : 0;
+  int64_t expectedNumResultHandles = getNumResultHandles();
+  if (numResultHandles != expectedNumResultHandles) {
+    // Failing case needs to propagate gracefully for both suppress and
+    // propagate modes.
+    for (int64_t idx = 0; idx < expectedNumResultHandles; ++idx)
+      results.set(getResults()[idx].cast<OpResult>(), {});
+    // Empty input handle corner case: always propagates empty handles in both
+    // suppress and propagate modes.
+    if (numResultHandles == 0)
+      return DiagnosedSilenceableFailure::success();
+    // If the input handle was not empty and the number of result handles does
+    // not match, this is a legit silenceable error.
+    return emitSilenceableError()
+           << getHandle() << " expected to contain " << expectedNumResultHandles
+           << " operation handles but it only contains " << numResultHandles
+           << " handles";
+  }
+  // Normal successful case.
+  for (auto en : llvm::enumerate(state.getPayloadOps(getHandle())))
+    results.set(getResults()[en.index()].cast<OpResult>(), en.value());
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::SplitHandlesOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  consumesHandle(getHandle(), effects);
+  producesHandle(getResults(), effects);
+  // There are no effects on the Payload IR as this is only a handle
+  // manipulation.
+}
+
+//===----------------------------------------------------------------------===//
 // PDLMatchOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PDLMatchOp::apply(transform::TransformResults &results,
                              transform::TransformState &state) {
   auto *extension = state.getExtension<PatternApplicatorExtension>();
   assert(extension &&

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

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %muli = transform.structured.match ops{["arith.muli"]} in %arg1
   // expected-error @below {{could not find a producer for operand number: 0 of}}
   %bbarg = get_producer_of_operand %muli[0]
 
 }
 
+// -----
+
+func.func @split_handles(%a: index, %b: index, %c: index) {
+  %0 = arith.muli %a, %b : index  
+  %1 = arith.muli %a, %c : index  
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%fun: !pdl.operation):
+  %muli = transform.structured.match ops{["arith.muli"]} in %fun
+  %h:2 = split_handles %muli in [2]
+  // expected-remark @below {{1}}
+  transform.test_print_number_of_associated_payload_ir_ops %h#0
+  %muli_2 = transform.structured.match ops{["arith.muli"]} in %fun
+  // expected-error @below {{expected to contain 3 operation handles but it only contains 2 handles}}
+  %h_2:3 = split_handles %muli_2 in [3]
+}
+
+// -----
+
+func.func @split_handles(%a: index, %b: index, %c: index) {
+  %0 = arith.muli %a, %b : index  
+  %1 = arith.muli %a, %c : index  
+  return
+}
+
+transform.sequence failures(suppress) {
+^bb1(%fun: !pdl.operation):
+  %muli = transform.structured.match ops{["arith.muli"]} in %fun
+  %h:2 = split_handles %muli in [2]
+  // expected-remark @below {{1}}
+  transform.test_print_number_of_associated_payload_ir_ops %h#0
+  %muli_2 = transform.structured.match ops{["arith.muli"]} in %fun
+  // Silenceable failure and all handles are now empty.
+  %h_2:3 = split_handles %muli_2 in [3]
+  // expected-remark @below {{0}}
+ transform.test_print_number_of_associated_payload_ir_ops %h_2#0
+}

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

mlir::test::TestMixedSuccessAndSilenceableOp::applyToOne(
   if (target->hasAttr("target_me"))
     return DiagnosedSilenceableFailure::success();
   return emitDefaultSilenceableFailure(target);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestPrintNumberOfAssociatedPayloadIROps::apply(
     transform::TransformResults &results, transform::TransformState &state) {
+  if (!getHandle())
+    emitRemark() << 0;
   emitRemark() << state.getPayloadOps(getHandle()).size();
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestPrintNumberOfAssociatedPayloadIROps::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::onlyReadsHandle(getHandle(), effects);
 }