[mlir] Use semantically readable functions for transform op effects

[ftynse]

Summary

A recent commit introduced helper functions with semantically meaningful names
to populate the lists of memory effects in transform ops, use them whenever
possible.

Depends On D129287

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

 /// cannot modify the Payload IR without reading it first). This is intended
 /// to disallow reordering of Transform IR operations that mutate the Payload IR
 /// while still allowing the reordering of those that only access it.
 struct PayloadIRResource
     : public SideEffects::Resource::Base<PayloadIRResource> {
   StringRef getName() override { return "transform.payload_ir"; }
 };
 
+/// 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);
+
 /// Trait implementing the MemoryEffectOpInterface for operations that "consume"
 /// their operands and produce new results.
 template <typename OpTy>
 class FunctionalStyleTransformOpTrait
     : public OpTrait::TraitBase<OpTy, FunctionalStyleTransformOpTrait> {
 public:
   /// This op "consumes" the operands by reading and freeing then, "produces"
   /// the results by allocating and writing it and reads/writes the payload IR
   /// in the process.
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-    for (Value operand : this->getOperation()->getOperands()) {
-      effects.emplace_back(MemoryEffects::Read::get(), operand,
-                           TransformMappingResource::get());
-      effects.emplace_back(MemoryEffects::Free::get(), operand,
-                           TransformMappingResource::get());
-    }
-    for (Value result : this->getOperation()->getResults()) {
-      effects.emplace_back(MemoryEffects::Allocate::get(), result,
-                           TransformMappingResource::get());
-      effects.emplace_back(MemoryEffects::Write::get(), result,
-                           TransformMappingResource::get());
-    }
-    effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
-    effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
+    consumesHandle(this->getOperation()->getOperands(), effects);
+    producesHandle(this->getOperation()->getResults(), effects);
+    modifiesPayload(effects);
   }
 
   /// Checks that the op matches the expectations of this trait.
   static LogicalResult verifyTrait(Operation *op) {
     if (!op->getName().getInterface<MemoryEffectOpInterface>()) {
       op->emitError()
           << "FunctionalStyleTransformOpTrait should only be attached to ops "
              "that implement MemoryEffectOpInterface";
     }
     return success();
   }
 };
 
 /// Trait implementing the MemoryEffectOpInterface for single-operand
 /// single-result operations that use their operand without consuming and
 /// without modifying the Payload IR to produce a new handle.
 template <typename OpTy>
 class NavigationTransformOpTrait
     : public OpTrait::TraitBase<OpTy, NavigationTransformOpTrait> {
 public:
   /// This op produces handles to the Payload IR without consuming the original
   /// handles and without modifying the IR itself.
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-    effects.emplace_back(MemoryEffects::Read::get(),
-                         this->getOperation()->getOperand(0),
-                         TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Allocate::get(),
-                         this->getOperation()->getResult(0),
-                         TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Write::get(),
-                         this->getOperation()->getResult(0),
-                         TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
+    onlyReadsHandle(this->getOperation()->getOperands(), effects);
+    producesHandle(this->getOperation()->getResults(), effects);
+    onlyReadsPayload(effects);
   }
 
   /// Checks that the op matches the expectation of this trait.
   static LogicalResult verifyTrait(Operation *op) {
     static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
                   "expected single-operand op");
     static_assert(OpTy::template hasTrait<OpTrait::OneResult>(),
                   "expected single-result op");
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/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
   results.push_back(lowTileSize);
   results.push_back(highTileSize);
   results.push_back(splitPoint);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::MultiTileSizesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
-                       transform::TransformMappingResource::get());
-  for (Value result : getResults()) {
-    effects.emplace_back(MemoryEffects::Allocate::get(), result,
-                         transform::TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Write::get(), result,
-                         transform::TransformMappingResource::get());
-  }
-
-  effects.emplace_back(MemoryEffects::Read::get(),
-                       transform::PayloadIRResource::get());
-  effects.emplace_back(MemoryEffects::Write::get(),
-                       transform::PayloadIRResource::get());
+  onlyReadsHandle(getTarget(), effects);
+  producesHandle(getResults(), effects);
+  modifiesPayload(effects);
 }
 
 //===---------------------------------------------------------------------===//
 // PadOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PadOp::applyToOne(linalg::LinalgOp target,
DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
 
   results.set(getFirst().cast<OpResult>(), first);
   results.set(getSecond().cast<OpResult>(), second);
   return DiagnosedSilenceableFailure::success();
 }
 
 void SplitOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  // The target handle is consumed.
