[Wave] Fix IREE benchmarking (again) (#31)

And add a test this time.

---------

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

Author: Hardcode84

tests/kernel/wave/attention/chained_gemm_test.py

@@ -166,7 +166,7 @@ def repeat(
             print(f"IR dumped to {filename}")
 
     iree_ref = device_zeros(batch, v_head_dim, q_seq_len, dtype=torch.float32)
-    generate_iree_ref("chain_mmt", [q, k, v], [iree_ref])
+    generate_iree_ref("chain_mmt", [q, k, v], [iree_ref], options)
     assert_close(output, iree_ref, check_device=False, atol=0, rtol=0)
 
     torch_qk = torch.matmul(q, k.transpose(-1, -2))
@@ -311,5 +311,5 @@ def repeat(
             f.write(asm)
 
     iree_ref = device_zeros(batch, v_head_dim, q_seq_len, dtype=torch.float32)
-    generate_iree_ref("chain_mmt_f8", [q, k, v], [iree_ref])
+    generate_iree_ref("chain_mmt_f8", [q, k, v], [iree_ref], options)
     assert_close(output, iree_ref, atol=7e-5, rtol=2e-3, check_device=False)

tests/kernel/wave/reordered_gemm_test.py

@@ -89,5 +89,5 @@ def testReorderedPureGemm(
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)

tests/kernel/wave/wave_e2e_test.py

@@ -1545,6 +1545,7 @@ def test_igemm_conv(
             "conv_2d_" + layout,
             [x, we],
             [iree_ref],
+            options,
         )
 
 

tests/kernel/wave/wave_gemm_test.py

@@ -68,6 +68,50 @@ def get_test_shapes(test_name: str) -> list[tuple[int]]:
     return default_test_shapes[test_name]
 
 
+@require_e2e
+def testGemmBench(tmp_path):
+    shape = (64, 64, 64)
+    perf_filename_tk = tmp_path / "wave_gemm_bench.txt"
+    perf_filename_iree = tmp_path / "iree_gemm_bench.txt"
+    enable_scheduling = SchedulingType.NONE
+    dynamic_dims = False
+    mfma_variant = MMAType.F32_16x16x16_F16
+    gemm, hyperparams, dynamic_symbols = get_gemm_kernel(
+        shape, dynamic_dims, mfma_variant, torch.float16
+    )
+
+    assert not perf_filename_tk.exists()
+
+    options = WaveCompileOptions(
+        subs=hyperparams,
+        canonicalize=True,
+        run_bench=True,
+        schedule=enable_scheduling,
+        use_scheduling_barriers=enable_scheduling_barriers,
+        dynamic_symbols=dynamic_symbols,
+        benchmark_batch_size=10,
+        benchmark_repetitions=3,
+        benchmark_results_file=perf_filename_tk,
+    )
+    options = set_default_run_config(options)
+    gemm = wave_compile(options, gemm)
+
+    a = device_randn(shape[0], shape[2], dtype=torch.float16)
+    b = device_randn(shape[1], shape[2], dtype=torch.float16)
+    c = device_zeros(shape[0], shape[1], dtype=torch.float32)
+    gemm(a, b, c)
+    assert perf_filename_tk.exists()
+    assert "real_time" in perf_filename_tk.read_text()
+
+    assert not perf_filename_iree.exists()
+    options.benchmark_results_file = perf_filename_iree
+
+    iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    assert perf_filename_iree.exists()
+    assert "real_time" in perf_filename_iree.read_text()
+
+
 @require_e2e
 @pytest.mark.parametrize("shape", get_test_shapes("test_gemm"))
 @pytest.mark.parametrize(
@@ -130,7 +174,7 @@ def testPureGemm(
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -202,7 +246,7 @@ def testGemmGatherToLDS(
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -336,7 +380,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -574,7 +618,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -627,7 +671,7 @@ def testGemmDumpOverrideSchedule(
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
     # Now reload the schedule and run the kernel again.
@@ -784,7 +828,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False, atol=1e-3, rtol=1e-3)
 
 
@@ -913,7 +957,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, atol=2e-4, rtol=3e-4, check_device=False)
 
 
@@ -1044,7 +1088,7 @@ def repeat(acc: tkl.Register[M, N, tkl.i32]) -> tkl.Register[M, N, tkl.i32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.int32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -1151,7 +1195,7 @@ def repeat(acc: tkl.Register[M, N, tkl.i32]) -> tkl.Register[M, N, tkl.i32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.int32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -1255,7 +1299,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt_f8", [a, b], [iree_ref])
+    generate_iree_ref("mmt_f8", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, atol=3e-5, rtol=3e-4, check_device=False)
 
 
@@ -1382,7 +1426,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -1516,7 +1560,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref])
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -1615,7 +1659,7 @@ def repeat(
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], shape[2], dtype=torch.float32)
-    generate_iree_ref("bmmt", [a, b], [iree_ref])
+    generate_iree_ref("bmmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
     torch_ref = torch.matmul(a, b.transpose(-2, -1))
@@ -1719,7 +1763,7 @@ def repeat(
         options.benchmark_results_file = perf_filename_iree
 
     iree_ref = device_zeros(shape[0], shape[1], shape[2], dtype=torch.float32)
-    generate_iree_ref("bmmt", [a, b], [iree_ref])
+    generate_iree_ref("bmmt", [a, b], [iree_ref], options)
     assert_close(c, iree_ref, check_device=False)
 
     torch_ref = torch.matmul(a, b.transpose(-2, -1))

wave_lang/kernel/wave/iree_utils.py

@@ -8,6 +8,10 @@
 
 from wave_lang.runtime.launch import Launchable
 from wave_lang.support.conversions import TORCH_DTYPE_TO_IREE_TYPE_ASM
+from .utils.run_utils import get_benchmark_flags, print_bench_result
+from .profiling import benchmark_module
+from .utils.compile_utils import compile_to_vmfb
+import iree.runtime as rt
 
 
 def get_chain_mmt_asm(
@@ -161,6 +165,7 @@ def generate_iree_ref(
     kernel_type: str,
     kernel_inputs: list[torch.Tensor],
     kernel_outputs: list[torch.Tensor],
+    options: "WaveCompileOptions",
 ):
     """
     Generate a reference output for the given kernel type and arguments.
@@ -211,10 +216,30 @@ def generate_iree_ref(
     else:
         raise ValueError(f"Unknown kernel type: {kernel_type}")
 
-    launchable = Launchable.jit_compile(asm, entry_point=func_name)
+    vmfb = compile_to_vmfb(asm, options)
+
+    def loader(device):
+        vm_instance = device.vm_instance
+        return rt.VmModule.copy_buffer(vm_instance, vmfb)
+
+    launchable = Launchable.from_vm_module(loader, entry_point=func_name)
     res = launchable(*kernel_inputs, outputs=kernel_outputs)
     if len(kernel_outputs) == 1:
         kernel_outputs[0][:] = res
     else:
         for r, k in zip(res, kernel_outputs):
             k[:] = r
+
+    if options.run_bench:
+        benchmark_flags = get_benchmark_flags(options)
+
+        benchmark_results = benchmark_module(
+            options,
+            kernel_inputs,
+            [],  # kernel_outputs,
+            [],  # dynamic_symbols,
+            vmfb,
+            func_name,
+            **benchmark_flags,
+        )
+        print_bench_result(benchmark_results, options.benchmark_results_file)

wave_lang/kernel/wave/profiling.py

@@ -165,7 +165,7 @@ def benchmark_module(
 
     err = err.decode()
     if "INVALID_ARGUMENT;" in err:
-        raise ValueError("Invalid inputs specified for benchmarking")
+        raise ValueError(f"Invalid inputs specified for benchmarking:\n{err}")
 
     # In the event benchmarking runs but encounteres an internal error,
     # return the internal error instead of benchmark results.

wave_lang/kernel/wave/utils/run_utils.py

@@ -87,7 +87,7 @@ def push_tensor_to_arg_list(arg_tensor: torch.Tensor):
         ) from e
 
 
-def _print_bench_result(result, filename):
+def print_bench_result(result, filename):
     import json
 
     res = json.dumps(result, sort_keys=True, indent=4)
@@ -97,6 +97,15 @@ def _print_bench_result(result, filename):
         print(res)
 
 
+def get_benchmark_flags(options: WaveCompileOptions):
+    benchmark_flags = {}
+    benchmark_flags["batch_size"] = options.benchmark_batch_size
+
+    if options.benchmark_repetitions is not None:
+        benchmark_flags["benchmark_repetitions"] = int(options.benchmark_repetitions)
+    return benchmark_flags
+
+
 def invoke_vmfb(
     vmfb: bytes,
     options: WaveCompileOptions,
@@ -120,16 +129,6 @@ def invoke_vmfb(
         return
 
     device = options.device
-    if options.run_bench:
-        benchmark_flags = {}
-        # If we use 1000 for bench_batch_size during compilation, and set this batch size to 1,
-        # then the latency is in milliseconds.
-        benchmark_flags["batch_size"] = 1
-
-        if options.benchmark_repetitions is not None:
-            benchmark_flags["benchmark_repetitions"] = int(
-                options.benchmark_repetitions
-            )
 
     # Select device as the GPU, where input tensors are coming from.
     device_list = tuple(
@@ -170,6 +169,7 @@ def invoke_vmfb(
     )
 
     if options.run_bench:
+        benchmark_flags = get_benchmark_flags(options)
         benchmark_results = benchmark_module(
             options,
             kernel_inputs,
@@ -179,7 +179,7 @@ def invoke_vmfb(
             options.func_name,
             **benchmark_flags,
         )
-        _print_bench_result(benchmark_results, options.benchmark_results_file)
+        print_bench_result(benchmark_results, options.benchmark_results_file)
 
 
 def invoke_with_wave_runtime(