Improve mlir-query tool by implementing `getBackwardSlice` and `getForwardSlice` matchers. As an addition `SetQuery` also needed to be added to enable custom configuration for each query. e.g: `inclusive`, `omitUsesFromAbove`, `omitBlockArguments`. Note: backwardSlice and forwardSlice algoritms are the same as the ones in `mlir/lib/Analysis/SliceAnalysis.cpp` Example of current matcher. The query was made to the file: `mlir/test/mlir-query/complex-test.mlir` ```mlir ./mlir-query /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir -c "match getDefinitions(hasOpName(\"arith.add f\"),2)" Match #1: /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir:5:8: %0 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) { ^ /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir:7:10: note: "root" binds here %2 = arith.addf %in, %in : f32 ^ Match #2: /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir:10:16: %collapsed = tensor.collapse_shape %0 [[0, 1]] : tensor<5x5xf32> into tensor<25xf32> ^ /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir:13:11: %c2 = arith.constant 2 : index ^ /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir:14:18: %extracted = tensor.extract %collapsed[%c2] : tensor<25xf32> ^ /home/dbudii/personal/llvm-project/mlir/test/mlir-query/complex-test.mlir:15:10: note: "root" binds here %2 = arith.addf %extracted, %extracted : f32 ^ 2 matches. ```
60 lines
2.1 KiB
C++
60 lines
2.1 KiB
C++
//===- MatchFinder.h - ------------------------------------------*- C++ -*-===//
|
|
//
|
|
// 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 contains the MatchFinder class, which is used to find operations
|
|
// that match a given matcher and print them.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
|
|
#define MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
|
|
|
|
#include "MatchersInternal.h"
|
|
#include "mlir/Query/Query.h"
|
|
#include "mlir/Query/QuerySession.h"
|
|
#include "llvm/ADT/SetVector.h"
|
|
|
|
namespace mlir::query::matcher {
|
|
|
|
/// A class that provides utilities to find operations in the IR.
|
|
class MatchFinder {
|
|
|
|
public:
|
|
/// A subclass which preserves the matching information. Each instance
|
|
/// contains the `rootOp` along with the matching environment.
|
|
struct MatchResult {
|
|
MatchResult() = default;
|
|
MatchResult(Operation *rootOp, std::vector<Operation *> matchedOps);
|
|
|
|
Operation *rootOp = nullptr;
|
|
/// Contains the matching environment.
|
|
std::vector<Operation *> matchedOps;
|
|
};
|
|
|
|
/// Traverses the IR and returns a vector of `MatchResult` for each match of
|
|
/// the `matcher`.
|
|
std::vector<MatchResult> collectMatches(Operation *root,
|
|
DynMatcher matcher) const;
|
|
|
|
/// Prints the matched operation.
|
|
void printMatch(llvm::raw_ostream &os, QuerySession &qs, Operation *op) const;
|
|
|
|
/// Labels the matched operation with the given binding (e.g., `"root"`) and
|
|
/// prints it.
|
|
void printMatch(llvm::raw_ostream &os, QuerySession &qs, Operation *op,
|
|
const std::string &binding) const;
|
|
|
|
/// Flattens a vector of `MatchResult` into a vector of operations.
|
|
std::vector<Operation *>
|
|
flattenMatchedOps(std::vector<MatchResult> &matches) const;
|
|
};
|
|
|
|
} // namespace mlir::query::matcher
|
|
|
|
#endif // MLIR_TOOLS_MLIRQUERY_MATCHER_MATCHERFINDER_H
|