mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp

   1 //===- mlir-src-sharder.cpp - A tool for sharder generated source files ---===//
   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 #include "mlir/Support/FileUtilities.h"
  10 #include "mlir/Support/LLVM.h"
  11 #include "llvm/Support/CommandLine.h"
  12 #include "llvm/Support/InitLLVM.h"
  13 #include "llvm/Support/MemoryBuffer.h"
  14 #include "llvm/Support/ToolOutputFile.h"
  15
  16 using namespace mlir;
  17
  18 /// Create a dependency file for `-d` option.
  19 ///
  20 /// This functionality is generally only for the benefit of the build system,
  21 /// and is modeled after the same option in TableGen.
  22 static LogicalResult createDependencyFile(StringRef outputFilename,
  23                                           StringRef dependencyFile) {
  24   if (outputFilename == "-") {
  25     llvm::errs() << "error: the option -d must be used together with -o\n";
  26     return failure();
  27   }
  28
  29   std::string errorMessage;
  30   std::unique_ptr<llvm::ToolOutputFile> outputFile =
  31       openOutputFile(dependencyFile, &errorMessage);
  32   if (!outputFile) {
  33     llvm::errs() << errorMessage << "\n";
  34     return failure();
  35   }
  36
  37   outputFile->os() << outputFilename << ":\n";
  38   outputFile->keep();
  39   return success();
  40 }
  41
  42 int main(int argc, char **argv) {
  43   // FIXME: This is necessary because we link in TableGen, which defines its
  44   // options as static variables.. some of which overlap with our options.
  45   llvm::cl::ResetCommandLineParser();
  46
  47   llvm::cl::opt<unsigned> opShardIndex(
  48       "op-shard-index", llvm::cl::desc("The current shard index"));
  49   llvm::cl::opt<std::string> inputFilename(llvm::cl::Positional,
  50                                            llvm::cl::desc("<input file>"),
  51                                            llvm::cl::init("-"));
  52   llvm::cl::opt<std::string> outputFilename(
  53       "o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"),
  54       llvm::cl::init("-"));
  55   llvm::cl::list<std::string> includeDirs(
  56       "I", llvm::cl::desc("Directory of include files"),
  57       llvm::cl::value_desc("directory"), llvm::cl::Prefix);
  58   llvm::cl::opt<std::string> dependencyFilename(
  59       "d", llvm::cl::desc("Dependency filename"),
  60       llvm::cl::value_desc("filename"), llvm::cl::init(""));
  61   llvm::cl::opt<bool> writeIfChanged(
  62       "write-if-changed",
  63       llvm::cl::desc("Only write to the output file if it changed"));
  64
  65   // `ResetCommandLineParser` at the above unregistered the "D" option
  66   // of `llvm-tblgen`, which caused `TestOps.cpp` to fail due to
  67   // "Unknnown command line argument '-D...`" when a macros name is
  68   // present. The following is a workaround to re-register it again.
  69   llvm::cl::list<std::string> macroNames(
  70       "D",
  71       llvm::cl::desc(
  72           "Name of the macro to be defined -- ignored by mlir-src-sharder"),
  73       llvm::cl::value_desc("macro name"), llvm::cl::Prefix);
  74
  75   llvm::InitLLVM y(argc, argv);
  76   llvm::cl::ParseCommandLineOptions(argc, argv);
  77
  78   // Open the input file.
  79   std::string errorMessage;
  80   std::unique_ptr<llvm::MemoryBuffer> inputFile =
  81       openInputFile(inputFilename, &errorMessage);
  82   if (!inputFile) {
  83     llvm::errs() << errorMessage << "\n";
  84     return 1;
  85   }
  86
  87   // Write the output to a buffer.
  88   std::string outputStr;
  89   llvm::raw_string_ostream os(outputStr);
  90   os << "#define GET_OP_DEFS_" << opShardIndex << "\n"
  91      << inputFile->getBuffer();
  92
  93   // Determine whether we need to write the output file.
  94   bool shouldWriteOutput = true;
  95   if (writeIfChanged) {
  96     // Only update the real output file if there are any differences. This
  97     // prevents recompilation of all the files depending on it if there aren't
  98     // any.
  99     if (auto existingOrErr =
 100             llvm::MemoryBuffer::getFile(outputFilename, /*IsText=*/true))
 101       if (std::move(existingOrErr.get())->getBuffer() == outputStr)
 102         shouldWriteOutput = false;
 103   }
 104
 105   // Populate the output file if necessary.
 106   if (shouldWriteOutput) {
 107     std::unique_ptr<llvm::ToolOutputFile> outputFile =
 108         openOutputFile(outputFilename, &errorMessage);
 109     if (!outputFile) {
 110       llvm::errs() << errorMessage << "\n";
 111       return 1;
 112     }
 113     outputFile->os() << os.str();
 114     outputFile->keep();
 115   }
 116
 117   // Always write the depfile, even if the main output hasn't changed. If it's
 118   // missing, Ninja considers the output dirty.
 119   if (!dependencyFilename.empty())
 120     if (failed(createDependencyFile(outputFilename, dependencyFilename)))
 121       return 1;
 122
 123   return 0;
 124 }