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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the general framework of the MLIR query tool. It
10 // parses the command line arguments, parses the MLIR file and outputs the query
11 // 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/SourceMgr.h"
25
26 //===----------------------------------------------------------------------===//
27 // Query Parser
28 //===----------------------------------------------------------------------===//
29
30 llvm::LogicalResult
mlirQueryMain(int argc,char ** argv,MLIRContext & context,const mlir::query::matcher::Registry & matcherRegistry)31 mlir::mlirQueryMain(int argc, char **argv, MLIRContext &context,
32 const mlir::query::matcher::Registry &matcherRegistry) {
33
34 // Override the default '-h' and use the default PrintHelpMessage() which
35 // won't print options in categories.
36 static llvm::cl::opt<bool> help("h", llvm::cl::desc("Alias for -help"),
37 llvm::cl::Hidden);
38
39 static llvm::cl::OptionCategory mlirQueryCategory("mlir-query options");
40
41 static llvm::cl::list<std::string> commands(
42 "c", llvm::cl::desc("Specify command to run"),
43 llvm::cl::value_desc("command"), llvm::cl::cat(mlirQueryCategory));
44
45 static llvm::cl::opt<std::string> inputFilename(
46 llvm::cl::Positional, llvm::cl::desc("<input file>"),
47 llvm::cl::cat(mlirQueryCategory));
48
49 static llvm::cl::opt<bool> noImplicitModule{
50 "no-implicit-module",
51 llvm::cl::desc(
52 "Disable implicit addition of a top-level module op during parsing"),
53 llvm::cl::init(false)};
54
55 static llvm::cl::opt<bool> allowUnregisteredDialects(
56 "allow-unregistered-dialect",
57 llvm::cl::desc("Allow operation with no registered dialects"),
58 llvm::cl::init(false));
59
60 llvm::cl::HideUnrelatedOptions(mlirQueryCategory);
61
62 llvm::InitLLVM y(argc, argv);
63
64 llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR test case query tool.\n");
65
66 if (help) {
67 llvm::cl::PrintHelpMessage();
68 return mlir::success();
69 }
70
71 // Set up the input file.
72 std::string errorMessage;
73 auto file = openInputFile(inputFilename, &errorMessage);
74 if (!file) {
75 llvm::errs() << errorMessage << "\n";
76 return mlir::failure();
77 }
78
79 auto sourceMgr = llvm::SourceMgr();
80 auto bufferId = sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
81
82 context.allowUnregisteredDialects(allowUnregisteredDialects);
83
84 // Parse the input MLIR file.
85 OwningOpRef<Operation *> opRef =
86 noImplicitModule ? parseSourceFile(sourceMgr, &context)
87 : parseSourceFile<mlir::ModuleOp>(sourceMgr, &context);
88 if (!opRef)
89 return mlir::failure();
90
91 mlir::query::QuerySession qs(opRef.get(), sourceMgr, bufferId,
92 matcherRegistry);
93 if (!commands.empty()) {
94 for (auto &command : commands) {
95 mlir::query::QueryRef queryRef = mlir::query::parse(command, qs);
96 if (mlir::failed(queryRef->run(llvm::outs(), qs)))
97 return mlir::failure();
98 }
99 } else {
100 llvm::LineEditor le("mlir-query");
101 le.setListCompleter([&qs](llvm::StringRef line, size_t pos) {
102 return mlir::query::complete(line, pos, qs);
103 });
104 while (std::optional<std::string> line = le.readLine()) {
105 mlir::query::QueryRef queryRef = mlir::query::parse(*line, qs);
106 (void)queryRef->run(llvm::outs(), qs);
107 llvm::outs().flush();
108 if (qs.terminate)
109 break;
110 }
111 }
112
113 return mlir::success();
114 }
115