[mlir] Add ReplicateOp to the Transform dialect

[ftynse]

Summary

This handle manipulation operation allows one to define a new handle that is
associated with a the same payload IR operations N times, where N can be driven
by the size of payload IR operation list associated with another handle. This
can be seen as a sort of broadcast that can be used to ensure the lists
associated with two handles have equal numbers of payload IR ops as expected by
many pairwise transform operations.

Introduce an additional "expensive" check that guards against consuming a
handle that is assocaited with the same payload IR operation more than once as
this is likely to lead to double-free or other undesired effects.

Depends On D129110

[ftynse]

Rebase.

[ftynse]

Rebase.

[nicolasvasilache]

LGTM conditioned on a more meaningful example / test (can be in a followup if too annoying now).

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

201

typo: Replicate

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

207

Here the key information is "consumed", can we make it stand out more ?
The linearity properties of the IR are not always trivial to make out from error messages and this error message is tricky, would be great to beef up.

371

Seems to me that consumed relates to free and not to read ?
Will it always be impossible to have free and !read ? (if so please add an assertion)

402

modifies seems unrelated to read, can you either rename the helper or split its uses further between reads and writes/modifies ?

mlir/test/Dialect/Transform/expensive-checks.mlir

57

yes, I would go overboard with the error message here as it is an exception to the linear IR we see in most place.

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

593

Any chance to have a more meaningful example in IR (either here or in the doc) ?
Would be great to link this to the example that prompted you to add this op.
As it stands right now it feels like it is useful, it is well tested in isolation but it is quite unclear how it really improves life.

[ftynse]

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

207

I wouldn't think about the model as global linearity with exceptions. Many transform ops consume their "main" operand handles and produce new handles because they transform the associated payload ops. They will not consume the handles pointing to options, such as tile sizes. A non-negligible amount of transform ops will not consume any of their operand handles. Without the linearity assumption, you need to pay attention to the side effects on each operand, but that makes it simpler IMO, just look at the documentation.

Added the operand number for a more helpful message, but there isn't much more that can be done here.

371

Consumed is read-and-free by convention. Just free sounds pretty much useless, it will mean the transformation drops a handle without doing anything with it.

402

It's hard for me to see how one can modify the payload IR without reading it, even if only for setting the right insertion point. The only situation where this happens is the complete replacement of the top-level payload IR operation (module), which is disallowed elsewhere, so the "write without reading" should never happen.

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

593

