mlir/examples/transform/Ch2/lib/MyExtension.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-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/TransformInterfaces.h"
  19 #include "mlir/Dialect/Transform/IR/TransformTypes.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.
  29 class MyExtension
  30     : public ::mlir::transform::TransformDialectExtension<MyExtension> {
  31 public:
  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
  41 void 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
  77 static 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
 127 void 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
 138 void registerMyExtension(::mlir::DialectRegistry &registry) {
 139   registry.addExtensions<MyExtension>();
 140 }