[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 ↗

(On Diff #466045)

Leftover debug?

[nicolasvasilache]

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

295–296 ↗

(On Diff #466045)

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 always succeeds and returns `num_result_handles`. 
+    The resulting handles may all be empty if the statically specified 
+    `num_result_handles` does not correspond to the dynamic number of operations
+    contained in the source `handle`. This behavior allows the propagation of
+    the "failed to match" information when the transform dialect is applied to
+    unexpected IR.
+  }];
+
+  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>]> {
   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/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 = state.getPayloadOps(getHandle()).size();
+  int64_t expectedNumResultHandles = getNumResultHandles();
+  // Propagate mode.
+  if (numResultHandles == 0) {
+    for (int64_t idx = 0; idx < expectedNumResultHandles; ++idx)
+      results.set(getResults()[idx].cast<OpResult>(), {});
+    return DiagnosedSilenceableFailure::success();
+  }
+  if (numResultHandles != expectedNumResultHandles) {
+    emitOpError("expects ")
+        << getHandle() << " to contain " << expectedNumResultHandles
+        << " operation handles but it only contains " << numResultHandles
+        << " handles";
+    return DiagnosedSilenceableFailure::definiteFailure();
+  }
+  int64_t idx = 0;
+  for (Operation *op : state.getPayloadOps(getHandle()))
+    results.set(getResults()[idx++].cast<OpResult>(), op);
+  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]
+  // Both form succeed, the second form has all empty handles.
+  %muli_2 = transform.structured.match ops{["arith.muli"]} in %fun
+  %h_2:3 = split_handles %muli in [3]
+}