167 lines · cpp
1//===- RuntimeOpVerification.cpp - Op Verification ------------------------===//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#include "mlir/Dialect/Linalg/Transforms/RuntimeOpVerification.h"10 11#include "mlir/Dialect/Affine/IR/AffineOps.h"12#include "mlir/Dialect/Arith/IR/Arith.h"13#include "mlir/Dialect/Arith/Utils/Utils.h"14#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"15#include "mlir/Dialect/Index/IR/IndexAttrs.h"16#include "mlir/Dialect/Index/IR/IndexDialect.h"17#include "mlir/Dialect/Index/IR/IndexOps.h"18#include "mlir/Dialect/Linalg/IR/Linalg.h"19#include "mlir/Dialect/MemRef/IR/MemRef.h"20#include "mlir/Dialect/SCF/IR/SCF.h"21#include "mlir/Dialect/Tensor/IR/Tensor.h"22#include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"23 24namespace mlir {25namespace linalg {26namespace {27/// Verify that the runtime sizes of the operands to linalg structured ops are28/// compatible with the runtime sizes inferred by composing the loop ranges with29/// the linalg op's indexing maps. This is similar to the verifier except that30/// here we insert IR to perform the verification at runtime.31template <typename T>32struct StructuredOpInterface33 : public RuntimeVerifiableOpInterface::ExternalModel<34 StructuredOpInterface<T>, T> {35 void36 generateRuntimeVerification(Operation *op, OpBuilder &builder, Location loc,37 function_ref<std::string(Operation *, StringRef)>38 generateErrorMessage) const {39 auto linalgOp = llvm::cast<LinalgOp>(op);40 41 SmallVector<Range> loopRanges = linalgOp.createLoopRanges(builder, loc);42 auto [starts, ends, _] = getOffsetsSizesAndStrides(loopRanges);43 44 auto zero = arith::ConstantIndexOp::create(builder, loc, 0);45 auto one = arith::ConstantIndexOp::create(builder, loc, 1);46 47 Value iterationDomainIsNonDegenerate;48 for (auto [start, end] : llvm::zip(starts, ends)) {49 auto startValue = getValueOrCreateConstantIndexOp(builder, loc, start);50 auto endValue = getValueOrCreateConstantIndexOp(builder, loc, end);51 52 // Loop Trip count > 0 iff start < end53 Value dimensionHasNonZeroTripCount = index::CmpOp::create(54 builder, loc, index::IndexCmpPredicate::SLT, startValue, endValue);55 56 if (!iterationDomainIsNonDegenerate) {57 iterationDomainIsNonDegenerate = dimensionHasNonZeroTripCount;58 } else {59 // Iteration domain is non-degenerate iff all dimensions have loop trip60 // count > 061 iterationDomainIsNonDegenerate =62 arith::AndIOp::create(builder, loc, iterationDomainIsNonDegenerate,63 dimensionHasNonZeroTripCount);64 }65 }66 67 if (!iterationDomainIsNonDegenerate)68 return;69 70 auto ifOp = scf::IfOp::create(builder, loc, iterationDomainIsNonDegenerate,71 /*withElseRegion=*/false);72 builder.setInsertionPointToStart(&ifOp.getThenRegion().front());73 74 // Subtract one from the loop ends before composing with the indexing map75 transform(ends, ends.begin(), [&](OpFoldResult end) {76 auto endValue = getValueOrCreateConstantIndexOp(builder, loc, end);77 return builder.createOrFold<index::SubOp>(loc, endValue, one);78 });79 80 for (OpOperand &opOperand : linalgOp->getOpOperands()) {81 AffineMap indexingMap = linalgOp.getMatchingIndexingMap(&opOperand);82 auto startIndices = affine::makeComposedFoldedMultiResultAffineApply(83 builder, loc, indexingMap, starts);84 auto endIndices = affine::makeComposedFoldedMultiResultAffineApply(85 builder, loc, indexingMap, ends);86 87 for (auto dim : llvm::seq(linalgOp.getRank(&opOperand))) {88 auto startIndex =89 getValueOrCreateConstantIndexOp(builder, loc, startIndices[dim]);90 auto endIndex =91 getValueOrCreateConstantIndexOp(builder, loc, endIndices[dim]);92 93 // Generate:94 // minIndex = min(startIndex, endIndex)95 // assert(minIndex >= 0)96 // To ensure we do not generate a negative index. We take the minimum of97 // the start and end indices in order to handle reverse loops such as98 // `affine_map<(i) -> (3 - i)>`99 auto min =100 builder.createOrFold<index::MinSOp>(loc, startIndex, endIndex);101 auto cmpOp = builder.createOrFold<index::CmpOp>(102 loc, index::IndexCmpPredicate::SGE, min, zero);103 auto msg = generateErrorMessage(104 linalgOp, "unexpected negative result on dimension #" +105 std::to_string(dim) + " of input/output operand #" +106 std::to_string(opOperand.getOperandNumber()));107 builder.createOrFold<cf::AssertOp>(loc, cmpOp, msg);108 109 // Generate:110 // inferredDimSize = max(startIndex, endIndex) + 1111 // actualDimSize = dim(operand)112 // assert(inferredDimSize <= actualDimSize)113 // To ensure that we do not index past the bounds of the operands.114 auto max =115 builder.createOrFold<index::MaxSOp>(loc, startIndex, endIndex);116 117 auto inferredDimSize =118 builder.createOrFold<index::AddOp>(loc, max, one);119 120 auto actualDimSize =121 createOrFoldDimOp(builder, loc, opOperand.get(), dim);122 123 // Similar to the verifier, when the affine expression in the indexing124 // map is complicated, we just check that the inferred dimension sizes125 // are in the boundary of the operands' size. Being more precise than126 // that is difficult.127 auto predicate = isa<AffineDimExpr>(indexingMap.getResult(dim))128 ? index::IndexCmpPredicate::EQ129 : index::IndexCmpPredicate::SLE;130 131 cmpOp = builder.createOrFold<index::CmpOp>(132 loc, predicate, inferredDimSize, actualDimSize);133 msg = generateErrorMessage(134 linalgOp, "dimension #" + std::to_string(dim) +135 " of input/output operand #" +136 std::to_string(opOperand.getOperandNumber()) +137 " is incompatible with inferred dimension size");138 builder.createOrFold<cf::AssertOp>(loc, cmpOp, msg);139 }140 }141 builder.setInsertionPointAfter(ifOp);142 }143};144 145template <typename... OpTs>146void attachInterface(MLIRContext *ctx) {147 (OpTs::template attachInterface<StructuredOpInterface<OpTs>>(*ctx), ...);148}149} // namespace150} // namespace linalg151} // namespace mlir152 153void mlir::linalg::registerRuntimeVerifiableOpInterfaceExternalModels(154 DialectRegistry ®istry) {155 registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *) {156 attachInterface<157#define GET_OP_LIST158#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"159 >(ctx);160 161 // Load additional dialects of which ops may get created.162 ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,163 cf::ControlFlowDialect, index::IndexDialect,164 tensor::TensorDialect>();165 });166}167