From 1eb6355182a5124639ce9d3ff165732a94ed9a70 Mon Sep 17 00:00:00 2001
From: Bing Xu <antinucleon@gmail.com>
Date: Sat, 2 Jul 2022 05:00:45 -0700
Subject: [PATCH] [activation] tanh (#550)

Co-authored-by: Bing Xu <bingxu@fb.com>
---
 include/cutlass/epilogue/thread/activation.h | 37 ++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/include/cutlass/epilogue/thread/activation.h b/include/cutlass/epilogue/thread/activation.h
index ca1f72c2..9763f5fc 100644
--- a/include/cutlass/epilogue/thread/activation.h
+++ b/include/cutlass/epilogue/thread/activation.h
@@ -98,6 +98,43 @@ struct ReLu<Array<T, N>> {
   }
 };
 
+// Tanh operator
+template <typename T>
+struct Tanh {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &scalar) const {
+    return fast_tanh(scalar);
+  }
+};
+
+template <typename T, int N>
+struct Tanh<Array<T, N> > {
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &rhs) const {
+    Array<T, N> y;
+    Tanh<T> tanh_op;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = tanh_op(rhs[i]);
+    }
+
+    return y;
+  }
+};
+
+template <int N>
+struct Tanh<Array<half_t, N>> {
+  using T = half_t;
+
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& z) const {
+    fast_tanh_op<Array<T, N>> tanh;
+    return tanh(z);
+
+  }
+};
+
 // Leaky Relu operator
 template <typename T>
 struct LeakyReLU {