llvm-external-projects/iree-dialects/lib/Transforms/TransformMatchers.cpp

// Copyright 2022 The IREE Authors
//
// Licensed 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

#include "iree-dialects/Transforms/TransformMatchers.h"

#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/Debug.h"

using namespace mlir;

#define DEBUG_TYPE "transform-matchers"
#define DBGS() llvm::dbgs() << "[" DEBUG_TYPE "] "

//===---------------------------------------------------------------------===//
// CapturingOpMatcher
//===---------------------------------------------------------------------===//

void transform_ext::CapturingOpMatcher::getAllNested(
    SmallVectorImpl<CapturingOpMatcher *> &nested) {
  int64_t start = nested.size();
  llvm::append_range(nested, nestedCapturingMatchers);
  for (int64_t position = start; position < nested.size(); ++position) {
    llvm::append_range(nested, nested[position]->nestedCapturingMatchers);
  }
}

void transform_ext::CapturingOpMatcher::getAllNestedValueMatchers(
    SmallVectorImpl<CapturingValueMatcher *> &nested) {
  llvm::append_range(nested, nestedCapturingValueMatchers);
}

//===---------------------------------------------------------------------===//
// ValueMatcher
//===---------------------------------------------------------------------===//

namespace {
struct DebugPrintValueWrapper {
  Value value;
};

llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                              const DebugPrintValueWrapper &wrapper) {
  if (auto opResult = wrapper.value.dyn_cast<OpResult>()) {
    return os << "op result #" << opResult.getResultNumber() << " in "
              << wrapper.value;
  }

  auto blockArg = wrapper.value.cast<BlockArgument>();
  os << "block argument #" << blockArg.getArgNumber();
  Block *parentBlock = blockArg.getParentBlock();
  Region *parentRegion = parentBlock->getParent();
  if (!parentRegion) {
    os << " of a detached block:\n";
    parentBlock->print(os);
    return os;
  }

  os << " of block #"
     << std::distance(parentRegion->begin(), parentBlock->getIterator());
  Operation *parentOp = parentRegion->getParentOp();
  if (!parentOp) {
    os << " of a detached region:\n";
    for (Block &b : *parentRegion)
      b.print(os);
    return os;
  }

  os << " in region #" << parentRegion->getRegionNumber() << " of "
     << *parentOp;
  return os;
}
} // namespace

bool transform_ext::ValueMatcher::match(Value value) {
  auto debugRAII =
      llvm::make_scope_exit([] { LLVM_DEBUG(DBGS() << "-------\n"); });
  LLVM_DEBUG(DBGS() << "matching " << DebugPrintValueWrapper{value} << "\n");

  if (getCaptured()) {
    LLVM_DEBUG(DBGS() << "found an already captured value: ");
    if (getCaptured() == value) {
      LLVM_DEBUG(llvm::dbgs() << "same\n");
      return true;
    } else {
      LLVM_DEBUG(llvm::dbgs() << "different\n");
      return false;
    }
  }

  for (const PredicateFn &fn : predicates) {
    bool result = fn(value);
    LLVM_DEBUG(llvm::dbgs() << ": " << result << "\n");
    if (!result)
      return false;
  }

  captured = value;
  return true;
}

//===---------------------------------------------------------------------===//
// StructuredOpMatcher and friends.
//===---------------------------------------------------------------------===//

void transform_ext::StructuredOpMatcher::debugOutputForCreate(
    ArrayRef<StringRef> opNames) {
  LLVM_DEBUG(DBGS() << "operation type is one of {";
             llvm::interleaveComma(opNames, llvm::dbgs()); llvm::dbgs() << "}");
}

bool transform_ext::StructuredOpMatcher::match(Operation *op) {
  auto debugRAII =
      llvm::make_scope_exit([] { LLVM_DEBUG(DBGS() << "-------\n"); });
  LLVM_DEBUG(DBGS() << "matching: " << *op << "\n");

  if (getCaptured()) {
    LLVM_DEBUG(DBGS() << "found an already captured op: ");
    if (getCaptured() == op) {
      LLVM_DEBUG(llvm::dbgs() << "same\n");
      return true;
    } else {
      LLVM_DEBUG(llvm::dbgs() << "different\n");
      return false;
    }
  }

  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
  if (!linalgOp) {
    LLVM_DEBUG(DBGS() << "not a structured op\n");
    return false;
  }

  if (!llvm::all_of(predicates, [linalgOp](const PredicateFn &fn) {
        bool result = fn(linalgOp);
        LLVM_DEBUG(llvm::dbgs() << ": " << result << "\n");
        return result;
      })) {
    return false;
  }

  captured = linalgOp;
  return true;
}

