[mlir] verify that transform ops have memory effects

[ftynse]

Summary

Add a verifier to the TransformOpInterface ensuring that operations
implementing the interface define memory effects on their operands and
results.

Add the missing effects to TileToForeachThreadOp, specifically for
operands that were added at a later version of the op without modifying
getEffects accordingly.

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

 /// to either the list of operations associated with its operand or the root of
 /// the payload IR, depending on what is available in the context.
 LogicalResult
 mapPossibleTopLevelTransformOpBlockArguments(TransformState &state,
                                              Operation *op, Region &region);
 
 /// Verification hook for PossibleTopLevelTransformOpTrait.
 LogicalResult verifyPossibleTopLevelTransformOpTrait(Operation *op);
+
+/// Verification hook for TransformOpInterface.
+LogicalResult verifyTransformOpInterface(Operation *op);
 } // namespace detail
 
 /// This trait is supposed to be attached to Transform dialect operations that
 /// can be standalone top-level transforms. Such operations typically contain
 /// other Transform dialect operations that can be executed following some
 /// control flow logic specific to the current operation. The operations with
 /// this trait are expected to have at least one single-block region with one
 /// argument of PDL Operation type. The operations are also expected to be valid

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

let extraSharedClassDeclaration = [{
     DiagnosedSilenceableFailure emitDefaultSilenceableFailure(
         Operation *target) {
       DiagnosedSilenceableFailure diag = emitSilenceableFailure($_op->getLoc());
       diag << $_op->getName() << " failed to apply";
       diag.attachNote(target->getLoc()) << "when applied to this op";
       return diag;
     }
   }];
+
+  let verify = [{
+    return ::mlir::transform::detail::verifyTransformOpInterface($_op);
+  }];
 }
 
 class TransformTypeInterfaceBase<string cppClass, string cppObjectType>
     : TypeInterface<cppClass> {
   let cppNamespace = "::mlir::transform";
 
   let methods = [
     InterfaceMethod<

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

DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply(
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::TileToForeachThreadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getTileSizes(), effects);
   onlyReadsHandle(getNumThreads(), effects);
+  onlyReadsHandle(getPackedNumThreads(), effects);
+  onlyReadsHandle(getPackedTileSizes(), effects);
   producesHandle(getResults(), effects);
 }
 
 SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedNumThreads() {
   Builder b(getContext());
   return getMixedValues(getStaticNumThreads(), getNumThreads(), b);
 }
 

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, typename ResourceTy = SideEffects::DefaultResource>
-static bool hasEffect(ArrayRef<MemoryEffects::EffectInstance> effects) {
+template <typename EffectTy, typename ResourceTy, typename Range>
+static bool hasEffect(Range &&effects) {
   return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
     return isa<EffectTy>(effect.getEffect()) &&
            isa<ResourceTy>(effect.getResource());
   });
 }
 
 bool transform::isHandleConsumed(Value handle,
                                  transform::TransformOpInterface transform) {

 }
 
 void transform::onlyReadsPayload(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
 }
 
 //===----------------------------------------------------------------------===//
+// Utilities for TransformOpInterface.
+//===----------------------------------------------------------------------===//
+
+LogicalResult transform::detail::verifyTransformOpInterface(Operation *op) {
+  auto iface = cast<MemoryEffectOpInterface>(op);
+  SmallVector<MemoryEffects::EffectInstance> effects;
+  iface.getEffects(effects);
+
+  auto effectsOn = [&](Value value) {
+    return llvm::make_filter_range(
+        effects, [value](const MemoryEffects::EffectInstance &instance) {
+          return instance.getValue() == value;
+        });
+  };
+
+  for (OpOperand &operand : op->getOpOperands()) {
+    auto range = effectsOn(operand.get());
+    if (range.empty()) {
+      InFlightDiagnostic diag =
+          op->emitError() << "TransformOpInterface requires memory effects "
+                             "on operands to be specified";
+      diag.attachNote() << "no effects specified for operand #"
+                        << operand.getOperandNumber();
+      return diag;
+    }
+  }
+  for (OpResult result : op->getResults()) {
+    auto range = effectsOn(result);
+    if (!::hasEffect<MemoryEffects::Allocate, TransformMappingResource>(
+            range)) {
+      InFlightDiagnostic diag =
+          op->emitError() << "TransformOpInterface requires 'allocate' memory "
+                             "effect to be specified for results";
+      diag.attachNote() << "no 'allocate' effect specified for result #"
+                        << result.getResultNumber();
+      return diag;
+    }
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // Entry point.
 //===----------------------------------------------------------------------===//
 
 LogicalResult transform::applyTransforms(Operation *payloadRoot,
                                          TransformOpInterface transform,
                                          const TransformOptions &options) {
 #ifndef NDEBUG
   if (!transform->hasTrait<PossibleTopLevelTransformOpTrait>() ||

mlir/test/Dialect/Transform/ops-invalid.mlir

^bb0(%arg0: !pdl.operation):
   // expected-note @below {{used here as operand #0}}
   transform.foreach %0 : !pdl.operation {
   ^bb1(%arg1: !pdl.operation):
     transform.test_consume_operand %arg1
   }
   // expected-note @below {{used here as operand #0}}
   transform.test_consume_operand %0
 }
+
+// -----
+
+transform.sequence failures(suppress) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{TransformOpInterface requires memory effects on operands to be specified}}
+  // expected-note @below {{no effects specified for operand #0}}
+  transform.test_required_memory_effects %arg0 : (!transform.any_op) -> !transform.any_op
+}
+
+// -----
+
+transform.sequence failures(suppress) {
+^bb0(%arg0: !transform.any_op):
+  // expected-error @below {{TransformOpInterface requires 'allocate' memory effect to be specified for results}}
+  // expected-note @below {{no 'allocate' effect specified for result #0}}
+  transform.test_required_memory_effects %arg0 {has_operand_effect} : (!transform.any_op) -> !transform.any_op
+}

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

 DiagnosedSilenceableFailure mlir::test::TestProduceNullPayloadOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   SmallVector<Operation *, 1> null({nullptr});
   results.set(getOut().cast<OpResult>(), null);
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestProduceNullParamOp::getEffects(
-    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::producesHandle(getOut(), effects);
+}
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceNullParamOp::apply(transform::TransformResults &results,
                                           transform::TransformState &state) {
   results.setParams(getOut().cast<OpResult>(), Attribute());
   return DiagnosedSilenceableFailure::success();
 }
 
+void mlir::test::TestRequiredMemoryEffectsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  if (getHasOperandEffect())
+    transform::consumesHandle(getIn(), effects);
+
+  if (getHasResultEffect())
+    transform::producesHandle(getOut(), effects);
+  else
+    transform::onlyReadsHandle(getOut(), effects);
+}
+
+DiagnosedSilenceableFailure mlir::test::TestRequiredMemoryEffectsOp::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  results.set(getOut().cast<OpResult>(), state.getPayloadOps(getIn()));
+  return DiagnosedSilenceableFailure::success();
+}
+
 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 TestProduceNullParamOp
   : Op<Transform_Dialect, "test_produce_null_param",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let results = (outs TransformParamTypeInterface:$out);
   let assemblyFormat = "attr-dict `:` type($out)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestRequiredMemoryEffectsOp
+  : Op<Transform_Dialect, "test_required_memory_effects",
+      [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins TransformHandleTypeInterface:$in,
+                       UnitAttr:$has_operand_effect,
+                       UnitAttr:$has_result_effect);
+  let results = (outs TransformHandleTypeInterface:$out);
+  let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
+  let cppNamespace = "::mlir::test";
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD