xref: /llvm-project/mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp (revision ffc80de8643969ffa0dbbd377c5b33e3a7488f5e)
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 }
125