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 |
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(Val: 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(Val: false)); |
59 | |
60 | llvm::cl::HideUnrelatedOptions(Category&: mlirQueryCategory); |
61 | |
62 | llvm::InitLLVM y(argc, argv); |
63 | |
64 | llvm::cl::ParseCommandLineOptions(argc, argv, Overview: "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: &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(F: std::move(file), IncludeLoc: SMLoc()); |
81 | |
82 | context.allowUnregisteredDialects(allow: allowUnregisteredDialects); |
83 | |
84 | // Parse the input MLIR file. |
85 | OwningOpRef<Operation *> opRef = |
86 | noImplicitModule ? parseSourceFile(sourceMgr, config: &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(line: command, qs); |
96 | if (mlir::failed(Result: queryRef->run(os&: 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: *line, qs); |
106 | (void)queryRef->run(os&: llvm::outs(), qs); |
107 | llvm::outs().flush(); |
108 | if (qs.terminate) |
109 | break; |
110 | } |
111 | } |
112 | |
113 | return mlir::success(); |
114 | } |
115 | |