Specialize.cpp

Go to the documentation of this file.

//===- Specialize.cpp - linalg generic ops to named ops  ------------------===//
//
// 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 implements a method to specialize generic operations to named
// operations. Conceptually it is the opposite of generalize.cpp.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/TypeSwitch.h"
 
namespace mlir {
#define GEN_PASS_DEF_LINALGSPECIALIZEGENERICOPSPASS
#include "mlir/Dialect/Linalg/Passes.h.inc"
} // namespace mlir
 
#define DEBUG_TYPE "linalg-specialization"
 
using namespace mlir;
using namespace mlir::linalg;
 
//===----------------------------------------------------------------------===//
// Specialize linalg generic to elementwise ops.
//===----------------------------------------------------------------------===//
 
// Given an elementwise single binary linalg generic op, checks whether the
// binary op accesses operands as swapped. e.g.
// this differentiates between a linalg-generic body that contains:
//    ^bb0(%a: f32, %b: f32, %c : f32):
//         %0 = arith.subf %a, %b : f32
//         linalg.yield %0: f32
// against:
//    ^bb0(%a: f32, %b: f32, %c : f32):
//         %0 = arith.subf %b, %a : f32
//         linalg.yield %0: f32
// Former is linalg.sub(a,b), latter is linalg.sub(b,a).
static bool areBinOpsSwapped(GenericOp genericOp, bool isTernary) {
  Block *body = genericOp.getBody();
  Operation *op = &body->front();
  bool swapped = false;
  if (op->getOpOperand(0 + isTernary).get() !=
      body->getArgument(0 + isTernary)) {
    swapped = true;
    assert(op->getOpOperand(0 + isTernary).get() ==
               body->getArgument(1 + isTernary) &&
           op->getOpOperand(1 + isTernary).get() ==
               body->getArgument(0 + isTernary) &&
           "binary op uses just one block arg");
  }
  return swapped;
}
 
// Given an elementwise single unary linalg generic op whose body operation is a
// binary operation, check if one of its operands is a scalar value defined
// outside the generic op, set its index, and return true. Otherwise return
// false. The index is unique because the block argument is used at
// least by one operand, as checked in `isaElemwiseSingleUnaryOpInterface`.
//
// Example:
//   %cst = arith.constant 3.14 : f32
//   %0 = linalg.generic { indexing_maps = [#mapA, #mapRes], ... }
//          ins(%A : tensor<?xf32>) outs(...) {
//   ^bb0(%a: f32, %out : f32):
//     %0 = arith.mulf %a, %cst : f32
//     linalg.yield %0: f32
//   } -> tensor<?xf32>
// Here, the returned index is 1, and the generic op can be represented as
//   %0 = linalg.elementwise <mul>
//          indexing_maps = [#mapA, affine_map<(d0) -> ()>, #mapRes]
//          ins(%A, %cst : tensor<?xf32>, f32) outs(...) -> tensor<?xf32>
static bool findIndexOfScalarOperand(GenericOp genericOp, int &index) {
  Block *body = genericOp.getBody();
  Operation *op = &body->front();
  for (auto [i, v] : llvm::enumerate(op->getOperands())) {
    if (auto blockArg = dyn_cast<BlockArgument>(v);
        blockArg && blockArg.getOwner() == body)
      continue; // not an outside value...
    index = i;
    return true;
  }
  return false;
}
 
