Set block sizes dynamically if not provided

Author: voltjia

src/ntops/kernels/addmm.py

@@ -14,10 +14,19 @@ def arrangement(
     beta,
     alpha,
     output,
-    block_size_m=mm.BLOCK_SIZE_M,
-    block_size_n=mm.BLOCK_SIZE_N,
-    block_size_k=mm.BLOCK_SIZE_K,
+    block_size_m=None,
+    block_size_n=None,
+    block_size_k=None,
 ):
+    if block_size_m is None:
+        block_size_m = mm.BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = mm.BLOCK_SIZE_N
+
+    if block_size_k is None:
+        block_size_k = mm.BLOCK_SIZE_K
+
     _, _, input_arranged = mm.arrangement(
         x,
         y,

src/ntops/kernels/bmm.py

@@ -7,13 +7,17 @@
 
 
 def arrangement(
-    input,
-    other,
-    output,
-    block_size_m=BLOCK_SIZE_M,
-    block_size_n=BLOCK_SIZE_N,
-    block_size_k=BLOCK_SIZE_K,
+    input, other, output, block_size_m=None, block_size_n=None, block_size_k=None
 ):
+    if block_size_m is None:
+        block_size_m = BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = BLOCK_SIZE_N
+
+    if block_size_k is None:
+        block_size_k = BLOCK_SIZE_K
+
     output_arranged = output.tile((1, block_size_m, block_size_n))
     output_arranged.dtype = output_arranged.dtype.squeeze(0)
 

src/ntops/kernels/element_wise.py

@@ -1,7 +1,10 @@
 import ninetoothed
 
 
-def arrangement(*tensors, block_size=ninetoothed.block_size()):
+def arrangement(*tensors, block_size=None):
+    if block_size is None:
+        block_size = ninetoothed.block_size()
+
     ndim = max(tensor.ndim for tensor in tensors)
 
     assert all(tensor.ndim == ndim or tensor.ndim == 0 for tensor in tensors)

src/ntops/kernels/mm.py

@@ -10,13 +10,17 @@
 
 
 def arrangement(
-    input,
-    other,
-    output,
-    block_size_m=BLOCK_SIZE_M,
-    block_size_n=BLOCK_SIZE_N,
-    block_size_k=BLOCK_SIZE_K,
+    input, other, output, block_size_m=None, block_size_n=None, block_size_k=None
 ):
+    if block_size_m is None:
+        block_size_m = BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = BLOCK_SIZE_N
+
+    if block_size_k is None:
+        block_size_k = BLOCK_SIZE_K
+
     output_arranged = output.tile((block_size_m, block_size_n))
 
     input_arranged = input.tile((block_size_m, block_size_k))

src/ntops/kernels/scaled_dot_product_attention.py

@@ -22,8 +22,8 @@ def arrangement(
     output,
     with_attn_mask,
     with_kv_cache,
-    block_size_m=BLOCK_SIZE_M,
-    block_size_n=BLOCK_SIZE_N,
+    block_size_m=None,
+    block_size_n=None,
 ):
     def arrange_query_or_output(input):
         arranged = input.tile((1, 1, block_size_m, -1)).tile(
@@ -58,6 +58,12 @@ def arrange_attn_mask(input):
 
         return arranged
 
+    if block_size_m is None:
+        block_size_m = BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = BLOCK_SIZE_N
+
     query_arranged = arrange_query_or_output(query)
     key_arranged = arrange_key_or_value(key)
     value_arranged = arrange_key_or_value(value)

src/ntops/kernels/softmax.py

@@ -7,7 +7,7 @@
 BLOCK_SIZE = ninetoothed.block_size()
 
 
-def arrangement(input, output, dim, block_size=BLOCK_SIZE):
+def arrangement(input, output, dim, block_size=None):
     assert input.ndim == output.ndim
 
     def create_axis_tile_shape(dim, dim_block):
@@ -28,6 +28,9 @@ def arrange(input):
         )
         return input_arranged
 
+    if block_size is None:
+        block_size = BLOCK_SIZE
+
     inner_block_shape = create_axis_tile_shape(dim, block_size)
     outer_block_shape = create_axis_tile_shape(dim, -1)