//===---------------------------------------------------------------------===//
// Constraints on op rank and dims.
//===---------------------------------------------------------------------===//

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::rank(NumGreaterEqualTo minRank) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "rank >= " << minRank.value);
    return linalgOp.getNumLoops() >= minRank.value;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::rank(NumLowerEqualTo maxRank) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "rank <= " << maxRank.value);
    return linalgOp.getNumLoops() <= maxRank.value;
  });
  return *this;
}

StringRef stringifyShapeKind(transform_ext::ShapeKind kind) {
  switch (kind) {
  case transform_ext::ShapeKind::Static:
    return "static";
  case transform_ext::ShapeKind::Dynamic:
    return "dynamic";
  }
  llvm_unreachable("unhandled shape kind");
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(SmallVector<int64_t> &&dimensions,
                                        ShapeKind kind) {
  predicates.push_back([dimensions = std::move(dimensions),
                        kind](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "dimensions [";
               llvm::interleaveComma(dimensions, llvm::dbgs());
               llvm::dbgs() << "] are " << stringifyShapeKind(kind));
    SmallVector<int64_t> shape = linalgOp.getStaticLoopRanges();
    for (auto dimension : dimensions) {
      int64_t transformedDimension =
          dimension >= 0 ? dimension : shape.size() + dimension;
      if (transformedDimension < 0 || transformedDimension >= shape.size())
        return false;
      if (ShapedType::isDynamic(shape[transformedDimension]) ^
          (kind == ShapeKind::Static))
        continue;
      return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(AllDims tag, ShapeKind kind) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all dimensions are " << stringifyShapeKind(kind));
    SmallVector<int64_t> shape = linalgOp.getStaticLoopRanges();
    return llvm::all_of(shape, [=](int64_t dimension) {
      return ShapedType::isDynamic(dimension) ^ (kind == ShapeKind::Static);
    });
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(SmallVector<int64_t> &&dimensions,
                                        utils::IteratorType kind) {
  predicates.push_back([dimensions = std::move(dimensions),
                        kind](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "dimensions [";
               llvm::interleaveComma(dimensions, llvm::dbgs());
               llvm::dbgs() << "] are " << utils::stringifyIteratorType(kind));
    int64_t rank = linalgOp.getNumLoops();
    for (auto dimension : dimensions) {
      int64_t transformedDimension =
          dimension >= 0 ? dimension : rank + dimension;
      if (transformedDimension < 0 || transformedDimension >= rank)
        return false;
      utils::IteratorType iteratorKind =
          linalgOp.getIteratorTypesArray()[transformedDimension];
      if (iteratorKind == kind)
        continue;
      return false;
    }
    return true;
  });
  return *this;
}
transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(AllDims tag, utils::IteratorType kind) {
  return dim(AllDimsExcept({}), kind);
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(AllDimsExcept &&dims,
                                        utils::IteratorType kind) {
  predicates.push_back([dimensions = std::move(dims),
                        kind](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all dimensions except [";
               llvm::interleaveComma(dimensions.getExcluded(), llvm::dbgs());
               llvm::dbgs() << "] are " << utils::stringifyIteratorType(kind));
    int64_t rank = linalgOp.getNumLoops();
    llvm::SmallDenseSet<int64_t> excludedDims;
    for (int64_t dim : dimensions.getExcluded()) {
      excludedDims.insert(dim >= 0 ? dim : rank + dim);
    }

    for (auto [index, type] :
         llvm::enumerate(linalgOp.getIteratorTypesArray())) {
      if (excludedDims.contains(index))
        continue;
      if (type == kind)
        continue;
      return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(int64_t dimension,
                                        DivisibleBy divisibleBy) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "dimension " << dimension << " is divisible by "
                      << divisibleBy.value);
    int64_t rank = linalgOp.getNumLoops();
    int64_t transformedDimension =
        dimension >= 0 ? dimension : rank + dimension;
    if (transformedDimension >= rank)
      return false;

    int64_t size = linalgOp.getStaticLoopRanges()[transformedDimension];
    return !ShapedType::isDynamic(size) && (size % divisibleBy.value == 0);
  });
  return *this;
}

//===---------------------------------------------------------------------===//
// Capture directives.
//===---------------------------------------------------------------------===//
transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::rank(CaptureRank capture) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "capture rank");
    capture.value = linalgOp.getNumLoops();
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(int64_t dimension, CaptureDim capture) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "capture dimension");
    int64_t rank = linalgOp.getNumLoops();
    int64_t transformedDimension =
        dimension >= 0 ? dimension : rank + dimension;
    if (transformedDimension >= rank)
      return false;

    capture.value = linalgOp.getStaticLoopRanges()[transformedDimension];
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::dim(AllDims tag, CaptureDims captures) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "capture all dimensions");
    captures.value = linalgOp.getStaticLoopRanges();
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher::StructuredOpMatcher(
    StructuredOpMatcher &A, StructuredOpMatcher &B) {

  predicates.push_back([&A, &B](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "start recursive lhs OR match {\n");
    {
      auto debugRAII = llvm::make_scope_exit(
          [] { LLVM_DEBUG(DBGS() << "} end recursive match"); });
      if (A.match(linalgOp))
        return true;
    }
    LLVM_DEBUG(DBGS() << "start recursive rhs OR match {\n");
    {
      auto debugRAII = llvm::make_scope_exit(
          [] { LLVM_DEBUG(DBGS() << "} end recursive match"); });
      if (B.match(linalgOp))
        return true;
    }
    return false;
  });
  recordNestedMatcher(A);
  recordNestedMatcher(B);
}