// Maps a linalg.generic body operation to its corresponding elementwise kind,
// if one exists. Handles the ops with a straightforward one-to-one mapping;
// ops needing extra context (e.g. boolean add/mul, reciprocal, square) are
// handled by the caller.
static std::optional<ElementwiseKind> getElementwiseKind(Operation *op) {
  return llvm::TypeSwitch<Operation *, std::optional<ElementwiseKind>>(op)
      .Case<arith::AddFOp, arith::AddIOp, complex::AddOp>(
          [](Operation *) { return ElementwiseKind::add; })
      .Case<arith::MulIOp, arith::MulFOp, complex::MulOp>(
          [](Operation *) { return ElementwiseKind::mul; })
      .Case([](math::ExpOp) { return ElementwiseKind::exp; })
      .Case([](math::AbsFOp) { return ElementwiseKind::abs; })
      .Case([](math::CeilOp) { return ElementwiseKind::ceil; })
      .Case([](math::FloorOp) { return ElementwiseKind::floor; })
      .Case([](arith::NegFOp) { return ElementwiseKind::negf; })
      .Case([](math::RoundOp) { return ElementwiseKind::round; })
      .Case([](math::SqrtOp) { return ElementwiseKind::sqrt; })
      .Case([](math::RsqrtOp) { return ElementwiseKind::rsqrt; })
      .Case([](math::TanhOp) { return ElementwiseKind::tanh; })
      .Case([](math::ErfOp) { return ElementwiseKind::erf; })
      .Case([](math::SinOp) { return ElementwiseKind::sin; })
      .Case([](math::CosOp) { return ElementwiseKind::cos; })
      .Case([](math::TanOp) { return ElementwiseKind::tan; })
      .Case([](math::AcosOp) { return ElementwiseKind::acos; })
      .Case([](math::AcoshOp) { return ElementwiseKind::acosh; })
      .Case([](math::AsinOp) { return ElementwiseKind::asin; })
      .Case([](math::AsinhOp) { return ElementwiseKind::asinh; })
      .Case([](math::AtanOp) { return ElementwiseKind::atan; })
      .Case([](math::AtanhOp) { return ElementwiseKind::atanh; })
      .Case([](math::LogOp) { return ElementwiseKind::log; })
      .Case([](math::Log10Op) { return ElementwiseKind::log10; })
      .Case([](math::Log1pOp) { return ElementwiseKind::log1p; })
      .Case([](math::Log2Op) { return ElementwiseKind::log2; })
      .Case<arith::SubIOp, arith::SubFOp, complex::SubOp>(
          [](Operation *) { return ElementwiseKind::sub; })
      .Case<arith::DivSIOp, arith::DivFOp, complex::DivOp>(
          [](Operation *) { return ElementwiseKind::div; })
      .Case([](arith::DivUIOp) { return ElementwiseKind::div_unsigned; })
      .Case<arith::MaxSIOp, arith::MaximumFOp>(
          [](Operation *) { return ElementwiseKind::max_signed; })
      .Case<arith::MinSIOp, arith::MinimumFOp>(
          [](Operation *) { return ElementwiseKind::min_signed; })
      .Case([](math::PowFOp) { return ElementwiseKind::powf; })
      .Case([](arith::MaxUIOp) { return ElementwiseKind::max_unsigned; })
      .Case([](arith::MinUIOp) { return ElementwiseKind::min_unsigned; })
      .Case([](arith::SelectOp) { return ElementwiseKind::select; })
      .Default([](Operation *) { return std::nullopt; });
}
 
