1325 lines · cpp
1//===- LinalgInterfaces.cpp - Linalg interfaces implementation ------------===//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/IR/LinalgInterfaces.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/Complex/IR/Complex.h"15#include "mlir/Dialect/Linalg/IR/Linalg.h"16#include "mlir/IR/AffineExpr.h"17#include "mlir/IR/AffineExprVisitor.h"18#include "mlir/IR/AffineMap.h"19#include "mlir/IR/BuiltinTypeInterfaces.h"20#include "mlir/IR/MLIRContext.h"21#include "mlir/IR/TypeUtilities.h"22#include "llvm/ADT/STLExtras.h"23#include "llvm/ADT/SetOperations.h"24#include "llvm/ADT/SmallBitVector.h"25#include "llvm/ADT/SmallVector.h"26#include "llvm/Support/Casting.h"27#include "llvm/Support/raw_ostream.h"28#include <optional>29 30using namespace mlir;31using namespace mlir::linalg;32 33/// Include the definitions of the copy operation interface.34#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.cpp.inc"35 36//===----------------------------------------------------------------------===//37// Interface utility functions38//===----------------------------------------------------------------------===//39 40bool linalg::detail::canOpOperandsBeDroppedImpl(41 linalg::LinalgOp linalgOp, ArrayRef<OpOperand *> droppedOperands) {42 SmallVector<AffineMap> indexingMaps;43 for (auto &opOperand : linalgOp->getOpOperands()) {44 if (llvm::is_contained(droppedOperands, &opOperand))45 continue;46 indexingMaps.push_back(linalgOp.getMatchingIndexingMap(&opOperand));47 }48 if (indexingMaps.empty()) {49 // If there are no indexing maps, the operand can only be dropped50 // if the op has no loops.51 return linalgOp.getNumLoops() == 0;52 }53 return inversePermutation(concatAffineMaps(54 indexingMaps, linalgOp.getContext())) != AffineMap();55}56 57//===----------------------------------------------------------------------===//58// CopyOpInterface implementation59//===----------------------------------------------------------------------===//60 61bool linalg::isaCopyOpInterface(LinalgOp op) {62 // Check all loops are parallel and linalgOp is single input and output.63 if (!op.isAllParallelLoops() || !op.isSingleInputOutput())64 return false;65 66 auto mapRange = op.getIndexingMapsArray();67 if (mapRange.size() != 2 || !mapRange.front().isIdentity() ||68 !mapRange.back().isIdentity()) {69 return false;70 }71 // Check yield first block argument.72 Block *body = op.getBlock();73 if (body->getOperations().size() != 1)74 return false;75 auto yieldOp = dyn_cast<linalg::YieldOp>(body->back());76 if (!yieldOp || yieldOp.getNumOperands() != 1)77 return false;78 return yieldOp->getOperand(0) == body->getArgument(0);79}80 81//===----------------------------------------------------------------------===//82// FillOpInterface implementation83//===----------------------------------------------------------------------===//84/// Detects if a linalg.generic operation represents a fill with an inlined85/// constant. If so, returns the constant value. Otherwise, returns86/// std::nullopt.87static std::optional<Value> isaInlinedFillOp(GenericOp op) {88 if (!op.isAllParallelLoops() || op.getNumDpsInits() != 1 ||89 op.getNumDpsInputs() != 0)90 return std::nullopt;91 92 // Init should not be referenced.93 if (op.payloadUsesValueFromOperand(op.getDpsInitOperand(0)))94 return std::nullopt;95 96 Block *body = op.getBody();97 if (body->getOperations().size() != 1)98 return std::nullopt;99 100 auto yieldOp = dyn_cast<linalg::YieldOp>(body->back());101 if (!yieldOp || yieldOp.getNumOperands() != 1)102 return std::nullopt;103 104 Value yieldOperand = yieldOp->getOperand(0);105 if (!yieldOperand.getDefiningOp<arith::ConstantOp>() &&106 !yieldOperand.getDefiningOp<complex::ConstantOp>())107 return std::nullopt;108 109 return yieldOperand;110}111 112/// Detects if a linalg.generic operation represents an external scalar input.113/// If so, returns the constant value. Otherwise, returns std::nullopt.114static std::optional<Value> isaExternalFillOp(GenericOp op) {115 // Structural.116 if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||117 !op.isSingleYieldOp())118 return std::nullopt;119 120 // Input should be referenced and init should not.121 if (!op.payloadUsesValueFromOperand(op.getDpsInputOperand(0)) ||122 op.payloadUsesValueFromOperand(op.getDpsInitOperand(0)))123 return std::nullopt;124 125 OpOperand *value = op.getDpsInputOperand(0);126 if (!op.isScalar(value))127 return std::nullopt;128 return value->get();129}130 131std::optional<Value> linalg::isaFillOpInterface(GenericOp op) {132 if (auto fillVal = isaInlinedFillOp(op))133 return fillVal;134 return isaExternalFillOp(op);135}136 137//===----------------------------------------------------------------------===//138// BroadcastOpInterface implementation139//===----------------------------------------------------------------------===//140std::optional<SmallVector<int64_t>>141linalg::isaBroadcastOpInterface(GenericOp op) {142 // Structural.143 if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||144 !op.isSingleYieldOp())145 return std::nullopt;146 147 auto srcTy = op.getDpsInputOperand(0)->get().getType();148 auto dstTy = op.getDpsInitOperand(0)->get().getType();149 if (!isa<MemRefType, RankedTensorType>(srcTy) ||150 !isa<MemRefType, RankedTensorType>(dstTy))151 return std::nullopt;152 153 // Check output is identity map. Broadcast could additionally be154 // employing permutation of indices and that would be expressible155 // in linalg.generic but is not expressible for named broadcast op.156 auto dstMap = op.getIndexingMapsArray()[1];157 if (!dstMap.isIdentity())158 return std::nullopt;159 160 SmallVector<int64_t> position;161 auto srcMap = op.getIndexingMapsArray()[0];162 163 if (srcMap.getResults().size() >= dstMap.getResults().size())164 return std::nullopt;165 166 // Check input map is monotonically increasing DimIds.167 for (unsigned i = 0; i < srcMap.getNumResults(); ++i) {168 auto expr = llvm::dyn_cast<AffineDimExpr>(srcMap.getResults()[i]);169 if (!expr)170 return std::nullopt;171 int64_t pos = expr.getPosition();172 if (i > 0 && pos <= position[i - 1])173 return std::nullopt;174 position.push_back(expr.getPosition());175 }176 177 SmallVector<int64_t> broadcastedDims;178 auto numDims = srcMap.getNumDims();179 // This is quadratic but number of items is generally small.180 for (auto dim : llvm::seq<int64_t>(0, numDims)) {181 if (!llvm::is_contained(position, dim))182 broadcastedDims.push_back(dim);183 }184 return broadcastedDims;185}186 187//===----------------------------------------------------------------------===//188// TransposeOpInterface implementation189//===----------------------------------------------------------------------===//190std::optional<SmallVector<int64_t>>191linalg::isaTransposeOpInterface(GenericOp op) {192 // To specialize as a transpose op, the genericOp must be193 // all parallel loops, single input, single output, and its body194 // should be just a yield op, yielding input as output as is (no compute).195 if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||196 !op.isSingleYieldOp())197 return std::nullopt;198 199 auto mapRange = op.getIndexingMapsArray();200 if (mapRange.size() != 2)201 return std::nullopt;202 203 auto mapOfInput = mapRange.front();204 auto mapOfResult = mapRange.back();205 206 // linalg.transpose permutes the dimensions of input using this207 // rule: dim(result, i) = dim(input, permutation[i])208 if (!mapOfResult.isIdentity() || !mapOfInput.isPermutation())209 return std::nullopt;210 211 SmallVector<int64_t> permutation(mapOfInput.getNumDims());212 for (unsigned i = 0; i < mapOfInput.getNumDims(); ++i) {213 auto expr = llvm::cast<AffineDimExpr>(mapOfInput.getResults()[i]);214 permutation[expr.getPosition()] = i;215 }216 return permutation;217}218 219//===----------------------------------------------------------------------===//220// Elementwise Single Unary/Binary-OpInterface implementation221//===----------------------------------------------------------------------===//222static bool isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,223 unsigned arity) {224 // Check all loops are parallel.225 if (!op.isAllParallelLoops() || op.getNumLoops() < 1)226 return false;227 228 // Check there are arity-inputs, 1-output and all are identity-maps.229 if (op.getNumDpsInputs() != arity || op.getNumDpsInits() != 1 ||230 !llvm::all_of(op.getIndexingMapsArray(),231 [](AffineMap map) { return map.isIdentity(); }))232 return false;233 234 // Init should not be referenced for elementwise operations.235 if (op.payloadUsesValueFromOperand(op.getDpsInitOperand(0)))236 return false;237 238 // A linalg.generic could be series of elementwise ops e.g. exp(neg(x)) such239 // as resulting from producer-consumer fusion. Here, we restrict to two ops in240 // the body, where the first is the elementwise single op and the second a241 // yield.242 Block *body = op.getBody();243 if (body->getOperations().size() != 2)244 return false;245 246 Operation *oper = &body->front();247 if (oper->getNumOperands() != arity || oper->getNumResults() != 1)248 return false;249 250 auto yieldOp = dyn_cast<linalg::YieldOp>(body->back());251 return !(!yieldOp || yieldOp.getNumOperands() != 1 ||252 yieldOp->getOperand(0).getDefiningOp() != oper);253}254 255bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op) {256 // All basic elemwise checks.257 if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 1))258 return false;259 260 // Check input is actully used.261 if (!op.payloadUsesValueFromOperand(op.getDpsInputOperand(0)))262 return false;263 return true;264}265 266bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {267 if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 2))268 return false;269 270 // Check both inputs are used (elementwise).271 OpOperand *inputOpOperand0 = op.getDpsInputOperand(0);272 OpOperand *inputOpOperand1 = op.getDpsInputOperand(1);273 return !(!op.payloadUsesValueFromOperand(inputOpOperand0) ||274 !op.payloadUsesValueFromOperand(inputOpOperand1));275}276 277//===----------------------------------------------------------------------===//278// ContractionOpInterface implementation279//===----------------------------------------------------------------------===//280 281/// If the value is defined by a chain of unary side effect-free, go up the282/// use-def chain until the first value that isn't defined by such an op.283// TODO: relax to multi-operands with constants, which are technically unary ops284// as needed (e.g. add5).285static Value getSourceSkipUnary(Value value) {286 Operation *op = value.getDefiningOp();287 while (op && op->getNumOperands() == 1) {288 auto iface = dyn_cast<MemoryEffectOpInterface>(op);289 if (!iface || !iface.hasNoEffect())290 break;291 value = op->getOperand(0);292 op = value.getDefiningOp();293 }294 return value;295}296 297bool mlir::linalg::detail::isContractionBody(298 Block &block, function_ref<bool(Operation *, Operation *)> isaPair,299 llvm::raw_ostream &errs) {300 if (block.empty() || !block.back().mightHaveTrait<OpTrait::IsTerminator>()) {301 errs << "no terminator in the block";302 return false;303 }304 305 if (block.getNumArguments() != 3) {306 errs << "expected block with 3 arguments";307 return false;308 }309 310 Operation *terminator = block.getTerminator();311 if (terminator->getNumOperands() != 1) {312 errs << "expected terminator with 1 operand";313 return false;314 }315 316 Value yielded = getSourceSkipUnary(terminator->getOperand(0));317 Operation *reductionOp = yielded.getDefiningOp();318 if (!reductionOp || reductionOp->getNumResults() != 1 ||319 reductionOp->getNumOperands() != 2) {320 errs << "expected reduction op to be binary";321 return false;322 }323 324 Value reductionLHS = getSourceSkipUnary(reductionOp->getOperand(0));325 Value reductionRHS = getSourceSkipUnary(reductionOp->getOperand(1));326 327 if (reductionLHS != block.getArgument(2) &&328 reductionRHS != block.getArgument(2)) {329 errs << "expected reduction to take block argument #2 as one of the "330 "operands (modulo unary casts)";331 return false;332 }333 334 Value contributed = getSourceSkipUnary(335 isa<BlockArgument>(reductionLHS) ? reductionRHS : reductionLHS);336 Operation *elementwiseOp = contributed.getDefiningOp();337 if (!elementwiseOp || elementwiseOp->getNumResults() != 1 ||338 elementwiseOp->getNumOperands() != 2) {339 errs << "expected elementwise op to be binary";340 return false;341 }342 343 if (!isaPair(elementwiseOp, reductionOp)) {344 errs << "expected reduction/elementwise op kind not satisfied";345 return false;346 }347 348 Value elementwiseLHS = getSourceSkipUnary(elementwiseOp->getOperand(0));349 Value elementwiseRHS = getSourceSkipUnary(elementwiseOp->getOperand(1));350 if ((elementwiseLHS == block.getArgument(0) &&351 elementwiseRHS == block.getArgument(1)) ||352 (elementwiseLHS == block.getArgument(1) &&353 elementwiseRHS == block.getArgument(0))) {354 return true;355 }356 357 errs << "expected elementwise op to apply to block arguments (modulo unary "358 "casts)";359 return false;360}361 362/// Returns true if the two operations are of the kinds specified by a pair of363/// consecutive template arguments.364template <typename AddOpTy, typename MulOpTy, typename... Args>365static bool isPairTemplateImpl(Operation *add, Operation *mul) {366 static_assert(sizeof...(Args) % 2 == 0,367 "expected an even number of template arguments");368 if (isa<AddOpTy>(add) && isa<MulOpTy>(mul))369 return true;370 371 if constexpr (sizeof...(Args) > 0)372 return isPairTemplateImpl<Args...>(add, mul);373 else374 return false;375}376 377/// Returns true if the block is a body of a contraction with the kinds of378/// operations given pairwise by template arguments.379template <typename... Args>380static bool isContractionBody(Block &block) {381 return linalg::detail::isContractionBody(block, &isPairTemplateImpl<Args...>);382}383 384/// Given an `indexingMap` and its corresponding `iterators`, returns385/// the positions of the iterators of type `iter` that are indexed by386/// the `indexingMap` as a permutation. This is useful to infer various387/// subcomputations on a `LinalgOp`. This is performed by looking up388/// each result in the `indexingMap` and determining whether:389/// - It is a single AffineDimExpr.390/// - It is the only result involving this AffineDimExpr.391static llvm::SmallDenseSet<int64_t>392findPermutationsIndexingOperand(AffineMap indexingMap,393 ArrayRef<utils::IteratorType> iterators,394 utils::IteratorType iter) {395 assert(iterators.size() == indexingMap.getNumDims());396 llvm::SmallDenseSet<int64_t> res;397 for (AffineExpr e : indexingMap.getResults()) {398 if (auto d = dyn_cast<AffineDimExpr>(e)) {399 if (iterators[d.getPosition()] == iter &&400 llvm::count_if(indexingMap.getResults(), [d](AffineExpr e) {401 return e.isFunctionOfDim(d.getPosition());402 }) == 1)403 res.insert(d.getPosition());404 }405 }406 return res;407}408 409namespace {410auto par = utils::IteratorType::parallel;411auto red = utils::IteratorType::reduction;412} // namespace413 414/// Infer the iterator types from the init affine map. This looks at which dims415/// are present in the map results, and returns an iterator types array with416/// parallel types for dims that are present, and reduction types for dims that417/// are not present.418static FailureOr<SmallVector<utils::IteratorType>>419inferIteratorsFromOutMap(AffineMap map) {420 if (!map.isProjectedPermutation())421 return failure();422 SmallVector<utils::IteratorType> iterators(map.getNumDims(), red);423 for (auto expr : map.getResults())424 if (auto dim = dyn_cast<AffineDimExpr>(expr))425 iterators[dim.getPosition()] = par;426 return iterators;427}428 429/// Find 2 parallel (m and n) and 1 reduction (k) dimension candidates that form430/// a matmul subcomputation within `linalgOp`. These dimensions are such that:431/// 1. The m dimension is involved in an outer-product along LHS432/// (i.e. it is a permutation on RES and LHS and does not appear in RHS).433/// 2. The n dimension is involved in an outer-product along RHS434/// (i.e. it is a permutation on RES and RHS and does not appear in LHS).435/// 3. The k dimension appears as a permutation on LHS and RHS.436/// 4. m, n and k appear only once in any given indexing.437/// 5. Optional batch dimensions that appear in all operands are captured.438/// This allows e.g. detecting that some contraction is embedded within439/// `linalgOp` with some orthogonal heuristic.440static FailureOr<ContractionDimensions>441inferContractionDimsImpl(ArrayRef<AffineMap> indexingMaps,442 ArrayRef<utils::IteratorType> iterators) {443 llvm::SmallDenseSet<int64_t> a =444 findPermutationsIndexingOperand(indexingMaps[0], iterators, par);445 llvm::SmallDenseSet<int64_t> b =446 findPermutationsIndexingOperand(indexingMaps[1], iterators, par);447 llvm::SmallDenseSet<int64_t> c =448 findPermutationsIndexingOperand(indexingMaps[2], iterators, par);449 450 // A & C - B are the iterators involved in an outer-product along A (the LHS).451 llvm::SmallDenseSet<int64_t> ac = a;452 llvm::set_intersect(ac, c);453 llvm::set_subtract(ac, b);454 // B & C - A are the iterators involved in an outer-product along B (the RHS).455 llvm::SmallDenseSet<int64_t> bc = b;456 llvm::set_intersect(bc, c);457 llvm::set_subtract(bc, a);458 // A & B & C are the "batch" dimensions.459 llvm::SmallDenseSet<int64_t> batches = a;460 llvm::set_intersect(batches, b);461 llvm::set_intersect(batches, c);462 463 // A & B red are the reduction dimensions.464 llvm::SmallDenseSet<int64_t> ra =465 findPermutationsIndexingOperand(indexingMaps[0], iterators, red);466 llvm::SmallDenseSet<int64_t> rb =467 findPermutationsIndexingOperand(indexingMaps[1], iterators, red);468 llvm::set_intersect(ra, rb);469 470 // Return each set in sorted order.471 ContractionDimensions dimensions{472 SmallVector<unsigned, 2>(batches.begin(), batches.end()),473 SmallVector<unsigned, 2>(ac.begin(), ac.end()),474 SmallVector<unsigned, 2>(bc.begin(), bc.end()),475 SmallVector<unsigned, 2>(ra.begin(), ra.end())};476 llvm::sort(dimensions.batch);477 llvm::sort(dimensions.m);478 llvm::sort(dimensions.n);479 llvm::sort(dimensions.k);480 return dimensions;481}482 483FailureOr<ContractionDimensions>484mlir::linalg::inferContractionDims(LinalgOp linalgOp) {485 if (linalgOp.getNumDpsInits() != 1 || linalgOp.getNumDpsInputs() != 2)486 return failure();487 return inferContractionDimsImpl(linalgOp.getIndexingMapsArray(),488 linalgOp.getIteratorTypesArray());489}490 491FailureOr<ContractionDimensions>492mlir::linalg::inferContractionDims(ArrayRef<AffineMap> indexingMaps) {493 if (indexingMaps.size() != 3)494 return failure();495 auto iterators = inferIteratorsFromOutMap(indexingMaps[2]);496 if (failed(iterators))497 return failure();498 return inferContractionDimsImpl(indexingMaps, iterators.value());499}500 501namespace mlir::linalg::detail {502enum class MatchContractionResult {503 Success = 0,504 NotLinalgOp,505 WrongNumOperands,506 NoReduction,507 NotProjectedPermutations,508 NotAddMul509};510} // namespace mlir::linalg::detail511 512mlir::linalg::detail::MatchContractionResult513mlir::linalg::detail::isContractionInterfaceImpl(514 Operation *op, mlir::linalg::ContractionDimensions *dimensions) {515 auto linalgOp = dyn_cast<linalg::LinalgOp>(op);516 if (!linalgOp)517 return MatchContractionResult::NotLinalgOp;518 if (linalgOp.getNumDpsInputs() != 2 || linalgOp.getNumDpsInits() != 1)519 return MatchContractionResult::WrongNumOperands;520 auto mapRange = linalgOp.getIndexingMapsArray();521 if (linalgOp.getNumReductionLoops() == 0)522 return MatchContractionResult::NoReduction;523 if (llvm::any_of(mapRange,524 [](AffineMap m) { return !m.isProjectedPermutation(); }))525 return MatchContractionResult::NotProjectedPermutations;526 // TODO: more fields than add/mul.527 // clang-format off528 if (!::isContractionBody<529 arith::MulFOp, arith::AddFOp,530 arith::MulIOp, arith::AddIOp,531 complex::MulOp, complex::AddOp,532 arith::AndIOp, arith::OrIOp>(533 *linalgOp.getBlock())) {534 return MatchContractionResult::NotAddMul;535 }536 // clang-format on537 538 if (dimensions) {539 FailureOr<ContractionDimensions> res = inferContractionDims(linalgOp);540 assert(succeeded(res) && "unexpected failure to infer contraction dims");541 *dimensions = *res;542 }543 return MatchContractionResult::Success;544}545 546StringRef547mlir::linalg::detail::getMatchContractionMessage(MatchContractionResult res) {548 switch (res) {549 case MatchContractionResult::NotLinalgOp:550 return "expected a LinalgOp";551 case MatchContractionResult::WrongNumOperands:552 return "expected op with 2 inputs and 1 output";553 case MatchContractionResult::NoReduction:554 return "expected at least 1 reduction";555 case MatchContractionResult::NotProjectedPermutations:556 return "expected indexing maps to be projected permutations";557 case MatchContractionResult::NotAddMul:558 return "expected add/mul op in the body";559 case MatchContractionResult::Success:560 return "";561 }562 llvm_unreachable("unhandled MatchContractionResult case");563}564 565bool mlir::linalg::isaContractionOpInterface(LinalgOp linalgOp) {566 if (!linalgOp)567 return false;568 Operation *op = linalgOp.getOperation();569 return isa<ContractionOpInterface>(op) ||570 (mlir::linalg::detail::isContractionInterfaceImpl(op) ==571 mlir::linalg::detail::MatchContractionResult::Success);572}573 574/// Verify that a LinalgOp `op` is a contraction.575/// A Linalg contraction is defined in general terms:576/// 1. Has 2 input and 1 output shapes.577/// 2. Has at least one reduction dimension.578/// 3. Has only projected permutation indexing maps.579/// 4. its body computes `u5(u1(c) + u2(u3(a) * u4(b)))` on some field580/// (AddOpType, MulOpType), where u1, u2, u3, u4 and u5 represent scalar unary581/// operations that may change the type (e.g. for mixed-precision).582/// As a consequence, when vectorization of such an op occurs, the only special583/// behavior is that the (unique) MulOpType is vectorized into a584/// `vector.contract`. All other ops are handled in a generic fashion.585/// In the future, we may wish to allow more input arguments and elementwise and586/// constant operations that do not involve the reduction dimension(s).587LogicalResult mlir::linalg::detail::verifyContractionInterface(Operation *op) {588 auto res = isContractionInterfaceImpl(op);589 if (res != MatchContractionResult::Success)590 return op->emitError(getMatchContractionMessage(res));591 return success();592}593 594//===----------------------------------------------------------------------===//595// ConvolutionOpInterface implementation596//===----------------------------------------------------------------------===//597 598/// Of the given two expressions returns one that is of type T (`lhs` gets599/// preference over `rhs`)600template <typename T>601static T getAffineExprOfType(AffineExpr lhs, AffineExpr rhs) {602 return isa<T>(lhs) ? cast<T>(lhs) : (isa<T>(rhs) ? cast<T>(rhs) : nullptr);603}604 605namespace {606/// Walk the indexing expressions for input of a convolution operation to verify607/// its of the right form, either608/// - AffineDimExpr609/// - AffineDimExpr (`*` (AffineSymbolExpr | AffineConstantExpr))?610/// (`+` AffineDimExpr (`*` (AffineSymbolExpr | AffineConstantExpr))?)*611///612/// classifies the AffineDimExpr as convolved dimensions or unconvolved613/// dimensions and verifies each dimension occurs only once.614struct ConvAccessExprWalker615 : public AffineExprVisitor<ConvAccessExprWalker, LogicalResult> {616 // Stores dimensions used in expressions of the above form.617 llvm::SmallDenseSet<int64_t> convolvedDims;618 // Stores the dual mapping between LHS and RHS of convolution exprs.619 llvm::SmallDenseMap<int64_t, int64_t> convolvedDimMapping;620 // Stores single use dimensions used by an AffineDimExpr.621 llvm::SmallDenseSet<int64_t> unConvolvedDims;622 // Stores a mapping from convolved dims to their coefficient.623 llvm::SmallDenseMap<int64_t, AffineExpr> strideAndDilationMapping;624 625 // Removes dims with multiple uses in the source input map from dimension626 // sets tracked by this walker.627 void clearMultiUseDims(AffineMap map) {628 for (int dimPos = 0, e = map.getNumDims(); dimPos < e; ++dimPos) {629 if (llvm::count_if(map.getResults(), [dimPos](AffineExpr e) {630 return e.isFunctionOfDim(dimPos);631 }) > 1) {632 convolvedDims.erase(dimPos);633 unConvolvedDims.erase(dimPos);634 // If a duplicate dim is marked as convolved, the pair of the duplicate635 // dim must be removed from the map as well.636 auto it = convolvedDimMapping.find(dimPos);637 if (it != convolvedDimMapping.end()) {638 int64_t pairedDim = it->second;639 convolvedDims.erase(pairedDim);640 unConvolvedDims.erase(pairedDim);641 strideAndDilationMapping.erase(pairedDim);642 convolvedDimMapping.erase(dimPos);643 convolvedDimMapping.erase(pairedDim);644 }645 }646 }647 }648 649 LogicalResult visitDimExpr(AffineDimExpr dimExpr) {650 unsigned position = dimExpr.getPosition();651 if (unConvolvedDims.count(position) || convolvedDims.count(position)) {652 return failure();653 }654 unConvolvedDims.insert(position);655 return success();656 }657 658 LogicalResult visitSymbolExpr(AffineSymbolExpr expr) { return failure(); }659 660 LogicalResult visitConstantExpr(AffineConstantExpr expr) { return failure(); }661 662 LogicalResult visitAffineBinaryOpExpr(AffineBinaryOpExpr binaryExpr) {663 // In pre-order visit, top level op has to be an add op.664 if (binaryExpr.getKind() != AffineExprKind::Add)665 return failure();666 auto lhsDimPos = getDimExprOrMulExprDimPos(binaryExpr.getLHS());667 auto rhsDimPos = getDimExprOrMulExprDimPos(binaryExpr.getRHS());668 if (failed(lhsDimPos) || failed(rhsDimPos))669 return failure();670 convolvedDimMapping[*lhsDimPos] = *rhsDimPos;671 convolvedDimMapping[*rhsDimPos] = *lhsDimPos;672 return success();673 }674 675 FailureOr<int64_t> getDimExprOrMulExprDimPos(AffineExpr expr) {676 if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {677 int64_t dim = dimExpr.getPosition();678 if (convolvedDims.count(dim) || unConvolvedDims.count(dim))679 return failure();680 // Stride/dilation for this dim is implicitly 1.681 strideAndDilationMapping[dim] =682 getAffineConstantExpr(1, expr.getContext());683 convolvedDims.insert(dim);684 return dim;685 }686 if (auto symbolMulExpr = dyn_cast<AffineBinaryOpExpr>(expr)) {687 if (symbolMulExpr.getKind() != AffineExprKind::Mul)688 return failure();689 auto lhsExpr = symbolMulExpr.getLHS();690 auto rhsExpr = symbolMulExpr.getRHS();691 // Check for symbol expression.692 AffineExpr mulExpr =693 getAffineExprOfType<AffineSymbolExpr>(lhsExpr, rhsExpr);694 // If there was no symbol expr, check for constant expression.695 if (!mulExpr) {696 mulExpr = getAffineExprOfType<AffineConstantExpr>(lhsExpr, rhsExpr);697 }698 auto dimExpr = getAffineExprOfType<AffineDimExpr>(lhsExpr, rhsExpr);699 if (!mulExpr || !dimExpr)700 return failure();701 int64_t dim = dimExpr.getPosition();702 if (convolvedDims.count(dim) || unConvolvedDims.count(dim))703 return failure();704 strideAndDilationMapping[dim] = mulExpr;705 convolvedDims.insert(dim);706 return dim;707 }708 return failure();709 }710};711} // namespace712 713static llvm::SmallDenseSet<int64_t> getPreservedDims(AffineMap map) {714 assert(map.isProjectedPermutation() &&715 "expected map to have projected permutations");716 llvm::SmallDenseSet<int64_t> preservedDims;717 for (auto expr : map.getResults())718 preservedDims.insert(cast<AffineDimExpr>(expr).getPosition());719 return preservedDims;720}721 722static SmallVector<int64_t, 2>723getConstantsFromExprList(const SmallVector<AffineExpr, 2> &exprs) {724 SmallVector<int64_t, 2> vals;725 for (auto e : exprs) {726 auto constantExpr = dyn_cast<AffineConstantExpr>(e);727 assert(constantExpr && "Found non-constant stride/dilation");728 vals.push_back(constantExpr.getValue());729 }730 return vals;731}732 733/// Classifies dimensions in the `linalgOp` used by a convolution734/// subcomputation, as captured by `inputExprWalker`. If735/// `allowEmptyConvolvedDims` is not set this this will fail if there is not736/// at least convolved dimension pair (output image + filter loop). Convolution737/// dimensions are specified in sorted order, and strides match the order of738/// the filter loop dimensions, while the dilations match the order of the739/// output image dimensions.740static FailureOr<ConvolutionDimensions>741inferConvolutionDimsImpl(LinalgOp linalgOp,742 ConvAccessExprWalker &inputExprWalker,743 bool allowEmptyConvolvedDims) {744 auto filterMap =745 linalgOp.getMatchingIndexingMap(linalgOp.getDpsInputOperand(1));746 auto outputMap =747 linalgOp.getMatchingIndexingMap(linalgOp.getDpsInitOperand(0));748 llvm::SmallDenseSet<int64_t> filterDims = findPermutationsIndexingOperand(749 filterMap, linalgOp.getIteratorTypesArray(), par);750 llvm::SmallDenseSet<int64_t> outputDims = findPermutationsIndexingOperand(751 outputMap, linalgOp.getIteratorTypesArray(), par);752 753 // unConvolvedDims & outputDims - filterDims are the batch iterators.754 llvm::SmallDenseSet<int64_t> batch = inputExprWalker.unConvolvedDims;755 llvm::set_intersect(batch, outputDims);756 llvm::set_subtract(batch, filterDims);757 758 // convolvedDims & outputDims are the output image iterators.759 llvm::SmallDenseSet<int64_t> oi = inputExprWalker.convolvedDims;760 llvm::set_intersect(oi, outputDims);761 762 // filterDims & outputDims - unConvolvedDims are the output channel iterators.763 llvm::SmallDenseSet<int64_t> oc = filterDims;764 llvm::set_intersect(oc, outputDims);765 llvm::set_subtract(oc, inputExprWalker.unConvolvedDims);766 767 // filterDims & outputDims & unConvolvedDims are the depth iterators.768 llvm::SmallDenseSet<int64_t> depth = filterDims;769 llvm::set_intersect(depth, outputDims);770 llvm::set_intersect(depth, inputExprWalker.unConvolvedDims);771 772 llvm::SmallDenseSet<int64_t> filterReducedDims =773 findPermutationsIndexingOperand(filterMap,774 linalgOp.getIteratorTypesArray(), red);775 776 // convolvedDims & filterReducedDims are the filter loop iterators.777 llvm::SmallDenseSet<int64_t> fl = inputExprWalker.convolvedDims;778 llvm::set_intersect(fl, filterReducedDims);779 780 // unConvolvedDims & filterReducedDims are the input channel iterators.781 llvm::SmallDenseSet<int64_t> ic = inputExprWalker.unConvolvedDims;782 llvm::set_intersect(ic, filterReducedDims);783 784 if (oi.empty() && !allowEmptyConvolvedDims)785 return failure();786 787 // Return each set in sorted order.788 ConvolutionDimensions dimensions{789 SmallVector<unsigned, 2>(batch.begin(), batch.end()),790 SmallVector<unsigned, 2>(oi.begin(), oi.end()),791 SmallVector<unsigned, 2>(oc.begin(), oc.end()),792 SmallVector<unsigned, 2>(fl.begin(), fl.end()),793 SmallVector<unsigned, 2>(ic.begin(), ic.end()),794 SmallVector<unsigned, 2>(depth.begin(), depth.end()),795 /*strides=*/SmallVector<int64_t, 2>{},796 /*dilations=*/SmallVector<int64_t, 2>{}};797 llvm::sort(dimensions.batch);798 llvm::sort(dimensions.outputImage);799 llvm::sort(dimensions.outputChannel);800 llvm::sort(dimensions.filterLoop);801 llvm::sort(dimensions.inputChannel);802 llvm::sort(dimensions.depth);803 804 // Use the op carried strides/dilations attribute if present.805 auto nativeStrides = linalgOp->getAttrOfType<DenseIntElementsAttr>("strides");806 if (!nativeStrides) {807 SmallVector<AffineExpr, 2> strideExprs;808 for (unsigned oiDim : dimensions.outputImage)809 strideExprs.push_back(inputExprWalker.strideAndDilationMapping[oiDim]);810 dimensions.strides = getConstantsFromExprList(strideExprs);811 } else {812 dimensions.strides = llvm::to_vector<2>(nativeStrides.getValues<int64_t>());813 }814 auto nativeDilations =815 linalgOp->getAttrOfType<DenseIntElementsAttr>("dilations");816 if (!nativeDilations) {817 SmallVector<AffineExpr, 2> dilationExprs;818 for (unsigned flDim : dimensions.filterLoop)819 dilationExprs.push_back(inputExprWalker.strideAndDilationMapping[flDim]);820 dimensions.dilations = getConstantsFromExprList(dilationExprs);821 } else {822 dimensions.dilations =823 llvm::to_vector<2>(nativeDilations.getValues<int64_t>());824 }825 return dimensions;826}827 828/// Find at least 1 parallel (output_image) and reduction (filter_loop)829/// dimension candidates that form a convolution subcomputation within830/// `linalgOp`. The LHS is assumed to be the convolution input while the831/// RHS is assumed as the filter.832/// These dimensions are such that:833/// 1. Optional batch dimensions that appear in the input and filter.834/// 2. The output_image dimension is involved in a cross-correlation along LHS835/// (i.e. it is a permutation on RES and LHS and has an associated836/// filter_loop in RHS).837/// 3. Optional output_channel dimension is involved in an outer-product along838/// RHS (i.e. it is a permutation on RES and RHS and does not appear in839/// LHS).840/// 4. Optional input_channel dimension appears as a permutation on LHS and841/// RHS.842/// 5. The filter_loop dimension appears as a permutation on the RHS and843/// represents the shape of the kernel cross-correlated along a844/// corresponding output_image dim.845/// 6. The input_channel dimension appears as a permutation on LHS and RHS.846/// 7. All dimensions appear only once in any given indexing map.847/// This allows e.g. detecting that some convolution is embedded within848/// `linalgOp` with some orthogonal heuristic.849/// When multiple dimension occurrences exist that match any classification850/// indices are returned in sorted order.851/// Returns a failure if `output_image` (and implicitly `filter_loop`) is empty.852FailureOr<ConvolutionDimensions>853mlir::linalg::inferConvolutionDims(LinalgOp linalgOp) {854 if (linalgOp.getNumDpsInits() != 1 || linalgOp.getNumDpsInputs() != 2)855 return failure();856 857 auto indexingMaps = linalgOp.getIndexingMapsArray();858 859 // Check the input indexing map has the right form.860 ConvAccessExprWalker inputExprWalker;861 for (AffineExpr expr : indexingMaps[0].getResults())862 (void)inputExprWalker.visit(expr);863 inputExprWalker.clearMultiUseDims(indexingMaps[0]);864 865 return inferConvolutionDimsImpl(linalgOp, inputExprWalker,866 /*allowEmptyConvolvedDims=*/false);867}868 869namespace mlir::linalg::detail {870enum class MatchConvolutionResult {871 Success = 0,872 NotLinalgOp,873 WrongNumOperands,874 WrongInputIndexingMap,875 NotProjectedPermutations,876 NonConvolutionLoop,877 OutputDimsNotParallel,878 NonOutputDimNotReduction,879 EmptyConvolvedDims880};881} // namespace mlir::linalg::detail882 883mlir::linalg::detail::MatchConvolutionResult884mlir::linalg::detail::isConvolutionInterfaceImpl(885 Operation *op, ConvolutionDimensions *dimensions,886 bool allowEmptyConvolvedDims) {887 auto linalgOp = dyn_cast<linalg::LinalgOp>(op);888 if (!linalgOp)889 return MatchConvolutionResult::NotLinalgOp;890 if (linalgOp.getNumDpsInputs() < 2 || linalgOp.getNumDpsInits() != 1)891 return MatchConvolutionResult::WrongNumOperands;892 893 auto indexingMaps = linalgOp.getIndexingMapsArray();894 895 // Check the input indexing map has the right form.896 ConvAccessExprWalker inputExprWalker;897 if (llvm::any_of(indexingMaps[0].getResults(),898 [&inputExprWalker](AffineExpr expr) {899 return failed(inputExprWalker.visit(expr));900 })) {901 return MatchConvolutionResult::WrongInputIndexingMap;902 }903 904 // Filter and output maps must be projected permutation.905 if (!indexingMaps[1].isProjectedPermutation() ||906 !indexingMaps.back().isProjectedPermutation())907 return MatchConvolutionResult::NotProjectedPermutations;908 909 auto iteratorTypes = linalgOp.getIteratorTypesArray();910 911 llvm::SmallDenseSet<int64_t> outputDims =912 getPreservedDims(indexingMaps.back());913 llvm::SmallDenseSet<int64_t> filterDims = getPreservedDims(indexingMaps[1]);914 // Make sure all loops are characterized as one of:915 // - Batch loop : present in output, as non-convolved in input, not present in916 // filter.917 // - Output image dimension : present in output, convolved dims in input, not918 // present in filter.919 // - Output channel dimension : present in output, not present in input,920 // present in filter.921 // - Filter loop dimension : present in filter, convolved in input, not922 // present in output.923 // - Input channel dimension : unconvolved in input, not present in output,924 // present in filter.925 // - Depth multiplier : unconvolved in input, present in output, present in926 // filter.927 llvm::SmallDenseSet<int64_t> allLoopDims;928 for (auto outputExpr : indexingMaps.back().getResults()) {929 int64_t outputDim = cast<AffineDimExpr>(outputExpr).getPosition();930 if (inputExprWalker.unConvolvedDims.count(outputDim) &&931 !filterDims.count(outputDim)) {932 // Batch dimension.933 if (iteratorTypes[outputDim] != utils::IteratorType::parallel)934 return MatchConvolutionResult::OutputDimsNotParallel;935 allLoopDims.insert(outputDim);936 continue;937 }938 if (inputExprWalker.convolvedDims.count(outputDim) &&939 !filterDims.count(outputDim)) {940 // Output image Loop dimension.941 if (iteratorTypes[outputDim] != utils::IteratorType::parallel)942 return MatchConvolutionResult::OutputDimsNotParallel;943 allLoopDims.insert(outputDim);944 continue;945 }946 if (!inputExprWalker.convolvedDims.count(outputDim) &&947 !inputExprWalker.unConvolvedDims.count(outputDim) &&948 filterDims.count(outputDim)) {949 // Output channel dimension.950 if (iteratorTypes[outputDim] != utils::IteratorType::parallel)951 return MatchConvolutionResult::OutputDimsNotParallel;952 allLoopDims.insert(outputDim);953 continue;954 }955 if (inputExprWalker.unConvolvedDims.count(outputDim) &&956 filterDims.count(outputDim)) {957 // Depth multiplier.958 if (iteratorTypes[outputDim] != utils::IteratorType::parallel)959 return MatchConvolutionResult::OutputDimsNotParallel;960 allLoopDims.insert(outputDim);961 continue;962 }963 return MatchConvolutionResult::NonConvolutionLoop;964 }965 for (auto filterExpr : indexingMaps[1].getResults()) {966 int64_t filterDim = cast<AffineDimExpr>(filterExpr).getPosition();967 if (outputDims.count(filterDim) &&968 !inputExprWalker.unConvolvedDims.count(filterDim) &&969 !inputExprWalker.convolvedDims.count(filterDim)) {970 // Output channel dimension. This is already seen, continue;971 continue;972 }973 if (inputExprWalker.convolvedDims.count(filterDim) &&974 !outputDims.count(filterDim)) {975 // Filter loop dimension.976 if (iteratorTypes[filterDim] != utils::IteratorType::reduction)977 return MatchConvolutionResult::NonOutputDimNotReduction;978 if (allLoopDims.count(filterDim))979 return MatchConvolutionResult::NonConvolutionLoop;980 allLoopDims.insert(filterDim);981 continue;982 }983 if (inputExprWalker.unConvolvedDims.count(filterDim) &&984 !outputDims.count(filterDim)) {985 // Input channel dimension.986 if (iteratorTypes[filterDim] != utils::IteratorType::reduction)987 return MatchConvolutionResult::NonOutputDimNotReduction;988 if (allLoopDims.count(filterDim))989 return MatchConvolutionResult::NonConvolutionLoop;990 allLoopDims.insert(filterDim);991 continue;992 }993 if (inputExprWalker.unConvolvedDims.count(filterDim) &&994 outputDims.count(filterDim)) {995 // Depthwise loop. Already seen.996 continue;997 }998 return MatchConvolutionResult::NonConvolutionLoop;999 }1000 // All loops must be covered now.1001 if (allLoopDims.size() != linalgOp.getNumLoops())1002 return MatchConvolutionResult::NonConvolutionLoop;1003 1004 if (!allowEmptyConvolvedDims && inputExprWalker.convolvedDims.empty())1005 return MatchConvolutionResult::EmptyConvolvedDims;1006 1007 if (dimensions) {1008 FailureOr<ConvolutionDimensions> res = inferConvolutionDimsImpl(1009 linalgOp, inputExprWalker, allowEmptyConvolvedDims);1010 assert(succeeded(res) && "unexpected failure to infer convolution dims");1011 *dimensions = *res;1012 }1013 1014 return MatchConvolutionResult::Success;1015}1016 1017StringRef1018mlir::linalg::detail::getMatchConvolutionMessage(MatchConvolutionResult res) {1019 switch (res) {1020 case MatchConvolutionResult::NotLinalgOp:1021 return "expected a LinalgOp";1022 case MatchConvolutionResult::WrongNumOperands:1023 return "expected op with 2 inputs and 1 output";1024 case MatchConvolutionResult::WrongInputIndexingMap:1025 return "unexpected input index map for convolutions";1026 case MatchConvolutionResult::NotProjectedPermutations:1027 return "expected output/filter indexing maps to be projected permutations";1028 case MatchConvolutionResult::NonConvolutionLoop:1029 return "unexpected loop dimension for convolution op";1030 case MatchConvolutionResult::OutputDimsNotParallel:1031 return "expected all iterators used to access outputs to be parallel";1032 case MatchConvolutionResult::NonOutputDimNotReduction:1033 return "expected all iterators not used to access outputs to be reduction";1034 case MatchConvolutionResult::EmptyConvolvedDims:1035 return "expected convolved dim to be non-empty";1036 case MatchConvolutionResult::Success:1037 return "";1038 }1039 llvm_unreachable("unhandled MatchConvolutionResult case");1040}1041 1042bool mlir::linalg::isaConvolutionOpInterface(LinalgOp linalgOp,1043 bool allowEmptyConvolvedDims) {1044 return linalg::detail::isConvolutionInterfaceImpl(1045 linalgOp.getOperation(), nullptr, allowEmptyConvolvedDims) ==1046 linalg::detail::MatchConvolutionResult::Success;1047}1048 1049LogicalResult mlir::linalg::detail::verifyConvolutionInterface(Operation *op) {1050 MatchConvolutionResult res = isConvolutionInterfaceImpl(op);1051 if (res != MatchConvolutionResult::Success)1052 return op->emitError(getMatchConvolutionMessage(res));1053 return success();1054}1055 1056//===----------------------------------------------------------------------===//1057// FillOpInterface implementation1058//===----------------------------------------------------------------------===//1059 1060namespace {1061enum class MatchFillResult {1062 Success = 0,1063 NotLinalgOp,1064 WrongNumOperands,1065 NotScalarInput,1066 TypeMismatch1067};1068} // namespace1069 1070static MatchFillResult isFillInterfaceImpl(Operation *op) {1071 auto linalgOp = dyn_cast<linalg::LinalgOp>(op);1072 if (!linalgOp)1073 return MatchFillResult::NotLinalgOp;1074 if (linalgOp.getNumDpsInputs() != 1 || linalgOp.getNumDpsInits() != 1)1075 return MatchFillResult::WrongNumOperands;1076 1077 OpOperand *value = linalgOp.getDpsInputOperand(0);1078 if (!linalgOp.isScalar(value))1079 return MatchFillResult::NotScalarInput;1080 1081 // Check that the scalar input type matches the output element type.1082 OpOperand *output = linalgOp.getDpsInitOperand(0);1083 Type scalarType = value->get().getType();1084 Type outputElementType = getElementTypeOrSelf(output->get().getType());1085 if (scalarType != outputElementType)1086 return MatchFillResult::TypeMismatch;1087 1088 return MatchFillResult::Success;1089}1090 1091LogicalResult mlir::linalg::detail::verifyFillInterface(Operation *op) {1092 MatchFillResult res = isFillInterfaceImpl(op);1093 if (res == MatchFillResult::NotLinalgOp)1094 return op->emitError("expected a LinalgOp");1095 if (res == MatchFillResult::WrongNumOperands)1096 return op->emitError("expected op with 1 input and 1 output");1097 if (res == MatchFillResult::NotScalarInput)1098 return op->emitError("expected op with scalar input");1099 if (res == MatchFillResult::TypeMismatch) {1100 auto linalgOp = cast<linalg::LinalgOp>(op);1101 Type scalarType = linalgOp.getDpsInputOperand(0)->get().getType();1102 Type outputElementType =1103 getElementTypeOrSelf(linalgOp.getDpsInitOperand(0)->get().getType());1104 return op->emitOpError("expected fill value type (")1105 << scalarType << ") to match output element type ("1106 << outputElementType << ")";1107 }1108 1109 return success();1110}1111 1112//===----------------------------------------------------------------------===//1113// StructuredOpInterface implementation1114//===----------------------------------------------------------------------===//1115 1116SmallVector<OpFoldResult> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,1117 Location loc) {1118 SmallVector<OpFoldResult> res;1119 for (OpOperand &opOperand : getOperation()->getOpOperands()) {1120 for (int64_t i = 0, e = getRank(&opOperand); i < e; ++i)1121 res.push_back(createFoldedDimOp(b, loc, opOperand.get(), i));1122 }1123 return res;1124}1125 1126SmallVector<int64_t, 4> LinalgOp::createFlatListOfOperandStaticDims() {1127 SmallVector<int64_t, 4> res;1128 assert(!hasDynamicShape() && "expected operands to have static shapes");1129 for (OpOperand &opOperand : getOperation()->getOpOperands())1130 llvm::append_range(res, getShape(&opOperand));1131 return res;1132}1133 1134SmallVector<Range, 4> LinalgOp::createLoopRanges(OpBuilder &b, Location loc) {1135 AffineMap map = getLoopsToShapesMap();1136 unsigned numDims = map.getNumDims(), numRes = map.getNumResults();1137 auto viewSizes = createFlatListOfOperandDims(b, loc);1138 SmallVector<Range, 4> res(numDims);1139 for (unsigned idx = 0; idx < numRes; ++idx) {1140 auto result = map.getResult(idx);1141 if (auto d = dyn_cast<AffineDimExpr>(result)) {1142 if (res[d.getPosition()].offset)1143 continue;1144 res[d.getPosition()] =1145 Range{b.getIndexAttr(0), viewSizes[idx], b.getIndexAttr(1)};1146 }1147 }1148 return res;1149}1150 1151/// Visitor to check if any of the given set of positions from AffineDimExprs1152/// are used within an AffineExpr.1153struct HasAffineDimExprVisitor1154 : public AffineExprVisitor<HasAffineDimExprVisitor, bool> {1155 HasAffineDimExprVisitor(llvm::SmallBitVector positions)1156 : positions(std::move(positions)) {}1157 1158 bool visitAffineBinaryOpExpr(AffineBinaryOpExpr binaryOpExpr) {1159 return visit(binaryOpExpr.getLHS()) || visit(binaryOpExpr.getRHS());1160 }1161 1162 bool visitDimExpr(AffineDimExpr dimExpr) {1163 return positions.test(dimExpr.getPosition());1164 }1165 1166 bool visitConstantExpr(AffineConstantExpr constExpr) { return false; }1167 1168 bool visitSymbolExpr(AffineSymbolExpr symbolExpr) { return false; }1169 1170private:1171 llvm::SmallBitVector positions;1172};1173 1174static std::pair<int64_t, int64_t>1175getResultsPositionInLoopsToShapeMap(LinalgOp &op) {1176 int64_t inputRankSum = 0;1177 int64_t outputRankSum = 0;1178 for (OpOperand *input : op.getDpsInputOperands())1179 inputRankSum += op.getRank(input);1180 for (OpOperand &output : op.getDpsInitsMutable())1181 outputRankSum += op.getRank(&output);1182 return {inputRankSum, inputRankSum + outputRankSum};1183}1184 1185LogicalResult1186LinalgOp::reifyResultShapes(OpBuilder &b,1187 ReifiedRankedShapedTypeDims &reifiedReturnShapes) {1188 // An example that helps understand the logic below.1189 // Consider the following expression O(i+j, j) += A(i,k) * B(k, j)1190 // We want to express the shape of dim 0 of O in terms of shape of the inputs.1191 // This is achieved as follows.1192 // loopsToShapesMap = (d0, d1, d2) -> (d0, d2, d2, d1, d0 + d1, d1)1193 // subMapOfResultShapes = (d0, d1, d2) -> (d0 + d1, d1)1194 // shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)1195 // resultShapesFromInputShapes = subMapOfResultDim.compose(shapesToLoopMap)1196 // = (d0, d1, d2, d3, d4, d5) -> (d0 + d1, d1)1197 AffineMap loopsToShapesMap = getLoopsToShapesMap();1198 1199 // Find the position in the above map that represents the shape of the1200 // result:dim being inferred.1201 auto resultShapesSubMapPos = getResultsPositionInLoopsToShapeMap(*this);1202 1203 /// From loopsToShapesMap extract the submap that represents the shape of the1204 /// (resultIdx, dim) needed.1205 AffineMap loopToResultsShapeMap = loopsToShapesMap.getSliceMap(1206 resultShapesSubMapPos.first,1207 resultShapesSubMapPos.second - resultShapesSubMapPos.first);1208 AffineMap resultShapesFromInputShapesMap =1209 loopToResultsShapeMap.compose(getShapesToLoopsMap());1210 1211 // Check that the result dim map does not contain the positions corresponding1212 // to the outputs.1213 llvm::SmallBitVector outputDims(resultShapesFromInputShapesMap.getNumDims());1214 outputDims.set(resultShapesSubMapPos.first, resultShapesSubMapPos.second);1215 HasAffineDimExprVisitor checkDimExpr(std::move(outputDims));1216 Location loc = getOperation()->getLoc();1217 IRRewriter rewriter(b);1218 SmallVector<OpFoldResult> allResultDimValues =1219 affine::makeComposedFoldedMultiResultAffineApply(1220 rewriter, loc, resultShapesFromInputShapesMap,1221 createFlatListOfOperandDims(b, loc));1222 int64_t pos = 0;1223 ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();1224 for (OpOperand &opOperand : getDpsInitsMutable()) {1225 SmallVector<OpFoldResult> shapes;1226 for (int64_t dim : llvm::seq<int64_t>(0, getRank(&opOperand))) {1227 auto shapedType = llvm::cast<ShapedType>(opOperand.get().getType());1228 if (!shapedType.isDynamicDim(dim)) {1229 // Static dim: Return IntegerAttr.1230 shapes.push_back(b.getIndexAttr(shapedType.getDimSize(dim)));1231 } else {1232 // Dynamic dim: Return Value.1233 OpFoldResult ofr = checkDimExpr.visit(shapeExprs[pos])1234 ? createOrFoldDimOp(b, loc, opOperand.get(), dim)1235 : allResultDimValues[pos];1236 shapes.push_back(getValueOrCreateConstantIndexOp(b, loc, ofr));1237 }1238 pos++;1239 }1240 reifiedReturnShapes.emplace_back(std::move(shapes));1241 }1242 return success();1243}1244 1245/// Return the index in the indexingMaps vector that corresponds to this1246/// `opOperand`.1247int64_t LinalgOp::getIndexingMapIndex(OpOperand *opOperand) {1248 auto operandNumber = opOperand->getOperandNumber();1249 auto dpsIface = cast<DestinationStyleOpInterface>(*this->getOperation());1250 if (!dpsIface.isDpsInput(opOperand))1251 return operandNumber;1252 unsigned start = dpsIface.getDpsInits().getBeginOperandIndex();1253 assert(!dpsIface.isDpsInit(opOperand));1254 // Account for potential inputs that are not DPS and may not appear in1255 // `indexingMaps`.1256 return cast<DestinationStyleOpInterface>(*this->getOperation())1257 .getNumDpsInputs() +1258 operandNumber - start;1259}1260 1261LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {1262 LinalgOp linalgOp = cast<LinalgOp>(op);1263 // Mixed tensor/buffer operands are not allowed.1264 if (!linalgOp.hasPureTensorSemantics() &&1265 !linalgOp.hasPureBufferSemantics() && op->getNumOperands() > 0)1266 return op->emitOpError("expected to have pure tensor or buffer semantics");1267 1268 // Before checking indexing maps, we need to make sure the attributes1269 // referenced by it are valid.1270 if (linalgOp.hasDynamicIndexingMaps())1271 if (failed(linalgOp.verifyIndexingMapRequiredAttributes()))1272 return failure();1273 1274 // Delayed calling of IndexingMapOpInterface::verifyImpl.1275 if (failed(cast<IndexingMapOpInterface>(op).verifyImpl()))1276 return failure();1277 1278 // Set this flag if this op has user defined maps. This is required to guard1279 // the below error condition which assume default indexing maps.1280 for (OpOperand &opOperand : linalgOp->getOpOperands()) {1281 AffineMap indexingMap = linalgOp.getMatchingIndexingMap(&opOperand);1282 // Domain must be consistent.1283 unsigned numLoops = linalgOp.getNumLoops();1284 if (indexingMap.getNumDims() != numLoops)1285 return op->emitOpError("expected indexing_map #")1286 << opOperand.getOperandNumber() << " to have " << numLoops1287 << " dim(s) to match the number of loops";1288 }1289 SmallVector<unsigned> redDims;1290 linalgOp.getReductionDims(redDims);1291 1292 if (!linalgOp.getShapesToLoopsMap())1293 return op->emitOpError("expected the shape-to-loops map to be non-null");1294 1295 // Check the region has exactly one block.1296 if (linalgOp->getNumRegions() != 1 || !linalgOp->getRegion(0).hasOneBlock())1297 return op->emitOpError("expects to have 1 region with 1 block");1298 1299 // Simplifying assumption: bbargs match 1-1 with shape operands elemental1300 // types.1301 // TODO: once ranked shape types are plugged in, we may want to drop the1302 // corresponding bbargs, that can never be read from. This will be subject to1303 // consistency discussions (i.e. what to do with output tensors whose bbarg is1304 // not used).1305 Block &block = linalgOp->getRegion(0).front();1306 1307 if (linalgOp.getOpOperandsMatchingBBargs().size() != block.getNumArguments())1308 return op->emitOpError("expected as many non-induction variable region "1309 "arguments as the number of input/output operands");1310 1311 for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {1312 Type elementType = opOperand->get().getType();1313 if (isa<MemRefType, RankedTensorType>(elementType))1314 elementType = getElementTypeOrSelf(opOperand->get().getType());1315 Type argType = block.getArgument(opOperand->getOperandNumber()).getType();1316 if (elementType != argType)1317 return op->emitOpError("expected type of bb argument #")1318 << opOperand->getOperandNumber() << " (" << argType << ")"1319 << " to match element or self type of the corresponding operand ("1320 << elementType << ")";1321 }1322 1323 return success();1324}1325