Use turbine.runtime for generate_iree_ref (iree-org#861)

`generate_iree_ref` only used to get reference values for Wave tests and
it also the last user of `inplace=False` flag in Wave.

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

Author: Hardcode84

iree/turbine/kernel/wave/iree_utils.py

@@ -1,19 +1,17 @@
-# Copyright 2024 The IREE Authors
+# Copyright 2025 The IREE Authors
 #
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import torch
-from typing import Any
-from .utils.run_utils import compile_and_invoke
 from ...support.conversions import TORCH_DTYPE_TO_IREE_TYPE_ASM
-from .compile import WaveCompileOptions
+from iree.turbine.runtime.launch import Launchable
 
 
 def get_chain_mmt_asm(
     query_type: str, key_type: str, value_type: str, output_type: str
-) -> str:
+) -> tuple[str, str]:
     B, M, K1, input_dtype = query_type.split("x")
     B, K2, K1, input_dtype = key_type.split("x")
     B, N, K2, input_dtype = value_type.split("x")
@@ -22,7 +20,8 @@ def get_chain_mmt_asm(
     intermediate_cast_type = f"{B}x{K2}x{M}x{input_dtype}"
     transposed_cast_type = f"{B}x{M}x{K2}x{input_dtype}"
     transposed_output_type = f"{B}x{M}x{N}x{output_dtype}"
-    return f"""
+    return (
+        f"""
     func.func @chain_mmt(%query: tensor<{query_type}>, %key: tensor<{key_type}>, %value: tensor<{value_type}>) -> tensor<{output_type}> {{
       %c0 = arith.constant 0.0 : f32
       %init = tensor.empty() : tensor<{intermediate_output_type}>
@@ -39,12 +38,14 @@ def get_chain_mmt_asm(
       %init4 = tensor.empty() : tensor<{output_type}>
       %transpose2 = linalg.transpose ins(%result2: tensor<{transposed_output_type}>) outs(%init4: tensor<{output_type}>) permutation=[0, 2, 1]
       return %transpose2 : tensor<{output_type}>
-    }}"""
+    }}""",
+        "chain_mmt",
+    )
 
 
 def get_chain_mmt_f8_asm(
     query_type: str, key_type: str, value_type: str, output_type: str
-) -> str:
+) -> tuple[str, str]:
     B, M, K1, input_dtype = query_type.split("x")
     B, K2, K1, input_dtype = key_type.split("x")
     B, N, K2, input_dtype = value_type.split("x")
