Fix usage of kernel-repeats in tuningRunner (#2126)

Author: pabloantoniom

mlir/test/rocmlir-driver/large-kernel-no-scavenge.mlir

@@ -1,4 +1,4 @@
-// RUN: rocmlir-gen --operation conv_bwd_data -t f16 --arch %arch --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 1 --in_channels 192 --in_h 64 --in_w 64 --out_channels 384 --fil_h 4 --fil_w 4 --dilation_h 1 --dilation_w 1 --conv_stride_h 2 --conv_stride_w 2 --padding_h 1 --padding_w 1 --groupsize 1 --kernel-repeats 10 --perf_config 'v3:128,64,8,128,64,16,1,1,2,1,1' | rocmlir-driver --kernel-pipeline=full | FileCheck %s
+// RUN: rocmlir-gen --operation conv_bwd_data -t f16 --arch %arch --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 1 --in_channels 192 --in_h 64 --in_w 64 --out_channels 384 --fil_h 4 --fil_w 4 --dilation_h 1 --dilation_w 1 --conv_stride_h 2 --conv_stride_w 2 --padding_h 1 --padding_w 1 --groupsize 1 --perf_config 'v3:128,64,8,128,64,16,1,1,2,1,1' | rocmlir-driver --kernel-pipeline=full | FileCheck %s
 
 // CHECK: gpu.binary {{.*}} rock.blocks_per_cu = {{.*}} : i32
 

mlir/test/rocmlir-driver/lds_barrier_gfx908.mlir

@@ -1,5 +1,5 @@
 // This is E2E compilation test to ensure gfx908 specific inlineAsm hack
 // does not blow up in the backend compiler.
 
-// RUN: rocmlir-gen -operation gemm -t f32 -out_datatype f32 --arch gfx908:sramecc+:xnack- --num_cu 120 -g 1 -m 24576 -k 768 -n 3072 -transA=False -transB=False --kernel-repeats 5 --perf_config= | rocmlir-driver -c | FileCheck %s
+// RUN: rocmlir-gen -operation gemm -t f32 -out_datatype f32 --arch gfx908:sramecc+:xnack- --num_cu 120 -g 1 -m 24576 -k 768 -n 3072 -transA=False -transB=False --perf_config= | rocmlir-driver -c | FileCheck %s
 // CHECK: gpu.binary

mlir/test/rocmlir-gen/options.mlir

@@ -0,0 +1,2 @@
+// RUN: not rocmlir-gen -operation gemm -t f32 -out_datatype f32 --arch %arch -g 1 -m 1024 -k 1024 -n 1024 -transA=False -transB=False --kernel-repeats=100 2>&1 | FileCheck %s --check-prefix=ERR_KERNEL_REPEATS
+// ERR_KERNEL_REPEATS: --kernel-repeats is only supported with host harness (-ph) or CPU validation (-pv).

mlir/test/rocmlir-gen/problem-key.mlir

@@ -19,13 +19,13 @@
 // RUN: rocmlir-gen --arch gfx942 --operation attention -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f16 -return_lse -split_kv 8 -g 1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_10
 // CHECK_10: -t f16 -transQ false -transK false -transV false -transO false -causal false -return_lse true -split_kv 8 -num_heads_q 1 -num_heads_kv 1 -g 1 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32
 
-// RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 256 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 256 --kernel-repeats 1 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_DEPTHWISE_CONV
+// RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 256 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 256 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_DEPTHWISE_CONV
 // CHECK_DEPTHWISE_CONV: convfp16 -F 1 -f GNC01 -I NGC01 -O NGC01 -n 64 -c 256 -H 20 -W 20 -k 256 -y 7 -x 7 -p 3 -q 3 -u 1 -v 1 -l 1 -j 1 -g 256
 
-// RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 256 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 128 --kernel-repeats 1 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_GROUP_CONV
+// RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 256 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 128 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_GROUP_CONV
 // CHECK_GROUP_CONV: convfp16 -F 1 -f GNC01 -I NGC01 -O NGC01 -n 64 -c 256 -H 20 -W 20 -k 256 -y 7 -x 7 -p 3 -q 3 -u 1 -v 1 -l 1 -j 1 -g 128
 
-// RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 512 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 128 --kernel-repeats 1 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_GROUP_CONV2
+// RUN: rocmlir-gen --arch gfx942 --operation conv -t f16 --fil_layout gkc01 --in_layout ngc01 --out_layout ngk01 --batchsize 64 --in_channels 256 --in_h 20 --in_w 20 --out_channels 512 --fil_h 7 --fil_w 7 --dilation_h 1 --dilation_w 1 --conv_stride_h 1 --conv_stride_w 1 --padding_h 3 --padding_w 3 --groupsize 128 --perf_config=v3:32,256,2,32,32,4,1,1,2,1,1 | rocmlir-gen --emit-tuning-key - | FileCheck %s  --check-prefixes=CHECK_GROUP_CONV2
 // CHECK_GROUP_CONV2: convfp16 -F 1 -f GNC01 -I NGC01 -O NGC01 -n 64 -c 256 -H 20 -W 20 -k 512 -y 7 -x 7 -p 3 -q 3 -u 1 -v 1 -l 1 -j 1 -g 128
 
 // Checking numCU

mlir/tools/rocmlir-gen/rocmlir-gen.cpp

@@ -5507,6 +5507,14 @@ int main(int argc, char **argv) {
     module = ModuleOp::create(UnknownLoc::get(&context));
   }
 
+  if (kernelRepeats.getNumOccurrences() > 0 && !genCPUValidation &&
+      !genHostHarness) {
+    llvm::errs()
+        << "--kernel-repeats is only supported with host harness (-ph) or "
+           "CPU validation (-pv).\n";
+    return EXIT_FAILURE;
+  }
+
   if (genCloneHarness.getValue()) {
     populateCloneHarnessLogic(*module);
   } else if (!hasUserKernel) {

mlir/utils/performance/perfRunner.py

@@ -499,8 +499,9 @@ def generate_mlir_driver_commandline(self, rocmlir_gen_flags, kernel_repeats=MLI
             str(self.conv_stride_w), '--padding_h',
             str(self.padding_h), '--padding_w',
             str(self.padding_w), '--groupsize',
-            str(self.group), '--kernel-repeats',
-            str(kernel_repeats), f"--perf_config={self.perfconfig}"
+            str(self.group),
+            *(['--kernel-repeats', str(kernel_repeats)] if kernel_repeats is not None else []),
+            f"--perf_config={self.perfconfig}"
         ])
         result += ' '
         if rocmlir_gen_flags != '':
@@ -696,7 +697,7 @@ def get_gemm_configurations(filename,
 
                 # Skip unsupported datatypes
                 if datatype == 'f4E2M1FN':
-                    ## TODO: use information from AMDArchDB when it becomes available to determine supported chips
+                    # TODO: use information from AMDArchDB when it becomes available to determine supported chips
                     supported_chips = {'gfx950'}
                     if get_chip() not in supported_chips:
                         continue
@@ -926,32 +927,28 @@ def set_perfconfig(self, perf_config):
         self.perfconfig = perf_config
 
     def generate_mlir_driver_commandline(self, rocmlir_gen_flags, kernel_repeats=MLIR_N_REPEATS):
-        cmd_parts = [
+        result = ' '.join([
             '-operation', 'gemm', '-t', self.datatype, '-out_datatype', self.out_dtype, '--arch',
             self.arch, '--num_cu',
             str(self.num_cu), '-g',
             str(self.g), '-m',
             str(self.m), '-k',
             str(self.k), '-n',
-            str(self.n), f"-transA={self.trans_a}", f"-transB={self.trans_b}"
-        ]
+            str(self.n), f"-transA={self.trans_a}", f"-transB={self.trans_b}",
+            *(['--kernel-repeats', str(kernel_repeats)] if kernel_repeats is not None else []),
+            f"--perf_config={self.perfconfig}"
+        ])
 
         if self.scaled_gemm:
-            cmd_parts.append('-scaledGemm')
+            result += ' -scaledGemm'
         if self.scale_a_dtype:
-            cmd_parts.extend(['-scale_a_dtype', self.scale_a_dtype])
+            result += f' -scale_a_dtype {self.scale_a_dtype}'
         if self.scale_b_dtype:
-            cmd_parts.extend(['-scale_b_dtype', self.scale_b_dtype])
+            result += f' -scale_b_dtype {self.scale_b_dtype}'
         if self.trans_scale_a:
-            cmd_parts.append(f"-transScaleA={self.trans_scale_a}")
+            result += f' -transScaleA {str(self.trans_scale_a)}'
         if self.trans_scale_b:
-            cmd_parts.append(f"-transScaleB={self.trans_scale_b}")
-
-        cmd_parts.extend(
-            ['--kernel-repeats',
-             str(kernel_repeats), f"--perf_config={self.perfconfig}"])
-
-        result = ' '.join(cmd_parts)
+            result += f' -transScaleB {str(self.trans_scale_b)}'
 
         result += ' '
         if rocmlir_gen_flags != '':
@@ -1194,7 +1191,8 @@ def generate_mlir_driver_commandline(self, rocmlir_gen_flags, kernel_repeats=MLI
             f'--dilation_w={self.dilation_w}', f'--conv_stride_h={self.conv_stride_h}',
             f'--conv_stride_w={self.conv_stride_w}', f'--padding_h={self.padding_h}',
             f'--padding_w={self.padding_w}', f'--groupsize={self.group}', f'--gemmO={self.o}',
-            f'--kernel-repeats={kernel_repeats}', f"--perf_config={self.perfconfig}"
+            *(['--kernel-repeats', str(kernel_repeats)] if kernel_repeats is not None else []),
+            f"--perf_config={self.perfconfig}"
         ])
         result += ' '
         if rocmlir_gen_flags != '':
@@ -1363,8 +1361,9 @@ def generate_mlir_driver_commandline(self, rocmlir_gen_flags, kernel_repeats=MLI
             str(self.k), '-n',
             str(self.n), '-gemmO',
             str(self.o), f"-transA={self.trans_a}", f"-transB={self.trans_b}",
-            f"-transC={self.trans_c}", f"-transO={self.trans_o}", '--kernel-repeats',
-            str(kernel_repeats), f"--perf_config={self.perfconfig}"
+            f"-transC={self.trans_c}", f"-transO={self.trans_o}",
+            *(['--kernel-repeats', str(kernel_repeats)] if kernel_repeats is not None else []),
+            f"--perf_config={self.perfconfig}"
         ])
         result += ' '
         if rocmlir_gen_flags != '':