// Attempt to specialize unary/binary/ternary linalg.generic ops
// to linalg.elementwise.
//
// Example:
//   %0 = linalg.generic {
//       indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
//                        affine_map<(d0, d1) -> (d0, d1)>],
//       iterator_types = ["parallel", "parallel"]
//     } ins(%In : tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
//     ^bb0(%in: f32, %out: f32):
//       %1 = math.exp %in : f32
//       linalg.yield %1 : f32
//     } -> tensor<?x?xf32>
//
// is specialized to
//   linalg.elementwise <exp> ...
//
// The category op can carry non-identity indexing maps; these are
// transferred verbatim from the `genericOp`.
//
// In addition to the canonical forms used by the generalization path, this
// function can handle the following variations:
//
// 1) Swapped operands in binary ops (see the `areBinOpsSwapped` helper)
// 2) Unary generic ops with a binary body op (see the
//    `findIndexOfScalarOperand` helper)
static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
                                                       GenericOp genericOp) {
  // Classify the generic op.
  unsigned arity = genericOp.getNumDpsInputs();
  bool isUnary = arity == 1;
  bool isBinary = arity == 2;
  bool isTernary = arity == 3;
 
  // Will inspect the body operation to determine named op or elementwise kind.
  Operation *op = &genericOp.getBody()->front();
 
  // Detect variations from canonical forms.
  bool hasSwappedOperands =
      (isBinary || isTernary) && areBinOpsSwapped(genericOp, isTernary);
  int scalarOprIdx = -1;
  bool hasScalarOperand = isUnary && op->getNumOperands() == 2 &&
                          findIndexOfScalarOperand(genericOp, scalarOprIdx);
 
  // Helper to dispatch between named op and `linalg.elementwise`.
  // Lambdas with explicit template parameter list are a C++20 feature, hence
  // the dummy op object.
  auto replaceOp = [&](ElementwiseKind kind,
                       bool mayHoistScalarOperand = true) -> LinalgOp {
    SmallVector<Value> inputs = genericOp.getDpsInputs();
    SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
    if (hasSwappedOperands) {
      // Ternary indices are +1, since the first is the boolean mask.
      // If new ternary with non-booleans as first argument are created,
      // we may need to calculate all combinations possible.
      std::swap(inputs[0 + isTernary], inputs[1 + isTernary]);
      std::swap(indexingMaps[0 + isTernary], indexingMaps[1 + isTernary]);
    }
 
    if (hasScalarOperand && mayHoistScalarOperand) {
      // Adjust inputs and indexing maps accordingly.
      inputs.insert(inputs.begin() + scalarOprIdx,
                    op->getOperand(scalarOprIdx));
      auto scalarBroadcastMap =
          AffineMap::get(genericOp.getNumParallelLoops(), /*symbolCount=*/0,
                         rewriter.getContext());
      indexingMaps.insert(indexingMaps.begin() + scalarOprIdx,
                          scalarBroadcastMap);
    }
    auto newOp = ElementwiseOp::create(
        rewriter, genericOp.getLoc(), inputs, genericOp.getDpsInits(), kind,
        rewriter.getAffineMapArrayAttr(indexingMaps));
 
    rewriter.replaceOp(genericOp, newOp);
    return newOp;
  };
 
  // Reciprocal
  if (auto divOp = dyn_cast<arith::DivFOp>(op)) {
    if (auto constOp = dyn_cast_if_present<arith::ConstantOp>(
            divOp.getLhs().getDefiningOp()))
      if (cast<FloatAttr>(constOp.getValue()).getValue().isExactlyValue(1.0))
        return replaceOp(ElementwiseKind::reciprocal,
                         /*mayHoistScalarOperand=*/false);
  }
 
  // Square
  if (auto mulOp = dyn_cast<arith::MulFOp>(op))
    if (mulOp.getLhs() == mulOp.getRhs())
      return replaceOp(ElementwiseKind::square);
 
  // Boolean-typed `add` and `mul`.
  if (isBinary && llvm::all_of(op->getOperands(), [](Value v) {
        return v.getType().isInteger(1);
      })) {
    if (isa<arith::OrIOp>(op))
      return replaceOp(ElementwiseKind::add);
    if (isa<arith::AndIOp>(op))
      return replaceOp(ElementwiseKind::mul);
  }
 
  // Table driven
  if (std::optional<ElementwiseKind> kind = getElementwiseKind(op)) {
    auto arityGroupAndKind = linalg::getArityGroupAndKind(*kind);
    // A hoisted scalar operand adds one input to the elementwise op.
    unsigned numInputs = arity + (hasScalarOperand ? 1 : 0);
    if (numInputs == static_cast<unsigned>(arityGroupAndKind.arityGroup))
      return replaceOp(*kind);
  }
 
  return rewriter.notifyMatchFailure(
      genericOp, "elementwise operation cannot be specialized to category op");
}
 
//===----------------------------------------------------------------------===//
// Specialize linalg generic to reduce op.
//===----------------------------------------------------------------------===//
 
// Specialize a `linalg.generic` that is semantically a reduction into a
// `linalg.reduce`, moving the generic's payload into the reduce's combiner.
static FailureOr<LinalgOp>
specializeLinalgReductions(RewriterBase &rewriter, GenericOp genericOp,
                           ArrayRef<int64_t> reductionDims) {
  auto reduceOp = ReduceOp::create(
      rewriter, genericOp.getLoc(), genericOp.getDpsInputs(),
      genericOp.getDpsInits(), reductionDims, /*bodyBuild=*/nullptr);
 
  // Move the generic's payload region into the reduce's combiner.
  rewriter.inlineRegionBefore(genericOp.getRegion(), reduceOp.getCombiner(),
                              reduceOp.getCombiner().end());
  rewriter.replaceOp(genericOp, reduceOp);
  return cast<LinalgOp>(reduceOp.getOperation());
}
 
