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 "mlir/Support/LogicalResult.h" |
22 | #include "llvm/LineEditor/LineEditor.h" |
23 | #include "llvm/Support/CommandLine.h" |
24 | #include "llvm/Support/InitLLVM.h" |
25 | #include "llvm/Support/SourceMgr.h" |
26 | |
27 | //===----------------------------------------------------------------------===// |
28 | // Query Parser |
29 | //===----------------------------------------------------------------------===// |
30 | |
31 | mlir::LogicalResult |
32 | mlir::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() which |
36 | // 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>" ), |
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(Val: 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(Val: false)); |
60 | |
61 | llvm::cl::HideUnrelatedOptions(Category&: mlirQueryCategory); |
62 | |
63 | llvm::InitLLVM y(argc, argv); |
64 | |
65 | llvm::cl::ParseCommandLineOptions(argc, argv, Overview: "MLIR test case query tool.\n" ); |
66 | |
67 | if (help) { |
68 | llvm::cl::PrintHelpMessage(); |
69 | return mlir::success(); |
70 | } |
71 | |
72 | // Set up the input file. |
73 | std::string errorMessage; |
74 | auto file = openInputFile(inputFilename, errorMessage: &errorMessage); |
75 | if (!file) { |
76 | llvm::errs() << errorMessage << "\n" ; |
77 | return mlir::failure(); |
78 | } |
79 | |
80 | auto sourceMgr = llvm::SourceMgr(); |
81 | auto bufferId = sourceMgr.AddNewSourceBuffer(F: std::move(file), IncludeLoc: SMLoc()); |
82 | |
83 | context.allowUnregisteredDialects(allow: allowUnregisteredDialects); |
84 | |
85 | // Parse the input MLIR file. |
86 | OwningOpRef<Operation *> opRef = |
87 | noImplicitModule ? parseSourceFile(sourceMgr, config: &context) |
88 | : parseSourceFile<mlir::ModuleOp>(sourceMgr, &context); |
89 | if (!opRef) |
90 | return mlir::failure(); |
91 | |
92 | mlir::query::QuerySession qs(opRef.get(), sourceMgr, bufferId, |
93 | matcherRegistry); |
94 | if (!commands.empty()) { |
95 | for (auto &command : commands) { |
96 | mlir::query::QueryRef queryRef = mlir::query::parse(line: command, qs); |
97 | if (mlir::failed(result: queryRef->run(os&: llvm::outs(), qs))) |
98 | return mlir::failure(); |
99 | } |
100 | } else { |
101 | llvm::LineEditor le("mlir-query" ); |
102 | le.setListCompleter([&qs](llvm::StringRef line, size_t pos) { |
103 | return mlir::query::complete(line, pos, qs); |
104 | }); |
105 | while (std::optional<std::string> line = le.readLine()) { |
106 | mlir::query::QueryRef queryRef = mlir::query::parse(line: *line, qs); |
107 | (void)queryRef->run(os&: llvm::outs(), qs); |
108 | llvm::outs().flush(); |
109 | if (qs.terminate) |
110 | break; |
111 | } |
112 | } |
113 | |
114 | return mlir::success(); |
115 | } |
116 | |