llvm/mlir/examples/transform/Ch3/lib/MyExtension.cpp

//===-- MyExtension.cpp - Transform dialect tutorial ----------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file defines Transform dialect extension operations used in the
// Chapter 3 of the Transform dialect tutorial.
//
//===----------------------------------------------------------------------===//

#include "MyExtension.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/IR/TransformTypes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "llvm/ADT/TypeSwitch.h"

#define GET_TYPEDEF_CLASSES
#include "MyExtensionTypes.cpp.inc"

#define GET_OP_CLASSES
#include "MyExtension.cpp.inc"

//===---------------------------------------------------------------------===//
// MyExtension
//===---------------------------------------------------------------------===//

// Define a new transform dialect extension. This uses the CRTP idiom to
// identify extensions.
class MyExtension
    : public ::mlir::transform::TransformDialectExtension<MyExtension> { … };

void MyExtension::init() { … }

//===---------------------------------------------------------------------===//
// ChangeCallTargetOp
//===---------------------------------------------------------------------===//

static void updateCallee(mlir::func::CallOp call, llvm::StringRef newTarget) { … }

// Implementation of our transform dialect operation.
// This operation returns a tri-state result that can be one of:
// - success when the transformation succeeded;
// - definite failure when the transformation failed in such a way that
// following
//   transformations are impossible or undesirable, typically it could have left
//   payload IR in an invalid state; it is expected that a diagnostic is emitted
//   immediately before returning the definite error;
// - silenceable failure when the transformation failed but following
// transformations
//   are still applicable, typically this means a precondition for the
//   transformation is not satisfied and the payload IR has not been modified.
// The silenceable failure additionally carries a Diagnostic that can be emitted
// to the user.
::mlir::DiagnosedSilenceableFailure
mlir::transform::ChangeCallTargetOp::applyToOne(
    // The rewriter that should be used when modifying IR.
    ::mlir::transform::TransformRewriter &rewriter,
    // The single payload operation to which the transformation is applied.
    ::mlir::func::CallOp call,
    // The payload IR entities that will be appended to lists associated with
    // the results of this transform operation. This list contains one entry per
    // result.
    ::mlir::transform::ApplyToEachResultList &results,
    // The transform application state. This object can be used to query the
    // current associations between transform IR values and payload IR entities.
    // It can also carry additional user-defined state.
    ::mlir::transform::TransformState &state) { … }

void mlir::transform::ChangeCallTargetOp::getEffects(
    ::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) { … }

//===---------------------------------------------------------------------===//
// CallToOp
//===---------------------------------------------------------------------===//

static mlir::Operation *replaceCallWithOp(mlir::RewriterBase &rewriter,
                                          mlir::CallOpInterface call) { … }

// See above for the signature description.
mlir::DiagnosedSilenceableFailure mlir::transform::CallToOp::applyToOne(
    mlir::transform::TransformRewriter &rewriter, mlir::CallOpInterface call,
    mlir::transform::ApplyToEachResultList &results,
    mlir::transform::TransformState &state) { … }

//===---------------------------------------------------------------------===//
// CallOpInterfaceHandleType
//===---------------------------------------------------------------------===//

// The interface declares this method to verify constraints this type has on
// payload operations. It returns the now familiar tri-state result.
mlir::DiagnosedSilenceableFailure
mlir::transform::CallOpInterfaceHandleType::checkPayload(
    // Location at which diagnostics should be emitted.
    mlir::Location loc,
    // List of payload operations that are about to be associated with the
    // handle that has this type.
    llvm::ArrayRef<mlir::Operation *> payload) const { … }

//===---------------------------------------------------------------------===//
// Extension registration
//===---------------------------------------------------------------------===//

void registerMyExtension(::mlir::DialectRegistry &registry) { … }