//===---------------------------------------------------------------------===//
// Constraints on input operands.
//===---------------------------------------------------------------------===//

void transform_ext::StructuredOpMatcher::addInputMatcher(
    int64_t position, std::function<bool(Operation *)> matcher,
    OptionalMatch optional) {
  addInputMatcher(
      position,
      // No need to handle optional inside the lambda, the wrapper will do that.
      [matcher = std::move(matcher)](Value value) {
        Operation *definingOp = value.getDefiningOp();
        return definingOp && matcher(definingOp);
      },
      optional);
}

void transform_ext::StructuredOpMatcher::addInputMatcher(
    int64_t position, std::function<bool(Value)> matcher,
    OptionalMatch optional) {
  predicates.push_back([position, optional, matcher = std::move(matcher)](
                           linalg::LinalgOp linalgOp) -> bool {
    int64_t transformedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInputs() + position;
    if (transformedPosition >= linalgOp.getNumDpsInputs()) {
      LLVM_DEBUG(DBGS() << "input operand #" << position
                        << " does not exist but match required");
      return false;
    }

    LLVM_DEBUG(DBGS() << "input operand #" << position
                      << (optional.value ? " (optional match) " : " ")
                      << "is\n");

    // We MUST run the matcher at this point, even if the match is optional,
    // to allow for capture.
    LLVM_DEBUG(DBGS() << "start recursive match {\n");
    auto debugRAII = llvm::make_scope_exit(
        [] { LLVM_DEBUG(DBGS() << "} end recursive match"); });
    if (matcher(linalgOp.getDpsInputOperand(transformedPosition)->get()))
      return true;
    return optional.value;
  });
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(AllOperands tag, IsPermutation) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all input operands have permutation maps");
    // all_of with a lambda requires const-casting dance, so using a loop.
    for (OpOperand *operand : linalgOp.getDpsInputOperands()) {
      if (!linalgOp.getMatchingIndexingMap(operand).isPermutation())
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(AllOperands tag,
                                          IsProjectedPermutation) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all input operands have projected permutation maps");
    // all_of with a lambda requires const-casting dance, so using a loop.
    for (OpOperand *operand : linalgOp.getDpsInputOperands()) {
      if (!linalgOp.getMatchingIndexingMap(operand).isProjectedPermutation())
        return false;
    }
    return true;
  });
  return *this;
}