//===----------------------------------------------------------------------===//
// Specialize linalg generic to matmul variants.
//===----------------------------------------------------------------------===//
/// Identifies linalg.generic that is essentially named op of the form:
//    ` linalg.{batch_}?matmul{_transpose_a | _transpose_b}? `
//
// It is possible that a linalg.generic may be implementing a matmul but not
// in a straight-forward way e.g. below is matrix multiply over some slice
// ```
//  %0 = linalg.generic {
//          indexing_maps = [affine_map<(d0, d1, d2) -> (3, d1, d0)>,
//                           affine_map<(d0, d1, d2) -> (d0, 5, d2)>,
//                           affine_map<(d0, d1, d2) -> (d2, d1, 13)>],
//          iterator_types = ["parallel", "parallel", "parallel"]}
//          ins(%A, %B : tensor<20x20x20xf32>,  tensor<20x20x20xf32>)
//          outs(%C : tensor<20x20x20xf32>) {
//             ^bb0(%a: f32, %b: f32, %c : f32):
//                %mul = arith.mulf %a, %b : f32
//                %add = arith.addf %mul, %c : f32
//                linalg.yield %add : f32
//       } -> tensor<20x20x20xf32>
// ```
// It is not possible to represent above as named op.
// e.g. linalg.batch_matmul(%A, %B :  tensor<20x20x20xf32>, ...) is
// not  the same as linalg.generic above.
namespace {
enum class IndexMatchResult {
  Match = 0,  // identity map.
  Transposed, // transposed map.
  Mismatch    // none of the above.
};
 
// Checks whether the input Affine `map` contains two consecutive dims that
// can be interpreted as accessing a 2D matrix. It is assumed that the row
// column dimension are adjacent axis (in this order) and start at
// `rowDimIdx` in the input map.
//
//  e.g. consider A matrix in `C[M,N] = A[M,K] * B[K,N]`. We will check
//  whether the map of A is identity (match), transposed, or something
//  completely different (mis-match). Similar for B and C.
static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,
                                        unsigned expectedPosOfRowDim,
                                        unsigned expectedPosOfColDim) {
  // Get the matrix multiply indices. They are past the batch indices.
  auto exprOfRowDim = map.getResults()[rowDimIdx];
  auto exprOfColDim = map.getResults()[rowDimIdx + 1];
 
  // They should be pure dimension ids.
  if (exprOfRowDim.getKind() != AffineExprKind::DimId ||
      exprOfColDim.getKind() != AffineExprKind::DimId)
    return IndexMatchResult::Mismatch;
 
  auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();
  auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();
 
  if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)
    return IndexMatchResult::Match;
 
  if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)
    return IndexMatchResult::Transposed;
 
  return IndexMatchResult::Mismatch;
}
 
// Replaces genericOp with `NamedOpTy` op, supplied as a template arg.
// All the variants expressed as pseudo regular expression:
// `linalg.{batch_}?matmul` have same number of ins/out, so it's easy to
// stamp different versions.
// `castTy` is an optional type function that indicates whether (and which) cast
// attribute is needed for the named matmul op variant.
template <typename NamedOpTy>
static LinalgOp replaceWithMatmulVariant(RewriterBase &rewriter, GenericOp op,
                                         std::optional<TypeFn> castTy,
                                         ArrayRef<AffineMap> indexingMaps) {
  SmallVector<NamedAttribute> attributes;
  // Only explicitly specify the cast attribute for unsigned cast; signed is
  // the default for linalg.matmul/linalg.batch_matmul.
  if (castTy.has_value() && *castTy == TypeFn::cast_unsigned) {
    auto castAttr = rewriter.getNamedAttr(
        "cast", TypeFnAttr::get(rewriter.getContext(), *castTy));
    attributes.push_back(castAttr);
  }
 
  // Set the original generic's maps to preserve operand indexing semantics like
  // transposition.
  SmallVector<Attribute, 3> indexingMapsAttrVal =
      llvm::map_to_vector(indexingMaps, [](AffineMap map) -> Attribute {
        return AffineMapAttr::get(map);
      });
  auto indexingMapsAttr = rewriter.getNamedAttr(
      "indexing_maps", rewriter.getArrayAttr(indexingMapsAttrVal));
  attributes.push_back(indexingMapsAttr);
 
  LinalgOp namedOp = rewriter.replaceOpWithNewOp<NamedOpTy>(
      op, ValueRange{op.getDpsInputs()[0], op.getDpsInputs()[1]},
      ValueRange{op.getDpsInits()[0]}, attributes);
 
  return namedOp;
}
 
