[mlir] add interchange, pad and scalarize to structured transform dialect

Add ops to the structured transform extension of the transform dialect that
perform interchange, padding and scalarization on structured ops. Along with
tiling that is already defined, this provides a minimal set of transformations
necessary to build vectorizable code for a single structured op.

Define two helper traits: one that implements TransformOpInterface by applying
a function to each payload op independently and another that provides a simple
"functional-style" producer/consumer list of memory effects for the transform
ops.

Reviewed By: nicolasvasilache

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

Author: ftynse

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h

@@ -9,9 +9,16 @@
 #ifndef MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H
 #define MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H
 
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/OpImplementation.h"
 
+namespace mlir {
+namespace linalg {
+class LinalgOp;
+} // namespace linalg
+} // namespace mlir
+
 //===----------------------------------------------------------------------===//
 // Linalg Transform Operations
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

@@ -16,6 +16,81 @@ include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
+def InterchangeOp : Op<Transform_Dialect, "structured.interchange",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+    TransformOpInterface, TransformEachOpTrait]> {
+  let description = [{
+    Interchanges the iterators of the operations pointed to by the target handle
+    using the iterator interchange attribute.
+  }];
+
+  let arguments =
+    (ins PDL_Operation:$target,
+         DefaultValuedAttr<I64ArrayAttr, "{}">:$iterator_interchange);
+  let results = (outs PDL_Operation:$transformed);
+
+  let assemblyFormat = "$target attr-dict";
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
+        ::mlir::linalg::LinalgOp target);
+  }];
+}
+
+def PadOp : Op<Transform_Dialect, "structured.pad",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     TransformOpInterface, TransformEachOpTrait]> {
+  let description = [{
+    Pads the operations pointed to by the target handle using the options
+    provides as operation attributes.
+  }];
+
+  let arguments =
+    (ins PDL_Operation:$target,
+         DefaultValuedAttr<ArrayAttr, "{}">:$padding_values,
+         DefaultValuedAttr<I64ArrayAttr, "{}">:$padding_dimensions,
+         DefaultValuedAttr<I64ArrayAttr, "{}">:$pack_paddings,
+         DefaultValuedAttr<I64ArrayAttr, "{}">:$hoist_paddings,
+         DefaultValuedAttr<
+          TypedArrayAttrBase<I64ArrayAttr, "array of arrays of i64">,
+          "{}">:$transpose_paddings);
+  let results = (outs PDL_Operation:$transformed);
+
+  let assemblyFormat = "$target attr-dict";
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
+        ::mlir::linalg::LinalgOp target);
+  }];
+}
+
+def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     TransformOpInterface, TransformEachOpTrait]> {
+  let description = [{
+    Indicates that ops of a specific kind in the given function should be
+    scalarized (i.e. their dynamic dimensions tiled by 1).
+
+    This operation returns the tiled op but not the loops.
+
+    We make this design choice because it is hard to know ahead of time the
+    number of loops that will be produced (it depends on the number of dynamic
+    dimensions after multiple transformations have been applied).
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
+        ::mlir::linalg::LinalgOp target);
+  }];
+}
+
 def TileOp : Op<Transform_Dialect, "structured.tile",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {

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

@@ -397,6 +397,31 @@ class PossibleTopLevelTransformOpTrait
   }
 };
 
+/// Trait implementing the TransformOpInterface for operations applying a
+/// transformation to a single operation handle and producing a single operation
+/// handle. The op must implement a method with one of the following signatures:
+///   - FailureOr<convertible-to-Operation*> applyToOne(OpTy)
+///   - LogicalResult applyToOne(OpTy)
+/// to perform a transformation that is applied in turn to all payload IR
+/// operations that correspond to the handle of the transform IR operation.
+/// In the functions above, OpTy is either Operation * or a concrete payload IR
+/// Op class that the transformation is applied to (NOT the class of the
+/// transform IR op). The op is expected to have one operand and zero or one
+/// results.
+template <typename OpTy>
+class TransformEachOpTrait
+    : public OpTrait::TraitBase<OpTy, TransformEachOpTrait> {
+public:
+  /// Calls `applyToOne` for every payload operation associated with the operand
+  /// of this transform IR op. If `applyToOne` returns ops, associates them with
+  /// the result of this transform op.
+  LogicalResult apply(TransformResults &transformResults,
+                      TransformState &state);
+
+  /// Checks that the op matches the expectations of this trait.
+  static LogicalResult verifyTrait(Operation *op);
+};
+
 /// Side effect resource corresponding to the mapping between Transform IR
 /// values and Payload IR operations. An Allocate effect from this resource
 /// means creating a new mapping entry, it is always accompanied by a Write
@@ -426,9 +451,150 @@ struct PayloadIRResource
   StringRef getName() override { return "transform.payload_ir"; }
 };
 