/// Helper to check if the map is an identity map with a projected dim.
static bool isProjectedMap(AffineMap map, int64_t projectedDim) {
  if (!map.isProjectedPermutation())
    return false;
  int64_t dimCounter = 0;
  for (unsigned i = 0, e = map.getNumResults(); i < e; i++) {
    // Skip the project dim.
    if (dimCounter == projectedDim)
      dimCounter++;
    if (map.getDimPosition(i) != dimCounter++) {
      return false;
    }
  }
  return true;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(SmallVector<int64_t> &&positions,
                                          IsProjected dim) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "operands ";
               llvm::interleaveComma(positions, llvm::dbgs());
               llvm::dbgs() << " have a permutation maps with " << dim.value
                            << " projected");
    int64_t updatedDim =
        dim.value >= 0 ? dim.value : linalgOp.getNumLoops() + dim.value;
    for (int64_t position : positions) {
      OpOperand *operand = linalgOp.getDpsInputOperand(position);
      if (!isProjectedMap(linalgOp.getMatchingIndexingMap(operand), updatedDim))
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(AllOperands tag, IsIdentity) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all input operands have identity maps");
    // all_of with a lambda requires const-casting dance, so using a loop.
    for (OpOperand *operand : linalgOp.getDpsInputOperands()) {
      if (!linalgOp.getMatchingIndexingMap(operand).isIdentity())
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(SmallVector<int64_t> &&positions,
                                          IsIdentity) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "input operands ";
               llvm::interleaveComma(positions, llvm::dbgs());
               llvm::dbgs() << " have identity maps");
    // all_of with a lambda requires const-casting dance, so using a loop.
    for (int64_t position : positions) {
      int64_t updatedPosition =
          position >= 0 ? position : linalgOp.getNumDpsInputs() + position;
      OpOperand *operand = linalgOp.getDpsInputOperand(updatedPosition);
      if (!linalgOp.getMatchingIndexingMap(operand).isIdentity())
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(int64_t position,
                                          ElementTypeBitWidth width) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "input operand #" << position
                      << " has elemental type with bit width " << width.value);
    int64_t updatedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInputs() + position;
    if (0 < updatedPosition || updatedPosition >= linalgOp.getNumDpsInputs())
      return false;
    auto shapedType = linalgOp.getDpsInputOperand(updatedPosition)
                          ->get()
                          .getType()
                          .dyn_cast<ShapedType>();
    return shapedType && shapedType.getElementType().isIntOrFloat() &&
           shapedType.getElementType().getIntOrFloatBitWidth() == width.value;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(int64_t position,
                                          CaptureElementTypeBitWidth width) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "input operand #" << position << " capture bitwidth");
    int64_t updatedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInputs() + position;
    if (0 < updatedPosition || updatedPosition >= linalgOp.getNumDpsInputs())
      return false;
    auto shapedType = linalgOp.getDpsInputOperand(updatedPosition)
                          ->get()
                          .getType()
                          .dyn_cast<ShapedType>();
    if (!shapedType || !shapedType.getElementType().isIntOrFloat())
      return false;
    width.value = shapedType.getElementType().getIntOrFloatBitWidth();
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(NumEqualsTo num) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "number of input operands == " << num.value);
    return linalgOp.getNumDpsInputs() == num.value;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(int64_t position,
                                          ConstantFloatMinOrMinusInf) {
  return input(position, [](llvm::APFloat f) {
    return (f.isLargest() || f.isInfinity()) && f.isNegative();
  });
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::input(int64_t position, ConstantFloatZero) {
  return input(position, [](llvm::APFloat f) { return f.isZero(); });
}

transform_ext::StructuredOpMatcher &transform_ext::StructuredOpMatcher::input(
    int64_t position, std::function<bool(llvm::APFloat)> floatValueFn) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "input operands " << position
                      << " is a special floating point constant");
    int64_t updatedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInputs() + position;
    if (0 > updatedPosition || updatedPosition >= linalgOp.getNumDpsInputs())
      return false;
    auto cstOp = linalgOp.getDpsInputOperand(updatedPosition)
                     ->get()
                     .getDefiningOp<arith::ConstantFloatOp>();
    if (!cstOp)
      return false;
    return floatValueFn(cstOp.value());
  });
  return *this;
}

//===---------------------------------------------------------------------===//
// Constraints on output operands.
//===---------------------------------------------------------------------===//

void transform_ext::StructuredOpMatcher::addOutputMatcher(
    int64_t position, std::function<bool(Operation *)> matcher,
    OptionalMatch optional) {
  predicates.push_back([position, optional, matcher = std::move(matcher)](
                           linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "output operand #" << position
                      << (optional.value ? " (optional match) " : " ")
                      << "is produced by\n");
    int64_t transformedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInits() + position;
    if (transformedPosition >= linalgOp.getNumDpsInits())
      return false;

    Operation *definingOp =
        linalgOp.getDpsInitOperand(transformedPosition)->get().getDefiningOp();
    if (!definingOp)
      return optional.value;
    // We MUST run the matcher at this point, even if the match is optional,
    // to allow for capture.
    LLVM_DEBUG(DBGS() << "start recursive match {\n");
    auto debugRAII = llvm::make_scope_exit(
        [] { LLVM_DEBUG(DBGS() << "} end recursive match"); });
    if (matcher(definingOp))
      return true;
    return optional.value;
  });
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(AllOperands tag, IsPermutation) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all output operands have permutation maps");
    for (OpOperand *operand : linalgOp.getDpsInitOperands()) {
      if (!linalgOp.getMatchingIndexingMap(operand).isPermutation())
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(AllOperands tag,
                                           IsProjectedPermutation) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all output operands have projected permutation maps");
    for (OpOperand *operand : linalgOp.getDpsInitOperands()) {
      if (!linalgOp.getMatchingIndexingMap(operand).isProjectedPermutation())
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(AllOperands tag, IsProjected dim) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all output operands have a maps with projected");
    int64_t updatedDim =
        dim.value >= 0 ? dim.value : linalgOp.getNumLoops() + dim.value;
    // all_of with a lambda requires const-casting dance, so using a loop.
    for (OpOperand *operand : linalgOp.getDpsInitOperands()) {
      if (!isProjectedMap(linalgOp.getMatchingIndexingMap(operand), updatedDim))
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(AllOperands tag, IsIdentity) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "all output operands have identity permutation maps");
    for (OpOperand *operand : linalgOp.getDpsInitOperands()) {
      if (!linalgOp.getMatchingIndexingMap(operand).isIdentity())
        return false;
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(int64_t position,
                                           ElementTypeBitWidth width) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "output operand #" << position
                      << " has elemental type with bit width " << width.value);
    int64_t updatedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInits() + position;
    if (0 < updatedPosition || updatedPosition >= linalgOp.getNumDpsInits())
      return false;
    auto shapedType = linalgOp.getDpsInitOperand(updatedPosition)
                          ->get()
                          .getType()
                          .dyn_cast<ShapedType>();
    return shapedType && shapedType.getElementType().isIntOrFloat() &&
           shapedType.getElementType().getIntOrFloatBitWidth() == width.value;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(int64_t position,
                                           CaptureElementTypeBitWidth width) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "output operand #" << position << " capture bitwidth");
    int64_t updatedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInits() + position;
    if (0 < updatedPosition || updatedPosition >= linalgOp.getNumDpsInits())
      return false;
    auto shapedType = linalgOp.getDpsInitOperand(updatedPosition)
                          ->get()
                          .getType()
                          .dyn_cast<ShapedType>();
    if (!shapedType || !shapedType.getElementType().isIntOrFloat())
      return false;
    width.value = shapedType.getElementType().getIntOrFloatBitWidth();
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(int64_t position,
                                           SingleCombinerReduction tag) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "output operand #" << position
                      << " is populated by a single-combiner reduction");
    int64_t updatedPosition =
        position >= 0 ? position : linalgOp.getNumDpsInits() + position;
    if (0 < updatedPosition || updatedPosition >= linalgOp.getNumDpsInits())
      return false;
    SmallVector<Operation *> combinerOps;
    return matchReduction(linalgOp.getRegionOutputArgs(), updatedPosition,
                          combinerOps) &&
           llvm::hasSingleElement(combinerOps);
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::output(NumEqualsTo num) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "number of output operands == " << num.value);
    return linalgOp.getNumDpsInits() == num.value;
  });
  return *this;
}

