[KUNLUNXIN] speed up 'exp, celu, relu' (#1284)

Co-authored-by: xuerui <xuerui06@baidu.com>

Author: xuexingtu

src/flag_gems/runtime/backend/_kunlunxin/ops/celu.py

@@ -19,16 +19,17 @@
 )
 
 
-@pointwise_dynamic(
-    is_tensor=[True, False], promotion_methods=[(0, "DEFAULT")], config=config_
-)
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, "DEFAULT")])
 @triton.jit
+# celu another way: max(0, x) + alpha * (exp(min(0, x) / alpha) - 1), getting smaller instrs.
 def celu_forward_kernel(x, alpha):
-    return tl.where(
-        x > 0,
-        x,
-        alpha * (tl.exp(x / alpha) - 1),
-    )
+    inv_alpha = 1.0 / alpha
+
+    pos_part = tl.maximum(0.0, x)
+
+    neg_part_input = x - pos_part
+
+    return pos_part + alpha * (tl.exp(neg_part_input * inv_alpha) - 1.0)
 
 
 def celu(A, alpha=1.0):

src/flag_gems/runtime/backend/_kunlunxin/ops/exp.py

@@ -16,6 +16,7 @@
     prefer_1d_tile=True,
     buffer_size_limit=4096,
     isCloseVectorization=True,
+    kunlunAutoGrid=True,
     unroll_num=8,
 )
 

src/flag_gems/runtime/backend/_kunlunxin/ops/relu.py

@@ -10,8 +10,10 @@
 
 @pointwise_dynamic(promotion_methods=[(0, "DEFAULT")])
 @triton.jit
+# relu another way: maximum(x, 0)
+# tl.maximum(x, 0) to one max_instr，but tl.where two instr compare and select
 def relu_forward(x):
-    return tl.where(x > 0, x, 0)
+    return tl.maximum(x, 0)
 
 
 @pointwise_dynamic(promotion_methods=[(0, "DEFAULT")])