mlir/lib/Tools/mlir-query/MlirQueryMain.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-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(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 }