[mlir] add transform dialect entry point

[ftynse]

Summary

Introduce transform::applyTransforms as a top-level entry point to the
Transform dialect-driven transformation infrastructure, by analogy with
applyFull/PartialConversion. Clients are expected to use this function
and no longer need to maintain the transformation state. Make the
constructor of the TransformState private for that purpose.

[springerm]

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

461

nit: Can be written as getNumOperands != 0

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

 
 include "mlir/IR/OpBase.td"
 
 def Transform_Dialect : Dialect {
   let summary = "Fine-grain transformation control dialect";
   let description = [{
     ## Disclaimer
 
-    ** Proceed with care: not ready for general use. **
+    **This dialect is actively developed and may change frequently.**
 
-    This dialect is evolving rapidly and may change on a very short notice. To
-    decrease the maintenance burden and churn, only a few in-tree use cases are
-    currently supported in the main tree:
+    To decrease the maintenance burden and churn, please post a description of
+    the intended use case on the MLIR forum. A few in-tree use cases are
+    currently supported:
 
       - high-level transformations on "structured ops" (i.e. ops that operate on
         chunks of data in a way that can be decomposed into operations on
         smaller chunks of data and control flow) in Linalg, Tensor and Vector
-        dialects.
+        dialects;
+      - loop transformations in the SCF dialect.
       
-    *Please post a description of the intended use case on the MLIR forum and
-    wait for confirmation.*
 
     ## Overview
 
     This dialect provides operations that can be used to control transformation
     of the IR using a different portion of the IR. It refers to the IR being
     transformed as payload IR, and to the IR guiding the transformation as
     transform IR.
 
let description = [{
 
     Overall, Transform IR ops are expected to be contained in a single top-level
     op. Such top-level ops specify how to apply the transformations described
     by the operations they contain, e.g., `transform.sequence` executes
     transformations one by one and fails if any of them fails. Such ops are
     expected to have the `PossibleTopLevelTransformOpTrait` and may be used
     without arguments.
 
+    A program transformation expressed using the Transform dialect can be
+    programmatically triggered by calling:
+
+    ```c++
+    LogicalResult transform::applyTransforms(Operation *payloadRoot,
+                                             TransformOpInterface transform,
+                                             const TransformOptions &options);
+    ```
+
+    that applies the transformations specified by the top-level `transform` to
+    payload IR contained in `payloadRoot`.
+
     ## Dialect Extension Mechanism
 
     This dialect is designed to be extensible, that is, clients of this dialect
     are allowed to inject additional operations into this dialect using the
     `TransformDialectExtension` mechanism. This allows the dialect to avoid a
     dependency on the implementation of the transformation as well as to avoid
     introducing dialect-specific transform dialects. In the example above,
     the operations may have been injected by a notional `loop` dialect rather

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

public:
 
   /// Returns true if the expensive checks are requested.
   bool getExpensiveChecksEnabled() const { return expensiveChecksEnabled; }
 
 private:
   bool expensiveChecksEnabled = true;
 };
 
+/// Entry point to the Transform dialect infrastructure. Applies the
+/// transformation specified by `transform` to payload IR contained in
+/// `payloadRoot`. The `transform` operation may contain other operations that
+/// will be executed following the internal logic of the operation. It must
+/// have the `PossibleTopLevelTransformOp` trait and not have any operands.
+/// This function internally keeps track of the transformation state.
+LogicalResult
+applyTransforms(Operation *payloadRoot, TransformOpInterface transform,
+                const TransformOptions &options = TransformOptions());
+
 /// The state maintained across applications of various ops implementing the
 /// TransformOpInterface. The operations implementing this interface and the
 /// surrounding structure are referred to as transform IR. The operations to
 /// which transformations apply are referred to as payload IR. The state thus
 /// contains the many-to-many mapping between values defined in the transform IR
 /// ops and payload IR ops. The "expensive-checks" option can be passed to
 /// the constructor at transformation execution time that transform IR values
 /// used as operands by a transform IR operation are not associated with
class TransformState {
 
   /// Bidirectional mappings between transform IR values and payload IR
   /// operations.
   struct Mappings {
     TransformOpMapping direct;
     TransformOpReverseMapping reverse;
   };
 
-public:
-  /// Creates a state for transform ops living in the given region. The parent
-  /// operation of the region. The second argument points to the root operation
-  /// in the payload IR being transformed, which may or may not contain the
-  /// region with transform ops. Additional options can be provided through the
-  /// trailing configuration object.
-  TransformState(Region &region, Operation *root,
-                 const TransformOptions &options = TransformOptions());
+  friend LogicalResult applyTransforms(Operation *payloadRoot,
+                                       TransformOpInterface transform,
+                                       const TransformOptions &options);
 
+public:
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
   Operation *getTopLevel() const;
 
   /// Returns the list of ops that the given transform IR value corresponds to.
   /// This is helpful for transformations that apply to a particular handle.
   ArrayRef<Operation *> getPayloadOps(Value value) const;
 
static_assert(
         "only an class derived from TransformState::Extension is allowed here");
     extensions.erase(TypeID::get<Ty>());
   }
 
 private:
   /// Identifier for storing top-level value in the `operations` mapping.
   static constexpr Value kTopLevelValue = Value();
 
+  /// Creates a state for transform ops living in the given region. The second
+  /// argument points to the root operation in the payload IR being transformed,
+  /// which may or may not contain the region with transform ops. Additional
+  /// options can be provided through the trailing configuration object.
+  TransformState(Region *region, Operation *payloadRoot,
+                 const TransformOptions &options = TransformOptions());
+
   /// Returns the mappings frame for the reigon in which the value is defined.
   const Mappings &getMapping(Value value) const {
     return const_cast<TransformState *>(this)->getMapping(value);
   }
   Mappings &getMapping(Value value) {
     auto it = mappings.find(value.getParentRegion());
     assert(it != mappings.end() &&
            "trying to find a mapping for a value from an unmapped region");

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

 //===- TransformInterfaces.cpp - Transform Dialect Interfaces -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #define DEBUG_TYPE "transform-dialect"
 #define DEBUG_PRINT_AFTER_ALL "transform-dialect-print-top-level-after-all"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TransformState
 //===----------------------------------------------------------------------===//
 
 constexpr const Value transform::TransformState::kTopLevelValue;
 
-transform::TransformState::TransformState(Region &region, Operation *root,
+transform::TransformState::TransformState(Region *region,
+                                          Operation *payloadRoot,
                                           const TransformOptions &options)
-    : topLevel(root), options(options) {
-  auto result = mappings.try_emplace(&region);
+    : topLevel(payloadRoot), options(options) {
+  auto result = mappings.try_emplace(region);
   assert(result.second && "the region scope is already present");
   (void)result;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
-  regionStack.push_back(&region);
+  regionStack.push_back(region);
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 }
 
 Operation *transform::TransformState::getTopLevel() const { return topLevel; }
 
 ArrayRef<Operation *>
 transform::TransformState::getPayloadOps(Value value) const {
   const TransformOpMapping &operationMapping = getMapping(value).direct;

 }
 
 void transform::onlyReadsPayload(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
 }
 
 //===----------------------------------------------------------------------===//
+// Entry point.
+//===----------------------------------------------------------------------===//
+
+LogicalResult transform::applyTransforms(Operation *payloadRoot,
+                                         TransformOpInterface transform,
+                                         const TransformOptions &options) {
+#ifndef NDEBUG
+  if (!transform->hasTrait<PossibleTopLevelTransformOpTrait>() ||
+      transform->getNumOperands() != 0) {
+    transform->emitError()
+        << "expected transform to start at the top-level transform op";
+    llvm::report_fatal_error("could not run transforms",
+                             /*gen_crash_diag=*/false);
+  }
+#endif // NDEBUG
+
+  TransformState state(transform->getParentRegion(), payloadRoot, options);
+  return state.applyTransform(transform).checkAndReport();
+}
+
+//===----------------------------------------------------------------------===//
 // Generated interface implementation.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.cpp.inc"

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

 // RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
   // expected-remark @below {{applying transformation}}
   transform.test_transform_op
+}
 
 // -----
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
   %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
   // expected-remark @below {{succeeded}}
   transform.test_consume_operand_if_matches_param_or_fail %0[42]
+}
 
 // -----
 
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
   %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
   // expected-error @below {{expected the operand to be associated with 21 got 42}}
   transform.test_consume_operand_if_matches_param_or_fail %0[21]
+}
 
 // -----
 
 // It is okay to have multiple handles to the same payload op as long
 // as only one of them is consumed. The expensive checks mode is necessary
 // to detect double-consumption.
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
   %0 = transform.test_produce_param_or_forward_operand 42 { foo = "bar" }
   %1 = transform.test_copy_payload %0
   // expected-remark @below {{succeeded}}
   transform.test_consume_operand_if_matches_param_or_fail %0[42]
+}
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     // expected-remark @below {{applying transformation "a"}}

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

public:
   }
 
   StringRef getDescription() const override {
     return "apply transform dialect operations one by one";
   }
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
-    transform::TransformState state(
-        module.getBodyRegion(), module,
-        transform::TransformOptions().enableExpensiveChecks(
-            enableExpensiveChecks));
     for (auto op :
          module.getBody()->getOps<transform::TransformOpInterface>()) {
-      if (failed(state.applyTransform(op).checkAndReport()))
+      if (failed(transform::applyTransforms(
+              module, op,
+              transform::TransformOptions().enableExpensiveChecks(
+                  enableExpensiveChecks))))
         return signalPassFailure();
     }
   }
 
   Option<bool> enableExpensiveChecks{
       *this, "enable-expensive-checks", llvm::cl::init(false),
       llvm::cl::desc("perform expensive checks to better report errors in the "
                      "transform IR")};