stream-k v0.4 (#689)

* add perfci tuning shapes and fall back configs

* make test_correctness work with multiple kernels

* change persistent_gemm kernel name back

* add persistent_gemm unit tests

* adapt tune_streamk to be able to switch kernel tunning from command line

* fix output file name issue and merge final TFLOPS and time into yaml file

* remove comments

Author: xiaohuguo2023

python/perf-kernels/streamk/03-matrix-multiplication-stream-k.py

@@ -4,7 +4,7 @@
 
 from streamk_kernel import streamk_gemm
 #from streamk_kernel_atomic import streamk_gemm
-#from persistent_gemm import streamk_gemm
+#from persistent_gemm import persistent_gemm
 
 torch.manual_seed(123)
 random.seed(123)

python/perf-kernels/streamk/README.md

@@ -1,3 +1,23 @@
+# streak gemm script v0.4
+
+### features added:
+- enable tuning with different kernels from command line, default for streamk_kernel
+- change persistent_gemm kernel back
+-
+
+### command line example
+
+tuning shapes using persistent_gemm
+```
+python tune_streamk.py --gemm_size_file input.yaml --kernel persistent_gemm,persistent_gemm --ngpus 8 --jobs 24
+```
+
+test correctness using streamk_kernel
+```
+python tune_streamk.py --gemm_size_file tuning_results.yaml --compare_wo_tuning --kernel streamk_kernel,streamk_gemm
+
+```
+
 # streamk gemm script v0.3
 
 ### features added:

python/perf-kernels/streamk/fallback_tuned_configs.yaml

@@ -3,7 +3,7 @@
 
 
 @triton.jit()
-def streamk_gemm(
+def persistent_gemm(
     A,
     B,
     C,

python/perf-kernels/streamk/persistent_gemm.py

@@ -9,8 +9,7 @@
 import torch
 import triton
 import triton.language as tl
-
-from streamk_kernel import streamk_gemm
+import importlib
 
 from datetime import datetime
 import multiprocessing
@@ -206,9 +205,9 @@ def profile_batch_kernels(M, N, K, gpuid, gpus, jobs, verbose):
     os.environ['ROCR_VISIBLE_DEVICES'] = str(gpuid)
     jobId = gpuIdx
     while jobId < jobs:
-        kernel_name = get_filename_profile_driver(M, N, K, jobId)
+        kernelname = get_filename_profile_driver(M, N, K, jobId)
         if verbose:
-            print(f"profiling {kernel_name} on GPU {gpuid}")
+            print(f"profiling {kernelname} on GPU {gpuid}")
         run_bash_command_wrapper(
             f"rocprof --stats -o results_{jobId}.csv python {get_filename_profile_driver(M, N, K, jobId)}",
             #            f"rocprofv2 --plugin file --plugin-version 1 --kernel-trace -o {jobId} python {get_filename_profile_driver(M, N, K, jobId)}",
@@ -217,15 +216,15 @@ def profile_batch_kernels(M, N, K, gpuid, gpus, jobs, verbose):
 
 
 def tune_gemm_config(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, dtype_p, dtype_lock, init_type, configs,
-                     run_bench, jobs, iters, skipWarmup, verbose=0, num_threads=32, gpus=[0], rotating_buffer_size=256,
-                     bias_size=0, icache_flush=False):
+                     run_bench, jobs, iters, skipWarmup, module_name, kernel_name, verbose=0, num_threads=32, gpus=[0],
+                     rotating_buffer_size=256, bias_size=0, icache_flush=False):
 
     # precompile the kernels in parallel
     start_time = datetime.now()
     if not skipWarmup:
         # Generate kernel out of all configs
         fname = generate_compile_driver(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, dtype_p, dtype_lock,
-                                        init_type, configs, rotating_buffer_size, bias_size)
+                                        init_type, configs, rotating_buffer_size, bias_size, kernel_name)
 
         run_bash_command(f"python {fname} -n {num_threads}", capture=(verbose < 2))
     compile_end = datetime.now()
@@ -235,7 +234,8 @@ def tune_gemm_config(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, dtype_p,
 
     # Generate kernels out of all configs
     generate_profile_tasks(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, dtype_p, dtype_lock, init_type, configs,
-                           jobs, iters, run_bench, rotating_buffer_size, bias_size, icache_flush)
+                           jobs, iters, run_bench, rotating_buffer_size, bias_size, icache_flush, module_name,
+                           kernel_name)
 
     # profile generated kernels
     running = [
@@ -377,8 +377,8 @@ def gen_rotating_tensors(M, N, K, dtype_a, need_Trans_a, dtype_b, need_Trans_b,
     return in_outs
 
 
-def matmul(a, b, c, bias, P, locks, num_sms, block_m, block_n, block_k, group_m, num_warps, num_stages, waves_per_eu,
-           mfmaInstrSize, kpack, use_bias):
+def matmul(kernel_func, a, b, c, bias, P, locks, num_sms, block_m, block_n, block_k, group_m, num_warps, num_stages,
+           waves_per_eu, mfmaInstrSize, kpack, use_bias):
     # Check constraints.
     assert a.shape[1] == b.shape[0], "Incompatible dimensions"
     #assert a.is_contiguous(), "Matrix A must be contiguous"
@@ -396,7 +396,7 @@ def matmul(a, b, c, bias, P, locks, num_sms, block_m, block_n, block_k, group_m,
     streamk_tiles = m_tiles * n_tiles % num_sms
     # change num_xcds = 1 if using MI250
     num_xcds = 8
-    streamk_gemm[
+    kernel_func[
         grid,
     ](a, b, c, bias, P, locks, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1),
       stride_bias=stride_bias, BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_K=block_k, GROUP_SIZE_M=group_m,
@@ -405,7 +405,8 @@ def matmul(a, b, c, bias, P, locks, num_sms, block_m, block_n, block_k, group_m,
     return c
 
 
-def test_correctness(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, init_type, config, bias_vector, verbose):
+def test_correctness(kernel_func, M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, init_type, config, bias_vector,
+                     verbose):
     block_m, block_n, block_k, group_m, num_sms, num_warps, num_stages, waves_per_eu, mfmaInstrSize, kpack = read_config(
         config)
     use_bias = bias_vector
@@ -423,8 +424,8 @@ def test_correctness(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, init_type
     c = torch.zeros((M, N), device=a.device, dtype=tl_to_torch_types[name_to_tl_types[dtype_c]])
     locks = torch.zeros((num_sms, ), device="cuda", dtype=torch.int32)
     P = torch.zeros((num_sms, block_m * block_n), device="cuda", dtype=torch.float32)
-    triton_output = matmul(a, b, c, bias, P, locks, num_sms, block_m, block_n, block_k, group_m, num_warps, num_stages,
-                           waves_per_eu, mfmaInstrSize, kpack, use_bias)
+    triton_output = matmul(kernel_func, a, b, c, bias, P, locks, num_sms, block_m, block_n, block_k, group_m, num_warps,
+                           num_stages, waves_per_eu, mfmaInstrSize, kpack, use_bias)
     torch_output = torch.matmul(a_fp16, b_fp16)
     if use_bias:
         torch_output += bias_fp16[:, None]
@@ -435,7 +436,7 @@ def test_correctness(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, init_type
     size_str = ''
     if verbose:
         size_str = f'SIZE M: {M}, N: {N}, K: {K}, trans: {row_a_str}{row_b_str}'
-    print(f'{size_str} correctness check')
+    print(f'{kernel_func} {size_str} correctness check')
     torch.testing.assert_close(triton_output.to(torch.float16), torch_output, atol=atol, rtol=rtol)
     print(f'{size_str} Correct✅')
 
@@ -446,6 +447,8 @@ def parse_args():
         allow_abbrev=False,
     )
 
+    parser.add_argument('--kernel', default='streamk_kernel, streamk_gemm',
+                        help='can specify different kernel file name')
     parser.add_argument("-m", type=int, default=0)
     parser.add_argument("-n", type=int, default=0)
     parser.add_argument("-k", type=int, default=0)
@@ -486,11 +489,6 @@ def parse_args():
     parser.add_argument("--hack_triton_compiler", action='store_true', default=False,
                         help="Modify the triton source to avoid backend query")
     args = parser.parse_args()
-    if not args.o:
-        if args.benchmark:
-            args.o = "benchmarking_results.csv"
-        else:
-            args.o = get_default_tuning_result_filename()
 
     return args
 
@@ -542,6 +540,19 @@ def get_rocm_version():
 
 def main():
     args = parse_args()
+    # parse kernel file and kernel function name
+    module_name, kernel_name = args.kernel.split(',')
+    module_name = module_name.strip()
+    kernel_name = kernel_name.strip()
+    module = importlib.import_module(module_name)
+    kernel_func = getattr(module, kernel_name)
+
+    if not args.o:
+        if args.benchmark:
+            args.o = f"benchmarking_results_{kernel_name}.csv"
+        else:
+            args.o = get_default_tuning_result_filename(kernel_name)
+
     matrix_size_file = args.gemm_size_file
     output_file = args.o
     keepTmp = args.keep
@@ -613,7 +624,8 @@ def main():
         for (M, N, K, col_a, col_b, myConfig) in mnks:
             if myConfig is None:
                 raise Exception("kernel config is None, need to provide a tuning config")
-            test_correctness(M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, init_type, myConfig, bias_vector, True)
+            test_correctness(kernel_func, M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, init_type, myConfig,
+                             bias_vector, True)
         return
 
     configs_full = get_full_tuning_space()
@@ -658,7 +670,7 @@ def main():
         configs += delta_configs
 
         ## Append new configs into the tuning space
-        generate_matmul_kernels(delta_configs)
+        generate_matmul_kernels(delta_configs, module_name, kernel_name)
 
         row_a_str = 'N' if col_a else 'T'
         row_b_str = 'N' if col_b else 'T'
@@ -679,8 +691,9 @@ def main():
         bias_size = M if bias_vector else 0
         minTime, bestConfig, compile_time, profile_time, post_time = tune_gemm_config(
             M, N, K, col_a, col_b, dtype_a, dtype_b, dtype_c, dtype_p, dtype_lock, init_type, pruned_configs, run_bench,
-            jobs, iters, skipWarmup, num_threads=args.num_threads, gpus=gpus, verbose=verbose_level,
-            rotating_buffer_size=rotating_buffer_size, bias_size=bias_size, icache_flush=icache_flush)
+            jobs, iters, skipWarmup, module_name, kernel_name, num_threads=args.num_threads, gpus=gpus,
+            verbose=verbose_level, rotating_buffer_size=rotating_buffer_size, bias_size=bias_size,
+            icache_flush=icache_flush)
 
         # post processing the numbers
         perf_tflops = lambda us: 2 * M * N * K * 1e-12 / (us * 1e-6)
@@ -701,9 +714,9 @@ def main():
 
         sizeDict = {'M': M, 'N': N, 'K': K, 'rowMajorA': row_a_str, 'rowMajorB': row_b_str}
         sizeDict.update(bestConfig)
+        sizeDict.update({'TFLOPS': formatted_tflops, 'time(us)': minTime})
         if not run_bench:
             f_results.write("- " + str(sizeDict) + " ")
-            f_results.write(f'# TFLOPS: {formatted_tflops} time(us): {minTime}\n')
 
         # remove generated files if asked to
         if not keepTmp:

python/perf-kernels/streamk/tune_streamk.py

@@ -0,0 +1,80 @@
+# If new gemm sizes/configs are added please provide a fallback tuned config in fallback_configs.yaml.
+# The fallback configs will be used in case the weekly tuning has not been run
+#
+- {'M': 4864, 'N': 4096, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4160, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4224, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4288, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+
+- {'M': 4864, 'N': 4096, 'K': 4097, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4098, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4100, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4104, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 4112, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+
+- {'M': 4864, 'N': 8192, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 8192, 'K': 4160, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 8192, 'K': 8192, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 8192, 'K': 8256, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+
+- {'M': 1024, 'N': 1024, 'K': 1024, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 1024, 'K': 1024, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 8192, 'N': 8192, 'K': 8192, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 8192, 'N': 8192, 'K': 8192, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+
+- {'M': 4864, 'N': 4096, 'K': 8256, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4864, 'N': 4096, 'K': 8256, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+
+- {'M': 1024, 'N': 8192, 'K': 28672, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 8192, 'K': 28672, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1024, 'N': 28672, 'K': 8192, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 28672, 'K': 8192, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1024, 'N': 14336, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 14336, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1, 'N': 8192, 'K': 28672, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1, 'N': 8192, 'K': 28672, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1, 'N': 14336, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1, 'N': 14336, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1024, 'N': 16384, 'K': 53248, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 16384, 'K': 53248, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1024, 'N': 53248, 'K': 16384, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 53248, 'K': 16384, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 32, 'N': 16384, 'K': 53248, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 32, 'N': 16384, 'K': 53248, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 32, 'N': 53248, 'K': 16384, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 32, 'N': 53248, 'K': 16384, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+
+- {'M': 2, 'N': 3584, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 2, 'N': 3584, 'K': 4096, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 2, 'N': 4096, 'K': 1792, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 2, 'N': 4096, 'K': 1792, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 4096, 'N': 13312, 'K': 8896, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 4096, 'N': 13312, 'K': 8896, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 2048, 'N': 17792, 'K': 13312, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 2048, 'N': 17792, 'K': 13312, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 1024, 'N': 13312, 'K': 1664, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 1024, 'N': 13312, 'K': 1664, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+
+- {'M': 8192, 'N': 1536, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 8192, 'N': 1536, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 8192, 'N': 5120, 'K': 1024, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 8192, 'N': 5120, 'K': 1024, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 8192, 'N': 1024, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 8192, 'N': 1024, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 32768, 'N': 5120, 'K': 512, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 32768, 'N': 5120, 'K': 512, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 512, 'N': 1536, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 512, 'N': 1536, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 512, 'N': 5120, 'K': 1024, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 512, 'N': 5120, 'K': 1024, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 512, 'N': 1024, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 512, 'N': 1024, 'K': 5120, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 2048, 'N': 5120, 'K': 512, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 2048, 'N': 5120, 'K': 512, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+
+- {'M': 20196, 'N': 512, 'K': 1536, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 20196, 'N': 512, 'K': 1536, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 171792, 'N': 512, 'K': 1536, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 171792, 'N': 512, 'K': 1536, 'rowMajorA': 'T', 'rowMajorB': 'T'}
+- {'M': 173318, 'N': 512, 'K': 1536, 'rowMajorA': 'T', 'rowMajorB': 'N'}
+- {'M': 173318, 'N': 512, 'K': 1536, 'rowMajorA': 'T', 'rowMajorB': 'T'}