[Wave] Migrate iree.turbine.runtime -> wave_lang.runtime (#49)

Signed-off-by: Harsh Menon <[email protected]>

Author: harsh-nod

docs/core/runtime.rst

@@ -1,15 +1,15 @@
-`iree.turbine.runtime`
+`wave_lang.runtime`
 ======================
 
-.. automodule:: iree.turbine.runtime
+.. automodule:: wave_lang.runtime
   :imported-members:
   :members:
   :undoc-members:
 
 op_reg
 --------------
 
-.. automodule:: iree.turbine.runtime.op_reg
+.. automodule:: wave_lang.runtime.op_reg
   :imported-members:
   :members:
   :undoc-members:

iree/turbine/kernel/gen/kernel.py

@@ -30,7 +30,7 @@
     IrType,
 )
 
-from ...runtime.op_reg import (
+from wave_lang.runtime.op_reg import (
     def_library,
     CustomOp,
     KernelBuilder,

iree/turbine/kernel/wave/iree_utils.py

@@ -6,7 +6,7 @@
 
 import torch
 from wave_lang.support.conversions import TORCH_DTYPE_TO_IREE_TYPE_ASM
-from iree.turbine.runtime.launch import Launchable
+from wave_lang.runtime.launch import Launchable
 
 
 def get_chain_mmt_asm(

lit_tests/kernel/wave/sharktank_integration.py

@@ -35,12 +35,12 @@
 from iree.turbine.kernel.wave.utils.run_utils import (
     set_default_run_config,
 )
-from iree.turbine.runtime.op_reg.base import (
+from wave_lang.runtime.op_reg.base import (
     CustomOp,
     KernelBuilder,
     KernelSelection,
 )
-from iree.turbine.runtime.op_reg.impl_helper import (
+from wave_lang.runtime.op_reg.impl_helper import (
     call_function,
 )
 from wave_lang.transforms.merger import Merger

iree/turbine/runtime/__init__.py renamed to wave_lang/runtime/__init__.py

@@ -7,4 +7,3 @@
 from .device import *
 from .invoke import *
 from .launch import *
-from . import op_reg

iree/turbine/runtime/device.py renamed to wave_lang/runtime/device.py

@@ -4,33 +4,30 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+import atexit
+import ctypes
+import platform
 from functools import lru_cache
 from hashlib import sha1
+from threading import Lock, local
 from typing import Any, Callable, Dict, Optional, Union
-from threading import local, Lock
-
-import warnings
-import platform
-import atexit
 
-import ctypes
 import torch
 
 from iree.runtime import (
     BufferUsage,
+    ExternalTimepointFlags,
     HalBufferView,
     HalDevice,
     HalDriver,
     HalExternalTimepoint,
     MemoryType,
+    SemaphoreCompatibility,
     VmInstance,
     VmModule,
-    SemaphoreCompatibility,
-    ExternalTimepointFlags,
     create_hal_module,
     get_driver,
 )
-
 from wave_lang.support.conversions import (
     dtype_to_element_type,
     torch_dtype_to_numpy,

iree/turbine/runtime/invoke.py renamed to wave_lang/runtime/invoke.py

@@ -5,20 +5,17 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from typing import (
-    Any,
     Callable,
-    Sequence,
 )
 
-from .device import Device
-
 from iree.runtime import (
+    HalFence,
     VmContext,
     VmFunction,
-    HalFence,
     VmVariantList,
 )
 
+from .device import Device
 
 __all__ = [
     "invoke_vm_function",
@@ -33,7 +30,7 @@ def invoke_vm_function(
     arg_list: VmVariantList,
     ret_list: VmVariantList,
     *,
-    timer: Callable[[], float] = (lambda: 0.0)
+    timer: Callable[[], float] = (lambda: 0.0),
 ):
     """Invokes a vm function on a device, adding async fences to the arg_list if is_async.
 

iree/turbine/runtime/launch.py renamed to wave_lang/runtime/launch.py

@@ -11,31 +11,27 @@
 from torch import Tensor
 
 from iree.compiler.api import (
+    Output,
     Session,
     Source,
-    Output,
 )
-
 from iree.runtime import (
-    create_io_parameters_module,
     HalBufferView,
     HalElementType,
-    HalFence,
     ParameterProvider,
     VmContext,
     VmFunction,
     VmModule,
     VmRef,
     VmVariantList,
+    create_io_parameters_module,
 )
-
 from wave_lang.support.logging import runtime_logger as logger
 
 from .device import (
-    get_device_from_torch,
     Device,
+    get_device_from_torch,
 )
-
 from .invoke import invoke_vm_function
 
 __all__ = [

iree/turbine/runtime/op_reg/__init__.py renamed to wave_lang/runtime/op_reg/__init__.py

@@ -5,4 +5,3 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .base import *
-from . import impl_helper

iree/turbine/runtime/op_reg/base.py renamed to wave_lang/runtime/op_reg/base.py

@@ -8,41 +8,38 @@
 dispatcher.
 """
 
-from typing import Any, Callable, List, Optional, Sequence, Type, Union, cast
-
-from abc import ABC, abstractmethod
 import functools
 import logging
 import re
 import textwrap
 import threading
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Optional, Sequence, Type, Union, cast
 
 import torch
 from torch import Tensor
 
+from wave_lang.support.conversions import (
+    TORCH_DTYPE_TO_IREE_TYPE_ASM,
+)
 from wave_lang.support.ir_imports import (
     Block,
     Context,
     FunctionType,
     IndexType,
     InsertionPoint,
     IntegerAttr,
+    IrType,
     Location,
     StringAttr,
     SymbolTable,
-    IrType,
     Value,
     arith_d,
     builtin_d,
     func_d,
 )
-
 from wave_lang.support.logging import runtime_logger as logger
 
-from wave_lang.support.conversions import (
-    TORCH_DTYPE_TO_IREE_TYPE_ASM,
-)
-
 __all__ = [
     "ArgDescriptor",
     "AttrArg",