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
31mlir::LogicalResult
32mlir::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

source code of mlir/lib/Tools/mlir-query/MlirQueryMain.cpp