// Returns the cast type to use for a matmul-like named op. If the generic
// contains casts that cannot be represented (e.g. output casts or mixed
// signedness), return std::nullopt.
static std::optional<TypeFn> getCastTypeForMatmulLikeOp(GenericOp genericOp) {
  bool foundCastForMatmulOutput = false;
  SmallVector<TypeFn> castTyFns;
  genericOp.getBody()->walk([&](CastOpInterface castOp) {
    // Collect forward slice of the cast op to check if it is for the matmul
    // output.
    SetVector<Operation *> forwardSlice;
    getForwardSlice(castOp, &forwardSlice);
 
    // If there is no multiplication op in the forward slice, then this cast
    // op is for the matmul output. Cast ops on matmul output cannot be
    // expressed by the matmul op variant.
    if (!llvm::any_of(forwardSlice, [](Operation *op) {
          // We check explicitly for these multiplication ops in
          // `specializeLinalgContractions()` to infer matmul-like ops.
          return isa<arith::MulIOp, arith::MulFOp, complex::MulOp>(op);
        })) {
      foundCastForMatmulOutput = true;
      return WalkResult::interrupt();
    }
 
    // Determine the cast type.
    if (isa<arith::ExtUIOp, arith::UIToFPOp, arith::FPToUIOp>(castOp))
      castTyFns.push_back(TypeFn::cast_unsigned);
    else if (isa<arith::ExtSIOp, arith::SIToFPOp, arith::FPToSIOp>(castOp))
      castTyFns.push_back(TypeFn::cast_signed);
 
    return WalkResult::advance();
  });
 
  if (foundCastForMatmulOutput)
    return std::nullopt;
 
  if (!castTyFns.empty()) {
    // If there were multiple different cast types found, then we can't express
    // them using matmul-like ops. They only allow a single cast type for all
    // inputs.
    if (!llvm::all_equal(castTyFns))
      return std::nullopt;
    return castTyFns.front();
  }
 
  // Default to signed cast for matmul-like ops.
  return TypeFn::cast_signed;
}
 
static FailureOr<LinalgOp> specializeLinalgMmt4D(RewriterBase &rewriter,
                                                 GenericOp genericOp,
                                                 std::optional<TypeFn> castTy,
                                                 ContractionDimensions &dims) {
  // Should all be rank 4 and dim 6
  auto indexingMaps = genericOp.getIndexingMapsArray();
  if (llvm::any_of(indexingMaps, [](AffineMap m) {
        return m.getResults().size() != 4 || m.getNumDims() != 6;
      }))
    return failure();
 
  auto aOuter = matchOperandMap(indexingMaps[0], 0, dims.m[0], dims.k[0]);
  auto aInner = matchOperandMap(indexingMaps[0], 2, dims.m[1], dims.k[1]);
 
  auto bOuter = matchOperandMap(indexingMaps[1], 0, dims.k[0], dims.n[0]);
  auto bInner = matchOperandMap(indexingMaps[1], 2, dims.k[1], dims.n[1]);
 
  auto cOuter = matchOperandMap(indexingMaps[2], 0, dims.m[0], dims.n[0]);
  auto cInner = matchOperandMap(indexingMaps[2], 2, dims.m[1], dims.n[1]);
 
  if (llvm::is_contained({aOuter, bOuter, cOuter}, IndexMatchResult::Mismatch))
    return failure();
  if (llvm::is_contained({aInner, bInner, cInner}, IndexMatchResult::Mismatch))
    return failure();
 
  SmallVector<AffineMap> namedOpMaps = {indexingMaps[0], indexingMaps[1],
                                        indexingMaps[2]};
 
  return replaceWithMatmulVariant<Mmt4DOp>(rewriter, genericOp, castTy,
                                           namedOpMaps);
}
 
