[KUNLUNXIN] Speed Up exp (#1219)

Author: Ason93

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

@@ -2,13 +2,25 @@
 
 import triton
 import triton.language as tl
+from _kunlunxin.utils.codegen_config_utils import CodeGenConfig
 
 from ..utils.pointwise_dynamic import pointwise_dynamic
 
 logger = logging.getLogger("flag_gems").getChild(__name__.lstrip("."))
 
+config_ = CodeGenConfig(
+    512,
+    (65536, 65536, 65536),
+    32,
+    True,
+    prefer_1d_tile=True,
+    buffer_size_limit=4096,
+    isCloseVectorization=True,
+    unroll_num=8,
+)
 
-@pointwise_dynamic(promotion_methods=[(0, "INT_TO_FLOAT")])
+
+@pointwise_dynamic(promotion_methods=[(0, "INT_TO_FLOAT")], config=config_)
 @triton.jit
 def exp_func(x):
     return tl.exp(x.to(tl.float32))

src/flag_gems/runtime/backend/_kunlunxin/utils/codegen_config_utils.py

@@ -46,6 +46,7 @@ class CodeGenConfig:
     isCloseDtypeConvert: bool = False
     isCloseMemoryAsync: bool = True
     buffer_size_limit: int = 0
+    unroll_num: int = 0
 
     def __post_init__(self):
         if self.prefer_1d_tile:

src/flag_gems/runtime/backend/_kunlunxin/utils/pointwise_dynamic.py

@@ -978,6 +978,8 @@ def gen_kernel_launch(
                     code.writeline("isCloseMemoryAsync=False,")
                 if os.getenv("XPU_cmp_nan") == "1":
                     code.writeline("isOpenCmpNan=True,")
+                if self.config.unroll_num:
+                    code.writeline(f"unroll_num={self.config.unroll_num},")
             code.writeline(")")
 
     def gen_kernel_launch_1d(
@@ -1043,6 +1045,8 @@ def gen_kernel_launch_1d(
                     code.writeline("isCloseMemoryAsync=False,")
                 if os.getenv("XPU_cmp_nan") == "1":
                     code.writeline("isOpenCmpNan=True,")
+                if self.config.unroll_num:
+                    code.writeline(f"unroll_num={self.config.unroll_num},")
             code.writeline(")")
 
     def gen_return(self, code: IndentedBuffer):