From 399dae666d5ec978c527d512eb81f2f5e7121fbb Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Tue, 15 Jun 2021 00:44:45 -0700
Subject: [PATCH] Add MLIR generated ExpOp GPU kernel for complex types.

We lower lmhlo::ExpOp to mlir::complex::ExpOp for complex types.

PiperOrigin-RevId: 379432147
---
 .../mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h  | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h b/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h
index 0a3c240..288e04c 100644
--- a/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h
+++ b/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h
@@ -323,7 +323,8 @@ inline Value MapLhloOpToStdScalarOp<lmhlo::ExpOp>(Location loc,
                                                   ArrayRef<Type> arg_types,
                                                   ArrayRef<Value> args,
                                                   OpBuilder* b) {
-  return MapLhloOpToScalarOpImpl<isFloatType, ::mlir::math::ExpOp>{}(
+  return MapLhloOpToScalarOpImpl<isFloatType, ::mlir::math::ExpOp,
+                                 isComplexType, ::mlir::complex::ExpOp>{}(
       loc, result_types, arg_types, args, b);
 }