[mlir][transform] Add transform op for loop-invariant subset hoisting

Author: matthias-springer

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

@@ -2247,56 +2247,6 @@ def ConvertConv2DToImg2ColOp : Op<Transform_Dialect,
  }];
 }
 
-//===----------------------------------------------------------------------===//
-// HoistRedundantTensorSubsetsOp
-//===----------------------------------------------------------------------===//
-
-def HoistRedundantTensorSubsetsOp :
- Op<Transform_Dialect, "structured.hoist_redundant_tensor_subsets",
- [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
- TransformEachOpTrait,
- TransformOpInterface,
- ReportTrackingListenerFailuresOpTrait]> {
- let description = [{
- Hoists supported tensor subset extract/insert operation pairs out of
- immediately enclosing loop iteratively, if the following conditions
- are true:
- 1. The 2 ops access the same tensor subset.
- 2. All operands are invariant under the enclosing loop.
-
- The supported subset extract/insert operation pairs currently comprise:
- - tensor.extract_slice / tensor.insert_slice
- - vector.transfer_read / vector.transfer_write on tensors
-
- Only scf.for loops are currently supported.
-
- When applied to:
- 1. an scf.for loop, hoist out of this loop only.
- 2. a non-loop op, apply hoisting to all the contained loop ops.
-
- #### Return modes:
-
- The operation always succeeds and returns nothing.
- }];
-
- let arguments = (ins TransformHandleTypeInterface:$target);
- let results = (outs);
-
- let assemblyFormat = [{
- $target
- attr-dict
- `:` functional-type(operands, results)
- }];
-
- let extraClassDeclaration = [{
- ::mlir::DiagnosedSilenceableFailure applyToOne(
- ::mlir::transform::TransformRewriter &rewriter,
- ::mlir::Operation *target,
- ::mlir::transform::ApplyToEachResultList &results,
- ::mlir::transform::TransformState &state);
- }];
-}
-
 //===----------------------------------------------------------------------===//
 // InsertSliceToCopyOp
 //===----------------------------------------------------------------------===//

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

@@ -1,3 +1,4 @@
 add_subdirectory(IR)
+add_subdirectory(LoopExtension)
 add_subdirectory(PDLExtension)
 add_subdirectory(Transforms)

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

@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS LoopExtensionOps.td)
+mlir_tablegen(LoopExtensionOps.h.inc -gen-op-decls)
+mlir_tablegen(LoopExtensionOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRTransformDialectLoopExtensionOpsIncGen)
+
+add_mlir_doc(LoopExtensionOps LoopExtensionOps Dialects/ -gen-op-doc)

mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtension.h

@@ -0,0 +1,16 @@
+//===- LoopExtension.h - Loop extension for Transform dialect ---*- 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
+//
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+class DialectRegistry;
+
+namespace transform {
+/// Registers the loop extension of the Transform dialect in the given registry.
+void registerLoopExtension(DialectRegistry &dialectRegistry);
+} // namespace transform
+} // namespace mlir

mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.h

@@ -0,0 +1,23 @@
+//===- LoopExtensionOps.h - Loop ext. for Transform dialect -----*- 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_LOOPEXTENSION_LOOPEXTENSIONOPS_H
+#define MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS_H
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS_H

mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.td