// Attempts to specialize `genericOp` to a specific named matmul variant
// (`matmul`, `batch_matmul`, or `mmt4d`). Returns failure without modifying the
// IR if no named variant matches. `castTy` is the cast type inferred for the
// contraction body.
static FailureOr<LinalgOp>
specializeToNamedContraction(RewriterBase &rewriter, GenericOp genericOp,
                             std::optional<TypeFn> castTy) {
  // Linalg generic contraction can be across multiple axis e.g.
  // ```
  //      linalg.generic
  //           {indexing_maps = [affine_map<(m, n, k1, k2) -> (m, k1, k2)>,
  //                             affine_map<(m, n, k1, k2) -> (k2, k1, n)>,
  //                             affine_map<(m, n, k1, k2) -> (m, n)>],
  //           iterator_types = ["parallel", "parallel",
  //                             "reduction", "reduction"]}
  //           ins(%A, %B : tensor<10x20x30xf32>, tensor<30x20x40xf32>)
  //           outs(%C : tensor<10x40xf32>) {
  //           ^bb0(%a: f32, %b: f32, %c: f32):
  //                 %1 = arith.mulf %a, %b : f32
  //                 %2 = arith.addf %c, %1 : f32
  //                 linalg.yield %2 : f32
  //      } -> tensor<10x40xf32>
  //  ```
  //  In above contraction, there are two reduction dimensions {k1, k2}
  //  and although a valid linalg contraction, it is not a named-op
  //  matrix multiply kind. Therefore, reject multi-dim reduction.
  auto res = inferContractionDims(genericOp);
  if (!succeeded(res))
    return failure();
  auto dims = *res;
  if (dims.m.size() == 2 && dims.n.size() == 2 && dims.k.size() == 2)
    return specializeLinalgMmt4D(rewriter, genericOp, castTy, dims);
  if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)
    return failure();
 
  // Check rank of operands
  auto indexingMaps = genericOp.getIndexingMapsArray();
  if (llvm::any_of(indexingMaps, [&dims](AffineMap m) {
        return m.getResults().size() !=
               dims.batch.size() + 2 /* any two of {m,n,k} */;
      }))
    return failure();
 
  auto numOfBatchDims = dims.batch.size();
  if (indexingMaps[0].getNumDims() != numOfBatchDims + 3)
    return failure();
 
  if (numOfBatchDims) {
    // Each operand in a linalg generic contraction  could express different
    // permutations for its batch dimension. But for named op it must be
    // identity since separate maps are not specified.
    if (llvm::any_of(indexingMaps, [numOfBatchDims](AffineMap m) {
          for (unsigned i = 0; i < numOfBatchDims; ++i) {
            auto expr = m.getResults()[i];
            if (expr.getKind() != AffineExprKind::DimId ||
                cast<AffineDimExpr>(expr).getPosition() != i)
              return true;
          }
          return false;
        }))
      return failure();
  }
 
  auto a =
      matchOperandMap(indexingMaps[0], numOfBatchDims, dims.m[0], dims.k[0]);
  auto b =
      matchOperandMap(indexingMaps[1], numOfBatchDims, dims.k[0], dims.n[0]);
  auto c =
      matchOperandMap(indexingMaps[2], numOfBatchDims, dims.m[0], dims.n[0]);
 
  if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))
    return failure();
 
  // Build indexing maps for the named op in its canonical dimension ordering
  auto *ctx = genericOp.getContext();
  unsigned numLoopDims = numOfBatchDims + 3;
  unsigned mIdx = numOfBatchDims;
  unsigned nIdx = mIdx + 1;
  unsigned kIdx = mIdx + 2;
 
  // TODO: add support for indexing_maps with broadcasts.
  auto makeMap = [&](IndexMatchResult match, unsigned rowIdx, unsigned colIdx) {
    SmallVector<unsigned> tensorDims;
    for (unsigned i = 0; i < numOfBatchDims; ++i)
      tensorDims.push_back(i);
    if (match == IndexMatchResult::Transposed)
      llvm::append_values(tensorDims, colIdx, rowIdx);
    else
      llvm::append_values(tensorDims, rowIdx, colIdx);
    return AffineMap::getMultiDimMapWithTargets(numLoopDims, tensorDims, ctx);
  };
 
  auto mapA = makeMap(a, mIdx, kIdx);
  auto mapB = makeMap(b, kIdx, nIdx);
  auto mapC = makeMap(c, mIdx, nIdx);
 
  SmallVector<AffineMap> namedOpMaps = {mapA, mapB, mapC};
 
  // Codegen the different matmul variants.
  if (numOfBatchDims) {
    return replaceWithMatmulVariant<BatchMatmulOp>(rewriter, genericOp, castTy,
                                                   namedOpMaps);
  }
  return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp, castTy,
                                            namedOpMaps);
}
 