//===---------------------------------------------------------------------===//
// Constraints on results.
//===---------------------------------------------------------------------===//

void transform_ext::StructuredOpMatcher::addResultMatcher(
    int64_t position, HasAnyUse tag, std::function<bool(Operation *)> matcher,
    OptionalMatch optional) {
  predicates.push_back([matcher = std::move(matcher), optional,
                        position](linalg::LinalgOp linalgOp) -> bool {
    LLVM_DEBUG(DBGS() << "result #" << position
                      << (optional.value ? " (optional match) " : " ")
                      << "has a use\n");
    int64_t transformedPosition =
        position >= 0 ? position : linalgOp->getNumResults() + position;
    if (transformedPosition >= linalgOp->getNumResults())
      return false;

    // We MUST run the matcher at this point, even if the match is optional,
    // to allow for capture.
    LLVM_DEBUG(DBGS() << "start recursive match {\n");
    auto debugRAII = llvm::make_scope_exit(
        [] { LLVM_DEBUG(DBGS() << "} end recursive match"); });
    if (llvm::any_of(linalgOp->getResult(transformedPosition).getUsers(),
                     [&matcher](Operation *op) { return matcher(op); })) {
      return true;
    }
    return optional.value;
  });
}

bool transform_ext::StructuredOpMatcher::checkAllTilableMatched(
    Operation *parent, linalg::LinalgOp linalgOp,
    ArrayRef<transform_ext::CapturingOpMatcher *> matchers) {
  LLVM_DEBUG(DBGS() << "all tilable ops captured");
  int64_t numTilableOps = 0;
  if (!parent)
    return false;
  parent->walk([&](TilingInterface Op) { ++numTilableOps; });

  llvm::SmallPtrSet<Operation *, 6> matched;
  for (CapturingOpMatcher *nested : matchers) {
    if (Operation *captured = nested->getCaptured()) {
      matched.insert(captured);
    }
  }

  // Don't forget to include the root matcher.
  matched.insert(linalgOp);
  return numTilableOps == matched.size();
}

