MLIR  21.0.0git
Query.cpp
Go to the documentation of this file.
1 //===---- Query.cpp - -----------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Query/Query.h"
10 #include "QueryParser.h"
12 #include "mlir/IR/IRMapping.h"
13 #include "mlir/IR/Verifier.h"
16 #include "llvm/ADT/SetVector.h"
17 #include "llvm/Support/SourceMgr.h"
18 #include "llvm/Support/raw_ostream.h"
19 
20 namespace mlir::query {
21 
22 QueryRef parse(llvm::StringRef line, const QuerySession &qs) {
23  return QueryParser::parse(line, qs);
24 }
25 
26 std::vector<llvm::LineEditor::Completion>
27 complete(llvm::StringRef line, size_t pos, const QuerySession &qs) {
28  return QueryParser::complete(line, pos, qs);
29 }
30 
31 // TODO: Extract into a helper function that can be reused outside query
32 // context.
33 static Operation *extractFunction(std::vector<Operation *> &ops,
34  MLIRContext *context,
35  llvm::StringRef functionName) {
36  context->loadDialect<func::FuncDialect>();
37  OpBuilder builder(context);
38 
39  // Collect data for function creation
40  std::vector<Operation *> slice;
41  std::vector<Value> values;
42  std::vector<Type> outputTypes;
43 
44  for (auto *op : ops) {
45  // Return op's operands are propagated, but the op itself isn't needed.
46  if (!isa<func::ReturnOp>(op))
47  slice.push_back(op);
48 
49  // All results are returned by the extracted function.
50  llvm::append_range(outputTypes, op->getResults().getTypes());
51 
52  // Track all values that need to be taken as input to function.
53  llvm::append_range(values, op->getOperands());
54  }
55 
56  // Create the function
57  FunctionType funcType =
58  builder.getFunctionType(TypeRange(ValueRange(values)), outputTypes);
59  auto loc = builder.getUnknownLoc();
60  func::FuncOp funcOp = func::FuncOp::create(loc, functionName, funcType);
61 
62  builder.setInsertionPointToEnd(funcOp.addEntryBlock());
63 
64  // Map original values to function arguments
65  IRMapping mapper;
66  for (const auto &arg : llvm::enumerate(values))
67  mapper.map(arg.value(), funcOp.getArgument(arg.index()));
68 
69  // Clone operations and build function body
70  std::vector<Operation *> clonedOps;
71  std::vector<Value> clonedVals;
72  // TODO: Handle extraction of operations with compute payloads defined via
73  // regions.
74  for (Operation *slicedOp : slice) {
75  Operation *clonedOp =
76  clonedOps.emplace_back(builder.clone(*slicedOp, mapper));
77  clonedVals.insert(clonedVals.end(), clonedOp->result_begin(),
78  clonedOp->result_end());
79  }
80  // Add return operation
81  builder.create<func::ReturnOp>(loc, clonedVals);
82 
83  // Remove unused function arguments
84  size_t currentIndex = 0;
85  while (currentIndex < funcOp.getNumArguments()) {
86  // Erase if possible.
87  if (funcOp.getArgument(currentIndex).use_empty())
88  if (succeeded(funcOp.eraseArgument(currentIndex)))
89  continue;
90  ++currentIndex;
91  }
92 
93  return funcOp;
94 }
95 
96 Query::~Query() = default;
97 
98 LogicalResult InvalidQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
99  os << errStr << "\n";
100  return mlir::failure();
101 }
102 
103 LogicalResult NoOpQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
104  return mlir::success();
105 }
106 
107 LogicalResult HelpQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
108  os << "Available commands:\n\n"
109  " match MATCHER, m MATCHER "
110  "Match the mlir against the given matcher.\n"
111  " quit "
112  "Terminates the query session.\n\n";
113  return mlir::success();
114 }
115 
116 LogicalResult QuitQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
117  qs.terminate = true;
118  return mlir::success();
119 }
120 
121 LogicalResult MatchQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {
122  Operation *rootOp = qs.getRootOp();
123  int matchCount = 0;
124  matcher::MatchFinder finder;
125  auto matches = finder.collectMatches(rootOp, std::move(matcher));
126 
127  // An extract call is recognized by considering if the matcher has a name.
128  // TODO: Consider making the extract more explicit.
129  if (matcher.hasFunctionName()) {
130  auto functionName = matcher.getFunctionName();
131  std::vector<Operation *> flattenedMatches =
132  finder.flattenMatchedOps(matches);
133  Operation *function =
134  extractFunction(flattenedMatches, rootOp->getContext(), functionName);
135  if (failed(verify(function)))
136  return mlir::failure();
137  os << "\n" << *function << "\n\n";
138  function->erase();
139  return mlir::success();
140  }
141 
142  os << "\n";
143  for (auto &results : matches) {
144  os << "Match #" << ++matchCount << ":\n\n";
145  for (auto op : results.matchedOps) {
146  if (op == results.rootOp) {
147  finder.printMatch(os, qs, op, "root");
148  } else {
149  finder.printMatch(os, qs, op);
150  }
151  }
152  }
153  os << matchCount << (matchCount == 1 ? " match.\n\n" : " matches.\n\n");
154  return mlir::success();
155 }
156 
157 } // namespace mlir::query
This is a utility class for mapping one set of IR entities to another.
Definition: IRMapping.h:26
void map(Value from, Value to)
Inserts a new mapping for 'from' to 'to'.
Definition: IRMapping.h:30
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
void loadDialect()
Load a dialect in the context.
Definition: MLIRContext.h:107
This class helps build Operations.
Definition: Builders.h:205
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
result_iterator result_begin()
Definition: Operation.h:413
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
static Operation * create(Location location, OperationName name, TypeRange resultTypes, ValueRange operands, NamedAttrList &&attributes, OpaqueProperties properties, BlockRange successors, unsigned numRegions)
Create a new Operation with the specific fields.
Definition: Operation.cpp:67
result_iterator result_end()
Definition: Operation.h:414
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:37
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
static QueryRef parse(llvm::StringRef line, const QuerySession &qs)
static std::vector< llvm::LineEditor::Completion > complete(llvm::StringRef line, size_t pos, const QuerySession &qs)
Finds and collects matches from the IR.
Definition: MatchFinder.h:27
void printMatch(llvm::raw_ostream &os, QuerySession &qs, Operation *op) const
Prints the matched operation.
Definition: MatchFinder.cpp:39
std::vector< MatchResult > collectMatches(Operation *root, DynMatcher matcher) const
Traverses the IR and returns a vector of MatchResult for each match of the matcher.
Definition: MatchFinder.cpp:21
std::vector< Operation * > flattenMatchedOps(std::vector< MatchResult > &matches) const
Flattens a vector of MatchResult into a vector of operations.
Definition: MatchFinder.cpp:59
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
QueryRef parse(llvm::StringRef line, const QuerySession &qs)
Definition: Query.cpp:22
static Operation * extractFunction(std::vector< Operation * > &ops, MLIRContext *context, llvm::StringRef functionName)
Definition: Query.cpp:33
llvm::IntrusiveRefCntPtr< Query > QueryRef
Definition: Query.h:36
std::vector< llvm::LineEditor::Completion > complete(llvm::StringRef line, size_t pos, const QuerySession &qs)
Definition: Query.cpp:27
LogicalResult verify(Operation *op, bool verifyRecursively=true)
Perform (potentially expensive) checks of invariants, used to detect compiler bugs,...
Definition: Verifier.cpp:423
llvm::LogicalResult run(llvm::raw_ostream &os, QuerySession &qs) const override
Definition: Query.cpp:107
llvm::LogicalResult run(llvm::raw_ostream &os, QuerySession &qs) const override
Definition: Query.cpp:98
std::string errStr
Definition: Query.h:50
llvm::LogicalResult run(llvm::raw_ostream &os, QuerySession &qs) const override
Definition: Query.cpp:121
const matcher::DynMatcher matcher
Definition: Query.h:97
llvm::LogicalResult run(llvm::raw_ostream &os, QuerySession &qs) const override
Definition: Query.cpp:103
llvm::LogicalResult run(llvm::raw_ostream &os, QuerySession &qs) const override
Definition: Query.cpp:116