@@ -1691,17 +1690,25 @@ def run_config_with_mlir(config: PerfConfiguration,
     # remove the result file generated by rocprof in previous benchmarking
     if os.path.exists(get_profiler_output_path(arch, BENCHMARKING_STATS_FILE_NAME)):
         os.remove(get_profiler_output_path(arch, BENCHMARKING_STATS_FILE_NAME))
-    commandline_options = config.generate_mlir_driver_commandline(rocmlir_gen_flags)
+    use_tuning_driver = (not use_rocprof) and bool(config.perfconfig)
+    use_host_harness = not use_tuning_driver
+
+    rocmlir_gen_flags = rocmlir_gen_flags + ' -ph' if use_host_harness else ''
+    # We want to use kernel_repeats only if we are passing ' -ph' to rocmlir-gen, otherwise we use None.
+    # This is because the kernel-repeats flag is only supported with host harness or CPU validation.
+    kernel_repeats = MLIR_N_REPEATS if use_host_harness else None
+
+    commandline_options = config.generate_mlir_driver_commandline(rocmlir_gen_flags, kernel_repeats)
+    rocmlir_gen_cmd = paths.mlir_paths.rocmlir_gen_path + ' ' + commandline_options
     if debug:
         print("Running MLIR Benchmark: ", repr(config))
 
     nanoseconds = np.nan
 
     # Use HIP timing via tuning-driver if rocprof is disabled and perfconfig is present
-    if not use_rocprof and config.perfconfig:
+    if use_tuning_driver:
         if debug:
             print("Using HIP timing for benchmarking")
-        rocmlir_gen_cmd = paths.mlir_paths.rocmlir_gen_path + ' ' + commandline_options
         tuning_driver_command = [
             paths.mlir_paths.rocmlir_tuning_driver_path, f'--benchmark-config={config.perfconfig}',
             f'--num-iterations={MLIR_N_REPEATS}', f'--warmup-iterations={WARMUP_ITERATIONS}',
@@ -1719,7 +1726,6 @@ def run_config_with_mlir(config: PerfConfiguration,
     else:
         if debug:
             print("Using rocprof for benchmarking")
-        rocmlir_gen_cmd = paths.mlir_paths.rocmlir_gen_path + ' -ph ' + commandline_options
         rocmlir_driver_cmd = [paths.mlir_paths.rocmlir_driver_path, '-c']
         mlir_cpu_runner_args = [
             f'--shared-libs={paths.mlir_paths.libmlir_rocm_runtime_path},{paths.mlir_paths.libconv_validation_wrappers_path},{paths.mlir_paths.libmlir_runtime_utils_path},{paths.mlir_paths.libmlir_c_runner_utils_path}',

mlir/utils/performance/tuningRunner.py

@@ -183,8 +183,10 @@ def tune_mlir_kernels(configs, conf_class, paths: Paths, options: Options):
             test_vector = config.to_command_line()
             print("Tuning:", test_vector, file=sys.stderr)
             command_line_options = config.generate_mlir_driver_commandline(
-                options.rocmlir_gen_flags, kernel_repeats=MLIR_N_REPEATS)
-            # Note, we don't need the -ph, this goes to the tuning driver
+                options.rocmlir_gen_flags, kernel_repeats=None)
+            # Note, we don't need the -ph, this goes to the tuning driver.
+            # Because we don't set -ph, kernel_repeats is set to None.
+            # This is because the kernel-repeats flag is only supported with host harness or CPU validation.
             kernel_gen_command = paths.mlir_paths.rocmlir_gen_path + ' ' + command_line_options
             kernel_gen = subprocess.Popen(kernel_gen_command.split(),
                                           stdout=subprocess.PIPE,