diff --git a/mlir/include/mlir/Dialect/Affine/CMakeLists.txt b/mlir/include/mlir/Dialect/Affine/CMakeLists.txt --- 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 --- /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 --- /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]> { + 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, + "1">:$num_loops); + let results = (outs TransformTypeInterface:$parent); + + let assemblyFormat = + "$target attr-dict `:` functional-type(operands, results)"; +} + + +def AffineLoopUnrollOp : Op { + 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:$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 --- /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 --- 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 @@ 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 --- 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 --- /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 parents; + for (Operation *target : state.getPayloadOps(getTarget())) { + AffineForOp loop; + Operation *current = target; + for (unsigned i = 0, e = getNumLoops(); i < e; ++i) { + loop = current->getParentOfType(); + if (!loop) { + DiagnosedSilenceableFailure diag = emitSilenceableError() + << "could not find an '" + << AffineForOp::getOperationName() + << "' parent"; + diag.attachNote(target->getLoc()) << "target op"; + results.set(getResult().cast(), {}); + return diag; + } + current = loop; + } + parents.insert(loop); + } + results.set(getResult().cast(), parents.getArrayRef()); + return DiagnosedSilenceableFailure::success(); +} + + +//===----------------------------------------------------------------------===// +// LoopUnrollOp +//===----------------------------------------------------------------------===// + +DiagnosedSilenceableFailure +transform::AffineLoopUnrollOp::applyToOne(AffineForOp target, + SmallVector &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(); + + 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(); +} diff --git a/mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt new file mode 100644 --- /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 --- /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"> +} +