Support generation of external dialects in mlir-tblgen-jl

JuliaLLVM · May 24, 2024 · 5c35af3 · 5c35af3
1 parent a48d146
commit 5c35af3
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/deps/tblgen/jl-generators.cc b/deps/tblgen/jl-generators.cc
@@ -158,6 +158,7 @@ namespace
 } // namespace
 
 extern bool disableModuleWrap;
+extern bool isExternal;
 
 bool emitOpTableDefs(const llvm::RecordKeeper &recordKeeper,
                      llvm::raw_ostream &os)
@@ -168,25 +169,34 @@ bool emitOpTableDefs(const llvm::RecordKeeper &recordKeeper,
   std::vector<llvm::Record *> opdefs = recordKeeper.getAllDerivedDefinitions("Op");
 #endif
 
+  const char *imports;
+  if (isExternal)
+  {
+    imports = R"(import MLIR.IR: IR, NamedAttribute, Value, Location, Block, Region, Attribute, context, IndexType
+import MLIR.Dialects: namedattribute, operandsegmentsizes)";
+  }
+  else
+  {
+    imports = R"(import ...IR: IR, NamedAttribute, Value, Location, Block, Region, Attribute, context, IndexType
+import ..Dialects: namedattribute, operandsegmentsizes)";
+  }
+
   const char *moduleTemplate;
   if (disableModuleWrap)
   {
-    moduleTemplate = R"(import ...IR: IR, NamedAttribute, Value, Location, Block, Region, Attribute, context, IndexType
-import ..Dialects: namedattribute, operandsegmentsizes
-import ...API
+    // 0: imports, 1: content
+    moduleTemplate = R"({0}
 
-{0}
+{1}
 )";
   }
   else
   {
+    // 0: module name, 1: imports, 2: content
     moduleTemplate = R"(module {0}
 
-import ...IR: IR, NamedAttribute, Value, Location, Block, Region, Attribute, context, IndexType
-import ..Dialects: namedattribute, operandsegmentsizes
-import ...API
-
 {1}
+{2}
 end # {0}
 )";
   }
@@ -439,11 +449,11 @@ end
 
   if (disableModuleWrap)
   {
-    os << llvm::formatv(moduleTemplate, modulecontents);
+    os << llvm::formatv(moduleTemplate, imports, modulecontents);
   }
   else
   {
-    os << llvm::formatv(moduleTemplate, modulename, modulecontents);
+    os << llvm::formatv(moduleTemplate, modulename, imports, modulecontents);
   }
 
   return false;

diff --git a/deps/tblgen/mlir-jl-tblgen.cc b/deps/tblgen/mlir-jl-tblgen.cc
@@ -43,11 +43,13 @@ static std::array<GeneratorInfo, 1> generators {{
 
 generator_function* generator;
 bool disableModuleWrap;
+bool isExternal;
 
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   llvm::cl::opt<std::string> generatorOpt("generator", llvm::cl::desc("Generator to run"), cl::Required);
   llvm::cl::opt<bool> disableModuleWrapOpt("disable-module-wrap", llvm::cl::desc("Disable module wrap"), cl::init(false));
+  llvm::cl::opt<bool> isExternalOpt("external", llvm::cl::desc("Mark the dialect as external and generate bindings accordingly"), cl::init(false));
   cl::ParseCommandLineOptions(argc, argv);
   for (const auto& spec : generators) {
     if (generatorOpt == spec.name) {
@@ -60,6 +62,7 @@ int main(int argc, char **argv) {
     abort();
   }
   disableModuleWrap = disableModuleWrapOpt;
+  isExternal = isExternalOpt;
 
   return TableGenMain(argv[0], [](raw_ostream& os, RecordKeeper &records) {
     return generator(records, os);