[Mlir-commits] [mlir] 9b63bdd - [mlir] Improve mlir-query tool by implementing `getBackwardSlice` and `getForwardSlice` matchers (#115670)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 13 04:18:18 PDT 2025
Author: Denzel-Brian Budii
Date: 2025-05-13T13:18:14+02:00
New Revision: 9b63bdd15418e90d0c6889d26e490db89ae744d1
URL: https://github.com/llvm/llvm-project/commit/9b63bdd15418e90d0c6889d26e490db89ae744d1
DIFF: https://github.com/llvm/llvm-project/commit/9b63bdd15418e90d0c6889d26e490db89ae744d1.diff
LOG: [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.
```
Added:
mlir/include/mlir/Query/Matcher/SliceMatchers.h
mlir/lib/Query/Matcher/MatchFinder.cpp
mlir/test/mlir-query/complex-test.mlir
Modified:
mlir/include/mlir/Query/Matcher/Marshallers.h
mlir/include/mlir/Query/Matcher/MatchFinder.h
mlir/include/mlir/Query/Matcher/MatchersInternal.h
mlir/include/mlir/Query/Matcher/VariantValue.h
mlir/lib/Query/Matcher/CMakeLists.txt
mlir/lib/Query/Matcher/Parser.cpp
mlir/lib/Query/Matcher/Parser.h
mlir/lib/Query/Matcher/RegistryManager.cpp
mlir/lib/Query/Matcher/VariantValue.cpp
mlir/lib/Query/Query.cpp
mlir/tools/mlir-query/mlir-query.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Query/Matcher/Marshallers.h b/mlir/include/mlir/Query/Matcher/Marshallers.h
index 6ed35ac0ddccc..012bf7b9ec4a9 100644
--- a/mlir/include/mlir/Query/Matcher/Marshallers.h
+++ b/mlir/include/mlir/Query/Matcher/Marshallers.h
@@ -50,6 +50,36 @@ struct ArgTypeTraits<llvm::StringRef> {
}
};
+template <>
+struct ArgTypeTraits<int64_t> {
+ static bool hasCorrectType(const VariantValue &value) {
+ return value.isSigned();
+ }
+
+ static unsigned get(const VariantValue &value) { return value.getSigned(); }
+
+ static ArgKind getKind() { return ArgKind::Signed; }
+
+ static std::optional<std::string> getBestGuess(const VariantValue &) {
+ return std::nullopt;
+ }
+};
+
+template <>
+struct ArgTypeTraits<bool> {
+ static bool hasCorrectType(const VariantValue &value) {
+ return value.isBoolean();
+ }
+
+ static unsigned get(const VariantValue &value) { return value.getBoolean(); }
+
+ static ArgKind getKind() { return ArgKind::Boolean; }
+
+ static std::optional<std::string> getBestGuess(const VariantValue &) {
+ return std::nullopt;
+ }
+};
+
template <>
struct ArgTypeTraits<DynMatcher> {
diff --git a/mlir/include/mlir/Query/Matcher/MatchFinder.h b/mlir/include/mlir/Query/Matcher/MatchFinder.h
index b008a21f53ae2..f8abf20ef60bb 100644
--- a/mlir/include/mlir/Query/Matcher/MatchFinder.h
+++ b/mlir/include/mlir/Query/Matcher/MatchFinder.h
@@ -7,7 +7,7 @@
//===----------------------------------------------------------------------===//
//
// This file contains the MatchFinder class, which is used to find operations
-// that match a given matcher.
+// that match a given matcher and print them.
//
//===----------------------------------------------------------------------===//
@@ -15,25 +15,43 @@
#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 {
-// MatchFinder is used to find all operations that match a given matcher.
+/// A class that provides utilities to find operations in the IR.
class MatchFinder {
+
public:
- // Returns all operations that match the given matcher.
- static std::vector<Operation *> getMatches(Operation *root,
- DynMatcher matcher) {
- std::vector<Operation *> matches;
-
- // Simple match finding with walk.
- root->walk([&](Operation *subOp) {
- if (matcher.match(subOp))
- matches.push_back(subOp);
- });
-
- return matches;
- }
+ /// 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
diff --git a/mlir/include/mlir/Query/Matcher/MatchersInternal.h b/mlir/include/mlir/Query/Matcher/MatchersInternal.h
index 117f7d4edef9e..183b2514e109f 100644
--- a/mlir/include/mlir/Query/Matcher/MatchersInternal.h
+++ b/mlir/include/mlir/Query/Matcher/MatchersInternal.h
@@ -8,8 +8,9 @@
//
// Implements the base layer of the matcher framework.
//
-// Matchers are methods that return a Matcher which provides a method
-// match(Operation *op)
+// Matchers are methods that return a Matcher which provides a method one of the
+// following methods: match(Operation *op), match(Operation *op,
+// SetVector<Operation *> &matchedOps)
//
// The matcher functions are defined in include/mlir/IR/Matchers.h.
// This file contains the wrapper classes needed to construct matchers for
@@ -25,6 +26,31 @@
namespace mlir::query::matcher {
+// Defaults to false if T has no match() method with the signature:
+// match(Operation* op).
+template <typename T, typename = void>
+struct has_simple_match : std::false_type {};
+
+// Specialized type trait that evaluates to true if T has a match() method
+// with the signature: match(Operation* op).
+template <typename T>
+struct has_simple_match<T, std::void_t<decltype(std::declval<T>().match(
+ std::declval<Operation *>()))>>
+ : std::true_type {};
+
+// Defaults to false if T has no match() method with the signature:
+// match(Operation* op, SetVector<Operation*>&).
+template <typename T, typename = void>
+struct has_bound_match : std::false_type {};
+
+// Specialized type trait that evaluates to true if T has a match() method
+// with the signature: match(Operation* op, SetVector<Operation*>&).
+template <typename T>
+struct has_bound_match<T, std::void_t<decltype(std::declval<T>().match(
+ std::declval<Operation *>(),
+ std::declval<SetVector<Operation *> &>()))>>
+ : std::true_type {};
+
// Generic interface for matchers on an MLIR operation.
class MatcherInterface
: public llvm::ThreadSafeRefCountedBase<MatcherInterface> {
@@ -32,6 +58,7 @@ class MatcherInterface
virtual ~MatcherInterface() = default;
virtual bool match(Operation *op) = 0;
+ virtual bool match(Operation *op, SetVector<Operation *> &matchedOps) = 0;
};
// MatcherFnImpl takes a matcher function object and implements
@@ -40,14 +67,25 @@ template <typename MatcherFn>
class MatcherFnImpl : public MatcherInterface {
public:
MatcherFnImpl(MatcherFn &matcherFn) : matcherFn(matcherFn) {}
- bool match(Operation *op) override { return matcherFn.match(op); }
+
+ bool match(Operation *op) override {
+ if constexpr (has_simple_match<MatcherFn>::value)
+ return matcherFn.match(op);
+ return false;
+ }
+
+ bool match(Operation *op, SetVector<Operation *> &matchedOps) override {
+ if constexpr (has_bound_match<MatcherFn>::value)
+ return matcherFn.match(op, matchedOps);
+ return false;
+ }
private:
MatcherFn matcherFn;
};
-// Matcher wraps a MatcherInterface implementation and provides a match()
-// method that redirects calls to the underlying implementation.
+// Matcher wraps a MatcherInterface implementation and provides match()
+// methods that redirect calls to the underlying implementation.
class DynMatcher {
public:
// Takes ownership of the provided implementation pointer.
@@ -62,12 +100,13 @@ class DynMatcher {
}
bool match(Operation *op) const { return implementation->match(op); }
+ bool match(Operation *op, SetVector<Operation *> &matchedOps) const {
+ return implementation->match(op, matchedOps);
+ }
- void setFunctionName(StringRef name) { functionName = name.str(); };
-
- bool hasFunctionName() const { return !functionName.empty(); };
-
- StringRef getFunctionName() const { return functionName; };
+ void setFunctionName(StringRef name) { functionName = name.str(); }
+ bool hasFunctionName() const { return !functionName.empty(); }
+ StringRef getFunctionName() const { return functionName; }
private:
llvm::IntrusiveRefCntPtr<MatcherInterface> implementation;
diff --git a/mlir/include/mlir/Query/Matcher/SliceMatchers.h b/mlir/include/mlir/Query/Matcher/SliceMatchers.h
new file mode 100644
index 0000000000000..5bb8251672eb7
--- /dev/null
+++ b/mlir/include/mlir/Query/Matcher/SliceMatchers.h
@@ -0,0 +1,141 @@
+//===- SliceMatchers.h - Matchers for slicing analysis ----------*- 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 provides matchers for MLIRQuery that peform slicing analysis
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_MLIRQUERY_MATCHERS_SLICEMATCHERS_H
+#define MLIR_TOOLS_MLIRQUERY_MATCHERS_SLICEMATCHERS_H
+
+#include "mlir/Analysis/SliceAnalysis.h"
+
+/// A matcher encapsulating `getBackwardSlice` method from SliceAnalysis.h.
+/// Additionally, it limits the slice computation to a certain depth level using
+/// a custom filter.
+///
+/// Example: starting from node 9, assuming the matcher
+/// computes the slice for the first two depth levels:
+/// ============================
+/// 1 2 3 4
+/// |_______| |______|
+/// | | |
+/// | 5 6
+/// |___|_____________|
+/// | |
+/// 7 8
+/// |_______________|
+/// |
+/// 9
+///
+/// Assuming all local orders match the numbering order:
+/// {5, 7, 6, 8, 9}
+namespace mlir::query::matcher {
+
+template <typename Matcher>
+class BackwardSliceMatcher {
+public:
+ BackwardSliceMatcher(Matcher innerMatcher, int64_t maxDepth, bool inclusive,
+ bool omitBlockArguments, bool omitUsesFromAbove)
+ : innerMatcher(std::move(innerMatcher)), maxDepth(maxDepth),
+ inclusive(inclusive), omitBlockArguments(omitBlockArguments),
+ omitUsesFromAbove(omitUsesFromAbove) {}
+
+ bool match(Operation *rootOp, SetVector<Operation *> &backwardSlice) {
+ BackwardSliceOptions options;
+ options.inclusive = inclusive;
+ options.omitUsesFromAbove = omitUsesFromAbove;
+ options.omitBlockArguments = omitBlockArguments;
+ return (innerMatcher.match(rootOp) &&
+ matches(rootOp, backwardSlice, options, maxDepth));
+ }
+
+private:
+ bool matches(Operation *rootOp, llvm::SetVector<Operation *> &backwardSlice,
+ BackwardSliceOptions &options, int64_t maxDepth);
+
+private:
+ // The outer matcher (e.g., BackwardSliceMatcher) relies on the innerMatcher
+ // to determine whether we want to traverse the IR or not. For example, we
+ // want to explore the IR only if the top-level operation name is
+ // `"arith.addf"`.
+ Matcher innerMatcher;
+ // `maxDepth` specifies the maximum depth that the matcher can traverse the
+ // IR. For example, if `maxDepth` is 2, the matcher will explore the defining
+ // operations of the top-level op up to 2 levels.
+ int64_t maxDepth;
+ bool inclusive;
+ bool omitBlockArguments;
+ bool omitUsesFromAbove;
+};
+
+template <typename Matcher>
+bool BackwardSliceMatcher<Matcher>::matches(
+ Operation *rootOp, llvm::SetVector<Operation *> &backwardSlice,
+ BackwardSliceOptions &options, int64_t maxDepth) {
+ backwardSlice.clear();
+ llvm::DenseMap<Operation *, int64_t> opDepths;
+ // Initializing the root op with a depth of 0
+ opDepths[rootOp] = 0;
+ options.filter = [&](Operation *subOp) {
+ // If the subOp hasn't been recorded in opDepths, it is deeper than
+ // maxDepth.
+ if (!opDepths.contains(subOp))
+ return false;
+ // Examine subOp's operands to compute depths of their defining operations.
+ for (auto operand : subOp->getOperands()) {
+ int64_t newDepth = opDepths[subOp] + 1;
+ // If the newDepth is greater than maxDepth, further computation can be
+ // skipped.
+ if (newDepth > maxDepth)
+ continue;
+
+ if (auto definingOp = operand.getDefiningOp()) {
+ // Registers the minimum depth
+ if (!opDepths.contains(definingOp) || newDepth < opDepths[definingOp])
+ opDepths[definingOp] = newDepth;
+ } else {
+ auto blockArgument = cast<BlockArgument>(operand);
+ Operation *parentOp = blockArgument.getOwner()->getParentOp();
+ if (!parentOp)
+ continue;
+
+ if (!opDepths.contains(parentOp) || newDepth < opDepths[parentOp])
+ opDepths[parentOp] = newDepth;
+ }
+ }
+ return true;
+ };
+ getBackwardSlice(rootOp, &backwardSlice, options);
+ return options.inclusive ? backwardSlice.size() > 1
+ : backwardSlice.size() >= 1;
+}
+
+/// Matches transitive defs of a top-level operation up to N levels.
+template <typename Matcher>
+inline BackwardSliceMatcher<Matcher>
+m_GetDefinitions(Matcher innerMatcher, int64_t maxDepth, bool inclusive,
+ bool omitBlockArguments, bool omitUsesFromAbove) {
+ assert(maxDepth >= 0 && "maxDepth must be non-negative");
+ return BackwardSliceMatcher<Matcher>(std::move(innerMatcher), maxDepth,
+ inclusive, omitBlockArguments,
+ omitUsesFromAbove);
+}
+
+/// Matches all transitive defs of a top-level operation up to N levels
+template <typename Matcher>
+inline BackwardSliceMatcher<Matcher> m_GetAllDefinitions(Matcher innerMatcher,
+ int64_t maxDepth) {
+ assert(maxDepth >= 0 && "maxDepth must be non-negative");
+ return BackwardSliceMatcher<Matcher>(std::move(innerMatcher), maxDepth, true,
+ false, false);
+}
+
+} // namespace mlir::query::matcher
+
+#endif // MLIR_TOOLS_MLIRQUERY_MATCHERS_SLICEMATCHERS_H
diff --git a/mlir/include/mlir/Query/Matcher/VariantValue.h b/mlir/include/mlir/Query/Matcher/VariantValue.h
index 449f8b3a01e02..98c0a18e25101 100644
--- a/mlir/include/mlir/Query/Matcher/VariantValue.h
+++ b/mlir/include/mlir/Query/Matcher/VariantValue.h
@@ -21,7 +21,7 @@
namespace mlir::query::matcher {
// All types that VariantValue can contain.
-enum class ArgKind { Matcher, String };
+enum class ArgKind { Boolean, Matcher, Signed, String };
// A variant matcher object to abstract simple and complex matchers into a
// single object type.
@@ -81,6 +81,8 @@ class VariantValue {
// Specific constructors for each supported type.
VariantValue(const llvm::StringRef string);
VariantValue(const VariantMatcher &matcher);
+ VariantValue(int64_t signedValue);
+ VariantValue(bool setBoolean);
// String value functions.
bool isString() const;
@@ -92,21 +94,36 @@ class VariantValue {
const VariantMatcher &getMatcher() const;
void setMatcher(const VariantMatcher &matcher);
+ // Signed value functions.
+ bool isSigned() const;
+ int64_t getSigned() const;
+ void setSigned(int64_t signedValue);
+
+ // Boolean value functions.
+ bool isBoolean() const;
+ bool getBoolean() const;
+ void setBoolean(bool booleanValue);
// String representation of the type of the value.
std::string getTypeAsString() const;
+ explicit operator bool() const { return hasValue(); }
+ bool hasValue() const { return type != ValueType::Nothing; }
private:
void reset();
// All supported value types.
enum class ValueType {
+ Boolean,
+ Matcher,
Nothing,
+ Signed,
String,
- Matcher,
};
// All supported value types.
union AllValues {
+ bool Boolean;
+ int64_t Signed;
llvm::StringRef *String;
VariantMatcher *Matcher;
};
diff --git a/mlir/lib/Query/Matcher/CMakeLists.txt b/mlir/lib/Query/Matcher/CMakeLists.txt
index 3adff9f99243f..629479bf7adc1 100644
--- a/mlir/lib/Query/Matcher/CMakeLists.txt
+++ b/mlir/lib/Query/Matcher/CMakeLists.txt
@@ -1,4 +1,5 @@
add_mlir_library(MLIRQueryMatcher
+ MatchFinder.cpp
Parser.cpp
RegistryManager.cpp
VariantValue.cpp
diff --git a/mlir/lib/Query/Matcher/MatchFinder.cpp b/mlir/lib/Query/Matcher/MatchFinder.cpp
new file mode 100644
index 0000000000000..1d4817e32417d
--- /dev/null
+++ b/mlir/lib/Query/Matcher/MatchFinder.cpp
@@ -0,0 +1,68 @@
+//===- MatchFinder.cpp - --------------------------------------------------===//
+//
+// 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 method definitions for the `MatchFinder` class
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Query/Matcher/MatchFinder.h"
+namespace mlir::query::matcher {
+
+MatchFinder::MatchResult::MatchResult(Operation *rootOp,
+ std::vector<Operation *> matchedOps)
+ : rootOp(rootOp), matchedOps(std::move(matchedOps)) {}
+
+std::vector<MatchFinder::MatchResult>
+MatchFinder::collectMatches(Operation *root, DynMatcher matcher) const {
+ std::vector<MatchResult> results;
+ llvm::SetVector<Operation *> tempStorage;
+ root->walk([&](Operation *subOp) {
+ if (matcher.match(subOp)) {
+ MatchResult match;
+ match.rootOp = subOp;
+ match.matchedOps.push_back(subOp);
+ results.push_back(std::move(match));
+ } else if (matcher.match(subOp, tempStorage)) {
+ results.emplace_back(subOp, std::vector<Operation *>(tempStorage.begin(),
+ tempStorage.end()));
+ }
+ tempStorage.clear();
+ });
+ return results;
+}
+
+void MatchFinder::printMatch(llvm::raw_ostream &os, QuerySession &qs,
+ Operation *op) const {
+ auto fileLoc = cast<FileLineColLoc>(op->getLoc());
+ SMLoc smloc = qs.getSourceManager().FindLocForLineAndColumn(
+ qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
+ llvm::SMDiagnostic diag =
+ qs.getSourceManager().GetMessage(smloc, llvm::SourceMgr::DK_Note, "");
+ diag.print("", os, true, false, true);
+}
+
+void MatchFinder::printMatch(llvm::raw_ostream &os, QuerySession &qs,
+ Operation *op, const std::string &binding) const {
+ auto fileLoc = cast<FileLineColLoc>(op->getLoc());
+ auto smloc = qs.getSourceManager().FindLocForLineAndColumn(
+ qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
+ qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,
+ "\"" + binding + "\" binds here");
+}
+
+std::vector<Operation *>
+MatchFinder::flattenMatchedOps(std::vector<MatchResult> &matches) const {
+ std::vector<Operation *> newVector;
+ for (auto &result : matches) {
+ newVector.insert(newVector.end(), result.matchedOps.begin(),
+ result.matchedOps.end());
+ }
+ return newVector;
+}
+
+} // namespace mlir::query::matcher
diff --git a/mlir/lib/Query/Matcher/Parser.cpp b/mlir/lib/Query/Matcher/Parser.cpp
index 3609e24f9939f..e392a885c511b 100644
--- a/mlir/lib/Query/Matcher/Parser.cpp
+++ b/mlir/lib/Query/Matcher/Parser.cpp
@@ -135,6 +135,18 @@ class Parser::CodeTokenizer {
case '\'':
consumeStringLiteral(&result);
break;
+ case '0':
+ case '1':
+ case '2':
+ case '3':
+ case '4':
+ case '5':
+ case '6':
+ case '7':
+ case '8':
+ case '9':
+ consumeNumberLiteral(&result);
+ break;
default:
parseIdentifierOrInvalid(&result);
break;
@@ -144,6 +156,18 @@ class Parser::CodeTokenizer {
return result;
}
+ void consumeNumberLiteral(TokenInfo *result) {
+ StringRef original = code;
+ unsigned value = 0;
+ if (!code.consumeInteger(0, value)) {
+ size_t numConsumed = original.size() - code.size();
+ result->text = original.take_front(numConsumed);
+ result->kind = TokenKind::Literal;
+ result->value = static_cast<int64_t>(value);
+ return;
+ }
+ }
+
// Consume a string literal, handle escape sequences and missing closing
// quote.
void consumeStringLiteral(TokenInfo *result) {
@@ -195,9 +219,22 @@ class Parser::CodeTokenizer {
break;
++tokenLength;
}
- result->kind = TokenKind::Ident;
- result->text = code.substr(0, tokenLength);
+ llvm::StringRef token = code.substr(0, tokenLength);
code = code.drop_front(tokenLength);
+ // Check if the identifier is a boolean literal
+ if (token == "true") {
+ result->text = "false";
+ result->kind = TokenKind::Literal;
+ result->value = true;
+ } else if (token == "false") {
+ result->text = "false";
+ result->kind = TokenKind::Literal;
+ result->value = false;
+ } else {
+ // Otherwise it is treated as a normal identifier
+ result->kind = TokenKind::Ident;
+ result->text = token;
+ }
} else {
result->kind = TokenKind::InvalidChar;
result->text = code.substr(0, 1);
@@ -257,13 +294,19 @@ bool Parser::parseIdentifierPrefixImpl(VariantValue *value) {
if (tokenizer->nextTokenKind() != TokenKind::OpenParen) {
// Parse as a named value.
- auto namedValue =
- namedValues ? namedValues->lookup(nameToken.text) : VariantValue();
+ if (auto namedValue = namedValues ? namedValues->lookup(nameToken.text)
+ : VariantValue()) {
- if (!namedValue.isMatcher()) {
- error->addError(tokenizer->peekNextToken().range,
- ErrorType::ParserNotAMatcher);
- return false;
+ if (tokenizer->nextTokenKind() != TokenKind::Period) {
+ *value = namedValue;
+ return true;
+ }
+
+ if (!namedValue.isMatcher()) {
+ error->addError(tokenizer->peekNextToken().range,
+ ErrorType::ParserNotAMatcher);
+ return false;
+ }
}
if (tokenizer->nextTokenKind() == TokenKind::NewLine) {
diff --git a/mlir/lib/Query/Matcher/Parser.h b/mlir/lib/Query/Matcher/Parser.h
index 58968023022d5..2199a2335ba9c 100644
--- a/mlir/lib/Query/Matcher/Parser.h
+++ b/mlir/lib/Query/Matcher/Parser.h
@@ -16,8 +16,11 @@
// provided to the parser.
//
// The grammar for the supported expressions is as follows:
-// <Expression> := <StringLiteral> | <MatcherExpression>
+// <Expression> := <Literal> | <MatcherExpression>
+// <Literal> := <StringLiteral> | <NumericLiteral> | <BooleanLiteral>
// <StringLiteral> := "quoted string"
+// <BooleanLiteral> := "true" | "false"
+// <NumericLiteral> := [0-9]+
// <MatcherExpression> := <MatcherName>(<ArgumentList>)
// <MatcherName> := [a-zA-Z]+
// <ArgumentList> := <Expression> | <Expression>,<ArgumentList>
diff --git a/mlir/lib/Query/Matcher/RegistryManager.cpp b/mlir/lib/Query/Matcher/RegistryManager.cpp
index 645db7109c2de..4b511c5f009e7 100644
--- a/mlir/lib/Query/Matcher/RegistryManager.cpp
+++ b/mlir/lib/Query/Matcher/RegistryManager.cpp
@@ -19,16 +19,15 @@
namespace mlir::query::matcher {
namespace {
-// This is needed because these matchers are defined as overloaded functions.
-using IsConstantOp = detail::constant_op_matcher();
-using HasOpAttrName = detail::AttrOpMatcher(llvm::StringRef);
-using HasOpName = detail::NameOpMatcher(llvm::StringRef);
-
// Enum to string for autocomplete.
static std::string asArgString(ArgKind kind) {
switch (kind) {
+ case ArgKind::Boolean:
+ return "Boolean";
case ArgKind::Matcher:
return "Matcher";
+ case ArgKind::Signed:
+ return "Signed";
case ArgKind::String:
return "String";
}
diff --git a/mlir/lib/Query/Matcher/VariantValue.cpp b/mlir/lib/Query/Matcher/VariantValue.cpp
index 65bd4bd77bcf8..1cb2d48f9d56f 100644
--- a/mlir/lib/Query/Matcher/VariantValue.cpp
+++ b/mlir/lib/Query/Matcher/VariantValue.cpp
@@ -56,6 +56,14 @@ VariantValue::VariantValue(const VariantMatcher &matcher)
value.Matcher = new VariantMatcher(matcher);
}
+VariantValue::VariantValue(int64_t signedValue) : type(ValueType::Signed) {
+ value.Signed = signedValue;
+}
+
+VariantValue::VariantValue(bool setBoolean) : type(ValueType::Boolean) {
+ value.Boolean = setBoolean;
+}
+
VariantValue::~VariantValue() { reset(); }
VariantValue &VariantValue::operator=(const VariantValue &other) {
@@ -69,6 +77,12 @@ VariantValue &VariantValue::operator=(const VariantValue &other) {
case ValueType::Matcher:
setMatcher(other.getMatcher());
break;
+ case ValueType::Signed:
+ setSigned(other.getSigned());
+ break;
+ case ValueType::Boolean:
+ setBoolean(other.getBoolean());
+ break;
case ValueType::Nothing:
type = ValueType::Nothing;
break;
@@ -85,12 +99,34 @@ void VariantValue::reset() {
delete value.Matcher;
break;
// Cases that do nothing.
+ case ValueType::Signed:
+ case ValueType::Boolean:
case ValueType::Nothing:
break;
}
type = ValueType::Nothing;
}
+// Signed
+bool VariantValue::isSigned() const { return type == ValueType::Signed; }
+
+int64_t VariantValue::getSigned() const { return value.Signed; }
+
+void VariantValue::setSigned(int64_t newValue) {
+ type = ValueType::Signed;
+ value.Signed = newValue;
+}
+
+// Boolean
+bool VariantValue::isBoolean() const { return type == ValueType::Boolean; }
+
+bool VariantValue::getBoolean() const { return value.Signed; }
+
+void VariantValue::setBoolean(bool newValue) {
+ type = ValueType::Boolean;
+ value.Signed = newValue;
+}
+
bool VariantValue::isString() const { return type == ValueType::String; }
const llvm::StringRef &VariantValue::getString() const {
@@ -123,6 +159,10 @@ std::string VariantValue::getTypeAsString() const {
return "String";
case ValueType::Matcher:
return "Matcher";
+ case ValueType::Signed:
+ return "Signed";
+ case ValueType::Boolean:
+ return "Boolean";
case ValueType::Nothing:
return "Nothing";
}
diff --git a/mlir/lib/Query/Query.cpp b/mlir/lib/Query/Query.cpp
index 73f313cd37fd0..803284d6df86a 100644
--- a/mlir/lib/Query/Query.cpp
+++ b/mlir/lib/Query/Query.cpp
@@ -12,6 +12,7 @@
#include "mlir/IR/IRMapping.h"
#include "mlir/Query/Matcher/MatchFinder.h"
#include "mlir/Query/QuerySession.h"
+#include "llvm/ADT/SetVector.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
@@ -26,15 +27,6 @@ complete(llvm::StringRef line, size_t pos, const QuerySession &qs) {
return QueryParser::complete(line, pos, qs);
}
-static void printMatch(llvm::raw_ostream &os, QuerySession &qs, Operation *op,
- const std::string &binding) {
- auto fileLoc = op->getLoc()->findInstanceOf<FileLineColLoc>();
- auto smloc = qs.getSourceManager().FindLocForLineAndColumn(
- qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
- qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,
- "\"" + binding + "\" binds here");
-}
-
// TODO: Extract into a helper function that can be reused outside query
// context.
static Operation *extractFunction(std::vector<Operation *> &ops,
@@ -126,28 +118,34 @@ LogicalResult QuitQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
LogicalResult MatchQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
Operation *rootOp = qs.getRootOp();
int matchCount = 0;
- std::vector<Operation *> matches =
- matcher::MatchFinder().getMatches(rootOp, matcher);
+ matcher::MatchFinder finder;
+ auto matches = finder.collectMatches(rootOp, std::move(matcher));
// An extract call is recognized by considering if the matcher has a name.
// TODO: Consider making the extract more explicit.
if (matcher.hasFunctionName()) {
auto functionName = matcher.getFunctionName();
+ std::vector<Operation *> flattenedMatches =
+ finder.flattenMatchedOps(matches);
Operation *function =
- extractFunction(matches, rootOp->getContext(), functionName);
+ extractFunction(flattenedMatches, rootOp->getContext(), functionName);
os << "\n" << *function << "\n\n";
function->erase();
return mlir::success();
}
os << "\n";
- for (Operation *op : matches) {
+ for (auto &results : matches) {
os << "Match #" << ++matchCount << ":\n\n";
- // Placeholder "root" binding for the initial draft.
- printMatch(os, qs, op, "root");
+ for (auto op : results.matchedOps) {
+ if (op == results.rootOp) {
+ finder.printMatch(os, qs, op, "root");
+ } else {
+ finder.printMatch(os, qs, op);
+ }
+ }
}
os << matchCount << (matchCount == 1 ? " match.\n\n" : " matches.\n\n");
-
return mlir::success();
}
diff --git a/mlir/test/mlir-query/complex-test.mlir b/mlir/test/mlir-query/complex-test.mlir
new file mode 100644
index 0000000000000..ad96f03747a43
--- /dev/null
+++ b/mlir/test/mlir-query/complex-test.mlir
@@ -0,0 +1,32 @@
+// RUN: mlir-query %s -c "m getAllDefinitions(hasOpName(\"arith.addf\"),2)" | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @slice_use_from_above(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) {
+ %0 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %2 = arith.addf %in, %in : f32
+ linalg.yield %2 : f32
+ } -> tensor<5x5xf32>
+ %collapsed = tensor.collapse_shape %0 [[0, 1]] : tensor<5x5xf32> into tensor<25xf32>
+ %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %c2 = arith.constant 2 : index
+ %extracted = tensor.extract %collapsed[%c2] : tensor<25xf32>
+ %2 = arith.addf %extracted, %extracted : f32
+ linalg.yield %2 : f32
+ } -> tensor<5x5xf32>
+ return
+}
+
+// CHECK: Match #1:
+
+// CHECK: %[[LINALG:.*]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME: ins(%arg0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>)
+// CHECK: %[[ADDF1:.*]] = arith.addf %in, %in : f32
+
+// CHECK: Match #2:
+
+// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %[[LINALG]] {{\[\[.*\]\]}} : tensor<5x5xf32> into tensor<25xf32>
+// CHECK: %[[C2:.*]] = arith.constant {{.*}} : index
+// CHECK: %[[EXTRACTED:.*]] = tensor.extract %[[COLLAPSED]][%[[C2]]] : tensor<25xf32>
+// CHECK: %[[ADDF2:.*]] = arith.addf %[[EXTRACTED]], %[[EXTRACTED]] : f32
diff --git a/mlir/tools/mlir-query/mlir-query.cpp b/mlir/tools/mlir-query/mlir-query.cpp
index 0ed4f94d5802b..78c0ec97c0cdf 100644
--- a/mlir/tools/mlir-query/mlir-query.cpp
+++ b/mlir/tools/mlir-query/mlir-query.cpp
@@ -16,6 +16,7 @@
#include "mlir/IR/Matchers.h"
#include "mlir/InitAllDialects.h"
#include "mlir/Query/Matcher/Registry.h"
+#include "mlir/Query/Matcher/SliceMatchers.h"
#include "mlir/Tools/mlir-query/MlirQueryMain.h"
using namespace mlir;
@@ -39,6 +40,12 @@ int main(int argc, char **argv) {
query::matcher::Registry matcherRegistry;
// Matchers registered in alphabetical order for consistency:
+ matcherRegistry.registerMatcher(
+ "getDefinitions",
+ query::matcher::m_GetDefinitions<query::matcher::DynMatcher>);
+ matcherRegistry.registerMatcher(
+ "getAllDefinitions",
+ query::matcher::m_GetAllDefinitions<query::matcher::DynMatcher>);
matcherRegistry.registerMatcher("hasOpAttrName",
static_cast<HasOpAttrName *>(m_Attr));
matcherRegistry.registerMatcher("hasOpName", static_cast<HasOpName *>(m_Op));
More information about the Mlir-commits
mailing list