225 lines · cpp
1//===-- MyExtension.cpp - Transform dialect tutorial ----------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file defines Transform dialect extension operations used in the10// Chapter 3 of the Transform dialect tutorial.11//12//===----------------------------------------------------------------------===//13 14#include "MyExtension.h"15#include "mlir/Dialect/Func/IR/FuncOps.h"16#include "mlir/Dialect/SCF/IR/SCF.h"17#include "mlir/Dialect/Transform/IR/TransformDialect.h"18#include "mlir/Dialect/Transform/IR/TransformTypes.h"19#include "mlir/IR/DialectImplementation.h"20#include "mlir/Interfaces/CallInterfaces.h"21#include "llvm/ADT/TypeSwitch.h"22 23#define GET_TYPEDEF_CLASSES24#include "MyExtensionTypes.cpp.inc"25 26#define GET_OP_CLASSES27#include "MyExtension.cpp.inc"28 29//===---------------------------------------------------------------------===//30// MyExtension31//===---------------------------------------------------------------------===//32 33// Define a new transform dialect extension. This uses the CRTP idiom to34// identify extensions.35class MyExtension36 : public ::mlir::transform::TransformDialectExtension<MyExtension> {37public:38 // The TypeID of this extension.39 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MyExtension)40 41 // The extension must derive the base constructor.42 using Base::Base;43 44 // This function initializes the extension, similarly to `initialize` in45 // dialect definitions. List individual operations and dependent dialects46 // here.47 void init();48};49 50void MyExtension::init() {51 // Similarly to dialects, an extension can declare a dependent dialect. This52 // dialect will be loaded along with the extension and, therefore, along with53 // the Transform dialect. Only declare as dependent the dialects that contain54 // the attributes or types used by transform operations. Do NOT declare as55 // dependent the dialects produced during the transformation.56 // declareDependentDialect<MyDialect>();57 58 // When transformations are applied, they may produce new operations from59 // previously unloaded dialects. Typically, a pass would need to declare60 // itself dependent on the dialects containing such new operations. To avoid61 // confusion with the dialects the extension itself depends on, the Transform62 // dialects differentiates between:63 // - dependent dialects, which are used by the transform operations, and64 // - generated dialects, which contain the entities (attributes, operations,65 // types) that may be produced by applying the transformation even when66 // not present in the original payload IR.67 // In the following chapter, we will be add operations that generate function68 // calls and structured control flow operations, so let's declare the69 // corresponding dialects as generated.70 declareGeneratedDialect<::mlir::scf::SCFDialect>();71 declareGeneratedDialect<::mlir::func::FuncDialect>();72 73 // Register the additional transform dialect types with the dialect. List all74 // types generated from ODS.75 registerTypes<76#define GET_TYPEDEF_LIST77#include "MyExtensionTypes.cpp.inc"78 >();79 80 // ODS generates these helpers for type printing and parsing, but the81 // Transform dialect provides its own support for types supplied by the82 // extension. Reference these functions to avoid a compiler warning.83 (void)&generatedTypeParser;84 (void)&generatedTypePrinter;85 86 // Finally, we register the additional transform operations with the dialect.87 // List all operations generated from ODS. This call will perform additional88 // checks that the operations implement the transform and memory effect89 // interfaces required by the dialect interpreter and assert if they do not.90 registerTransformOps<91#define GET_OP_LIST92#include "MyExtension.cpp.inc"93 >();94}95 96//===---------------------------------------------------------------------===//97// ChangeCallTargetOp98//===---------------------------------------------------------------------===//99 100static void updateCallee(mlir::func::CallOp call, llvm::StringRef newTarget) {101 call.setCallee(newTarget);102}103 104// Implementation of our transform dialect operation.105// This operation returns a tri-state result that can be one of:106// - success when the transformation succeeded;107// - definite failure when the transformation failed in such a way that108// following109// transformations are impossible or undesirable, typically it could have left110// payload IR in an invalid state; it is expected that a diagnostic is emitted111// immediately before returning the definite error;112// - silenceable failure when the transformation failed but following113// transformations114// are still applicable, typically this means a precondition for the115// transformation is not satisfied and the payload IR has not been modified.116// The silenceable failure additionally carries a Diagnostic that can be emitted117// to the user.118::mlir::DiagnosedSilenceableFailure119mlir::transform::ChangeCallTargetOp::applyToOne(120 // The rewriter that should be used when modifying IR.121 ::mlir::transform::TransformRewriter &rewriter,122 // The single payload operation to which the transformation is applied.123 ::mlir::func::CallOp call,124 // The payload IR entities that will be appended to lists associated with125 // the results of this transform operation. This list contains one entry per126 // result.127 ::mlir::transform::ApplyToEachResultList &results,128 // The transform application state. This object can be used to query the129 // current associations between transform IR values and payload IR entities.130 // It can also carry additional user-defined state.131 ::mlir::transform::TransformState &state) {132 133 // Dispatch to the actual transformation.134 updateCallee(call, getNewTarget());135 136 // If everything went well, return success.137 return DiagnosedSilenceableFailure::success();138}139 140void mlir::transform::ChangeCallTargetOp::getEffects(141 ::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {142 // Indicate that the `call` handle is only read by this operation because the143 // associated operation is not erased but rather modified in-place, so the144 // reference to it remains valid.145 onlyReadsHandle(getCallMutable(), effects);146 147 // Indicate that the payload is modified by this operation.148 modifiesPayload(effects);149}150 151//===---------------------------------------------------------------------===//152// CallToOp153//===---------------------------------------------------------------------===//154 155static mlir::Operation *replaceCallWithOp(mlir::RewriterBase &rewriter,156 mlir::CallOpInterface call) {157 // Construct an operation from an unregistered dialect. This is discouraged158 // and is only used here for brevity of the overall example.159 mlir::OperationState state(call.getLoc(), "my.mm4");160 state.types.assign(call->result_type_begin(), call->result_type_end());161 state.operands.assign(call->operand_begin(), call->operand_end());162 163 mlir::Operation *replacement = rewriter.create(state);164 rewriter.replaceOp(call, replacement->getResults());165 return replacement;166}167 168// See above for the signature description.169mlir::DiagnosedSilenceableFailure mlir::transform::CallToOp::applyToOne(170 mlir::transform::TransformRewriter &rewriter, mlir::CallOpInterface call,171 mlir::transform::ApplyToEachResultList &results,172 mlir::transform::TransformState &state) {173 174 // Dispatch to the actual transformation.175 Operation *replacement = replaceCallWithOp(rewriter, call);176 177 // Associate the payload operation produced by the rewrite with the result178 // handle of this transform operation.179 results.push_back(replacement);180 181 // If everything went well, return success.182 return DiagnosedSilenceableFailure::success();183}184 185//===---------------------------------------------------------------------===//186// CallOpInterfaceHandleType187//===---------------------------------------------------------------------===//188 189// The interface declares this method to verify constraints this type has on190// payload operations. It returns the now familiar tri-state result.191mlir::DiagnosedSilenceableFailure192mlir::transform::CallOpInterfaceHandleType::checkPayload(193 // Location at which diagnostics should be emitted.194 mlir::Location loc,195 // List of payload operations that are about to be associated with the196 // handle that has this type.197 llvm::ArrayRef<mlir::Operation *> payload) const {198 199 // All payload operations are expected to implement CallOpInterface, check200 // this.201 for (Operation *op : payload) {202 if (llvm::isa<mlir::CallOpInterface>(op))203 continue;204 205 // By convention, these verifiers always emit a silenceable failure since206 // they are checking a precondition.207 DiagnosedSilenceableFailure diag =208 emitSilenceableError(loc)209 << "expected the payload operation to implement CallOpInterface";210 diag.attachNote(op->getLoc()) << "offending operation";211 return diag;212 }213 214 // If everything is okay, return success.215 return DiagnosedSilenceableFailure::success();216}217 218//===---------------------------------------------------------------------===//219// Extension registration220//===---------------------------------------------------------------------===//221 222void registerMyExtension(::mlir::DialectRegistry ®istry) {223 registry.addExtensions<MyExtension>();224}225