//===-------------------------------------------------------------------===//
// Constraints on op region.
//===-------------------------------------------------------------------===//

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::singleOpWithCanonicaleArgs(
    StringRef opcode, bool commutative) {
  predicates.push_back([=](linalg::LinalgOp linalgOp) {
    if (linalgOp.getBlock()->getOperations().size() != 2)
      return false;
    Operation *innerOp = &(*linalgOp.getBlock()->getOperations().begin());
    if (innerOp->getName().getStringRef() != opcode ||
        innerOp->getNumResults() != 1)
      return false;
    Operation *yieldOp = linalgOp.getBlock()->getTerminator();
    if (yieldOp->getNumOperands() != 1)
      return false;
    if (yieldOp->getOperand(0).getDefiningOp() != innerOp)
      return false;
    if (commutative && innerOp->getNumOperands() == 2) {
      auto arg0 = dyn_cast<BlockArgument>(innerOp->getOperand(0));
      auto arg1 = dyn_cast<BlockArgument>(innerOp->getOperand(1));
      if (!arg0 || !arg1)
        return false;
      if (arg0.getParentBlock() != linalgOp.getBlock() ||
          arg1.getParentBlock() != linalgOp.getBlock())
        return false;
      if (!((arg0.getArgNumber() == 0 && arg1.getArgNumber() == 1) ||
            (arg1.getArgNumber() == 0 && arg0.getArgNumber() == 1)))
        return false;
    } else {
      for (auto [index, operand] : llvm::enumerate(innerOp->getOperands())) {
        auto arg = dyn_cast<BlockArgument>(operand);
        if (!arg || arg.getParentBlock() != linalgOp.getBlock() ||
            arg.getArgNumber() != index)
          return false;
      }
    }
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::isFloatReciprocal() {
  predicates.push_back([=](linalg::LinalgOp linalgOp) {
    LLVM_DEBUG(DBGS() << "op region represents a reciprocal operation");
    if (linalgOp.getBlock()->getOperations().size() != 2)
      return false;
    Operation *innerOp = &(*linalgOp.getBlock()->getOperations().begin());
    if (!isa<arith::DivFOp>(innerOp) || innerOp->getNumResults() != 1)
      return false;
    Operation *yieldOp = linalgOp.getBlock()->getTerminator();
    if (yieldOp->getNumOperands() != 1)
      return false;
    if (yieldOp->getOperand(0).getDefiningOp() != innerOp)
      return false;
    auto cst = innerOp->getOperand(0).getDefiningOp<arith::ConstantFloatOp>();
    if (!cst || cst.value().convertToDouble() != 1.0)
      return false;
    auto arg = dyn_cast<BlockArgument>(innerOp->getOperand(1));
    if (!arg || arg.getParentBlock() != linalgOp.getBlock() ||
        arg.getArgNumber() != 0)
      return false;
    return true;
  });
  return *this;
}

transform_ext::StructuredOpMatcher &
transform_ext::StructuredOpMatcher::passThroughOp() {
  predicates.push_back([=](linalg::LinalgOp linalgOp) {
    if (linalgOp.getBlock()->getOperations().size() != 1)
      return false;
    Operation *yieldOp = linalgOp.getBlock()->getTerminator();
    for (auto [index, operand] : llvm::enumerate(yieldOp->getOperands())) {
      auto arg = dyn_cast<BlockArgument>(operand);
      if (!arg || arg.getParentBlock() != linalgOp.getBlock() ||
          arg.getArgNumber() != index)
        return false;
    }
    return true;
  });
  return *this;
}

//===---------------------------------------------------------------------===//
// MatchCallbackResult.
//===---------------------------------------------------------------------===//

ArrayRef<Operation *>
transform_ext::MatchCallbackResult::getPayloadGroup(int64_t position) const {
  assert(position < payloadGroupLengths.size());
  int64_t start = 0;
  for (int64_t i = 0; i < position; ++i) {
    start += payloadGroupLengths[i];
  }
  return llvm::ArrayRef(payloadOperations)
      .slice(start, payloadGroupLengths[position]);
}

//===---------------------------------------------------------------------===//
// Case-specific matcher builders.
//===---------------------------------------------------------------------===//

static constexpr int64_t kCudaWarpSize = 32;

void transform_ext::makeReductionMatcher(
    transform_ext::MatcherContext &matcherContext,
    transform_ext::StructuredOpMatcher *&reductionCapture,
    transform_ext::StructuredOpMatcher *&fillCapture,
    transform_ext::StructuredOpMatcher *&leadingCapture,
    transform_ext::StructuredOpMatcher *&trailingCapture,
    MatchedReductionCaptures &captures) {
  // The core part of the matcher is anchored on a particular reduction op.
  auto &reduction =
      m_StructuredOp(matcherContext)
          // Op has at least a parallel a reduction dimension and at
          // most 3 parallel dimensions.
          // TODO: relax once we have global collapse/expand_shape.
          //
          .rank(NumGreaterEqualTo(2))
          .rank(NumLowerEqualTo(4))
          .rank(CaptureRank(captures.reductionRank))
          // Op has a single most-minor reduction.
          .dim(-1, utils::IteratorType::reduction)
          // Capture op sizes.
          .dim(AllDims(), CaptureDims(captures.reductionOpSizes))
          // All other dimensions are parallel.
          .dim(AllDimsExcept({-1}), utils::IteratorType::parallel)
          // Single input for now, can be arbitrary projected permutations.
          // TODO: Multiple inputs, can be arbitrary projected permutations.
          // TODO: Watch out for multiple inputs though as a reduction turns
          //       into a contraction when mixed with projected
          //       permutations. A reduction is often bandwidth bound but
          //       contraction is a different beast that is compute bound
          //       and has a very different schedule.
          //
          .input(NumEqualsTo(1))
          .input(AllOperands(), IsProjectedPermutation())
          // Single output supported atm.
          // TODO: Multiple outputs.
          //
          .output(NumEqualsTo(1))
          // A reduction output must be a projected permutation, match it but we
          // could also drop this technically.
          .output(AllOperands(), IsProjectedPermutation())
          // Only single combiner for now due to reduction warp
          // distribution.
          // TODO: relax this once reduction distribution is more powerful.
          //
          .output(0, CaptureElementTypeBitWidth(
                         captures.reductionOutputElementalTypeBitWidth))
          .output(0, SingleCombinerReduction());
  reductionCapture = &reduction;

  // Mandatory FillOp must create the unique output of the reduction.
  // TODO: Relax this, as any map, broadcast, transpose should also work.
  //
  auto &fill = m_StructuredOp<linalg::FillOp>(matcherContext);
  reduction = reduction.output(NumEqualsTo(1)).output(0, fill);
  fillCapture = &fill;

  // Optional leading or trailing op can be any map, transpose, broadcast but
  // not reduce or windowing operation for now.
  // It must create the unique input for the reduction.
  // TODO: match more optional leading ops, one per input of the reduction.
  // TODO: careful about multi-output and turning into a contraction.
  //
  transform_ext::StructuredOpMatcher commonLeadingOrTrailing =
      m_StructuredOp<linalg::GenericOp>(matcherContext)
          // All parallel dimensions.
          .dim(AllDims(), utils::IteratorType::parallel)
          // All inputs are any projected permutation.
          .input(AllOperands(), IsProjectedPermutation())
          .output(AllOperands(), IsPermutation())
          // leading and trailing may have 0, 1 or more input as long as they do
          // not come from unmatched ops. This extra constraint is taken care of
          // separately. This is also a noop but we document it.
          // TODO: Base and derived classes, atm this does not compile.
          // .input(NumGreaterEqualTo(0))
          // Single output supported atm.
          // TODO: extend this.
          //
          .output(NumEqualsTo(1));
  // TODO: match more optional leading ops, one per input of the reduction.
  // TODO: careful about multi-output and turning into a contraction.
  //
  auto &leading =
      m_StructuredOp(matcherContext, commonLeadingOrTrailing)
          .rank(CaptureRank(captures.maybeLeadingRank))
          // Capture op sizes.
          .dim(AllDims(), CaptureDims(captures.leadingOpSizes))
          // Capture output elemental type.
          .output(0, CaptureElementTypeBitWidth(
                         captures.maybeLeadingOutputElementalTypeBitWidth));
  reduction = reduction.input(0, leading, OptionalMatch());
  leadingCapture = &leading;

  // Optional trailing can be any map, transpose, broadcast but not reduce or
  // windowing operation for now.
  // It must be fed by the unique input for the reduction.
  // TODO: match more optional leading ops, one per input of the reduction.
  // TODO: careful about multi-output and turning into a contraction.
  //
  auto &trailing =
      m_StructuredOp(matcherContext, commonLeadingOrTrailing)
          .rank(CaptureRank(captures.maybeTrailingRank))
          // Capture op sizes.
          .dim(AllDims(), CaptureDims(captures.trailingOpSizes))
          // Capture output elemental type.
          .output(0, CaptureElementTypeBitWidth(
                         captures.maybeTrailingOutputElementalTypeBitWidth));
  reduction = reduction.result(0, HasAnyUse(), trailing, OptionalMatch())
                  .allTilableOpsCaptured<func::FuncOp>();
  trailingCapture = &trailing;
}

void transform_ext::makeReductionMatcher(transform_ext::MatcherContext &context,
                                         StructuredOpMatcher *&reductionCapture,
                                         MatchedReductionCaptures &captures) {
  StructuredOpMatcher *fill;
  StructuredOpMatcher *leading;
  StructuredOpMatcher *trailing;
  makeReductionMatcher(context, reductionCapture, fill, leading, trailing,
                       captures);
}

/// Match sum(%src, broadcast(%reduction))
static void matchSubBroadcat(transform_ext::MatcherContext &matcherContext,
                             transform_ext::StructuredOpMatcher &maxReduction,
                             transform_ext::ValueMatcher &softmaxSourceOperand,
                             transform_ext::StructuredOpMatcher *&sub) {
  using namespace transform_ext;

  auto &broadcast =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .passThroughOp()
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(1))
          .input(0, IsProjected(-1))
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());
  broadcast = broadcast.input(0, maxReduction);

  auto &subParallel =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::SubFOp>()
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(2))
          .input(0, IsIdentity())
          .input(1, IsIdentity())
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());
  subParallel = subParallel.input(0, softmaxSourceOperand);
  subParallel = subParallel.input(1, broadcast);

  auto &subBroadcast =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::SubFOp>()
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(2))
          .input(0, IsIdentity())
          .input(1, IsProjected(-1))
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());
  subBroadcast = subBroadcast.input(0, softmaxSourceOperand);
  subBroadcast = subBroadcast.input(1, maxReduction);
  auto &subOr = transform_ext::m_StructuredOp_Or(matcherContext, subBroadcast,
                                                 subParallel);
  sub = &subOr;
}

