414 lines · cpp
1//===- Specialize.cpp - linalg generic ops to named ops ------------------===//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 implements a method to specialize generic operations to named10// operations. Conceptually it is the opposite of generalize.cpp.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Complex/IR/Complex.h"15#include "mlir/Dialect/Linalg/IR/Linalg.h"16#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"17#include "mlir/Dialect/Linalg/Passes.h"18#include "mlir/Dialect/Linalg/Transforms/Transforms.h"19#include "mlir/Dialect/Math/IR/Math.h"20#include "mlir/IR/PatternMatch.h"21#include "mlir/Transforms/GreedyPatternRewriteDriver.h"22 23namespace mlir {24#define GEN_PASS_DEF_LINALGSPECIALIZEGENERICOPSPASS25#include "mlir/Dialect/Linalg/Passes.h.inc"26} // namespace mlir27 28#define DEBUG_TYPE "linalg-specialization"29 30#define REPLACE_BINARY_OP(NEWOP, OPERANDS_SWAP) \31 (rewriter.replaceOpWithNewOp<NEWOP>( \32 genericOp, \33 ValueRange{genericOp.getDpsInputs()[(OPERANDS_SWAP) ? 1 : 0], \34 genericOp.getDpsInputs()[(OPERANDS_SWAP) ? 0 : 1]}, \35 ValueRange{genericOp.getDpsInits()[0]}))36 37#define REPLACE_UNARY_OP(NEWOP) \38 (rewriter.replaceOpWithNewOp<NEWOP>(genericOp, \39 ValueRange{genericOp.getDpsInputs()[0]}, \40 ValueRange{genericOp.getDpsInits()[0]}))41 42using namespace mlir;43using namespace mlir::linalg;44 45// Given a elementwise single binary linalg generic op, checks whether the46// binary op accesses operands as swapped. e.g.47// this differentiates between a linalg-generic body that contains:48// ^bb0(%a: f32, %b: f32, %c : f32):49// %0 = arith.subf %a, %b : f3250// linalg.yield %0: f3251// against:52// ^bb0(%a: f32, %b: f32, %c : f32):53// %0 = arith.subf %b, %a : f3254// linalg.yield %0: f3255// Former is linalg.sub(a,b), latter is linalg.sub(b,a).56static bool areBinOpsSwapped(GenericOp genericOp) {57 Block *body = genericOp.getBody();58 Operation *op = &body->front();59 bool swapped = false;60 if (op->getOpOperand(0).get() != body->getArgument(0)) {61 swapped = true;62 assert(op->getOpOperand(0).get() == body->getArgument(1) &&63 op->getOpOperand(1).get() == body->getArgument(0) &&64 "binary op uses just one block arg");65 }66 return swapped;67}68 69//===----------------------------------------------------------------------===//70// Specialize linalg generic to matmul variants.71//===----------------------------------------------------------------------===//72/// Identifies linalg.generic that is essentially named op of the form:73// ` linalg.{batch_}?matmul{_transpose_a | _transpose_b}? `74//75// It is possible that a linalg.generic may be implementing a matmul but not76// in a straight-forward way e.g. below is matrix multiply over some slice77// ```78// %0 = linalg.generic {79// indexing_maps = [affine_map<(d0, d1, d2) -> (3, d1, d0)>,80// affine_map<(d0, d1, d2) -> (d0, 5, d2)>,81// affine_map<(d0, d1, d2) -> (d2, d1, 13)>],82// iterator_types = ["parallel", "parallel", "parallel"]}83// ins(%A, %B : tensor<20x20x20xf32>, tensor<20x20x20xf32>)84// outs(%C : tensor<20x20x20xf32>) {85// ^bb0(%a: f32, %b: f32, %c : f32):86// %mul = arith.mulf %a, %b : f3287// %add = arith.addf %mul, %c : f3288// linalg.yield %add : f3289// } -> tensor<20x20x20xf32>90// ```91// It is not possible to represent above as named op.92// e.g. linalg.batch_matmul(%A, %B : tensor<20x20x20xf32>, ...) is93// not the same as linalg.generic above.94namespace {95enum class IndexMatchResult {96 Match = 0, // identity map.97 Transposed, // transposed map.98 Mismatch // none of the above.99};100 101// Checks whether the input Affine `map` contains two consecutive dims that102// can be interpreted as accessing a 2D matrix. It is assumed that the row103// column dimension are adjacent axis (in this order) and start at104// `rowDimIdx` in the input map.105//106// e.g. consider A matrix in `C[M,N] = A[M,K] * B[K,N]`. We will check107// whether the map of A is identity (match), transposed, or something108// completely different (mis-match). Similar for B and C.109static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,110 unsigned expectedPosOfRowDim,111 unsigned expectedPosOfColDim) {112 // Get the matrix multiply indices. They are past the batch indices.113 auto exprOfRowDim = map.getResults()[rowDimIdx];114 auto exprOfColDim = map.getResults()[rowDimIdx + 1];115 116 // They should be pure dimension ids.117 if (exprOfRowDim.getKind() != AffineExprKind::DimId ||118 exprOfColDim.getKind() != AffineExprKind::DimId)119 return IndexMatchResult::Mismatch;120 121 auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();122 auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();123 124 if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)125 return IndexMatchResult::Match;126 127 if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)128 return IndexMatchResult::Transposed;129 130 return IndexMatchResult::Mismatch;131}132 133// Replaces genericOp with `NamedOpTy` op, supplied as a template arg.134// All the variants expressed as pseudo regular expression:135// `linalg.{batch_}?matmul{_transpose_a | _transpose_b}?`136// have same number of ins/out, so its easy to stamp different versions.137template <typename NamedOpTy>138static LinalgOp replaceWithMatmulVariant(RewriterBase &rewriter, GenericOp op) {139 LinalgOp namedOp = rewriter.replaceOpWithNewOp<NamedOpTy>(140 op, ValueRange{op.getDpsInputs()[0], op.getDpsInputs()[1]},141 ValueRange{op.getDpsInits()[0]});142 return namedOp;143}144 145// Converts linalg.generic to named linalg.*matmul* where possible.146static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,147 GenericOp genericOp) {148 if (genericOp.getNumDpsInputs() != 2 || genericOp.getNumDpsInits() != 1)149 return failure();150 151 // Early exit if not projected permutations.152 auto mapRange = genericOp.getIndexingMapsArray();153 if (llvm::any_of(mapRange,154 [](AffineMap m) { return !m.isProjectedPermutation(); }))155 return failure();156 157 // Linalg generic contraction can be across multiple axis e.g.158 // ```159 // linalg.generic160 // {indexing_maps = [affine_map<(m, n, k1, k2) -> (m, k1, k2)>,161 // affine_map<(m, n, k1, k2) -> (k2, k1, n)>,162 // affine_map<(m, n, k1, k2) -> (m, n)>],163 // iterator_types = ["parallel", "parallel",164 // "reduction", "reduction"]}165 // ins(%A, %B : tensor<10x20x30xf32>, tensor<30x20x40xf32>)166 // outs(%C : tensor<10x40xf32>) {167 // ^bb0(%a: f32, %b: f32, %c: f32):168 // %1 = arith.mulf %a, %b : f32169 // %2 = arith.addf %c, %1 : f32170 // linalg.yield %2 : f32171 // } -> tensor<10x40xf32>172 // ```173 // In above contraction, there are two reduction dimensions {k1, k2}174 // and although a valid linalg contraction, it is not a named-op175 // matrix multiply kind. Therefore, reject multi-dim reduction.176 auto res = inferContractionDims(genericOp);177 if (!succeeded(res))178 return failure();179 auto dims = *res;180 if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)181 return failure();182 183 if (!mlir::linalg::detail::isContractionBody(184 *genericOp.getBlock(), [](Operation *first, Operation *second) {185 return (isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||186 (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)) ||187 (isa<complex::MulOp>(first) && isa<complex::AddOp>(second));188 }))189 return failure();190 191 // Check rank of operands192 auto indexingMaps = genericOp.getIndexingMapsArray();193 if (llvm::any_of(indexingMaps, [&dims](AffineMap m) {194 return m.getResults().size() !=195 dims.batch.size() + 2 /* any two of {m,n,k} */;196 }))197 return failure();198 199 auto numOfBatchDims = dims.batch.size();200 if (indexingMaps[0].getNumDims() != numOfBatchDims + 3)201 return failure();202 203 if (numOfBatchDims) {204 // Each operand in a linalg generic contraction could express different205 // permutations for its batch dimension. But for named op it must be206 // identity since separate maps are not specified.207 if (llvm::any_of(indexingMaps, [numOfBatchDims](AffineMap m) {208 for (unsigned i = 0; i < numOfBatchDims; ++i) {209 auto expr = m.getResults()[i];210 if (expr.getKind() != AffineExprKind::DimId ||211 cast<AffineDimExpr>(expr).getPosition() != i)212 return true;213 }214 return false;215 }))216 return failure();217 }218 219 auto a =220 matchOperandMap(indexingMaps[0], numOfBatchDims, dims.m[0], dims.k[0]);221 auto b =222 matchOperandMap(indexingMaps[1], numOfBatchDims, dims.k[0], dims.n[0]);223 auto c =224 matchOperandMap(indexingMaps[2], numOfBatchDims, dims.m[0], dims.n[0]);225 226 if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))227 return failure();228 229 if (c != IndexMatchResult::Match ||230 (a == IndexMatchResult::Transposed && b == IndexMatchResult::Transposed))231 return failure();232 233 /// Codegen the different matmul variants.234 if (numOfBatchDims) {235 return replaceWithMatmulVariant<BatchMatmulOp>(rewriter, genericOp);236 }237 return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp);238}239 240/// Utility to specialize a `genericOp` with a convolution op of type `ConvOpTy`241/// with `dilations` and `strides`.242template <typename ConvOpTy>243static FailureOr<LinalgOp>244specializeToConvOp(RewriterBase &rewriter, GenericOp genericOp,245 ArrayRef<int64_t> dilations, ArrayRef<int64_t> strides) {246 SmallVector<Value> inputs = genericOp.getDpsInputs();247 ValueRange outputs = genericOp.getDpsInits();248 SmallVector<Type> resultTypes = genericOp.hasPureTensorSemantics()249 ? TypeRange(ValueRange(outputs))250 : TypeRange{};251 LinalgOp namedOp;252 // Ops with no dilations and no strides.253 if constexpr (std::is_same_v<ConvOpTy, linalg::Conv1DOp> ||254 std::is_same_v<ConvOpTy, linalg::Conv2DOp> ||255 std::is_same_v<ConvOpTy, linalg::Conv3DOp>) {256 namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(genericOp, resultTypes,257 inputs, outputs);258 } else {259 Attribute stridesAttr = rewriter.getI64TensorAttr(strides);260 Attribute dilationsAttr = rewriter.getI64TensorAttr(dilations);261 namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(262 genericOp, resultTypes, inputs, outputs, stridesAttr, dilationsAttr);263 }264 return namedOp;265}266 267/// Converts linalg.generic to named linalg.*conv/pooling* where possible.268static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,269 GenericOp genericOp) {270 SmallVector<int64_t> dilations, strides;271#define CONV_OP_SPECIALIZER(ConvOpTy) \272 if (isaConvolutionOpOfType<ConvOpTy>(genericOp, &dilations, &strides)) \273 return specializeToConvOp<ConvOpTy>(rewriter, genericOp, dilations, \274 strides); \275 // -----------------------------276 // Convolution ops.277 // -----------------------------278 CONV_OP_SPECIALIZER(linalg::Conv1DOp);279 CONV_OP_SPECIALIZER(linalg::Conv1DNwcWcfOp);280 CONV_OP_SPECIALIZER(linalg::Conv1DNcwFcwOp);281 CONV_OP_SPECIALIZER(linalg::Conv2DOp);282 CONV_OP_SPECIALIZER(linalg::Conv3DOp);283 // -----------------------------284 // Depthwise Convolution ops.285 // -----------------------------286 CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNcwCwOp);287 CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNwcWcOp);288 CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNwcWcmOp);289 CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNchwChwOp);290 CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNdhwcDhwcmOp);291 // -----------------------------292 // Pooling ops.293 // -----------------------------294 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMaxOp);295 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMinOp);296 CONV_OP_SPECIALIZER(linalg::PoolingNhwcSumOp);297 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMaxUnsignedOp);298 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMinUnsignedOp);299#undef CONV_OP_SPECIALIZER300 return failure();301}302 303} // namespace304 305//===----------------------------------------------------------------------===//306// Categorize linalg generic to named op where possible.307//===----------------------------------------------------------------------===//308FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(RewriterBase &rewriter,309 GenericOp genericOp) {310 // Copy311 if (isaCopyOpInterface(genericOp)) {312 LinalgOp namedOp = rewriter.replaceOpWithNewOp<CopyOp>(313 genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0]);314 return namedOp;315 }316 317 // Fill318 if (std::optional<Value> fillValue = isaFillOpInterface(genericOp)) {319 // Always use the detected fill value, regardless of pattern320 LinalgOp namedOp = rewriter.replaceOpWithNewOp<FillOp>(321 genericOp, *fillValue, genericOp.getDpsInits()[0]);322 return namedOp;323 }324 325 // Broadcast326 std::optional<SmallVector<int64_t>> equivalentToBroadcast =327 isaBroadcastOpInterface(genericOp);328 if (equivalentToBroadcast) {329 auto dims = *equivalentToBroadcast;330 LinalgOp namedOp = rewriter.replaceOpWithNewOp<BroadcastOp>(331 genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0],332 dims);333 return namedOp;334 }335 336 // Transpose337 std::optional<SmallVector<int64_t>> equivalentToTranspose =338 isaTransposeOpInterface(genericOp);339 if (equivalentToTranspose) {340 auto permutation = *equivalentToTranspose;341 LinalgOp namedOp = rewriter.replaceOpWithNewOp<TransposeOp>(342 genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0],343 permutation);344 return namedOp;345 }346 347 // Elementwise Unary348 if (isaElemwiseSingleUnaryOpInterface(genericOp)) {349 Operation *op = &genericOp.getBody()->front();350 if (isa<math::ExpOp>(op)) {351 LinalgOp namedOp = REPLACE_UNARY_OP(ExpOp);352 return namedOp;353 }354 }355 356 // Elementwise Binary357 if (isaElemwiseSingleBinaryOpInterface(genericOp)) {358 bool swap = areBinOpsSwapped(genericOp);359 Operation *op = &genericOp.getBody()->front();360 if (isa<arith::AddFOp>(op)) {361 LinalgOp namedOp = REPLACE_BINARY_OP(AddOp, swap);362 return namedOp;363 }364 if (isa<arith::SubFOp>(op)) {365 LinalgOp namedOp = REPLACE_BINARY_OP(SubOp, swap);366 return namedOp;367 }368 if (isa<arith::MulFOp>(op)) {369 LinalgOp namedOp = REPLACE_BINARY_OP(MulOp, swap);370 return namedOp;371 }372 if (isa<arith::DivFOp>(op)) {373 LinalgOp namedOp = REPLACE_BINARY_OP(DivOp, swap);374 return namedOp;375 }376 }377 378 // Contraction - e.g. matmul379 if (isaContractionOpInterface(genericOp)) {380 return specializeLinalgContractions(rewriter, genericOp);381 }382 383 // Convolution - e.g. *conv/pooling*384 if (isaConvolutionOpInterface(genericOp)) {385 return specializeLinalgConvolutions(rewriter, genericOp);386 }387 return failure();388}389 390namespace {391struct LinalgSpecializeGenericOpsPass392 : public impl::LinalgSpecializeGenericOpsPassBase<393 LinalgSpecializeGenericOpsPass> {394 395 using impl::LinalgSpecializeGenericOpsPassBase<396 LinalgSpecializeGenericOpsPass>::LinalgSpecializeGenericOpsPassBase;397 void runOnOperation() override;398};399} // namespace400 401void LinalgSpecializeGenericOpsPass::runOnOperation() {402 RewritePatternSet patterns(&getContext());403 populateLinalgGenericOpsSpecializationPatterns(patterns);404 populateDecomposeProjectedPermutationPatterns(patterns);405 406 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))407 signalPassFailure();408}409 410void mlir::linalg::populateLinalgGenericOpsSpecializationPatterns(411 RewritePatternSet &patterns) {412 patterns.add<LinalgSpecializationPattern>(patterns.getContext());413}414