From 2ab16024cf112214ae1655b75262d10d8e408b1a Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Thu, 17 Jun 2021 06:33:13 -0700
Subject: [PATCH] Add tests for lowering HLO_ExpOp for complex types to Linalg.

PiperOrigin-RevId: 379944871
---
 tests/hlo-legalize-to-linalg.mlir  | 11 +++++++++++
 tests/lhlo-legalize-to-linalg.mlir | 14 ++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/tests/hlo-legalize-to-linalg.mlir b/tests/hlo-legalize-to-linalg.mlir
index 0a08562..775fd3a 100644
--- a/tests/hlo-legalize-to-linalg.mlir
+++ b/tests/hlo-legalize-to-linalg.mlir
@@ -169,6 +169,17 @@ func @float_exp(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 
 // -----
 
+// CHECK-LABEL: func @complex_exp
+func @complex_exp(%arg0: tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>> {
+  // CHECK: linalg.generic
+  // CHECK: complex.exp
+  %0 = "mhlo.exponential"(%arg0) : (tensor<2x2xcomplex<f32>>)
+                                 -> tensor<2x2xcomplex<f32>>
+  return %0 : tensor<2x2xcomplex<f32>>
+}
+
+// -----
+
 // CHECK-LABEL: func @float_expm1
 func @float_expm1(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
   // CHECK: linalg.generic
diff --git a/tests/lhlo-legalize-to-linalg.mlir b/tests/lhlo-legalize-to-linalg.mlir
index f8d7dea..0902ea1 100644
--- a/tests/lhlo-legalize-to-linalg.mlir
+++ b/tests/lhlo-legalize-to-linalg.mlir
@@ -135,6 +135,20 @@ func @exp(%input: memref<2x2xf32>, %result: memref<2x2xf32>) {
 
 // -----
 
+// CHECK-LABEL: func @complex_exp
+func @complex_exp(%input: memref<2x2xcomplex<f32>>,
+                  %result: memref<2x2xcomplex<f32>>) {
+  "lmhlo.exponential"(%input, %result)
+      : (memref<2x2xcomplex<f32>>, memref<2x2xcomplex<f32>>) -> ()
+  return
+}
+// CHECK: linalg.generic
+// CHECK-NEXT: ^bb0(%[[OPERAND_IN:.*]]: complex<f32>, %[[RESULT_OUT:.*]]):
+// CHECK-NEXT:   %[[RESULT:.*]] = complex.exp %[[OPERAND_IN]] : complex<f32>
+// CHECK-NEXT:   linalg.yield %[[RESULT]] : complex<f32>
+
+// -----
+
 // CHECK-LABEL: func @log
 func @log(%input: memref<2x2xf32>, %result: memref<2x2xf32>) {
   "lmhlo.log"(%input, %result) : (memref<2x2xf32>, memref<2x2xf32>) -> ()