// Converts linalg.generic to named linalg.*matmul* where possible, falling
// back to the generic `linalg.contract` op otherwise.
static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
                                                        GenericOp genericOp,
                                                        bool emitCategoryOp) {
  // Early exit if not projected permutations.
  auto mapRange = genericOp.getIndexingMapsArray();
  if (llvm::any_of(mapRange,
                   [](AffineMap m) { return !m.isProjectedPermutation(); }))
    return failure();
 
  // Determine the cast type for the named matmul op, or bail out if casts
  // cannot be represented by the named op.
  std::optional<TypeFn> castTy = getCastTypeForMatmulLikeOp(genericOp);
  if (!castTy)
    return rewriter.notifyMatchFailure(
        genericOp, "contains invalid cast ops for the named matmul op");
 
  // TODO: When `emitCategoryOp` is set, skip the named-variant matching below
  // and go straight to the `linalg.contract` category op.
 
  // Try to specialize to a specific named matmul variant first.
  if (!emitCategoryOp) {
    FailureOr<LinalgOp> namedOp =
        specializeToNamedContraction(rewriter, genericOp, castTy);
    if (succeeded(namedOp))
      return namedOp;
  }
 
  // No named variant matched; fall back to the generic `linalg.contract` op,
  // which supports a wider range of variants.
  return replaceWithMatmulVariant<ContractOp>(rewriter, genericOp, castTy,
                                              genericOp.getIndexingMapsArray());
}
 
/// Utility to specialize a `genericOp` with a convolution op of type `ConvOpTy`
/// with `dilations` and `strides`.
template <typename ConvOpTy>
static FailureOr<LinalgOp>
specializeToConvOp(RewriterBase &rewriter, GenericOp genericOp,
                   ArrayRef<int64_t> dilations, ArrayRef<int64_t> strides) {
  SmallVector<Value> inputs = genericOp.getDpsInputs();
  ValueRange outputs = genericOp.getDpsInits();
  SmallVector<Type> resultTypes = genericOp.hasPureTensorSemantics()
                                      ? TypeRange(ValueRange(outputs))
                                      : TypeRange{};
  LinalgOp namedOp;
  // Ops with no dilations and no strides.
  if constexpr (std::is_same_v<ConvOpTy, linalg::Conv1DOp> ||
                std::is_same_v<ConvOpTy, linalg::Conv2DOp> ||
                std::is_same_v<ConvOpTy, linalg::Conv3DOp>) {
    namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(genericOp, resultTypes,
                                                    inputs, outputs);
  } else {
    Attribute stridesAttr = rewriter.getI64TensorAttr(strides);
    Attribute dilationsAttr = rewriter.getI64TensorAttr(dilations);
    namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(
        genericOp, resultTypes, inputs, outputs, stridesAttr, dilationsAttr);
  }
  return namedOp;
}
 
/// Converts linalg.generic to named linalg.*conv/pooling* where possible.
static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,
                                                        GenericOp genericOp) {
#define CONV_OP_SPECIALIZER(ConvOpTy)                                          \
  if (std::optional<DilationsAndStrides> convParams =                          \
          matchConvolutionOpOfType<ConvOpTy>(genericOp))                       \
    return specializeToConvOp<ConvOpTy>(                                       \
        rewriter, genericOp, convParams->dilations, convParams->strides);      \
  // -----------------------------
  // Convolution ops.
  // -----------------------------
  CONV_OP_SPECIALIZER(linalg::Conv1DOp);
  CONV_OP_SPECIALIZER(linalg::Conv1DNwcWcfOp);
  CONV_OP_SPECIALIZER(linalg::Conv1DNcwFcwOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNhwcHwcfOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNhwcHwcfQOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNhwcFhwcOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNhwcFhwcQOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNchwFchwOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNchwFchwQOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNgchwFgchwOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNgchwGfchwOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNgchwGfchwQOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNhwgcGfhwcOp);
  CONV_OP_SPECIALIZER(linalg::Conv2DNhwgcGfhwcQOp);
  CONV_OP_SPECIALIZER(linalg::Conv3DOp);
  CONV_OP_SPECIALIZER(linalg::Conv3DNdhwcDhwcfOp);
  CONV_OP_SPECIALIZER(linalg::Conv3DNdhwcDhwcfQOp);
  CONV_OP_SPECIALIZER(linalg::Conv3DNcdhwFcdhwOp);
  // -----------------------------
  // Depthwise Convolution ops.
  // -----------------------------
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNcwCwOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNwcWcOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNwcWcmOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNchwChwOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcQOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcmOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcmQOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNdhwcDhwcOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNcdhwCdhwOp);
  CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNdhwcDhwcmOp);
  // -----------------------------
  // Pooling ops.
  // -----------------------------
  CONV_OP_SPECIALIZER(linalg::PoolingNhwcMaxOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNhwcMinOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNhwcSumOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNhwcMaxUnsignedOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNhwcMinUnsignedOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNchwSumOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNchwMaxOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNwcSumOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNcwSumOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNwcMaxOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNwcMaxUnsignedOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNcwMaxOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNwcMinOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNwcMinUnsignedOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNdhwcSumOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNdhwcMaxOp);
  CONV_OP_SPECIALIZER(linalg::PoolingNdhwcMinOp);