There are 7 follow-up patches...

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

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
+
+/// Populates `effects` with the memory effects indicating the operation on the
+/// given handle value:
+///   - consumes = Read + Free,
+///   - produces = Allocate + Write,
+///   - onlyReads = Read.
+void consumesHandle(ValueRange handles,
+                    SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
+void producesHandle(ValueRange handles,
+                    SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
+void onlyReadsHandle(ValueRange handles,
+                     SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
+
+/// Checks whether the transform op consumes the given handle.
+bool isHandleConsumed(Value handle, transform::TransformOpInterface transform);
+
+/// Populates `effects` with the memory effects indicating the access to payload
+/// IR resource.
+void modifiesPayload(SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
+void onlyReadsPayload(SmallVectorImpl<MemoryEffects::EffectInstance> &effects);
+
 } // namespace transform
 } // namespace mlir
 
 template <typename OpTy>
 mlir::DiagnosedSilenceableFailure
 mlir::transform::TransformEachOpTrait<OpTy>::apply(
     TransformResults &transformResults, TransformState &state) {
   using TransformOpType = typename llvm::function_traits<

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

let arguments = (ins
     SymbolRefAttr:$pattern_name);
   let results = (outs
     Res<PDL_Operation, "Handle to the matched Payload IR ops",
         [TransformMappingAlloc, TransformMappingWrite]>:$matched);
 
   let assemblyFormat = "$pattern_name `in` $root attr-dict";
 }
 
+def ReplicateOp : TransformDialectOp<"replicate",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Lists payload ops multiple times in the new handle";
+  let description = [{
+    Produces a new handle associated with a list of payload IR ops that is
+    computed by repeating the list of payload IR ops associated with the
+    operand handle as many times as the "pattern" handle has associated
+    operations. For example, if pattern is associated with [op1, op2] and the
+    operand handle is associated with [op3, op4, op5], the resulting handle
+    will be associated with [op3, op4, op5, op3, op4, op5].
+
+    This transformation is useful to "align" the sizes of payload IR lists
+    before a transformation that expects, e.g., identically-sized lists. For
+    example, a transformation may be parameterized by same notional per-target 
+    size computed at runtime and supplied as another handle, the replication
+    allows this size to be computed only once and used for every target instead
+    of replicating the computation itself.
+
+    Note that it is undesirable to pass a handle with duplicate operations to
+    an operation that consumes the handle. Handle consumption often indicates
+    that the associated payload IR ops are destroyed, so having the same op
+    listed more than once will lead to double-free. Single-operand
+    MergeHandlesOp may be used to deduplicate the associated list of payload IR
+    ops when necessary. Furthermore, a combination of ReplicateOp and
+    MergeHandlesOp can be used to construct arbitrary lists with repetitions.
+  }];
+
+  let arguments = (ins PDL_Operation:$pattern,
+                       Variadic<PDL_Operation>:$handles);
+  let results = (outs Variadic<PDL_Operation>:$replicated);
+  let assemblyFormat =
+    "`num` `(` $pattern `)` $handles "
+    "custom<PDLOpTypedResults>(type($replicated), ref($handles)) attr-dict";
+}
+
 def SequenceOp : TransformDialectOp<"sequence",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getSuccessorEntryOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      OpAsmOpInterface, PossibleTopLevelTransformOpTrait,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {

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

if (Value handle = mapping.reverse.lookup(op))
       return handle;
   }
   return Value();
 }
 
 LogicalResult transform::TransformState::tryEmplaceReverseMapping(
     Mappings &map, Operation *operation, Value handle) {
   auto insertionResult = map.reverse.insert({operation, handle});
-  if (!insertionResult.second) {
+  if (!insertionResult.second && insertionResult.first->second != handle) {
     InFlightDiagnostic diag = operation->emitError()
                               << "operation tracked by two handles";
     diag.attachNote(handle.getLoc()) << "handle";
     diag.attachNote(insertionResult.first->second.getLoc()) << "handle";
     return diag;
   }
   return success();
 }
if (llvm::find_if(effects, consumesTarget) != effects.end())
       recordHandleInvalidation(target);
   }
   return success();
 }
 
 DiagnosedSilenceableFailure
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   LLVM_DEBUG(DBGS() << "applying: " << transform << "\n");
-  if (options.getExpensiveChecksEnabled() &&
-      failed(checkAndRecordHandleInvalidation(transform))) {
+  if (options.getExpensiveChecksEnabled()) {
+    if (failed(checkAndRecordHandleInvalidation(transform)))
       return DiagnosedSilenceableFailure::definiteFailure();
+
+    for (OpOperand &operand : transform->getOpOperands()) {
+      if (!isHandleConsumed(operand.get(), transform))
+        continue;
+
+      DenseSet<Operation *> seen;
+      for (Operation *op : getPayloadOps(operand.get())) {
+        if (!seen.insert(op).second) {
+          DiagnosedSilenceableFailure diag =
+              transform.emitSilenceableError()
+              << "a handle passed as operand #" << operand.getOperandNumber()
+              << " and consumed by this operation points to a payload "
+                 "operation more than once";
+          diag.attachNote(op->getLoc()) << "repeated target op";
+          return diag;
+        }
+      }
+    }
   }
 
   transform::TransformResults results(transform->getNumResults());
   DiagnosedSilenceableFailure result(transform.apply(results, *this));
   if (!result.succeeded())
     return result;
 
   // Remove the mapping for the operand if it is consumed by the operation. This
if (op->getNumOperands() == 0) {
       return diag;
     }
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// Memory effects.
+//===----------------------------------------------------------------------===//
+
+void transform::consumesHandle(
+    ValueRange handles,
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  for (Value handle : handles) {
+    effects.emplace_back(MemoryEffects::Read::get(), handle,
+                         TransformMappingResource::get());
+    effects.emplace_back(MemoryEffects::Free::get(), handle,
+                         TransformMappingResource::get());
+  }
+}
+
+/// Returns `true` if the given list of effects instances contains an instance
+/// with the effect type specified as template parameter.
+template <typename EffectTy>
+static bool hasEffect(ArrayRef<MemoryEffects::EffectInstance> effects) {
+  return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
+    return isa<EffectTy>(effect.getEffect());
+  });
+}
+
+bool transform::isHandleConsumed(Value handle,
+                                 transform::TransformOpInterface transform) {
+  auto iface = cast<MemoryEffectOpInterface>(transform.getOperation());
+  SmallVector<MemoryEffects::EffectInstance> effects;
+  iface.getEffectsOnValue(handle, effects);
+  return hasEffect<MemoryEffects::Read>(effects) &&
+         hasEffect<MemoryEffects::Free>(effects);
+}
+
+void transform::producesHandle(
+    ValueRange handles,
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  for (Value handle : handles) {
+    effects.emplace_back(MemoryEffects::Allocate::get(), handle,
+                         TransformMappingResource::get());
+    effects.emplace_back(MemoryEffects::Write::get(), handle,
+                         TransformMappingResource::get());
+  }
+}
+
+void transform::onlyReadsHandle(
+    ValueRange handles,
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  for (Value handle : handles) {
+    effects.emplace_back(MemoryEffects::Read::get(), handle,
+                         TransformMappingResource::get());
+  }
+}
+
+void transform::modifiesPayload(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
+  effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
+}
+
+void transform::onlyReadsPayload(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
+}
+
+//===----------------------------------------------------------------------===//
 // Generated interface implementation.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.cpp.inc"

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;
 
