From 73b5a44f338e670b062853aec61ded0a2271a0b7 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <aminim@google.com>
Date: Mon, 24 Aug 2020 10:58:19 -0700
Subject: [PATCH] Update
 third_party/tensorflow/compiler/mlir/tensorflow/utils/... to not depend on
 the global Dialect Registry (NFC)

PiperOrigin-RevId: 328171679
---
 include/mlir-hlo/Dialect/mhlo/IR/register.h | 3 +++
 lib/Dialect/mhlo/IR/init.cc                 | 8 ++++++++
 2 files changed, 11 insertions(+)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/register.h b/include/mlir-hlo/Dialect/mhlo/IR/register.h
index 5773901..90ff6c9 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/register.h
+++ b/include/mlir-hlo/Dialect/mhlo/IR/register.h
@@ -17,10 +17,13 @@ limitations under the License.
 #define MLIR_HLO_DIALECT_MHLO_IR_REGISTER_H_
 
 namespace mlir {
+class DialectRegistry;
 namespace mhlo {
 
 void registerAllDialects();
 
+// Add chlo, mhlo, lmhlo dialects to the provided registry.
+void registerAllMhloDialects(DialectRegistry &registry);
 }
 }  // namespace mlir
 
diff --git a/lib/Dialect/mhlo/IR/init.cc b/lib/Dialect/mhlo/IR/init.cc
index 9fffeae..cf8bd25 100644
--- a/lib/Dialect/mhlo/IR/init.cc
+++ b/lib/Dialect/mhlo/IR/init.cc
@@ -31,3 +31,11 @@ void mlir::mhlo::registerAllDialects() {
 
   // Dependent dialects
 }
+
+void mlir::mhlo::registerAllMhloDialects(mlir::DialectRegistry &registry) {
+  // clang-format off
+  registry.insert<mlir::chlo::HloClientDialect,
+                  mlir::lmhlo::LmhloDialect,
+                  mlir::mhlo::MhloDialect>();
+  // clang-format on
+}