#undef CONV_OP_SPECIALIZER
  return failure();
}
 
} // namespace
 
//===----------------------------------------------------------------------===//
// Categorize linalg generic to named op where possible.
//===----------------------------------------------------------------------===//
FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(RewriterBase &rewriter,
                                                      GenericOp genericOp,
                                                      bool emitCategoryOps) {
  // Elementwise - e.g. exp, add are always category ops
  if (isaElemwiseSingleUnaryOpInterface(genericOp) ||
      isaElemwiseSingleBinaryOpInterface(genericOp) ||
      isaElemwiseSingleTernaryOpInterface(genericOp)) {
    return specializeLinalgElementwise(rewriter, genericOp);
  }
 
  // Contraction - e.g. matmul
  if (isaContractionOpInterface(genericOp)) {
    return specializeLinalgContractions(rewriter, genericOp, emitCategoryOps);
  }
 
  // Early exit in case of category specialization.
  // TODO: Remove when matches for other ops account for both named and
  // category.
  if (emitCategoryOps)
    return rewriter.notifyMatchFailure(
        genericOp, "no matching category op specialization");
 
  // Copy
  if (isaCopyOpInterface(genericOp)) {
    LinalgOp namedOp = rewriter.replaceOpWithNewOp<CopyOp>(
        genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0]);
    return namedOp;
  }
 
  // Fill
  if (std::optional<Value> fillValue = isaFillOpInterface(genericOp)) {
    // Always use the detected fill value, regardless of pattern
    LinalgOp namedOp = rewriter.replaceOpWithNewOp<FillOp>(
        genericOp, *fillValue, genericOp.getDpsInits()[0]);
    return namedOp;
  }
 
  // Broadcast
  std::optional<SmallVector<int64_t>> equivalentToBroadcast =
      isaBroadcastOpInterface(genericOp);
  if (equivalentToBroadcast) {
    auto dims = *equivalentToBroadcast;
    LinalgOp namedOp = rewriter.replaceOpWithNewOp<BroadcastOp>(
        genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0],
        dims);
    return namedOp;
  }
 
  // Transpose
  std::optional<SmallVector<int64_t>> equivalentToTranspose =
      isaTransposeOpInterface(genericOp);
  if (equivalentToTranspose) {
    auto permutation = *equivalentToTranspose;
    LinalgOp namedOp = rewriter.replaceOpWithNewOp<TransposeOp>(
        genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0],
        permutation);
    return namedOp;
  }
 
  // Reduction - e.g. linalg.reduce
  if (std::optional<SmallVector<int64_t>> reductionDims =
          isaReductionOpInterface(genericOp)) {
    return specializeLinalgReductions(rewriter, genericOp, *reductionDims);
  }
 
  // Convolution - e.g. *conv/pooling*
  if (isaConvolutionOpInterface(genericOp))
    return specializeLinalgConvolutions(rewriter, genericOp);
 
  return rewriter.notifyMatchFailure(genericOp,
                                     "no matching named op specialization");
}
 
namespace {
struct LinalgSpecializeGenericOpsPass
    : public impl::LinalgSpecializeGenericOpsPassBase<
          LinalgSpecializeGenericOpsPass> {
 
  using impl::LinalgSpecializeGenericOpsPassBase<
      LinalgSpecializeGenericOpsPass>::LinalgSpecializeGenericOpsPassBase;
  void runOnOperation() override;
};
} // namespace
 
void LinalgSpecializeGenericOpsPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  populateLinalgGenericOpsSpecializationPatterns(patterns);
  populateDecomposeProjectedPermutationPatterns(patterns);
 
  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
    signalPassFailure();
}
 
void mlir::linalg::populateLinalgGenericOpsSpecializationPatterns(
    RewritePatternSet &patterns, bool emitCategoryOps) {
  patterns.add<LinalgSpecializationPattern>(patterns.getContext(),
                                            emitCategoryOps);
}