207 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 4 of the Transform dialect tutorial.11//12//===----------------------------------------------------------------------===//13 14#include "MyExtension.h"15#include "mlir/Dialect/Transform/IR/TransformDialect.h"16#include "llvm/Support/DebugLog.h"17 18#define DEBUG_TYPE "transform-matcher"19 20#define GET_OP_CLASSES21#include "MyExtension.cpp.inc"22 23//===---------------------------------------------------------------------===//24// MyExtension25//===---------------------------------------------------------------------===//26 27// Define a new transform dialect extension. This uses the CRTP idiom to28// identify extensions.29class MyExtension30 : public ::mlir::transform::TransformDialectExtension<MyExtension> {31public:32 // The TypeID of this extension.33 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MyExtension)34 35 // The extension must derive the base constructor.36 using Base::Base;37 38 // This function initializes the extension, similarly to `initialize` in39 // dialect definitions. List individual operations and dependent dialects40 // here.41 void init();42};43 44void MyExtension::init() {45 // Register the additional match operations with the dialect similarly to46 // other transform operations. List all operations generated from ODS. This47 // call will perform additional checks that the operations implement the48 // transform and memory effect interfaces required by the dialect interpreter49 // and assert if they do not.50 registerTransformOps<51#define GET_OP_LIST52#include "MyExtension.cpp.inc"53 >();54}55 56//===---------------------------------------------------------------------===//57// HasOperandSatisfyingOp58//===---------------------------------------------------------------------===//59 60/// Returns `true` if both types implement one of the interfaces provided as61/// template parameters.62template <typename... Tys>63static bool implementSameInterface(mlir::Type t1, mlir::Type t2) {64 return ((llvm::isa<Tys>(t1) && llvm::isa<Tys>(t2)) || ... || false);65}66 67/// Returns `true` if both types implement one of the transform dialect68/// interfaces.69static bool implementSameTransformInterface(mlir::Type t1, mlir::Type t2) {70 return implementSameInterface<71 mlir::transform::TransformHandleTypeInterface,72 mlir::transform::TransformParamTypeInterface,73 mlir::transform::TransformValueHandleTypeInterface>(t1, t2);74}75 76// Matcher ops implement `apply` similarly to other transform ops. They are not77// expected to modify payload, but use the tri-state result to signal failure or78// success to match, as well as potential irrecoverable errors.79mlir::DiagnosedSilenceableFailure80mlir::transform::HasOperandSatisfyingOp::apply(81 mlir::transform::TransformRewriter &rewriter,82 mlir::transform::TransformResults &results,83 mlir::transform::TransformState &state) {84 // For simplicity, only handle a single payload op. Actual implementations85 // can use `SingleOpMatcher` trait to simplify implementation and document86 // this expectation.87 auto payloadOps = state.getPayloadOps(getOp());88 if (!llvm::hasSingleElement(payloadOps))89 return emitSilenceableError() << "expected single payload";90 91 // Iterate over all operands of the payload op to see if they can be matched92 // using the body of this op.93 Operation *payload = *payloadOps.begin();94 for (OpOperand &operand : payload->getOpOperands()) {95 // Create a scope for transform values defined in the body. This corresponds96 // to the syntactic scope of the region attached to this op. Any values97 // associated with payloads from now on will be automatically dissociated98 // when this object is destroyed, i.e. at the end of the iteration.99 // Associate the block argument handle with the operand.100 auto matchScope = state.make_region_scope(getBody());101 if (failed(state.mapBlockArgument(getBody().getArgument(0),102 {operand.get()}))) {103 return DiagnosedSilenceableFailure::definiteFailure();104 }105 106 // Iterate over all nested matchers with the current mapping and see if they107 // succeed.108 bool matchSucceeded = true;109 for (Operation &matcher : getBody().front().without_terminator()) {110 // Matcher ops are applied similarly to any other transform op.111 DiagnosedSilenceableFailure diag =112 state.applyTransform(cast<TransformOpInterface>(matcher));113 114 // Definite failures are immediately propagated as they are irrecoverable.115 if (diag.isDefiniteFailure())116 return diag;117 118 // On success, keep checking the remaining conditions.119 if (diag.succeeded())120 continue;121 122 // Report failure-to-match for debugging purposes and stop matching this123 // operand.124 assert(diag.isSilenceableFailure());125 LDBG() << "failed to match operand #" << operand.getOperandNumber()126 << ": " << diag.getMessage();127 (void)diag.silence();128 matchSucceeded = false;129 break;130 }131 // If failed to match this operand, try other operands.132 if (!matchSucceeded)133 continue;134 135 // If we reached this point, the matching succeeded for the current operand.136 // Remap the values associated with terminator operands to be associated137 // with op results, and also map the parameter result to the operand's138 // position. Note that it is safe to do here despite the end of the scope139 // as `results` are integrated into `state` by the interpreter after `apply`140 // returns rather than immediately.141 SmallVector<SmallVector<MappedValue>> yieldedMappings;142 transform::detail::prepareValueMappings(143 yieldedMappings, getBody().front().getTerminator()->getOperands(),144 state);145 results.setParams(cast<OpResult>(getPosition()),146 {rewriter.getI32IntegerAttr(operand.getOperandNumber())});147 for (auto &&[result, mapping] : llvm::zip(getResults(), yieldedMappings))148 results.setMappedValues(result, mapping);149 return DiagnosedSilenceableFailure::success();150 }151 152 // If we reached this point, none of the operands succeeded the match.153 return emitSilenceableError()154 << "none of the operands satisfied the conditions";155}156 157// By convention, operations implementing MatchOpInterface must not modify158// payload IR and must therefore specify that they only read operand handles and159// payload as their effects.160void mlir::transform::HasOperandSatisfyingOp::getEffects(161 llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &effects) {162 onlyReadsPayload(effects);163 onlyReadsHandle(getOpMutable(), effects);164 producesHandle(getOperation()->getOpResults(), effects);165}166 167// Verify well-formedness of the operation and emit diagnostics if it is168// ill-formed.169llvm::LogicalResult mlir::transform::HasOperandSatisfyingOp::verify() {170 mlir::Block &bodyBlock = getBody().front();171 if (bodyBlock.getNumArguments() != 1 ||172 !isa<TransformValueHandleTypeInterface>(173 bodyBlock.getArgument(0).getType())) {174 return emitOpError()175 << "expects the body to have one value handle argument";176 }177 if (bodyBlock.getTerminator()->getNumOperands() != getNumResults() - 1) {178 return emitOpError() << "expects the body to yield "179 << (getNumResults() - 1) << " values, got "180 << bodyBlock.getTerminator()->getNumOperands();181 }182 for (auto &&[i, operand, result] :183 llvm::enumerate(bodyBlock.getTerminator()->getOperands().getTypes(),184 getResults().getTypes())) {185 if (implementSameTransformInterface(operand, result))186 continue;187 return emitOpError() << "expects terminator operand #" << i188 << " and result #" << (i + 1)189 << " to implement the same transform interface";190 }191 192 for (Operation &op : bodyBlock.without_terminator()) {193 if (!isa<TransformOpInterface>(op) || !isa<MatchOpInterface>(op)) {194 InFlightDiagnostic diag = emitOpError()195 << "expects body to contain match ops";196 diag.attachNote(op.getLoc()) << "non-match operation";197 return diag;198 }199 }200 201 return success();202}203 204void registerMyExtension(::mlir::DialectRegistry ®istry) {205 registry.addExtensions<MyExtension>();206}207