[KUNLUNXIN] speed up compare ops

Author: dongjibin1996

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

@@ -16,6 +16,8 @@
     True,
     prefer_1d_tile=True,
     isCloseMemoryAsync=False,
+    KunlunAtuoGrid=True,
+    unroll_num=8,
 )
 
 

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

@@ -17,6 +17,8 @@
     True,
     prefer_1d_tile=True,
     isCloseMemoryAsync=False,
+    KunlunAtuoGrid=True,
+    unroll_num=8,
 )
 
 

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

@@ -17,6 +17,8 @@
     True,
     prefer_1d_tile=True,
     isCloseMemoryAsync=False,
+    KunlunAtuoGrid=True,
+    unroll_num=8,
 )
 
 

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

@@ -16,6 +16,8 @@
     True,
     prefer_1d_tile=True,
     isCloseMemoryAsync=False,
+    unroll_num=8,
+    KunlunAtuoGrid=True,
 )
 
 

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

@@ -46,6 +46,7 @@ class CodeGenConfig:
     isCloseDtypeConvert: bool = False
     isCloseMemoryAsync: bool = True
     isCloseInterleave: bool = False
+    KunlunAtuoGrid: bool = False
     buffer_size_limit: int = 0
     unroll_num: int = 0
 

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

@@ -867,9 +867,19 @@ def gen_task_partition_1d(self, code: IndentedBuffer):
             code.writeline("num_tiles = triton.cdiv(num_tasks, tile_size)")
             # max_grid_size0 = self.config.max_grid_size[0]
             # code.writeline(f"num_ctas = min({max_grid_size0}, num_tiles)")
-
-            code.writeline("num_ctas = 12 # XPU BLOCK_NUM")
-            code.writeline("num_tiles = 12 # XPU BLOCK_NUM")
+            determine_num_ctas_and_tiles = [
+                "if sum(out0.shape) <= 2048*64:",
+                "   num_ctas = 1 # XPU BLOCK_NUM",
+                "   num_tiles = 1 # XPU BLOCK_NUM",
+                "else:",
+                "   num_ctas = 12 # XPU BLOCK_NUM",
+                "   num_tiles = 12 # XPU BLOCK_NUM",
+            ]
+            if self.config.KunlunAtuoGrid is True:
+                code.writelines(determine_num_ctas_and_tiles)
+            else:
+                code.writeline("num_ctas = 12 # XPU BLOCK_NUM")
+                code.writeline("num_tiles = 12 # XPU BLOCK_NUM")
             code.writeline(
                 "tile_size = triton.next_power_of_2(triton.cdiv(num_tasks, num_tiles)) # XPU BLOCK_NUM"
             )