[Hardware][Power] Add FAST_EXP for Power (#49571)

Signed-off-by: Akash kaothalkar <[email protected]>
Co-authored-by: Akash kaothalkar <[email protected]>

benchmarks/kernels/cpu/benchmark_cpu_attn.py

@@ -154,7 +154,7 @@ def main(
                 scale=scale,
                 causal=True,
                 alibi_slopes=None,
-                sliding_window=window_size,
+                sliding_window=window_size if sliding_window is not None else -1,
                 block_table=block_tables,
                 softcap=0,
                 scheduler_metadata=metadata,

csrc/cpu/cpu_arch_macros.h

@@ -172,4 +172,15 @@
 
 #endif  // __riscv_v
 
+// Power VSX
+#ifdef __powerpc__
+  // FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which
+  // implements a vectorised 5-term minimax polynomial using VSX intrinsics.
+  #define DEFINE_FAST_EXP                                                     \
+    auto fast_exp = [&](const vec_op::FP32Vec16& vec)                         \
+                        __attribute__((always_inline)) { return vec.exp(); }; \
+    auto fast_exp_f16 = fast_exp;
+
+#endif  // __powerpc__
+
 #endif

csrc/cpu/cpu_types_vsx.hpp

@@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
 
   explicit FP32Vec4(__vector float data) : reg(data) {}
 
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
+  FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
 };
 
 struct FP32Vec8 : public Vec<FP32Vec8> {
@@ -316,7 +316,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
 
   explicit FP32Vec8(f32x4x2_t data) : reg(data) {}
 
-  explicit FP32Vec8(const FP32Vec8& data) {
+  FP32Vec8(const FP32Vec8& data) {
     reg.val[0] = data.reg.val[0];
     reg.val[1] = data.reg.val[1];
   }
@@ -593,7 +593,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
   explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
 
-  explicit FP32Vec16(const FP32Vec16& data) {
+  FP32Vec16(const FP32Vec16& data) {
     reg.val[0] = data.reg.val[0];
     reg.val[1] = data.reg.val[1];
     reg.val[2] = data.reg.val[2];
@@ -747,6 +747,15 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
                                 vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
   }
 
+  FP32Vec16 exp() const {
+    FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
+    FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
+    auto lo_e = lo.exp();
+    auto hi_e = hi.exp();
+    return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
+                               hi_e.reg.val[0], hi_e.reg.val[1]});
+  }
+
   float reduce_max() {
     __vector float max01 = vec_max(reg.val[0], reg.val[1]);
     __vector float max23 = vec_max(reg.val[2], reg.val[3]);