xref: /llvm-project/mlir/lib/Tools/mlir-translate/Translation.cpp (revision 18546ff8dd45a81e72c0a2ed0561b5aec8c15ca3)
1 //===- Translation.cpp - Translation registry -----------------------------===//
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 // Definitions of the translation registry.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Tools/mlir-translate/Translation.h"
14 #include "mlir/IR/AsmState.h"
15 #include "mlir/IR/BuiltinOps.h"
16 #include "mlir/IR/Dialect.h"
17 #include "mlir/IR/Verifier.h"
18 #include "mlir/Parser/Parser.h"
19 #include "mlir/Tools/ParseUtilities.h"
20 #include "llvm/Support/SourceMgr.h"
21 
22 using namespace mlir;
23 
24 //===----------------------------------------------------------------------===//
25 // Translation CommandLine Options
26 //===----------------------------------------------------------------------===//
27 
28 struct TranslationOptions {
29   llvm::cl::opt<bool> noImplicitModule{
30       "no-implicit-module",
31       llvm::cl::desc("Disable the parsing of an implicit top-level module op"),
32       llvm::cl::init(false)};
33 };
34 
35 static llvm::ManagedStatic<TranslationOptions> clOptions;
36 
37 void mlir::registerTranslationCLOptions() { *clOptions; }
38 
39 //===----------------------------------------------------------------------===//
40 // Translation Registry
41 //===----------------------------------------------------------------------===//
42 
43 /// Get the mutable static map between registered file-to-file MLIR
44 /// translations.
45 static llvm::StringMap<Translation> &getTranslationRegistry() {
46   static llvm::StringMap<Translation> translationBundle;
47   return translationBundle;
48 }
49 
50 /// Register the given translation.
51 static void registerTranslation(StringRef name, StringRef description,
52                                 Optional<llvm::Align> inputAlignment,
53                                 const TranslateFunction &function) {
54   auto &registry = getTranslationRegistry();
55   if (registry.count(name))
56     llvm::report_fatal_error(
57         "Attempting to overwrite an existing <file-to-file> function");
58   assert(function &&
59          "Attempting to register an empty translate <file-to-file> function");
60   registry[name] = Translation(function, description, inputAlignment);
61 }
62 
63 TranslateRegistration::TranslateRegistration(
64     StringRef name, StringRef description, const TranslateFunction &function) {
65   registerTranslation(name, description, /*inputAlignment=*/std::nullopt,
66                       function);
67 }
68 
69 //===----------------------------------------------------------------------===//
70 // Translation to MLIR
71 //===----------------------------------------------------------------------===//
72 
73 // Puts `function` into the to-MLIR translation registry unless there is already
74 // a function registered for the same name.
75 static void registerTranslateToMLIRFunction(
76     StringRef name, StringRef description, Optional<llvm::Align> inputAlignment,
77     const TranslateSourceMgrToMLIRFunction &function) {
78   auto wrappedFn = [function](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
79                               raw_ostream &output, MLIRContext *context) {
80     OwningOpRef<Operation *> op = function(sourceMgr, context);
81     if (!op || failed(verify(*op)))
82       return failure();
83     op.get()->print(output);
84     return success();
85   };
86   registerTranslation(name, description, inputAlignment, wrappedFn);
87 }
88 
89 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
90     StringRef name, StringRef description,
91     const TranslateSourceMgrToMLIRFunction &function,
92     Optional<llvm::Align> inputAlignment) {
93   registerTranslateToMLIRFunction(name, description, inputAlignment, function);
94 }
95 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
96     StringRef name, StringRef description,
97     const TranslateRawSourceMgrToMLIRFunction &function,
98     Optional<llvm::Align> inputAlignment) {
99   registerTranslateToMLIRFunction(
100       name, description, inputAlignment,
101       [function](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
102                  MLIRContext *ctx) { return function(*sourceMgr, ctx); });
103 }
104 /// Wraps `function` with a lambda that extracts a StringRef from a source
105 /// manager and registers the wrapper lambda as a to-MLIR conversion.
106 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
107     StringRef name, StringRef description,
108     const TranslateStringRefToMLIRFunction &function,
109     Optional<llvm::Align> inputAlignment) {
110   registerTranslateToMLIRFunction(
111       name, description, inputAlignment,
112       [function](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
113                  MLIRContext *ctx) {
114         const llvm::MemoryBuffer *buffer =
115             sourceMgr->getMemoryBuffer(sourceMgr->getMainFileID());
116         return function(buffer->getBuffer(), ctx);
117       });
118 }
119 
120 //===----------------------------------------------------------------------===//
121 // Translation from MLIR
122 //===----------------------------------------------------------------------===//
123 
124 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
125     StringRef name, StringRef description,
126     const TranslateFromMLIRFunction &function,
127     const std::function<void(DialectRegistry &)> &dialectRegistration) {
128   registerTranslation(
129       name, description, /*inputAlignment=*/std::nullopt,
130       [function,
131        dialectRegistration](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
132                             raw_ostream &output, MLIRContext *context) {
133         DialectRegistry registry;
134         dialectRegistration(registry);
135         context->appendDialectRegistry(registry);
136         bool implicitModule =
137             (!clOptions.isConstructed() || !clOptions->noImplicitModule);
138         OwningOpRef<Operation *> op =
139             parseSourceFileForTool(sourceMgr, context, implicitModule);
140         if (!op || failed(verify(*op)))
141           return failure();
142         return function(op.get(), output);
143       });
144 }
145 
146 //===----------------------------------------------------------------------===//
147 // Translation Parser
148 //===----------------------------------------------------------------------===//
149 
150 TranslationParser::TranslationParser(llvm::cl::Option &opt)
151     : llvm::cl::parser<const Translation *>(opt) {
152   for (const auto &kv : getTranslationRegistry())
153     addLiteralOption(kv.first(), &kv.second, kv.second.getDescription());
154 }
155 
156 void TranslationParser::printOptionInfo(const llvm::cl::Option &o,
157                                         size_t globalWidth) const {
158   TranslationParser *tp = const_cast<TranslationParser *>(this);
159   llvm::array_pod_sort(tp->Values.begin(), tp->Values.end(),
160                        [](const TranslationParser::OptionInfo *lhs,
161                           const TranslationParser::OptionInfo *rhs) {
162                          return lhs->Name.compare(rhs->Name);
163                        });
164   llvm::cl::parser<const Translation *>::printOptionInfo(o, globalWidth);
165 }
166