/// Match sum(%exp, broadcast(%sum))
static void matchdivBroadcast(transform_ext::MatcherContext &matcherContext,
                              transform_ext::StructuredOpMatcher &expOperand,
                              transform_ext::StructuredOpMatcher &sum,
                              transform_ext::StructuredOpMatcher *&div) {
  using namespace transform_ext;

  auto &broadcast =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .passThroughOp()
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(1))
          .input(0, IsProjected(-1))
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());
  broadcast = broadcast.input(0, sum);

  auto &divNoBroadcast =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::DivFOp>()
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(2))
          .input(0, IsIdentity())
          .input(1, IsIdentity())
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());

  divNoBroadcast = divNoBroadcast.input(0, expOperand);
  divNoBroadcast = divNoBroadcast.input(1, broadcast);

  auto &divBroadcast =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::DivFOp>()
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(2))
          .input(0, IsIdentity())
          .input(1, IsProjected(-1))
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());

  divBroadcast = divBroadcast.input(0, expOperand);
  divBroadcast = divBroadcast.input(1, sum);

  auto &divMerge = transform_ext::m_StructuredOp_Or(
      matcherContext, divNoBroadcast, divBroadcast);
  div = &divMerge;
}

void transform_ext::makeSoftmaxMatcher(
    transform_ext::MatcherContext &matcherContext,
    transform_ext::StructuredOpMatcher *&maxReductionCapture,
    transform_ext::StructuredOpMatcher *&softmaxRootCapture) {
  auto &softmaxSourceOperand = m_Value(matcherContext);

  auto &fillMinusInf = m_StructuredOp<linalg::FillOp>(matcherContext)
                           .input(0, ConstantFloatMinOrMinusInf());
  auto &maxReduction =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::MaxFOp>(/*commutative=*/true)
          // Only handle most inner reduction for now.
          .dim(-1, utils::IteratorType::reduction)
          .dim(AllDimsExcept({-1}), utils::IteratorType::parallel)
          .input(NumEqualsTo(1))
          .input(AllOperands(), IsIdentity())
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsProjected(-1));
  maxReduction = maxReduction.input(0, softmaxSourceOperand);
  maxReduction = maxReduction.output(0, fillMinusInf);
  maxReductionCapture = &maxReduction;

  transform_ext::StructuredOpMatcher *subOperand;
  matchSubBroadcat(matcherContext, maxReduction, softmaxSourceOperand,
                   subOperand);

  auto &expOperand = m_StructuredOp<linalg::GenericOp>(matcherContext)
                         .singleOpWithCanonicaleArgs<math::ExpOp>()
                         .dim(AllDims(), utils::IteratorType::parallel)
                         .input(NumEqualsTo(1))
                         .input(AllOperands(), IsIdentity())
                         .output(AllOperands(), IsIdentity())
                         .output(NumEqualsTo(1));
  expOperand = expOperand.input(0, *subOperand);

  auto &fillZero = m_StructuredOp<linalg::FillOp>(matcherContext)
                       .input(0, ConstantFloatZero());
  auto &sum =
      m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::AddFOp>(/*commutative=*/true)
          // Only handle most inner reduction for now.
          .dim(-1, utils::IteratorType::reduction)
          .dim(AllDimsExcept({-1}), utils::IteratorType::parallel)
          .input(NumEqualsTo(1))
          .input(AllOperands(), IsIdentity())
          .output(AllOperands(), IsProjected(-1))
          .output(NumEqualsTo(1));
  sum = sum.input(0, expOperand);
  sum = sum.output(0, fillZero);

  auto &rcpOperand = m_StructuredOp<linalg::GenericOp>(matcherContext)
                         .isFloatReciprocal()
                         .dim(AllDims(), utils::IteratorType::parallel)
                         .input(NumEqualsTo(1))
                         .input(AllOperands(), IsIdentity())
                         .output(AllOperands(), IsIdentity())
                         .output(NumEqualsTo(1));
  rcpOperand = rcpOperand.input(0, sum);

  auto &mulOperand =
      transform_ext::m_StructuredOp<linalg::GenericOp>(matcherContext)
          .singleOpWithCanonicaleArgs<arith::MulFOp>(/*commutative=*/true)
          .dim(AllDims(), utils::IteratorType::parallel)
          .input(NumEqualsTo(2))
          .input(0, IsIdentity())
          .input(1, IsProjected(-1))
          .output(NumEqualsTo(1))
          .output(AllOperands(), IsIdentity());

  mulOperand = mulOperand.input(0, expOperand);
  mulOperand = mulOperand.input(1, rcpOperand);

  transform_ext::StructuredOpMatcher *divOperand;
  matchdivBroadcast(matcherContext, expOperand, sum, divOperand);

  auto &softmaxRoot =
      transform_ext::m_StructuredOp_Or(matcherContext, mulOperand, *divOperand);
  softmaxRootCapture = &softmaxRoot;
}