[Mlir-commits] [mlir] 9a78d06 - [MLIR][Transform] Transform dialect support for affine
Lorenzo Chelini
llvmlistbot at llvm.org
Sun Nov 20 23:56:55 PST 2022
Author: Amy Wang
Date: 2022-11-21T08:53:59+01:00
New Revision: 9a78d06f84f5b6c810613e1fef44101ccc38b553
URL: https://github.com/llvm/llvm-project/commit/9a78d06f84f5b6c810613e1fef44101ccc38b553
DIFF: https://github.com/llvm/llvm-project/commit/9a78d06f84f5b6c810613e1fef44101ccc38b553.diff
LOG: [MLIR][Transform] Transform dialect support for affine
Create the directory and file structures to support transform ops for
the affine dialect in the future. Enable affine.unroll and affine.get_parent_for
transform operations in this first batch of check in.
Reviewed By: ftynse
Differential Revision: https://reviews.llvm.org/D137997
Added:
mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.h
mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.td
mlir/include/mlir/Dialect/Affine/TransformOps/CMakeLists.txt
mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt
mlir/test/Dialect/Affine/transform-ops.mlir
Modified:
mlir/include/mlir/Dialect/Affine/CMakeLists.txt
mlir/include/mlir/InitAllDialects.h
mlir/lib/Dialect/Affine/CMakeLists.txt
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/CMakeLists.txt b/mlir/include/mlir/Dialect/Affine/CMakeLists.txt
index 131101878c7e0..fe1b372f6e03d 100644
--- a/mlir/include/mlir/Dialect/Affine/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/Affine/CMakeLists.txt
@@ -1,4 +1,5 @@
add_subdirectory(IR)
+add_subdirectory(TransformOps)
set(LLVM_TARGET_DEFINITIONS Passes.td)
mlir_tablegen(Passes.h.inc -gen-pass-decls -name Affine)
diff --git a/mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.h b/mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.h
new file mode 100644
index 0000000000000..c1d67c7869646
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.h
@@ -0,0 +1,37 @@
+//===- AffineTransformOps.h - Affine transformation ops ---------*- 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_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
+#define MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/IR/OpImplementation.h"
+
+namespace mlir {
+class AffineForOp;
+namespace func {
+class FuncOp;
+} // namespace func
+namespace affine {
+class ForOp;
+} // namespace affine
+} // namespace mlir
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h.inc"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace affine {
+void registerTransformDialectExtension(DialectRegistry ®istry);
+} // namespace affine
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
diff --git a/mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.td b/mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.td
new file mode 100644
index 0000000000000..dc59435a7896d
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.td
@@ -0,0 +1,78 @@
+//===- AffineTransformOps.td - Affine transformation ops ---*- 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 AFFINE_TRANSFORM_OPS
+#define AFFINE_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformEffects.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+def Transform_AffineForOp : Transform_ConcreteOpType<"affine.for">;
+
+def AffineGetParentForOp : Op<Transform_Dialect, "affine.get_parent_for", [
+ NavigationTransformOpTrait, MemoryEffectsOpInterface,
+ DeclareOpInterfaceMethods<TransformOpInterface>]> {
+ let summary =
+ "Gets a handle to the parent 'affine.for' loop of the given operation";
+ let description = [{
+ Produces a handle to the n-th (default 1) parent `affine.for` loop for each
+ Payload IR operation associated with the operand. Fails if such a loop
+ cannot be found. The list of operations associated with the handle contains
+ parent operations in the same order as the list associated with the operand,
+ except for operations that are parents to more than one input which are only
+ present once.
+ }];
+
+ let arguments =
+ (ins TransformTypeInterface:$target,
+ DefaultValuedAttr<ConfinedAttr<I64Attr, [IntPositive]>,
+ "1">:$num_loops);
+ let results = (outs TransformTypeInterface:$parent);
+
+ let assemblyFormat =
+ "$target attr-dict `:` functional-type(operands, results)";
+}
+
+
+def AffineLoopUnrollOp : Op<Transform_Dialect, "affine.unroll", [
+ FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+ TransformOpInterface, TransformEachOpTrait]> {
+ let summary = "Unrolls the given loop with the given unroll factor";
+ let description = [{
+ Unrolls each loop associated with the given handle to have up to the given
+ number of loop body copies per iteration. If the unroll factor is larger
+ than the loop trip count, the latter is used as the unroll factor instead.
+
+ #### Return modes
+
+ This operation ignores non-affine::For ops and drops them in the return.
+ If all the operations referred to by the `target` PDLOperation unroll
+ properly, the transform succeeds. Otherwise the transform silently fails.
+
+ Does not return handles as the operation may result in the loop being
+ removed after a full unrolling.
+ }];
+
+ let arguments = (ins Transform_AffineForOp:$target,
+ ConfinedAttr<I64Attr, [IntPositive]>:$factor);
+
+ let assemblyFormat = "$target attr-dict `:` type($target)";
+
+ let extraClassDeclaration = [{
+ ::mlir::DiagnosedSilenceableFailure applyToOne(
+ ::mlir::AffineForOp target,
+ ::llvm::SmallVector<::mlir::Operation *> & results,
+ ::mlir::transform::TransformState & state);
+ }];
+}
+
+#endif // Affine_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/Affine/TransformOps/CMakeLists.txt b/mlir/include/mlir/Dialect/Affine/TransformOps/CMakeLists.txt
new file mode 100644
index 0000000000000..c743f5c0766ba
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Affine/TransformOps/CMakeLists.txt
@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS AffineTransformOps.td)
+mlir_tablegen(AffineTransformOps.h.inc -gen-op-decls)
+mlir_tablegen(AffineTransformOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRAffineTransformOpsIncGen)
+
+add_mlir_doc(AffineTransformOps AffineLoopTransformOps Dialects/ -gen-op-doc)
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index c82a6ac2f3f75..c8d742e1be21a 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -17,6 +17,7 @@
#include "mlir/Dialect/AMDGPU/AMDGPUDialect.h"
#include "mlir/Dialect/AMX/AMXDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
@@ -119,6 +120,7 @@ inline void registerAllDialects(DialectRegistry ®istry) {
memref::registerTransformDialectExtension(registry);
scf::registerTransformDialectExtension(registry);
gpu::registerTransformDialectExtension(registry);
+ affine::registerTransformDialectExtension(registry);
// Register all external models.
arith::registerBufferizableOpInterfaceExternalModels(registry);
diff --git a/mlir/lib/Dialect/Affine/CMakeLists.txt b/mlir/lib/Dialect/Affine/CMakeLists.txt
index 35c4201f21454..a981040c761da 100644
--- a/mlir/lib/Dialect/Affine/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/CMakeLists.txt
@@ -1,4 +1,5 @@
add_subdirectory(Analysis)
add_subdirectory(IR)
add_subdirectory(Transforms)
+add_subdirectory(TransformOps)
add_subdirectory(Utils)
diff --git a/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp b/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
new file mode 100644
index 0000000000000..7c32166fd6ad0
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
@@ -0,0 +1,99 @@
+//=== AffineTransformOps.cpp - Implementation of Affine transformation ops ===//
+//
+// 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/Affine/TransformOps/AffineTransformOps.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+
+using namespace mlir;
+
+namespace {
+/// A simple pattern rewriter that implements no special logic.
+class SimpleRewriter : public PatternRewriter {
+public:
+ SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// AffineGetParentForOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::AffineGetParentForOp::apply(transform::TransformResults &results,
+ transform::TransformState &state) {
+ SetVector<Operation *> parents;
+ for (Operation *target : state.getPayloadOps(getTarget())) {
+ AffineForOp loop;
+ Operation *current = target;
+ for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
+ loop = current->getParentOfType<AffineForOp>();
+ if (!loop) {
+ DiagnosedSilenceableFailure diag = emitSilenceableError()
+ << "could not find an '"
+ << AffineForOp::getOperationName()
+ << "' parent";
+ diag.attachNote(target->getLoc()) << "target op";
+ results.set(getResult().cast<OpResult>(), {});
+ return diag;
+ }
+ current = loop;
+ }
+ parents.insert(loop);
+ }
+ results.set(getResult().cast<OpResult>(), parents.getArrayRef());
+ return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
+// LoopUnrollOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::AffineLoopUnrollOp::applyToOne(AffineForOp target,
+ SmallVector<Operation *> &results,
+ transform::TransformState &state) {
+ if (failed(loopUnrollByFactor(target, getFactor()))) {
+ Diagnostic diag(target->getLoc(), DiagnosticSeverity::Note);
+ diag << "op failed to unroll";
+ return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+ }
+ return DiagnosedSilenceableFailure(success());
+}
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+class AffineTransformDialectExtension
+ : public transform::TransformDialectExtension<
+ AffineTransformDialectExtension> {
+public:
+ using Base::Base;
+
+ void init() {
+ declareGeneratedDialect<AffineDialect>();
+
+ registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.cpp.inc"
+ >();
+ }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.cpp.inc"
+
+void mlir::affine::registerTransformDialectExtension(
+ DialectRegistry ®istry) {
+ registry.addExtensions<AffineTransformDialectExtension>();
+}
diff --git a/mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt
new file mode 100644
index 0000000000000..740563af5afed
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt
@@ -0,0 +1,20 @@
+add_mlir_dialect_library(MLIRAffineTransformOps
+ AffineTransformOps.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine/TransformOps
+
+ DEPENDS
+ MLIRAffineTransformOpsIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRAffineDialect
+ MLIRFuncDialect
+ MLIRIR
+ MLIRPDLDialect
+ MLIRAffineDialect
+ MLIRAffineTransforms
+ MLIRAffineUtils
+ MLIRTransformDialect
+ MLIRVectorDialect
+)
diff --git a/mlir/test/Dialect/Affine/transform-ops.mlir b/mlir/test/Dialect/Affine/transform-ops.mlir
new file mode 100644
index 0000000000000..0a122092ac5e4
--- /dev/null
+++ b/mlir/test/Dialect/Affine/transform-ops.mlir
@@ -0,0 +1,67 @@
+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: @get_parent_for_op
+func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
+ // expected-remark @below {{first loop}}
+ affine.for %i = %arg0 to %arg1 {
+ // expected-remark @below {{second loop}}
+ affine.for %j = %arg0 to %arg1 {
+ // expected-remark @below {{third loop}}
+ affine.for %k = %arg0 to %arg1 {
+ arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+ %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+ // CHECK: = transform.affine.get_parent_for
+ %1 = transform.affine.get_parent_for %0 : (!pdl.operation) -> !transform.op<"affine.for">
+ %2 = transform.affine.get_parent_for %0 { num_loops = 2 } : (!pdl.operation) -> !transform.op<"affine.for">
+ %3 = transform.affine.get_parent_for %0 { num_loops = 3 } : (!pdl.operation) -> !transform.op<"affine.for">
+ transform.test_print_remark_at_operand %1, "third loop" : !transform.op<"affine.for">
+ transform.test_print_remark_at_operand %2, "second loop" : !transform.op<"affine.for">
+ transform.test_print_remark_at_operand %3, "first loop" : !transform.op<"affine.for">
+}
+
+// -----
+
+func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
+ // expected-note @below {{target op}}
+ arith.addi %arg0, %arg1 : index
+ return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+ %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+ // expected-error @below {{could not find an 'affine.for' parent}}
+ %1 = transform.affine.get_parent_for %0 : (!pdl.operation) -> !transform.op<"affine.for">
+}
+
+// -----
+
+func.func @loop_unroll_op() {
+ %c0 = arith.constant 0 : index
+ %c42 = arith.constant 42 : index
+ %c5 = arith.constant 5 : index
+ // CHECK: affine.for %[[I:.+]] =
+ // expected-remark @below {{affine for loop}}
+ affine.for %i = %c0 to %c42 {
+ // CHECK-COUNT-4: arith.addi
+ arith.addi %i, %i : index
+ }
+ return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+ %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+ %1 = transform.affine.get_parent_for %0 : (!pdl.operation) -> !transform.op<"affine.for">
+ transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
+ transform.affine.unroll %1 { factor = 4 } : !transform.op<"affine.for">
+}
+
More information about the Mlir-commits
mailing list