From 3c33fe4b9e83abe83b59840f76eb8016dd7f73f3 Mon Sep 17 00:00:00 2001
From: Tim Shen <timshen@google.com>
Date: Wed, 9 Dec 2020 08:42:11 -0800
Subject: [PATCH] [XLA/GPU] Migrate all unnested elementwise emitters.

PiperOrigin-RevId: 346559170
---
 include/mlir-hlo/utils/hlo_utils.h | 5 +++++
 lib/utils/hlo_utils.cc             | 8 ++++++++
 2 files changed, 13 insertions(+)

diff --git a/include/mlir-hlo/utils/hlo_utils.h b/include/mlir-hlo/utils/hlo_utils.h
index 74ea9c9..5513dc6 100644
--- a/include/mlir-hlo/utils/hlo_utils.h
+++ b/include/mlir-hlo/utils/hlo_utils.h
@@ -83,6 +83,11 @@ enum ScalarLimit {
 // Requires `ty` to be either FloatType or IntegerType.
 DenseElementsAttr GetScalarLimitOfType(Type ty, ScalarLimit limit);
 
+// Given `op_name` from LMHLO, returns the corresponding op name in MHLO.
+// Returns empty string if no such op exists.
+std::string LmhloToMhloOpName(llvm::StringRef op_name,
+                              mlir::MLIRContext* context);
+
 }  // namespace hlo
 }  // namespace mlir
 
diff --git a/lib/utils/hlo_utils.cc b/lib/utils/hlo_utils.cc
index 0bbd91e..8ff1ce3 100644
--- a/lib/utils/hlo_utils.cc
+++ b/lib/utils/hlo_utils.cc
@@ -132,5 +132,13 @@ DenseElementsAttr GetScalarLimitOfType(Type ty, ScalarLimit limit) {
   llvm_unreachable("unsupported type");
 }
 
+std::string LmhloToMhloOpName(llvm::StringRef op_name,
+                              mlir::MLIRContext *context) {
+  assert(op_name.startswith("lmhlo.") && "Expected an LMHLO op");
+  std::string mhlo_op_name(op_name.drop_front(1));
+  if (context->isOperationRegistered(mhlo_op_name)) return mhlo_op_name;
+  return "";
+}
+
 }  // namespace hlo
 }  // namespace mlir