[KUNLUNXIN] speed up 'any, div, neg, silu'

Author: xuerui

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

@@ -144,29 +144,31 @@ def any_kernel_2(mid, out, MID_SIZE, BLOCK_MID: tl.constexpr):
 def any(inp):
     logger.debug("GEMS ANY")
     n_elements = inp.numel()
-    block_size = min(
+    block_size = max(
         triton.cdiv(get_block(n_elements), cluster_num),
         triton.cdiv(buf_len_per_core * core_num, 4),
     )
+
     mid_size = triton.cdiv(n_elements, block_size)
     block_mid = triton.next_power_of_2(mid_size)
-
+    
     if n_elements >= vector_size * thread_num:
-        # according to api, op == any, use max to calculate
-        inpf = inp.to(torch.float)
-        midf = torch.empty((mid_size,), dtype=torch.float, device=inp.device)
-        outf = torch.empty([], dtype=torch.float, device=inp.device)
+        inp_uint8 = inp.view(torch.uint8)
+        
+        mid = torch.empty((mid_size,), dtype=torch.uint8, device=inp.device)
+        out = torch.empty([], dtype=torch.uint8, device=inp.device)
 
         with torch_device_fn.device(inp.device):
             max_kernel_1[(mid_size, 1)](
-                inpf, midf, n_elements, block_size, buffer_size_limit=2048
+                inp_uint8, mid, n_elements, block_size, buffer_size_limit=2048
             )
             if mid_size == 1:
-                return midf.to(torch.bool).reshape([])
+                return mid.view(torch.bool).reshape([])
+
             max_kernel_2[(1, 1)](
-                midf, outf, mid_size, block_mid, buffer_size_limit=2048
+                mid, out, mid_size, block_mid, buffer_size_limit=2048
             )
-        out = outf.to(torch.bool)
+        out = out.view(torch.bool)
     else:
         mid = torch.empty((mid_size,), dtype=torch.bool, device=inp.device)
         out = torch.empty([], dtype=torch.bool, device=inp.device)

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

@@ -7,6 +7,7 @@
 from flag_gems.utils import tl_extra_shim
 
 from ..utils.pointwise_dynamic import pointwise_dynamic
+from _kunlunxin.utils.codegen_config_utils import CodeGenConfig
 
 logger = logging.getLogger("flag_gems").getChild(__name__.lstrip("."))
 div_rn = tl_extra_shim.div_rn
@@ -15,8 +16,18 @@
 trunc = tl_extra_shim.trunc
 xpu_trunc_div = tl_extra_shim.xpu_trunc_div  # use it if we need to cmp result with xpu
 
-
-@pointwise_dynamic(promotion_methods=[(0, 1, "INT_TO_FLOAT")])
+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, 1, "INT_TO_FLOAT")], config=config_)
 @triton.jit
 def true_div_func(x, y):
     return x / y

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

@@ -3,11 +3,22 @@
 import triton
 
 from ..utils.pointwise_dynamic import pointwise_dynamic
+from _kunlunxin.utils.codegen_config_utils import CodeGenConfig
 
 logger = logging.getLogger("flag_gems").getChild(__name__.lstrip("."))
 
-
-@pointwise_dynamic(promotion_methods=[(0, "DEFAULT")])
+config_ = CodeGenConfig(
+    512,
+    (65536, 65536, 65536),
+    32,
+    True,
+    prefer_1d_tile=True,
+    buffer_size_limit=4096,
+    isCloseVectorization=False,
+    unroll_num=8,
+)
+
+@pointwise_dynamic(promotion_methods=[(0, "DEFAULT")], config=config_)
 @triton.jit
 def neg_func(x):
     return -x

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

@@ -6,12 +6,23 @@
 from flag_gems.utils import tl_extra_shim
 
 from ..utils.pointwise_dynamic import pointwise_dynamic
+from _kunlunxin.utils.codegen_config_utils import CodeGenConfig
 
 logger = logging.getLogger("flag_gems").getChild(__name__.lstrip("."))
 div_rn = tl_extra_shim.div_rn
 
-
-@pointwise_dynamic(promotion_methods=[(0, "DEFAULT")])
+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, "DEFAULT")], config=config_)
 @triton.jit
 def silu_forward(x):
     x_fp32 = x.to(tl.float32)