-  effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
-                       TransformMappingResource::get());
-  effects.emplace_back(MemoryEffects::Free::get(), getTarget(),
-                       TransformMappingResource::get());
-
-  // The dynamic split point handle is not consumed.
-  if (getDynamicSplitPoint()) {
-    effects.emplace_back(MemoryEffects::Read::get(), getDynamicSplitPoint(),
-                         TransformMappingResource::get());
-  }
-
-  // The resulting handles are produced.
-  for (Value result : getResults()) {
-    effects.emplace_back(MemoryEffects::Allocate::get(), result,
-                         TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Write::get(), result,
-                         TransformMappingResource::get());
-  }
-
-  effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
-  effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
+  consumesHandle(getTarget(), effects);
+  if (getDynamicSplitPoint())
+    onlyReadsHandle(getDynamicSplitPoint(), effects);
+  producesHandle(getResults(), effects);
+  modifiesPayload(effects);
 }
 
 ParseResult SplitOp::parse(OpAsmParser &parser, OperationState &result) {
   OpAsmParser::UnresolvedOperand target, dynamicSplitPoint;
   IntegerAttr staticSplitPoint;
   auto pdlOperationType =
       pdl::OperationType::get(parser.getBuilder().getContext());
   if (parser.parseOperand(target) ||

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

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>
+template <typename EffectTy, typename ResourceTy = SideEffects::DefaultResource>
 static bool hasEffect(ArrayRef<MemoryEffects::EffectInstance> effects) {
   return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
-    return isa<EffectTy>(effect.getEffect());
+    return isa<EffectTy>(effect.getEffect()) &&
+           isa<ResourceTy>(effect.getResource());
   });
 }
 
 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);
+  return hasEffect<MemoryEffects::Read, TransformMappingResource>(effects) &&
+         hasEffect<MemoryEffects::Free, TransformMappingResource>(effects);
 }
 
 void transform::producesHandle(
     ValueRange handles,
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   for (Value handle : handles) {
     effects.emplace_back(MemoryEffects::Allocate::get(), handle,
                          TransformMappingResource::get());

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

transform::MergeHandlesOp::apply(transform::TransformResults &results,
 
   SetVector<Operation *> uniqued(operations.begin(), operations.end());
   results.set(getResult().cast<OpResult>(), uniqued.getArrayRef());
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::MergeHandlesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  for (Value operand : getHandles()) {
-    effects.emplace_back(MemoryEffects::Read::get(), operand,
-                         transform::TransformMappingResource::get());
-    effects.emplace_back(MemoryEffects::Free::get(), operand,
-                         transform::TransformMappingResource::get());
-  }
-  effects.emplace_back(MemoryEffects::Allocate::get(), getResult(),
-                       transform::TransformMappingResource::get());
-  effects.emplace_back(MemoryEffects::Write::get(), getResult(),
-                       transform::TransformMappingResource::get());
+  consumesHandle(getHandles(), effects);
+  producesHandle(getResult(), effects);
 
   // There are no effects on the Payload IR as this is only a handle
   // manipulation.
 }
 
 OpFoldResult transform::MergeHandlesOp::fold(ArrayRef<Attribute> operands) {
   if (getDeduplicate() || getHandles().size() != 1)
     return {};
transform::SequenceOp::apply(transform::TransformResults &results,
   forwardTerminatorOperands(getBodyBlock(), state, results);
   return DiagnosedSilenceableFailure::success();
 }
 
 /// Returns `true` if the given op operand may be consuming the handle value in
 /// the Transform IR. That is, if it may have a Free effect on it.
 static bool isValueUsePotentialConsumer(OpOperand &use) {
   // Conservatively assume the effect being present in absence of the interface.
-  auto memEffectInterface = dyn_cast<MemoryEffectOpInterface>(use.getOwner());
-  if (!memEffectInterface)
+  auto iface = dyn_cast<transform::TransformOpInterface>(use.getOwner());
+  if (!iface)
     return true;
 
-  SmallVector<MemoryEffects::EffectInstance, 2> effects;
-  memEffectInterface.getEffectsOnValue(use.get(), effects);
-  return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
-    return isa<transform::TransformMappingResource>(effect.getResource()) &&
-           isa<MemoryEffects::Free>(effect.getEffect());
-  });
+  return isHandleConsumed(use.get(), iface);
 }
 
 LogicalResult
 checkDoubleConsume(Value value,
                    function_ref<InFlightDiagnostic()> reportError) {
   OpOperand *potentialConsumer = nullptr;
   for (OpOperand &use : value.getUses()) {
     if (!isValueUsePotentialConsumer(use))

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

def TestConsumeOperand : Op<Transform_Dialect, "test_consume_operand",
   let cppNamespace = "::mlir::test";
 }
 
 def TestConsumeOperandIfMatchesParamOrFail
   : Op<Transform_Dialect, "test_consume_operand_if_matches_param_or_fail",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let arguments = (ins 
     Arg<PDL_Operation, "",
-        [TransformMappingWrite, TransformMappingFree]>:$operand,
+        [TransformMappingRead, TransformMappingFree]>:$operand,
     I64Attr:$parameter);
   let assemblyFormat = "$operand `[` $parameter `]` attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestPrintRemarkAtOperandOp
   : Op<Transform_Dialect, "test_print_remark_at_operand",
        [DeclareOpInterfaceMethods<TransformOpInterface>]> {