[mlir][transform] Introduce transform.sequence op

Sequence is an important transform combination primitive that just indicates
transform ops being applied in a row. The simplest version requires fails
immediately if any transformation in the sequence fails. Introducing this
operation allows one to start placing transform IR within other IR.

Depends On D123135

Reviewed By: Mogball, rriddle

Differential Revision: https://reviews.llvm.org/D123664

Author: ftynse

mlir/include/mlir/Dialect/Transform/IR/CMakeLists.txt

@@ -1,8 +1,13 @@
-# The dialect does not have its own ops, so just generate the dialect files.
+# Generate the dialect files from the dialect .td.
+#
+# TODO: Make it possible to use XDialect instead of XOpsDialect in
+# add_mlir_dialect.
 set(LLVM_TARGET_DEFINITIONS TransformDialect.td)
 mlir_tablegen(TransformDialect.h.inc -gen-dialect-decls -dialect=transform)
 mlir_tablegen(TransformDialect.cpp.inc -gen-dialect-defs -dialect=transform)
 add_public_tablegen_target(MLIRTransformDialectIncGen)
 add_dependencies(mlir-headers MLIRTransformDialectIncGen)
 
+add_mlir_dialect(TransformOps transform)
+
 add_mlir_interface(TransformInterfaces)

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