@@ -0,0 +1,76 @@
+//===- LoopExtensionOps.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_LOOPEXTENSION_LOOPEXTENSIONOPS
+#define MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def HoistLoopInvariantSubsetsOp
+ : TransformDialectOp<"loop.hoist_loop_invariant_subsets",
+ [TransformOpInterface, TransformEachOpTrait,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+ ReportTrackingListenerFailuresOpTrait]> {
+ let summary = "Hoist loop invariant subset ops";
+ let description = [{
+ This transform hoists loop-invariant subset ops out of the targeted
+ loop-like op. It looks for matching subset extraction/insertion op pairs and
+ hoists them. The loop body operates on a newly introduced region iter_arg.
+
+ Subset ops are hoisted only from the targeted op. If subset ops should be
+ hoisted from an entire loop nest, this transformation must be applied to
+ each loop-like op of the loop nest, starting with the innermost loop and
+ ending with the outermost loop.
+
+ Example:
+ ```
+ %r = scf.for ... iter_args(%t = %a) -> (tensor<?xf32>) {
+ %0 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
+ %1 = "test.foo"(%0) : (tensor<5xf32>) -> (tensor<5xf32>)
+ %2 = tensor.insert_slice %1 into %t[0][5][1]
+ : tensor<5xf32> into tensor<?xf32>
+ scf.yield %2 : tensor<?xf32>
+ }
+ ```
+ Is transformed to:
+ ```
+ %0 = tensor.extract_slice %a[0][5][1] : tensor<?xf32> to tensor<5xf32>
+ %new_loop:2 = scf.for ... iter_args(%t = %a, %h = %0) -> (tensor<?xf32>) {
+ %1 = "test.foo"(%h) : (tensor<5xf32>) -> (tensor<5xf32>)
+ scf.yield %t, %2 : tensor<?xf32>, tensor<5xf32>
+ }
+ %r = tensor.insert_slice %new_loop#1 into %new_loop#0
+ : tensor<5xf32> into tensor<?xf32>
+ ```
+
+ Subset ops are hoisted only if there are no conflicting subset ops. E.g.,
+ if there were a second overlapping extraction in the above example, no ops
+ could be hoisted safely.
+
+ This transform reads the target handle and modifies the payload. This
+ transform does not invalidate any handles, but loop-like ops are replaced
+ with new loop-like ops when a subset op is hoisted. The transform rewriter
+ updates all handles accordingly.
+ }];
+
+ let arguments = (ins TransformHandleTypeInterface:$target);
+ let results = (outs);
+ let assemblyFormat = "$target attr-dict `:` type($target)";
+
+ let extraClassDeclaration = [{
+ ::mlir::DiagnosedSilenceableFailure applyToOne(
+ ::mlir::transform::TransformRewriter &rewriter,
+ ::mlir::LoopLikeOpInterface loopLikeOp,
+ ::mlir::transform::ApplyToEachResultList &results,
+ ::mlir::transform::TransformState &state);
+ }];
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS

mlir/include/mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.td

@@ -1,4 +1,4 @@
-//===- TransformOps.td - Transform dialect operations ------*- tablegen -*-===//
+//===- PDLExtensionOps.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.

mlir/include/mlir/InitAllExtensions.h

@@ -34,6 +34,7 @@
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+#include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
 #include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
@@ -74,6 +75,7 @@ inline void registerAllExtensions(DialectRegistry &registry) {
  scf::registerTransformDialectExtension(registry);
  sparse_tensor::registerTransformDialectExtension(registry);
  tensor::registerTransformDialectExtension(registry);
+ transform::registerLoopExtension(registry);
  transform::registerPDLExtension(registry);
  vector::registerTransformDialectExtension(registry);
 

mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h

@@ -18,6 +18,7 @@ namespace mlir {
 class LoopLikeOpInterface;
 class Operation;
 class Region;
+class RewriterBase;
 class Value;
 
 /// Given a list of regions, perform loop-invariant code motion. An operation is
@@ -108,7 +109,8 @@ size_t moveLoopInvariantCode(LoopLikeOpInterface loopLike);
 /// %r = tensor.insert_slice %new_loop#1 into %new_loop#0
 /// : tensor<5xf32> into tensor<?xf32>
 /// ```
-LoopLikeOpInterface hoistLoopInvariantSubsets(LoopLikeOpInterface loopLike);
+LoopLikeOpInterface hoistLoopInvariantSubsets(RewriterBase &rewriter,
+ LoopLikeOpInterface loopLike);
 
 } // end namespace mlir
 

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

@@ -3163,35 +3163,6 @@ DiagnosedSilenceableFailure transform::ConvertConv2DToImg2ColOp::applyToOne(
  return DiagnosedSilenceableFailure::success();
 }
 
-//===----------------------------------------------------------------------===//
-// HoistRedundantTensorSubsetsOp
-//===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure
-transform::HoistRedundantTensorSubsetsOp::applyToOne(
- transform::TransformRewriter &rewriter, Operation *target,
- transform::ApplyToEachResultList &results,
- transform::TransformState &state) {
- auto forOp = dyn_cast<scf::ForOp>(target);
- if (forOp) {
- linalg::hoistRedundantSubsetExtractInsert(rewriter, forOp);
- return DiagnosedSilenceableFailure::success();
- }
-
- // TODO: walking in some reverse / inside-out order would be more efficient
- // and would capture more cases.
- target->walk([&](scf::ForOp forOp) {
- hoistRedundantSubsetExtractInsert(rewriter, forOp);
- });
- return DiagnosedSilenceableFailure::success();
-}
-
-void transform::HoistRedundantTensorSubsetsOp::getEffects(
- SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
- transform::onlyReadsHandle(getTarget(), effects);
- transform::modifiesPayload(effects);
-}
-
 //===----------------------------------------------------------------------===//
 // InsertSliceToCopyOp
 //===----------------------------------------------------------------------===//