[Wave] Migrate iree.turbine.kernel -> wave_lang.kernel (#50)

Also remove iree.turbine from setup.py and consolidate ir_imports.

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

Author: harsh-nod

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
     -   id: end-of-file-fixer
     -   id: check-yaml
     -   id: check-added-large-files
-# Fix trivial whitespace, line endings, YAML, large files
+# Format code with Black
 -   repo: https://github.com/psf/black
     rev: 25.1.0
     hooks:
@@ -19,3 +19,5 @@ repos:
     rev: v0.12.3
     hooks:
     -   id: ruff
+        args: [--select=F401,I001, --fix]
+        name: ruff-unused-imports-and-sorting

docs/conf.py

@@ -13,10 +13,10 @@
 import sys
 
 try:
-    import iree.turbine  # noqa
+    import wave_lang  # noqa
 except ImportError:
     sys.path.insert(0, os.path.abspath("../"))
-    import iree.turbine  # noqa
+    import wave_lang  # noqa
 
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

docs/kernel/wave.rst

@@ -1,7 +1,7 @@
 TKW (`turbine.kernel.wave`)
 ===========================
 
-.. py:module:: iree.turbine.kernel.wave.wave
+.. py:module:: wave_lang.kernel.wave.wave
 
 .. autoclass:: LaunchableWave
   :members:

docs/kernel/wave_codegen.rst

@@ -1,6 +1,6 @@
 Wave Codegen
 ============
 
-.. automodule:: iree.turbine.kernel.wave.codegen
+.. automodule:: wave_lang.kernel.wave.codegen
   :members:
   :undoc-members:

docs/kernel/wave_ops.rst

@@ -1,6 +1,6 @@
 Wave Ops
 ===================
 
-.. automodule:: iree.turbine.kernel.ops.wave_ops
+.. automodule:: wave_lang.kernel.ops.wave_ops
   :members:
   :undoc-members:

docs/wave/gemm_tutorial.rst

@@ -21,14 +21,14 @@ First, we need to import the necessary modules and define our symbolic dimension
 
 .. code-block:: python
 
-    from iree.turbine.kernel._support.indexing import sym
-    from iree.turbine.kernel._support.dtype import f16, f32
-    from iree.turbine.kernel.lang.wave_types import *
-    from iree.turbine.kernel.lang.global_symbols import *
-    from iree.turbine.kernel.wave.utils.run_utils import set_default_run_config
-    import iree.turbine.kernel as tkl
-    import iree.turbine.kernel.wave as tkw
-    from iree.turbine.kernel.wave.compile import WaveCompileOptions, wave_compile
+    from wave_lang.kernel._support.indexing import sym
+    from wave_lang.kernel._support.dtype import f16, f32
+    from wave_lang.kernel.lang.wave_types import *
+    from wave_lang.kernel.lang.global_symbols import *
+    from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
+    import wave_lang.kernel as tkl
+    import wave_lang.kernel.wave as tkw
+    from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile
     import torch
 
     # Define symbolic dimensions for our matrices

docs/wave/optimize_schedule.rst

@@ -35,11 +35,11 @@ Here's a complete example of how to use `tune_attention`:
 
 .. code-block:: python
 
-    from iree.turbine.kernel.wave.tuner.tune_attention import (
+    from wave_lang.kernel.wave.tuner.tune_attention import (
         tune_attention, AttentionConfig
     )
-    from iree.turbine.kernel.wave.constraints import MMAType
-    from iree.turbine.kernel.wave.scheduling import SchedulingType
+    from wave_lang.kernel.wave.constraints import MMAType
+    from wave_lang.kernel.wave.scheduling import SchedulingType
 
     # Configure the attention kernel
     config = AttentionConfig(

examples/jupyter/wave_gemm_example.ipynb

@@ -35,13 +35,13 @@
    "source": [
     "import torch\n",
     "\n",
-    "import iree.turbine.kernel.wave as tkw\n",
-    "from iree.turbine.kernel._support.dtype import f16, f32\n",
-    "from iree.turbine.kernel._support.indexing import sym\n",
-    "from iree.turbine.kernel.lang.global_symbols import *\n",
-    "from iree.turbine.kernel.lang.wave_types import *\n",
-    "from iree.turbine.kernel.wave.compile import WaveCompileOptions, wave_compile\n",
-    "from iree.turbine.kernel.wave.utils.run_utils import set_default_run_config\n",
+    "import wave_lang.kernel.wave as tkw\n",
+    "from wave_lang.kernel._support.dtype import f16, f32\n",
+    "from wave_lang.kernel._support.indexing import sym\n",
+    "from wave_lang.kernel.lang.global_symbols import *\n",
+    "from wave_lang.kernel.lang.wave_types import *\n",
+    "from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile\n",
+    "from wave_lang.kernel.wave.utils.run_utils import set_default_run_config\n",
     "\n",
     "# Define symbolic dimensions for our matrices\n",
     "M = sym.M  # Rows of A and C\n",