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/Translation.h"
14 #include "mlir/IR/Module.h"
15 #include "mlir/IR/Verifier.h"
16 #include "mlir/Parser.h"
17 #include "mlir/Support/LLVM.h"
18 #include "llvm/Support/SourceMgr.h"
19
20 using namespace mlir;
21
22 //===----------------------------------------------------------------------===//
23 // Translation Registry
24 //===----------------------------------------------------------------------===//
25
26 /// Get the mutable static map between registered file-to-file MLIR translations
27 /// and the TranslateFunctions that perform those translations.
getTranslationRegistry()28 static llvm::StringMap<TranslateFunction> &getTranslationRegistry() {
29 static llvm::StringMap<TranslateFunction> translationRegistry;
30 return translationRegistry;
31 }
32
33 /// Register the given translation.
registerTranslation(StringRef name,const TranslateFunction & function)34 static void registerTranslation(StringRef name,
35 const TranslateFunction &function) {
36 auto &translationRegistry = getTranslationRegistry();
37 if (translationRegistry.find(name) != translationRegistry.end())
38 llvm::report_fatal_error(
39 "Attempting to overwrite an existing <file-to-file> function");
40 assert(function &&
41 "Attempting to register an empty translate <file-to-file> function");
42 translationRegistry[name] = function;
43 }
44
TranslateRegistration(StringRef name,const TranslateFunction & function)45 TranslateRegistration::TranslateRegistration(
46 StringRef name, const TranslateFunction &function) {
47 registerTranslation(name, function);
48 }
49
50 //===----------------------------------------------------------------------===//
51 // Translation to MLIR
52 //===----------------------------------------------------------------------===//
53
54 // Puts `function` into the to-MLIR translation registry unless there is already
55 // a function registered for the same name.
registerTranslateToMLIRFunction(StringRef name,const TranslateSourceMgrToMLIRFunction & function)56 static void registerTranslateToMLIRFunction(
57 StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
58 auto wrappedFn = [function](llvm::SourceMgr &sourceMgr, raw_ostream &output,
59 MLIRContext *context) {
60 OwningModuleRef module = function(sourceMgr, context);
61 if (!module || failed(verify(*module)))
62 return failure();
63 module->print(output);
64 return success();
65 };
66 registerTranslation(name, wrappedFn);
67 }
68
TranslateToMLIRRegistration(StringRef name,const TranslateSourceMgrToMLIRFunction & function)69 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
70 StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
71 registerTranslateToMLIRFunction(name, function);
72 }
73
74 /// Wraps `function` with a lambda that extracts a StringRef from a source
75 /// manager and registers the wrapper lambda as a to-MLIR conversion.
TranslateToMLIRRegistration(StringRef name,const TranslateStringRefToMLIRFunction & function)76 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
77 StringRef name, const TranslateStringRefToMLIRFunction &function) {
78 registerTranslateToMLIRFunction(
79 name, [function](llvm::SourceMgr &sourceMgr, MLIRContext *ctx) {
80 const llvm::MemoryBuffer *buffer =
81 sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
82 return function(buffer->getBuffer(), ctx);
83 });
84 }
85
86 //===----------------------------------------------------------------------===//
87 // Translation from MLIR
88 //===----------------------------------------------------------------------===//
89
TranslateFromMLIRRegistration(StringRef name,const TranslateFromMLIRFunction & function)90 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
91 StringRef name, const TranslateFromMLIRFunction &function) {
92 registerTranslation(name, [function](llvm::SourceMgr &sourceMgr,
93 raw_ostream &output,
94 MLIRContext *context) {
95 auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
96 if (!module)
97 return failure();
98 return function(module.get(), output);
99 });
100 }
101
102 //===----------------------------------------------------------------------===//
103 // Translation Parser
104 //===----------------------------------------------------------------------===//
105
TranslationParser(llvm::cl::Option & opt)106 TranslationParser::TranslationParser(llvm::cl::Option &opt)
107 : llvm::cl::parser<const TranslateFunction *>(opt) {
108 for (const auto &kv : getTranslationRegistry())
109 addLiteralOption(kv.first(), &kv.second, kv.first());
110 }
111
printOptionInfo(const llvm::cl::Option & o,size_t globalWidth) const112 void TranslationParser::printOptionInfo(const llvm::cl::Option &o,
113 size_t globalWidth) const {
114 TranslationParser *tp = const_cast<TranslationParser *>(this);
115 llvm::array_pod_sort(tp->Values.begin(), tp->Values.end(),
116 [](const TranslationParser::OptionInfo *lhs,
117 const TranslationParser::OptionInfo *rhs) {
118 return lhs->Name.compare(rhs->Name);
119 });
120 llvm::cl::parser<const TranslateFunction *>::printOptionInfo(o, globalWidth);
121 }
122