@@ -57,7 +58,8 @@ def get_chain_mmt_f8_asm(
     query_f8_type = "x".join([B, M, K1, f8_dtype])
     key_f8_type = "x".join([B, K2, K1, f8_dtype])
     value_f8_type = "x".join([B, N, K2, f8_dtype])
-    return f"""
+    return (
+        f"""
     func.func @chain_mmt_f8(%query: tensor<{query_type}>, %key: tensor<{key_type}>, %value: tensor<{value_type}>) -> tensor<{output_type}> {{
       %c0 = arith.constant 0.0 : f32
       %init = tensor.empty() : tensor<{intermediate_output_type}>
@@ -77,7 +79,9 @@ def get_chain_mmt_f8_asm(
       %init4 = tensor.empty() : tensor<{output_type}>
       %transpose2 = linalg.transpose ins(%result2: tensor<{transposed_output_type}>) outs(%init4: tensor<{output_type}>) permutation=[0, 2, 1]
       return %transpose2 : tensor<{output_type}>
-    }}"""
+    }}""",
+        "chain_mmt_f8",
+    )
 
 
 def get_mmt_asm(
@@ -86,7 +90,7 @@ def get_mmt_asm(
     acc_type: str,
     batch: bool = False,
     cast_fp8: bool = False,
-) -> str:
+) -> tuple[str, str]:
     acc_dtype = acc_type.split("x")[-1]
     operator = "batch_matmul_transpose_b" if batch else "matmul_transpose_b"
     func_name = "bmmt" if batch else "mmt"
@@ -118,14 +122,15 @@ def get_mmt_asm(
                      outs(%inital_result: tensor<{acc_type}>) -> tensor<{acc_type}>
           return %result : tensor<{acc_type}>
         }}"""
-    return matmul_function
+    return matmul_function, func_name
 
 
 def get_conv_asm(
     conv_type: str, lhs_type: str, rhs_type: str, res_type: str, stride: int
-) -> str:
+) -> tuple[str, str]:
     res_dtype = res_type.split("x")[-1]
-    return f"""
+    return (
+        f"""
     func.func @conv_{conv_type}(%lhs: tensor<{lhs_type}>, %rhs: tensor<{rhs_type}>) -> tensor<{res_type}> {{
       %c0 = arith.constant 0.0 : {res_dtype}
       %init = tensor.empty() : tensor<{res_type}>
@@ -135,7 +140,9 @@ def get_conv_asm(
                 ins(%lhs, %rhs : tensor<{lhs_type}>, tensor<{rhs_type}>)
                 outs(%inital_result : tensor<{res_type}>) -> tensor<{res_type}>
       return %result : tensor<{res_type}>
-    }}"""
+    }}""",
+        f"conv_{conv_type}",
+    )
 
 
 def dtype_str(dtype: torch.dtype) -> str:
@@ -153,7 +160,6 @@ 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.
@@ -165,7 +171,7 @@ def generate_iree_ref(
         lhs_type = get_type_str(kernel_inputs[0].shape, kernel_inputs[0].dtype)
         rhs_type = get_type_str(kernel_inputs[1].shape, kernel_inputs[1].dtype)
         acc_type = get_type_str(kernel_outputs[0].shape, kernel_outputs[0].dtype)
-        asm = get_mmt_asm(
+        asm, func_name = get_mmt_asm(
             lhs_type,
             rhs_type,
             acc_type,
@@ -176,37 +182,38 @@ def generate_iree_ref(
         lhs_type = get_type_str(kernel_inputs[0].shape, kernel_inputs[0].dtype)
         rhs_type = get_type_str(kernel_inputs[1].shape, kernel_inputs[1].dtype)
         acc_type = get_type_str(kernel_outputs[0].shape, kernel_outputs[0].dtype)
-        asm = get_mmt_asm(lhs_type, rhs_type, acc_type, batch=True)
+        asm, func_name = get_mmt_asm(lhs_type, rhs_type, acc_type, batch=True)
     elif kernel_type == "chain_mmt":
         query_type = get_type_str(kernel_inputs[0].shape, kernel_inputs[0].dtype)
         key_type = get_type_str(kernel_inputs[1].shape, kernel_inputs[1].dtype)
         value_type = get_type_str(kernel_inputs[2].shape, kernel_inputs[2].dtype)
         output_type = get_type_str(kernel_outputs[0].shape, kernel_outputs[0].dtype)
-        asm = get_chain_mmt_asm(query_type, key_type, value_type, output_type)
+        asm, func_name = get_chain_mmt_asm(
+            query_type, key_type, value_type, output_type
+        )
     elif kernel_type == "chain_mmt_f8":
         query_type = get_type_str(kernel_inputs[0].shape, kernel_inputs[0].dtype)
         key_type = get_type_str(kernel_inputs[1].shape, kernel_inputs[1].dtype)
         value_type = get_type_str(kernel_inputs[2].shape, kernel_inputs[2].dtype)
         output_type = get_type_str(kernel_outputs[0].shape, kernel_outputs[0].dtype)
-        asm = get_chain_mmt_f8_asm(query_type, key_type, value_type, output_type)
+        asm, func_name = get_chain_mmt_f8_asm(
+            query_type, key_type, value_type, output_type
+        )
     elif kernel_type.startswith(conv_str):
         lhs_type = get_type_str(kernel_inputs[0].shape, kernel_inputs[0].dtype)
         rhs_type = get_type_str(kernel_inputs[1].shape, kernel_inputs[1].dtype)
         acc_type = get_type_str(kernel_outputs[0].shape, kernel_outputs[0].dtype)
         conv_type = kernel_type[len(conv_str) :]
-        asm = get_conv_asm(
+        asm, func_name = get_conv_asm(
             conv_type, lhs_type, rhs_type, acc_type, int(kwargs["stride"])
         )
     else:
         raise ValueError(f"Unknown kernel type: {kernel_type}")
 
-    options.func_name = kernel_type
-    options.inplace = False
-    options.kernel_hash = None
-    options.dynamic_symbols_map = {}
-    compile_and_invoke(
-        asm,
-        kernel_inputs,
-        kernel_outputs,
-        options,
-    )
+    launchable = Launchable.jit_compile(asm, 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

iree/turbine/kernel/wave/nn/linear.py

@@ -12,7 +12,6 @@
 import iree.turbine.kernel.lang as tkl
 import iree.turbine.kernel.wave as tkw
 from iree.turbine.kernel.lang.global_symbols import *
-from iree.turbine.kernel.wave.iree_utils import generate_iree_ref
 from iree.turbine.kernel.wave.utils.general_utils import (
     get_default_scheduling_params,
 )

iree/turbine/runtime/launch.py

@@ -199,7 +199,9 @@ def _resolve_target_binary(self, turbine_device: Device) -> _TargetBinary:
             f"Could not load a target binary for device {turbine_device}"
         )
 
-    def __call__(self, *args, device: Optional[torch.device] = None):
+    def __call__(
+        self, *args, device: Optional[torch.device] = None, outputs: Sequence[Any] = ()
+    ):
         turbine_device: Optional[Device] = (
             None if device is None else get_device_from_torch(device)
         )
@@ -238,7 +240,13 @@ def __call__(self, *args, device: Optional[torch.device] = None):
 
         vm_context, vm_function = self._resolve_target_binary(turbine_device)
 
-        ret_list = VmVariantList(1)
+        ret_list = VmVariantList(len(outputs))
+        for output in outputs:
+            if isinstance(output, Tensor):
+                assert output.is_contiguous(), "Outputs must be contiguous"
+                ret_list.push_ref(turbine_device.import_torch_tensor(arg))
+            else:
+                raise ValueError(f"Unsupported output type: {type(output)}")
 
         invoke_vm_function(
             turbine_device, self._is_async, vm_context, vm_function, arg_list, ret_list

tests/kernel/wave/attention/chained_gemm_test.py

@@ -171,7 +171,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], options)
+    generate_iree_ref("chain_mmt", [q, k, v], [iree_ref])
     assert_close(output, iree_ref, check_device=False, atol=0, rtol=0)
 
     torch_qk = torch.matmul(q, k.transpose(-1, -2))
@@ -322,5 +322,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], options)
+    generate_iree_ref("chain_mmt_f8", [q, k, v], [iree_ref])
     assert_close(output, iree_ref, atol=7e-5, rtol=2e-3, check_device=False)

tests/kernel/wave/runtime/cache_test.py

@@ -11,7 +11,6 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import copy
-import tempfile
 import pytest
 import torch
 from torch.testing import assert_close
@@ -27,10 +26,8 @@
     reset_cache_manager,
 )
 from iree.turbine.kernel.lang.global_symbols import *
-from iree.turbine.kernel.wave.iree_utils import generate_iree_ref
 from iree.turbine.kernel.wave.utils.run_utils import (
     set_default_run_config,
-    get_default_arch,
 )
 from iree.turbine.kernel.wave.utils.general_utils import (
     get_default_scheduling_params,

tests/kernel/wave/wave_e2e_test.py

@@ -1557,9 +1557,6 @@ def test_igemm_conv(
             "conv_2d_" + layout,
             [x, we],
             [iree_ref],
-            options,
-            stride=stride,
-            run_bench=True,
         )
 
 

tests/kernel/wave/wave_gemm_test.py

@@ -210,7 +210,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
                 dump_perf, "iree_" + perf_filename
             )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -328,7 +328,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
                 dump_perf, "iree_" + perf_filename
             )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -478,7 +478,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
                 dump_perf, "iree_" + perf_filename
             )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False, atol=1e-3, rtol=1e-3)
 
 
@@ -621,7 +621,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
             dump_perf, "iree_" + request.node.name + ".json"
         )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, atol=2e-4, rtol=3e-4, check_device=False)
 
 
@@ -766,7 +766,7 @@ def repeat(acc: tkl.Register[M, N, tkl.i32]) -> tkl.Register[M, N, tkl.i32]:
             dump_perf, "iree_" + request.node.name + ".json"
         )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.int32)
-    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -879,7 +879,7 @@ def repeat(acc: tkl.Register[M, N, tkl.i32]) -> tkl.Register[M, N, tkl.i32]:
             dump_perf, "iree_" + request.node.name + ".json"
         )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.int32)
-    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False)
 
 
@@ -989,7 +989,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
             dump_perf, "iree_" + request.node.name + ".json"
         )
     iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
-    generate_iree_ref("mmt_f8", [a, b], [iree_ref], options)
+    generate_iree_ref("mmt_f8", [a, b], [iree_ref])
     assert_close(c, iree_ref, atol=3e-5, rtol=3e-4, check_device=False)
 
 
@@ -1094,7 +1094,7 @@ def repeat(
             dump_perf, "iree_" + request.node.name + ".json"
         )
     iree_ref = device_zeros(shape[0], shape[1], shape[2], dtype=torch.float32)
-    generate_iree_ref("bmmt", [a, b], [iree_ref], options)
+    generate_iree_ref("bmmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False)
 
     torch_ref = torch.matmul(a, b.transpose(-2, -1))
@@ -1206,7 +1206,7 @@ def repeat(
             dump_perf, "iree_" + request.node.name + ".json"
         )
     iree_ref = device_zeros(shape[0], shape[1], shape[2], dtype=torch.float32)
-    generate_iree_ref("bmmt", [a, b], [iree_ref], options)
+    generate_iree_ref("bmmt", [a, b], [iree_ref])
     assert_close(c, iree_ref, check_device=False)
 
     torch_ref = torch.matmul(a, b.transpose(-2, -1))