NestedMatcher.h

Go to the documentation of this file.

//===- NestedMacher.h - Nested matcher for Function -------------*- C++ -*-===//
//
// Part of the MLIR 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
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_ANALYSIS_MLFUNCTIONMATCHER_H_
#define MLIR_ANALYSIS_MLFUNCTIONMATCHER_H_

#include "mlir/IR/Function.h"
#include "mlir/IR/Operation.h"
#include "llvm/Support/Allocator.h"

namespace mlir {

class NestedPattern;
class Operation;

class NestedMatch {
public:
  static NestedMatch build(Operation *operation,
                           ArrayRef<NestedMatch> nestedMatches);
  NestedMatch(const NestedMatch &) = default;
  NestedMatch &operator=(const NestedMatch &) = default;

  explicit operator bool() { return matchedOperation != nullptr; }

  Operation *getMatchedOperation() { return matchedOperation; }
  ArrayRef<NestedMatch> getMatchedChildren() { return matchedChildren; }

private:
  friend class NestedPattern;
  friend class NestedPatternContext;

  static llvm::BumpPtrAllocator *&allocator();

  NestedMatch() = default;

  Operation *matchedOperation;
  ArrayRef<NestedMatch> matchedChildren;
};

using FilterFunctionType = std::function<bool(Operation &)>;
inline bool defaultFilterFunction(Operation &) { return true; }
class NestedPattern {
public:
  NestedPattern(ArrayRef<NestedPattern> nested,
                FilterFunctionType filter = defaultFilterFunction);
  NestedPattern(const NestedPattern &) = default;
  NestedPattern &operator=(const NestedPattern &) = default;

  void match(FuncOp func, SmallVectorImpl<NestedMatch> *matches) {
    func.walk([&](Operation *op) { matchOne(op, matches); });
  }

  void match(Operation *op, SmallVectorImpl<NestedMatch> *matches) {
    op->walk([&](Operation *child) { matchOne(child, matches); });
  }

  unsigned getDepth() const;

private:
  friend class NestedPatternContext;
  friend class NestedMatch;
  friend struct State;

  static llvm::BumpPtrAllocator *&allocator();

  void matchOne(Operation *op, SmallVectorImpl<NestedMatch> *matches);

  ArrayRef<NestedPattern> nestedPatterns;

  FilterFunctionType filter;

  Operation *skip;
};

class NestedPatternContext {
public:
  NestedPatternContext() {
    assert(NestedMatch::allocator() == nullptr &&
           "Only a single NestedPatternContext is supported");
    assert(NestedPattern::allocator() == nullptr &&
           "Only a single NestedPatternContext is supported");
    NestedMatch::allocator() = &allocator;
    NestedPattern::allocator() = &allocator;
  }
  ~NestedPatternContext() {
    NestedMatch::allocator() = nullptr;
    NestedPattern::allocator() = nullptr;
  }
  llvm::BumpPtrAllocator allocator;
};

namespace matcher {
// Syntactic sugar NestedPattern builder functions.
NestedPattern Op(FilterFunctionType filter = defaultFilterFunction);
NestedPattern If(NestedPattern child);
NestedPattern If(FilterFunctionType filter, NestedPattern child);
NestedPattern If(ArrayRef<NestedPattern> nested = {});
NestedPattern If(FilterFunctionType filter,
                 ArrayRef<NestedPattern> nested = {});
NestedPattern For(NestedPattern child);
NestedPattern For(FilterFunctionType filter, NestedPattern child);
NestedPattern For(ArrayRef<NestedPattern> nested = {});
NestedPattern For(FilterFunctionType filter,
                  ArrayRef<NestedPattern> nested = {});

bool isParallelLoop(Operation &op);
bool isReductionLoop(Operation &op);
bool isLoadOrStore(Operation &op);

} // end namespace matcher
} // end namespace mlir

#endif // MLIR_ANALYSIS_MLFUNCTIONMATCHER_H_