124 lines · cpp
1//===- MlirQueryMain.cpp - MLIR Query main --------------------------------===//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// This file implements the general framework of the MLIR query tool. It10// parses the command line arguments, parses the MLIR file and outputs the query11// results.12//13//===----------------------------------------------------------------------===//14 15#include "mlir/Tools/mlir-query/MlirQueryMain.h"16#include "mlir/IR/BuiltinOps.h"17#include "mlir/Parser/Parser.h"18#include "mlir/Query/Query.h"19#include "mlir/Query/QuerySession.h"20#include "mlir/Support/FileUtilities.h"21#include "llvm/LineEditor/LineEditor.h"22#include "llvm/Support/CommandLine.h"23#include "llvm/Support/InitLLVM.h"24#include "llvm/Support/Process.h"25#include "llvm/Support/SourceMgr.h"26 27//===----------------------------------------------------------------------===//28// Query Parser29//===----------------------------------------------------------------------===//30 31llvm::LogicalResult32mlir::mlirQueryMain(int argc, char **argv, MLIRContext &context,33 const mlir::query::matcher::Registry &matcherRegistry) {34 35 // Override the default '-h' and use the default PrintHelpMessage() which36 // won't print options in categories.37 static llvm::cl::opt<bool> help("h", llvm::cl::desc("Alias for -help"),38 llvm::cl::Hidden);39 40 static llvm::cl::OptionCategory mlirQueryCategory("mlir-query options");41 42 static llvm::cl::list<std::string> commands(43 "c", llvm::cl::desc("Specify command to run"),44 llvm::cl::value_desc("command"), llvm::cl::cat(mlirQueryCategory));45 46 static llvm::cl::opt<std::string> inputFilename(47 llvm::cl::Positional, llvm::cl::desc("<input file>"), llvm::cl::init("-"),48 llvm::cl::cat(mlirQueryCategory));49 50 static llvm::cl::opt<bool> noImplicitModule{51 "no-implicit-module",52 llvm::cl::desc(53 "Disable implicit addition of a top-level module op during parsing"),54 llvm::cl::init(false)};55 56 static llvm::cl::opt<bool> allowUnregisteredDialects(57 "allow-unregistered-dialect",58 llvm::cl::desc("Allow operation with no registered dialects"),59 llvm::cl::init(false));60 61 llvm::cl::HideUnrelatedOptions(mlirQueryCategory);62 63 llvm::InitLLVM y(argc, argv);64 65 llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR test case query tool.\n");66 67 if (help) {68 llvm::cl::PrintHelpMessage();69 return mlir::success();70 }71 72 // When reading from stdin and the input is a tty, it is often a user mistake73 // and the process "appears to be stuck". Print a message to let the user74 // know!75 if (inputFilename == "-" &&76 llvm::sys::Process::FileDescriptorIsDisplayed(fileno(stdin)))77 llvm::errs() << "(processing input from stdin now, hit ctrl-c/ctrl-d to "78 "interrupt)\n";79 80 // Set up the input file.81 std::string errorMessage;82 auto file = openInputFile(inputFilename, &errorMessage);83 if (!file) {84 llvm::errs() << errorMessage << "\n";85 return mlir::failure();86 }87 88 auto sourceMgr = llvm::SourceMgr();89 auto bufferId = sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());90 91 context.allowUnregisteredDialects(allowUnregisteredDialects);92 93 // Parse the input MLIR file.94 OwningOpRef<Operation *> opRef =95 noImplicitModule ? parseSourceFile(sourceMgr, &context)96 : parseSourceFile<mlir::ModuleOp>(sourceMgr, &context);97 if (!opRef)98 return mlir::failure();99 100 mlir::query::QuerySession qs(opRef.get(), sourceMgr, bufferId,101 matcherRegistry);102 if (!commands.empty()) {103 for (auto &command : commands) {104 mlir::query::QueryRef queryRef = mlir::query::parse(command, qs);105 if (mlir::failed(queryRef->run(llvm::outs(), qs)))106 return mlir::failure();107 }108 } else {109 llvm::LineEditor le("mlir-query");110 le.setListCompleter([&qs](llvm::StringRef line, size_t pos) {111 return mlir::query::complete(line, pos, qs);112 });113 while (std::optional<std::string> line = le.readLine()) {114 mlir::query::QueryRef queryRef = mlir::query::parse(*line, qs);115 (void)queryRef->run(llvm::outs(), qs);116 llvm::outs().flush();117 if (qs.terminate)118 break;119 }120 }121 122 return mlir::success();123}124