+static ParseResult parsePDLOpTypedResults(
+    OpAsmParser &parser, SmallVectorImpl<Type> &types,
+    const SmallVectorImpl<OpAsmParser::UnresolvedOperand> &handles) {
+  types.resize(handles.size(), pdl::OperationType::get(parser.getContext()));
+  return success();
+}
+
+static void printPDLOpTypedResults(OpAsmPrinter &, Operation *, TypeRange,
+                                   ValueRange) {}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // PatternApplicatorExtension
 //===----------------------------------------------------------------------===//
 
 namespace {
if (failed(extension->findAllMatches(
       return DiagnosedSilenceableFailure::definiteFailure();
     }
   }
   results.set(getResult().cast<OpResult>(), targets);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
+// ReplicateOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::ReplicateOp::apply(transform::TransformResults &results,
+                              transform::TransformState &state) {
+  unsigned numRepetitions = state.getPayloadOps(getPattern()).size();
+  for (const auto &en : llvm::enumerate(getHandles())) {
+    Value handle = en.value();
+    ArrayRef<Operation *> current = state.getPayloadOps(handle);
+    SmallVector<Operation *> payload;
+    payload.reserve(numRepetitions * current.size());
+    for (unsigned i = 0; i < numRepetitions; ++i)
+      llvm::append_range(payload, current);
+    results.set(getReplicated()[en.index()].cast<OpResult>(), payload);
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::ReplicateOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getPattern(), effects);
+  consumesHandle(getHandles(), effects);
+  producesHandle(getReplicated(), effects);
+}
+
+//===----------------------------------------------------------------------===//
 // SequenceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::SequenceOp::apply(transform::TransformResults &results,
                              transform::TransformState &state) {
   // Map the entry block argument to the list of operations.
   auto scope = state.make_region_scope(*getBodyBlock()->getParent());

mlir/python/mlir/dialects/_transform_ops_ext.py

def __init__(self,
     super().__init__(
         pdl.OperationType.get(),
         _get_op_result_or_value(target),
         _get_symbol_ref_attr(pattern_name),
         loc=loc,
         ip=ip)
 
 
+class ReplicateOp:
+
+  def __init__(self,
+               pattern: Union[Operation, Value],
+               handles: Sequence[Union[Operation, Value]],
+               *,
+               loc=None,
+               ip=None):
+    super().__init__(
+        [pdl.OperationType.get()] * len(handles),
+        _get_op_result_or_value(pattern),
+        [_get_op_result_or_value(h) for h in handles],
+        loc=loc,
+        ip=ip)
+
+
 class SequenceOp:
 
   @overload
   def __init__(self, resultsOrRoot: Sequence[Type],
                optionalRoot: Optional[Union[Operation, Value]]):
     ...
 
   @overload

mlir/test/Dialect/Transform/expensive-checks.mlir

^bb1(%arg1: !pdl.operation):
     // expected-note @below {{other handle}}
     %0 = pdl_match @return in %arg1
     %1 = get_closest_isolated_parent %0
     // expected-error @below {{invalidated the handle to payload operations nested in the payload operation associated with its operand #0}}
     test_consume_operand %1
     test_print_remark_at_operand %0, "remark"
   }
 }
