Denzel-Brian Budii 9b63bdd154
[mlir] Improve mlir-query tool by implementing getBackwardSlice and getForwardSlice matchers (#115670)
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.
```
2025-05-13 13:18:14 +02:00

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