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-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file defines Transform dialect extension operations used in the
10// Chapter 2 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/Dialect/Transform/Interfaces/TransformInterfaces.h"
20#include "mlir/IR/DialectRegistry.h"
21#include "mlir/IR/Operation.h"
22#include "mlir/Interfaces/SideEffectInterfaces.h"
23#include "mlir/Support/LLVM.h"
24#include "llvm/ADT/SmallVector.h"
25#include "llvm/ADT/StringRef.h"
26
27// Define a new transform dialect extension. This uses the CRTP idiom to
28// identify extensions.
29class MyExtension
30 : public ::mlir::transform::TransformDialectExtension<MyExtension> {
31public:
32 // The extension must derive the base constructor.
33 using Base::Base;
34
35 // This function initializes the extension, similarly to `initialize` in
36 // dialect definitions. List individual operations and dependent dialects
37 // here.
38 void init();
39};
40
41void MyExtension::init() {
42 // Similarly to dialects, an extension can declare a dependent dialect. This
43 // dialect will be loaded along with the extension and, therefore, along with
44 // the Transform dialect. Only declare as dependent the dialects that contain
45 // the attributes or types used by transform operations. Do NOT declare as
46 // dependent the dialects produced during the transformation.
47 // declareDependentDialect<MyDialect>();
48
49 // When transformations are applied, they may produce new operations from
50 // previously unloaded dialects. Typically, a pass would need to declare
51 // itself dependent on the dialects containing such new operations. To avoid
52 // confusion with the dialects the extension itself depends on, the Transform
53 // dialects differentiates between:
54 // - dependent dialects, which are used by the transform operations, and
55 // - generated dialects, which contain the entities (attributes, operations,
56 // types) that may be produced by applying the transformation even when
57 // not present in the original payload IR.
58 // In the following chapter, we will be add operations that generate function
59 // calls and structured control flow operations, so let's declare the
60 // corresponding dialects as generated.
61 declareGeneratedDialect<::mlir::scf::SCFDialect>();
62 declareGeneratedDialect<::mlir::func::FuncDialect>();
63
64 // Finally, we register the additional transform operations with the dialect.
65 // List all operations generated from ODS. This call will perform additional
66 // checks that the operations implement the transform and memory effect
67 // interfaces required by the dialect interpreter and assert if they do not.
68 registerTransformOps<
69#define GET_OP_LIST
70#include "MyExtension.cpp.inc"
71 >();
72}
73
74#define GET_OP_CLASSES
75#include "MyExtension.cpp.inc"
76
77static void updateCallee(mlir::func::CallOp call, llvm::StringRef newTarget) {
78 call.setCallee(newTarget);
79}
80
81// Implementation of our transform dialect operation.
82// This operation returns a tri-state result that can be one of:
83// - success when the transformation succeeded;
84// - definite failure when the transformation failed in such a way that
85// following transformations are impossible or undesirable, typically it could
86// have left payload IR in an invalid state; it is expected that a diagnostic
87// is emitted immediately before returning the definite error;
88// - silenceable failure when the transformation failed but following
89// transformations are still applicable, typically this means a precondition
90// for the transformation is not satisfied and the payload IR has not been
91// modified. The silenceable failure additionally carries a Diagnostic that
92// can be emitted to the user.
93::mlir::DiagnosedSilenceableFailure mlir::transform::ChangeCallTargetOp::apply(
94 // The rewriter that should be used when modifying IR.
95 ::mlir::transform::TransformRewriter &rewriter,
96 // The list of payload IR entities that will be associated with the
97 // transform IR values defined by this transform operation. In this case, it
98 // can remain empty as there are no results.
99 ::mlir::transform::TransformResults &results,
100 // The transform application state. This object can be used to query the
101 // current associations between transform IR values and payload IR entities.
102 // It can also carry additional user-defined state.
103 ::mlir::transform::TransformState &state) {
104
105 // First, we need to obtain the list of payload operations that are associated
106 // with the operand handle.
107 auto payload = state.getPayloadOps(getCall());
108
109 // Then, we iterate over the list of operands and call the actual IR-mutating
110 // function. We also check the preconditions here.
111 for (Operation *payloadOp : payload) {
112 auto call = dyn_cast<::mlir::func::CallOp>(payloadOp);
113 if (!call) {
114 DiagnosedSilenceableFailure diag =
115 emitSilenceableError() << "only applies to func.call payloads";
116 diag.attachNote(payloadOp->getLoc()) << "offending payload";
117 return diag;
118 }
119
120 updateCallee(call, getNewTarget());
121 }
122
123 // If everything went well, return success.
124 return DiagnosedSilenceableFailure::success();
125}
126
127void mlir::transform::ChangeCallTargetOp::getEffects(
128 ::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {
129 // Indicate that the `call` handle is only read by this operation because the
130 // associated operation is not erased but rather modified in-place, so the
131 // reference to it remains valid.
132 onlyReadsHandle(getCall(), effects);
133
134 // Indicate that the payload is modified by this operation.
135 modifiesPayload(effects);
136}
137
138void registerMyExtension(::mlir::DialectRegistry &registry) {
139 registry.addExtensions<MyExtension>();
140}
141

source code of mlir/examples/transform/Ch2/lib/MyExtension.cpp