@@ -161,6 +161,7 @@ def Transform_Dialect : Dialect {
 
   let name = "transform";
   let cppNamespace = "::mlir::transform";
+  let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 
   let extraClassDeclaration = [{
     // Make addOperations available to the TransformDialectExtension class.
@@ -172,4 +173,9 @@ def Transform_Dialect : Dialect {
   }];
 }
 
+// Base class for ops that belong to the tranfsorm dialect. Ops defined in
+// extensions of this dialect may also use this.
+class TransformDialectOp<string mnemonic, list<Trait> traits = []>
+    : Op<Transform_Dialect, mnemonic, traits>;
+
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMDIALECT

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

@@ -33,6 +33,14 @@ class TransformOpInterface;
 /// expected to populate the `TransformResults` class instance in order to
 /// update the mapping. The `applyTransform` method takes care of propagating
 /// the state of `TransformResults` into the instance of this class.
+///
+/// When applying transform IR operations with regions, the client is expected
+/// to create a RegionScope RAII object to create a new "stack frame" for
+/// values defined inside the region. The mappings from and to these values will
+/// be automatically dropped when the object goes out of scope, typically at the
+/// end of the "apply" function of the parent operation. If a region contains
+/// blocks with arguments, the client can map those arguments to payload IR ops
+/// using "mapBlockArguments".
 class TransformState {
   /// Mapping between a Value in the transform IR and the corresponding set of
   /// operations in the payload IR.
@@ -42,9 +50,19 @@ class TransformState {
   /// currently associated with.
   using TransformOpReverseMapping = DenseMap<Operation *, Value>;
 
+  /// Bidirectional mappings between transform IR values and payload IR
+  /// operations.
+  struct Mappings {
+    TransformOpMapping direct;
+    TransformOpReverseMapping reverse;
+  };
+
 public:
-  /// Creates a state for the transformation rooted at the given op.
-  explicit TransformState(Operation *root);
+  /// 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 beind transformed, which may or may not contain the
+  /// region with transform ops.
+  TransformState(Region &region, Operation *root);
 
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
@@ -58,10 +76,96 @@ class TransformState {
   /// the state accordingly.
   LogicalResult applyTransform(TransformOpInterface transform);
 
+  /// Records the mapping between a block argument in the transform IR and a
+  /// list of operations in the payload IR. The arguments must be defined in
+  /// blocks of the currently processed transform IR region, typically after a
+  /// region scope is defined.
+  LogicalResult mapBlockArguments(BlockArgument argument,
+                                  ArrayRef<Operation *> operations) {
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+    assert(argument.getParentRegion() == regionStack.back() &&
+           "mapping block arguments from a region other than the active one");
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
+    return setPayloadOps(argument, operations);
+  }
+
+  // Forward declarations to support limited visibility.
+  class RegionScope;
+
+  /// Creates a new region scope for the given region. The region is expected to
+  /// be nested in the currently processed region.
+  // Implementation note: this method is inline but implemented outside of the
+  // class body to comply with visibility and full-declaration requirements.
+  inline RegionScope make_region_scope(Region &region);
+
+  /// A RAII object maintaining a "stack frame" for a transform IR region. When
+  /// applying a transform IR operation that contains a region, the caller is
+  /// expected to create a RegionScope before applying the ops contained in the
+  /// region. This ensures that the mappings between values defined in the
+  /// transform IR region and payload IR operations are cleared when the region
+  /// processing ends; such values cannot be accessed outside the region.
+  class RegionScope {
+  public:
+    /// Forgets the mapping from or to values defined in the associated
+    /// transform IR region.
+    ~RegionScope() {
+      state.mappings.erase(region);
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+      state.regionStack.pop_back();
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
+    }
+
+  private:
+    /// Creates a new scope for mappings between values defined in the given
+    /// transform IR region and payload IR operations.
+    RegionScope(TransformState &state, Region &region)
+        : state(state), region(&region) {
+      auto res = state.mappings.try_emplace(this->region);
+      assert(res.second && "the region scope is already present");
+      (void)res;
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+      assert(state.regionStack.back()->isProperAncestor(&region) &&
+             "scope started at a non-nested region");
+      state.regionStack.push_back(&region);
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
+    }
+
+    /// Back-reference to the transform state.
+    TransformState &state;
+
+    /// The region this scope is associated with.
+    Region *region;
+
+    friend RegionScope TransformState::make_region_scope(Region &);
+  };
+  friend class RegionScope;
+
 private:
   /// Identifier for storing top-level value in the `operations` mapping.
   static constexpr Value kTopLevelValue = Value();
 
+  /// 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");
+    return it->second;
+  }
+
+  /// Returns the mappings frame for the region in which the operation resides.
+  const Mappings &getMapping(Operation *operation) const {
+    return const_cast<TransformState *>(this)->getMapping(operation);
+  }
+  Mappings &getMapping(Operation *operation) {
+    auto it = mappings.find(operation->getParentRegion());
+    assert(it != mappings.end() &&
+           "trying to find a mapping for an operation from an unmapped region");
+    return it->second;
+  }
+
   /// Sets the payload IR ops associated with the given transform IR value.
   /// Fails if this would result in multiple transform IR values with uses
   /// corresponding to the same payload IR ops. For example, a hypothetical
@@ -88,9 +192,19 @@ class TransformState {
   void updatePayloadOps(Value value,
                         function_ref<Operation *(Operation *)> callback);
 
-  /// The mapping between payload IR values and transform IR ops.
-  TransformOpMapping operationMapping;
-  TransformOpReverseMapping reverseMapping;
+  /// The mappings between transform IR values and payload IR ops, aggregated by
+  /// the region in which the transform IR values are defined.
+  llvm::SmallDenseMap<Region *, Mappings> mappings;
+
+  /// The top-level operation that contains all payload IR, typically a module.
+  Operation *topLevel;
+
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+  /// A stack of nested regions that are being processed in the transform IR.
+  /// Each region must be an ancestor of the following regions in this list.
+  /// These are also the keys for "mappings".
+  SmallVector<Region *> regionStack;
+#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 };
 
 /// Local mapping between values defined by a specific op implementing the
@@ -123,6 +237,10 @@ class TransformResults {
   SmallVector<Operation *> operations;
 };
 
+TransformState::RegionScope TransformState::make_region_scope(Region &region) {
+  return RegionScope(*this, region);
+}
+
 } // namespace transform
 } // namespace mlir
 

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

@@ -0,0 +1,20 @@
+//===- TransformDialect.h - Transform dialect operations --------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
+#define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
+
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/IR/TransformOps.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H

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

@@ -0,0 +1,78 @@
+//===- TransformOps.td - Transform dialect operations ------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS
+#define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS
+
+include "mlir/IR/OpAsmInterface.td"
+include "mlir/Dialect/PDL/IR/PDLTypes.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+
+def SequenceOp : TransformDialectOp<"sequence",
+    [DeclareOpInterfaceMethods<TransformOpInterface>, OpAsmOpInterface,
+     SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
+  let summary = "Contains a sequence of other transform ops to apply";
+  let description = [{
+    The transformations indicated by the sequence are applied in order of their
+    appearance. Each value produced by a transformation within the sequence
+    corresponds to an operation or a group of operations in the payload IR.
+    Each value may be used at most once by another transformation operation as
+    the transformation is likely to replace the transformed operation with
+    another operation or a group thereof. In such cases, the transformation
+    operation is expected to produce a new value to denote the newly produced
+    operations that can be transformed further. During application, if any
+    transformation in the sequence fails, the entire sequence fails immediately
+    leaving the payload IR in potentially invalid state, i.e., this operation
+    offers no transformation rollback capabilities.
+
+    The entry block of this operation has a single argument that maps to either
+    the operand if provided or the top-level container operation of the payload
+    IR, typically the root operation of the pass interpreting the transform
+    dialect. Operand omission is only allowed for sequences not contained in
+    another sequence.
+  }];
+
+  let arguments = (ins Optional<PDL_Operation>:$root);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat =
+    "($root^)? attr-dict-with-keyword regions (`:` type($results)^)?";
+
+  let extraClassDeclaration = [{
+    /// Allow the dialect prefix to be omitted.
+    static StringRef getDefaultDialect() { return "transform"; }
+
+    Block *getBodyBlock() {
+      return &getBody().front();
+    }
+  }];
+
+  let hasVerifier = 1;
+}
+
+def YieldOp : TransformDialectOp<"yield", [Terminator]> {
+  let summary = "Yields operation handles from a transform IR region";
+  let description = [{
+    This terminator operation yields operation handles from regions of the
+    transform IR ops back to the containing op. It is not itself associated with
+    any transformation on the payload IR and is used for flow purposes only.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+  let assemblyFormat = "operands attr-dict (`:` type($operands)^)?";
+
+  let builders = [
+    OpBuilder<(ins), [{
+      return build($_builder, $_state, ::mlir::ValueRange());
+    }]>
+  ];
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS

mlir/lib/Dialect/Transform/IR/CMakeLists.txt

@@ -1,11 +1,14 @@
 add_mlir_dialect_library(MLIRTransformDialect
   TransformDialect.cpp
   TransformInterfaces.cpp
+  TransformOps.cpp
 
   DEPENDS
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRPDL
+  MLIRPDLInterp
   )

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

@@ -7,9 +7,15 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
-
-#include "mlir/Dialect/Transform/IR/TransformDialect.cpp.inc"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
 
 using namespace mlir;
 
-void transform::TransformDialect::initialize() {}
+#include "mlir/Dialect/Transform/IR/TransformDialect.cpp.inc"
+
+void transform::TransformDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
+      >();
+}