Parameter Sweeps for Attention: Check all outputs, log failures, avoid kernel repeats (#1914)

Check for all outputs to be [1 1 1], write failing configs to file every time one is detected for both attn and conv, change the usage of kernel_repeats to be 1 when used in parameter sweeps, include currentSeqLen when printing attention configuration.

---------

Signed-off-by: Djordje Antic <[email protected]>

Author: dorde-antic

mlir/utils/jenkins/Jenkinsfile

@@ -393,7 +393,7 @@ void check_randomE2ETests(String codepath) {
 void parameterSweep(String CONFIG, String codepath) {
     timeout(time: 300, activity: true, unit: 'MINUTES') {
         dir('build') {
-            sh """python3 ./bin/parameterSweeps.py -j 5 ${CONFIG}"""
+            sh """python3 ./bin/parameterSweeps.py -j 5 ${CONFIG} --log-failures"""
         }
     }
 }

mlir/utils/performance/attentionSweeps.py

@@ -32,7 +32,6 @@
 # GLOBAL VARIABLES
 DATA_TYPES_ATTENTION = initializeDataTypesAttention()
 BOOLS = [True, False]
-LOGFILE = 'failing_configs.csv'
 
 # Week number is used as seed to make sure weekly CI is reproducible
 seed = datetime.utcnow().isocalendar()[1]
@@ -156,7 +155,7 @@ def logFailingConfigs(configs: List[AttentionConfiguration], filename: str):
         writer = csv.writer(csvfile)
         writer.writerow(['CommandLine'])
         for config in configs:
-            writer.writerow([' '.join(config.generateMlirDriverCommandLine(''))])
+            writer.writerow([config.generateMlirDriverCommandLine('', kernel_repeats=None)])
 
 def main():
     parser = argparse.ArgumentParser(
@@ -181,7 +180,8 @@ def main():
         arch=arch,
         flags=[],
         concurrent_tests=args.jobs,
-        numCu=getNumCU(chip)
+        numCu=getNumCU(chip),
+        logFailures=args.log_failures
     )
 
 
@@ -206,8 +206,6 @@ def main():
         print(f"{'Failing Configurations':^80}\n")
         for fail in failing:
             print(multilineRepr(fail))
-        if args.log_failures:
-            logFailingConfigs(failing, LOGFILE)
 
     print(f"\nPassed: {passed}, Invalid: {invalid}, Failed: {len(failing)}")
 

mlir/utils/performance/parameterSweeps.py

@@ -36,6 +36,7 @@ class Options:
     flags: list
     concurrent_tests: int
     numCu: int
+    logFailures: bool = False
 
 class PerfConfig:
     class Version(enum.Enum):
@@ -191,7 +192,7 @@ async def testConfig(config, options: Options, paths: Paths) -> TestResult:
     if isinstance(config, MLIROnlyConfig):
         rocmlirGenOpts = config.generateMlirDriverCommandLine(options.flags)
     else:
-        rocmlirGenOpts = config.generateMlirDriverCommandLine(' '.join(options.flags)).split()
+        rocmlirGenOpts = config.generateMlirDriverCommandLine(' '.join(options.flags), kernel_repeats=None).split()
         if getattr(config, "currentSeqLen") is not None:
             rocmlirGenOpts.append(f"--current_seq_len={','.join(map(str, config.currentSeqLen))}")
     rocmlirGenOpts.append('-pv')
@@ -264,13 +265,17 @@ async def testConfig(config, options: Options, paths: Paths) -> TestResult:
 Return code = {runner.returncode}""", file=sys.stderr)
         return TestResult.FAIL
 
-    if not CORRECT_RESULT_RE.search(runnerOut):
+    output_lines = [line.strip() for line in runnerOut.splitlines() if len(line.strip()) > 0]
+    expected_output = "[1 1 1]"
+    all_correct = all(line == expected_output for line in output_lines)
+    if not all_correct:
         print(f"""Config returned incorrect result
 Output = {runnerOut}
 Errors = {runnerErrs.decode('utf-8')}""", file=sys.stderr)
         return TestResult.FAIL
     return TestResult.PASS
 
+
 IterType = TypeVar('IterType')
 def grouper(iterable: Iterable[IterType], n: int):
     it = iter(iterable)
@@ -288,8 +293,16 @@ async def dropGoodConfig(config, options: Options, paths: Paths):
         if isinstance(config, MLIROnlyConfig):
             print(f"{result.name}: {config!r}")
         else:
+            print("-" * 100)
             print(f"{result.name}: {multilineRepr(config)}")
     if result == TestResult.FAIL:
+        if options.logFailures:
+            if isinstance(config, perfRunner.AttentionConfiguration):
+                with open("failing_attn_configs.txt", "a") as f:
+                    f.write(multilineRepr(config) + "\n")
+            else:
+                with open("failing_conv_configs.txt", "a") as f:
+                    f.write(multilineRepr(config) + "\n")
         return config
     return result
 
@@ -459,7 +472,7 @@ async def runConfig(paramIter: Iterable[IterType],
     if len(failures) != 0:
         print("*** Summary of failures ***")
         for c in failures:
-            print(' '.join(c.generateMlirDriverCommandLine(options.flags)))
+            print(' '.join(c.generateMlirDriverCommandLine(options.flags, kernel_repeats=None)))
     print(f"Passed: {n_passes}, Invalid: {n_invalids}, Failed: {len(failures)}")
     return len(failures) == 0
 
@@ -482,6 +495,8 @@ def main() -> bool:
         help='Use xdlops when generating kernels (default off)')
     parser.add_argument('--no-xdlops', '-X', dest='xdlops', action='store_false',
         help='Explicitly disable xdlops usage')
+    parser.add_argument('--log-failures', '-L', action='store_true', default=False,
+        help='Save failures to file')
     parser.add_argument(
         '--codepath',
         type=str,
@@ -527,7 +542,7 @@ def main() -> bool:
             # unknow arch info
             print(f"""Unknown arch {arch}""", file=sys.stderr)
 
-    options = Options(debug=args.debug, quiet=args.quiet,
+    options = Options(debug=args.debug, quiet=args.quiet, logFailures=args.log_failures,
         arch=arch, flags=rocmlir_gen_flags, concurrent_tests=args.jobs, numCu=getNumCU(perfRunner.getChip()))
 
     paths = perfRunner.create_paths(None, args.mlir_build_dir)

mlir/utils/performance/perfRunner.py

@@ -1337,7 +1337,7 @@ def tableEntry(self, nanoSeconds):
     def setPerfConfig(self, perf_config):
         self.perfConfig = perf_config
 
-    def generateMlirDriverCommandLine(self, rocmlir_gen_flags):
+    def generateMlirDriverCommandLine(self, rocmlir_gen_flags, kernel_repeats=MLIR_N_REPEATS):
         result = ' '.join(['-operation', 'attention',
                            '-t', self.dataType,
                            '--arch', self.arch,
@@ -1357,7 +1357,7 @@ def generateMlirDriverCommandLine(self, rocmlir_gen_flags):
                            f"-transO={self.transO}",
                            f"-causal={self.causal}",
                            f"-return_lse={self.return_lse}",
-                           '--kernel-repeats', str(MLIR_N_REPEATS),
+                        *(['--kernel-repeats', str(kernel_repeats)] if kernel_repeats is not None else []),
                            f"--perf_config={self.perfConfig}"])
         result += ' '
         if rocmlir_gen_flags != '':