158 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Query/Query.h"10#include "QueryParser.h"11#include "mlir/Dialect/Func/IR/FuncOps.h"12#include "mlir/IR/IRMapping.h"13#include "mlir/IR/Verifier.h"14#include "mlir/Query/Matcher/MatchFinder.h"15#include "mlir/Query/QuerySession.h"16#include "llvm/Support/SourceMgr.h"17#include "llvm/Support/raw_ostream.h"18 19namespace mlir::query {20 21QueryRef parse(llvm::StringRef line, const QuerySession &qs) {22 return QueryParser::parse(line, qs);23}24 25std::vector<llvm::LineEditor::Completion>26complete(llvm::StringRef line, size_t pos, const QuerySession &qs) {27 return QueryParser::complete(line, pos, qs);28}29 30// TODO: Extract into a helper function that can be reused outside query31// context.32static Operation *extractFunction(std::vector<Operation *> &ops,33 MLIRContext *context,34 llvm::StringRef functionName) {35 context->loadDialect<func::FuncDialect>();36 OpBuilder builder(context);37 38 // Collect data for function creation39 std::vector<Operation *> slice;40 std::vector<Value> values;41 std::vector<Type> outputTypes;42 43 for (auto *op : ops) {44 // Return op's operands are propagated, but the op itself isn't needed.45 if (!isa<func::ReturnOp>(op))46 slice.push_back(op);47 48 // All results are returned by the extracted function.49 llvm::append_range(outputTypes, op->getResults().getTypes());50 51 // Track all values that need to be taken as input to function.52 llvm::append_range(values, op->getOperands());53 }54 55 // Create the function56 FunctionType funcType =57 builder.getFunctionType(TypeRange(ValueRange(values)), outputTypes);58 auto loc = builder.getUnknownLoc();59 func::FuncOp funcOp = func::FuncOp::create(loc, functionName, funcType);60 61 builder.setInsertionPointToEnd(funcOp.addEntryBlock());62 63 // Map original values to function arguments64 IRMapping mapper;65 for (const auto &arg : llvm::enumerate(values))66 mapper.map(arg.value(), funcOp.getArgument(arg.index()));67 68 // Clone operations and build function body69 std::vector<Operation *> clonedOps;70 std::vector<Value> clonedVals;71 // TODO: Handle extraction of operations with compute payloads defined via72 // regions.73 for (Operation *slicedOp : slice) {74 Operation *clonedOp =75 clonedOps.emplace_back(builder.clone(*slicedOp, mapper));76 clonedVals.insert(clonedVals.end(), clonedOp->result_begin(),77 clonedOp->result_end());78 }79 // Add return operation80 func::ReturnOp::create(builder, loc, clonedVals);81 82 // Remove unused function arguments83 size_t currentIndex = 0;84 while (currentIndex < funcOp.getNumArguments()) {85 // Erase if possible.86 if (funcOp.getArgument(currentIndex).use_empty())87 if (succeeded(funcOp.eraseArgument(currentIndex)))88 continue;89 ++currentIndex;90 }91 92 return funcOp;93}94 95Query::~Query() = default;96 97LogicalResult InvalidQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {98 os << errStr << "\n";99 return mlir::failure();100}101 102LogicalResult NoOpQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {103 return mlir::success();104}105 106LogicalResult HelpQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {107 os << "Available commands:\n\n"108 " match MATCHER, m MATCHER "109 "Match the mlir against the given matcher.\n"110 " quit "111 "Terminates the query session.\n\n";112 return mlir::success();113}114 115LogicalResult QuitQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {116 qs.terminate = true;117 return mlir::success();118}119 120LogicalResult MatchQuery::run(llvm::raw_ostream &os, QuerySession &qs) const {121 Operation *rootOp = qs.getRootOp();122 int matchCount = 0;123 matcher::MatchFinder finder;124 125 StringRef functionName = matcher.getFunctionName();126 auto matches = finder.collectMatches(rootOp, std::move(matcher));127 128 // An extract call is recognized by considering if the matcher has a name.129 // TODO: Consider making the extract more explicit.130 if (!functionName.empty()) {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 (Operation *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::query158