From 8973b98c0ba8aabe86fad8a8be5a75643ccd0258 Mon Sep 17 00:00:00 2001
From: aminim <aminim@google.com>
Date: Mon, 24 Aug 2020 01:47:40 -0700
Subject: [PATCH] Fix removal global dialect registry

PiperOrigin-RevId: 328099318
---
 sair_opt.cc          | 5 ++---
 sair_registration.cc | 4 +++-
 sair_registration.h  | 6 ++++--
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/sair_opt.cc b/sair_opt.cc
index 343eb3e6..10f966d0 100644
--- a/sair_opt.cc
+++ b/sair_opt.cc
@@ -60,9 +60,7 @@ int main(int argc, char **argv) {
   llvm::InitLLVM init(argc, argv);
 
   // Register any pass manager command line options.
-  mlir::registerAllDialects();
   mlir::registerAllPasses();
-  sair::RegisterSairDialect();
   sair::RegisterSairPasses();
   mlir::registerPassManagerCLOptions();
   mlir::PassPipelineCLParser passPipeline("", "MLIR passes to run");
@@ -84,7 +82,8 @@ int main(int argc, char **argv) {
   }
 
   mlir::DialectRegistry registry;
-  mlir::getGlobalDialectRegistry().appendTo(registry);
+  sair::RegisterSairDialect(registry);
+  mlir::registerAllDialects(registry);
   return mlir::failed(
       mlir::MlirOptMain(outputFile->os(), std::move(inputFile), passPipeline,
                         registry, split_input_file, verify_diagnostics,
diff --git a/sair_registration.cc b/sair_registration.cc
index 32e462a0..9a235be6 100644
--- a/sair_registration.cc
+++ b/sair_registration.cc
@@ -64,4 +64,6 @@ void sair::RegisterSairPasses() {
       [](llvm::function_ref<void(const mlir::detail::PassOptions &)>) {});
 }
 
-void sair::RegisterSairDialect() { mlir::registerDialect<sair::SairDialect>(); }
+void sair::RegisterSairDialect(mlir::DialectRegistry &registry) {
+  registry.insert<sair::SairDialect>();
+}
diff --git a/sair_registration.h b/sair_registration.h
index e9394669..1f8431f1 100644
--- a/sair_registration.h
+++ b/sair_registration.h
@@ -14,10 +14,12 @@
 
 #ifndef SAIR_SAIR_REGISTRATION_H_
 #define SAIR_SAIR_REGISTRATION_H_
-
+namespace mlir {
+class DialectRegistry;
+}
 namespace sair {
 // Registers Sair dialect with MLIR.
-void RegisterSairDialect();
+void RegisterSairDialect(mlir::DialectRegistry &registry);
 
 // Registers all Sair passes with MLIR.
 void RegisterSairPasses();