Limit TK conv to f16xf32 (#41)

We doesn't support other datatypes yet.

Signed-off-by: Ivan Butygin <[email protected]>

Author: Hardcode84

convbench/conv_bench.py

@@ -9,7 +9,7 @@
 import sys
 from utils import *
 from conv_utils import *
-from problems import get_conv_configs, get_conv_test_configs
+from problems import get_conv_configs, get_tk_conv_configs, get_conv_test_configs
 
 from wave_conv_utils import compile_wave_conv_config
 
@@ -56,7 +56,7 @@ def compile_conv_wave(tag, config, kernel_dir, vmfb_dir, extra_compiler_args):
         sys.exit()
 
     # configs = get_conv_test_configs()
-    configs = get_conv_configs()
+    configs = get_tk_conv_configs() if args.tk else get_conv_configs()
     print(f"Generated {len(configs)} conv configs.")
 
     num_cpus = max(1, cpu_count() - 20)

convbench/problems.py

@@ -81,6 +81,14 @@ def get_conv_configs() -> list[tuple[str, ConvConfig]]:
 
     return configs
 
+def get_tk_conv_configs() -> list[tuple[str, ConvConfig]]:
+    def check(config_tuple: tuple[str, ConvConfig]) -> bool:
+        config = config_tuple[1]
+        return config.input_dtype == "f16" and config.output_dtype == "f32"
+
+    return list(filter(check, get_conv_configs()))
+
+
 # Test function to run only a few chosen shapes
 def get_conv_test_configs() -> list[tuple[str, ConvConfig]]:
     configs: list[tuple[str, ConvConfig]] = []