+/// Trait implementing the MemoryEffectOpInterface for single-operand
+/// single-result operations that "consume" their operand and produce a new
+/// result.
+template <typename OpTy>
+class FunctionalStyleTransformOpTrait
+    : public OpTrait::TraitBase<OpTy, FunctionalStyleTransformOpTrait> {
+public:
+  /// This op "consumes" the operand by reading and freeing it, "produces" the
+  /// result by allocating and writing it and reads/writes the payload IR in the
+  /// process.
+  void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+    effects.emplace_back(MemoryEffects::Read::get(),
+                         this->getOperation()->getOperand(0),
+                         TransformMappingResource::get());
+    effects.emplace_back(MemoryEffects::Free::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());
+    effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
+  }
+
+  /// Checks that the op matches the expectations 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");
+    if (!op->getName().getInterface<MemoryEffectOpInterface>()) {
+      op->emitError()
+          << "FunctionalStyleTransformOpTrait should only be attached to ops "
+             "that implement MemoryEffectOpInterface";
+    }
+    return success();
+  }
+};
+
 } // namespace transform
 } // namespace mlir
 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h.inc"
 
+namespace mlir {
+namespace transform {
+namespace detail {
+/// Appends `result` to the vector assuming it corresponds to the success state
+/// in `FailureOr<convertible-to-Operation*>`. If `result` is just a
+/// `LogicalResult`, does nothing.
+template <typename Ty>
+std::enable_if_t<std::is_same<Ty, LogicalResult>::value, LogicalResult>
+appendTransformResultToVector(Ty result,
+                              SmallVectorImpl<Operation *> &results) {
+  return result;
+}
+template <typename Ty>
+std::enable_if_t<!std::is_same<Ty, LogicalResult>::value, LogicalResult>
+appendTransformResultToVector(Ty result,
+                              SmallVectorImpl<Operation *> &results) {
+  static_assert(
+      std::is_convertible<typename Ty::value_type, Operation *>::value,
+      "expected transform function to return operations");
+  if (failed(result))
+    return failure();
+
+  results.push_back(*result);
+  return success();
+}
+
+/// Applies a one-to-one transform to each of the given targets. Puts the
+/// results of transforms, if any, in `results` in the same order. Fails if any
+/// of the application fails. Individual transforms must be callable with
+/// one of the following signatures:
+///   - FailureOr<convertible-to-Operation*>(OpTy)
+///   - LogicalResult(OpTy)
+/// where OpTy is either
+///   - Operation *, in which case the transform is always applied;
+///   - a concrete Op class, in which case a check is performed whether
+///   `targets` contains operations of the same class and a failure is reported
+///   if it does not.
+template <typename FnTy>
+LogicalResult applyTransformToEach(ArrayRef<Operation *> targets,
+                                   SmallVectorImpl<Operation *> &results,
+                                   FnTy transform) {
+  using OpTy = typename llvm::function_traits<FnTy>::template arg_t<0>;
+  static_assert(std::is_convertible<OpTy, Operation *>::value,
+                "expected transform function to take an operation");
+  using RetTy = typename llvm::function_traits<FnTy>::result_t;
+  static_assert(std::is_convertible<RetTy, LogicalResult>::value,
+                "expected transform function to return LogicalResult or "
+                "FailureOr<convertible-to-Operation*>");
+  for (Operation *target : targets) {
+    auto specificOp = dyn_cast<OpTy>(target);
+    if (!specificOp)
+      return failure();
+
+    auto result = transform(specificOp);
+    if (failed(appendTransformResultToVector(result, results)))
+      return failure();
+  }
+  return success();
+}
+} // namespace detail
+} // namespace transform
+} // namespace mlir
+
+template <typename OpTy>
+mlir::LogicalResult mlir::transform::TransformEachOpTrait<OpTy>::apply(
+    TransformResults &transformResults, TransformState &state) {
+  using TransformOpType = typename llvm::function_traits<
+      decltype(&OpTy::applyToOne)>::template arg_t<0>;
+  ArrayRef<Operation *> targets =
+      state.getPayloadOps(this->getOperation()->getOperand(0));
+  SmallVector<Operation *> results;
+  if (failed(detail::applyTransformToEach(
+          targets, results, [&](TransformOpType specificOp) {
+            return static_cast<OpTy *>(this)->applyToOne(specificOp);
+          })))
+    return failure();
+  if (OpTy::template hasTrait<OpTrait::OneResult>()) {
+    transformResults.set(
+        this->getOperation()->getResult(0).template cast<OpResult>(), results);
+  }
+  return success();
+}
+
+template <typename OpTy>
+mlir::LogicalResult
+mlir::transform::TransformEachOpTrait<OpTy>::verifyTrait(Operation *op) {
+  static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
+                "expected single-operand op");
+  static_assert(OpTy::template hasTrait<OpTrait::OneResult>() ||
+                    OpTy::template hasTrait<OpTrait::ZeroResults>(),
+                "expected zero- or single-result op");
+  if (!op->getName().getInterface<TransformOpInterface>()) {
+    return op->emitError() << "TransformEachOpTrait should only be attached to "
+                              "ops that implement TransformOpInterface";
+  }
+
+  return success();
+}
+
 #endif // DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H

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

@@ -49,4 +49,13 @@ def TransformOpInterface : OpInterface<"TransformOpInterface"> {
   ];
 }
 
+def FunctionalStyleTransformOpTrait
+    : NativeOpTrait<"FunctionalStyleTransformOpTrait"> {
+  let cppNamespace = "::mlir::transform";
+}
+
+def TransformEachOpTrait : NativeOpTrait<"TransformEachOpTrait"> {
+  let cppNamespace = "::mlir::transform";
+}
+
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_INTERFACES_TD