+
+// -----
+
+func.func @func1() {
+  // expected-note @below {{repeated target op}}
+  return
+}
+func.func private @func2()
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @func : benefit(1) {
+    %0 = operands
+    %1 = types
+    %2 = operation "func.func"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+    rewrite %2 with "transform.dialect"
+  }
+  pdl.pattern @return : benefit(1) {
+    %0 = operands
+    %1 = types
+    %2 = operation "func.return"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+    rewrite %2 with "transform.dialect"
+  }
+
+  sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @func in %arg1
+    %1 = pdl_match @return in %arg1
+    %2 = replicate num(%0) %1
+    // expected-error @below {{a handle passed as operand #0 and consumed by this operation points to a payload operation more than once}}
+    test_consume_operand %2
+    test_print_remark_at_operand %0, "remark"
+  }
+}

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

^bb0(%arg0: !pdl.operation):
   }
 
   transform.sequence %arg0 {
   ^bb0(%arg1: !pdl.operation):
     %0 = pdl_match @some in %arg1
     transform.test_mixed_sucess_and_silenceable %0
   }
 }
+
+// -----
+
+module {
+  func.func private @foo()
+  func.func private @bar()
+
+  transform.with_pdl_patterns {
+  ^bb0(%arg0: !pdl.operation):
+    pdl.pattern @func : benefit(1) {
+      %0 = pdl.operands
+      %1 = pdl.types
+      %2 = pdl.operation "func.func"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+      pdl.rewrite %2 with "transform.dialect"
+    }
+
+    transform.sequence %arg0 {
+    ^bb0(%arg1: !pdl.operation):
+      %0 = pdl_match @func in %arg1
+      %1 = replicate num(%0) %arg1
+      // expected-remark @below {{2}}
+      test_print_number_of_associated_payload_ir_ops %1
+      %2 = replicate num(%0) %1
+      // expected-remark @below {{4}}
+      test_print_number_of_associated_payload_ir_ops %2
+    }
+  }
+}

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

 mlir::test::TestMixedSuccessAndSilenceableOp::applyToOne(
     Operation *target, SmallVectorImpl<Operation *> &results,
     transform::TransformState &state) {
   if (target->hasAttr("target_me"))
     return DiagnosedSilenceableFailure::success();
   return emitDefaultSilenceableFailure(target);
 }
 
+DiagnosedSilenceableFailure
+mlir::test::TestPrintNumberOfAssociatedPayloadIROps::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  emitRemark() << state.getPayloadOps(getHandle()).size();
+  return DiagnosedSilenceableFailure::success();
+}
+
+void mlir::test::TestPrintNumberOfAssociatedPayloadIROps::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getHandle(), effects);
+}
+
 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 TestMixedSuccessAndSilenceableOp
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation * target, 
         ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def TestPrintNumberOfAssociatedPayloadIROps
+  : Op<Transform_Dialect, "test_print_number_of_associated_payload_ir_ops",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let arguments = (ins PDL_Operation:$handle);
+  let assemblyFormat = "$handle attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD

mlir/test/python/dialects/transform.py

def testMergeHandlesOp():
   sequence = transform.SequenceOp()
   with InsertionPoint(sequence.body):
     transform.MergeHandlesOp([sequence.bodyTarget])
     transform.YieldOp()
   # CHECK-LABEL: TEST: testMergeHandlesOp
   # CHECK: transform.sequence
   # CHECK: ^{{.*}}(%[[ARG1:.+]]: !pdl.operation):
   # CHECK:   = merge_handles %[[ARG1]]
+
+
+@run
+def testReplicateOp():
+  with_pdl = transform.WithPDLPatternsOp()
+  with InsertionPoint(with_pdl.body):
+    sequence = transform.SequenceOp(with_pdl.bodyTarget)
+    with InsertionPoint(sequence.body):
+      m1 = transform.PDLMatchOp(sequence.bodyTarget, "first")
+      m2 = transform.PDLMatchOp(sequence.bodyTarget, "second")
+      transform.ReplicateOp(m1, [m2])
+      transform.YieldOp()
+  # CHECK-LABEL: TEST: testReplicateOp
+  # CHECK: %[[FIRST:.+]] = pdl_match
+  # CHECK: %[[SECOND:.+]] = pdl_match
+  # CHECK: %{{.*}} = replicate num(%[[FIRST]]) %[[SECOND]]