Fully upstream lowering pipeline

setup.py

@@ -282,6 +282,12 @@ def initialize_options(self):
 
 if BUILD_WATER:
     ext_modules += [
+        CMakeExtension(
+            "wave_execution_engine",
+            "wave_lang/kernel/wave/execution_engine",
+            install_dir="wave_lang/kernel/wave/execution_engine",
+            need_llvm=True,
+        ),
         CMakeExtension(
             "water",
             "water",

tests/kernel/wave/common/utils.py

@@ -65,4 +65,17 @@ def param_bool(name, shortname=None, values=None):
     shortname = shortname or name
     values = values or [False, True]
     ids = [f"{shortname}" if v else f"no_{shortname}" for v in values]
-    return pytest.mark.parametrize(name, [pytest.param(v) for v in values], ids=ids)
+    return pytest.mark.parametrize(name, values, ids=ids)
+
+
+def _is_water_and_ee_available() -> bool:
+    from wave_lang.kernel.wave.water import is_water_available
+    from wave_lang.kernel.wave.execution_engine import is_execution_engine_available
+
+    return is_water_available() and is_execution_engine_available()
+
+
+require_water_and_ee = pytest.mark.skipif(
+    not _is_water_and_ee_available(),
+    reason="Water or execution engine are not available.",
+)

tests/kernel/wave/test_execution_engine_wrapper.py

@@ -0,0 +1,177 @@
+# 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
+
+"""
+Tests for the ExecutionEngine wrapper with weak reference caching.
+"""
+
+import gc
+import os
+import pytest
+import weakref
+
+from wave_lang.kernel.wave.execution_engine import (
+    clear_engine_cache,
+    get_execution_engine,
+    is_engine_cached,
+    is_execution_engine_available,
+)
+
+pytestmark = pytest.mark.skipif(
+    not is_execution_engine_available(),
+    reason="ExecutionEngine not available (C++ extension not built)",
+)
+
+
+ENV_VARS = [
+    "WAVE_ENABLE_OBJECT_CACHE",
+    "WAVE_ENABLE_GDB_LISTENER",
+    "WAVE_ENABLE_PERF_LISTENER",
+]
+
+
+@pytest.fixture
+def clean_env():
+    """Fixture to clean cache and environment variables before/after each test."""
+    clear_engine_cache()
+
+    # Save original env vars
+    orig_env = {}
+    for key in ENV_VARS:
+        orig_env[key] = os.environ.get(key)
+        if key in os.environ:
+            del os.environ[key]
+
+    yield
+
+    # Restore original env vars
+    for key, value in orig_env.items():
+        if value is not None:
+            os.environ[key] = value
+        elif key in os.environ:
+            del os.environ[key]
+
+    clear_engine_cache()
+    gc.collect()
+
+
+def test_basic_creation(clean_env):
+    """Test basic creation of execution engine."""
+    engine = get_execution_engine()
+    assert engine is not None
+    assert is_engine_cached()
+
+
+def test_cache_returns_same_instance(clean_env):
+    """Test that cache returns the same instance."""
+    engine1 = get_execution_engine()
+    engine2 = get_execution_engine()
+
+    # Should be the exact same object
+    assert engine1 is engine2
+
+
+def test_weak_reference_cleanup(clean_env):
+    """Test that engine is cleaned up when references are released."""
+    # Create engine and get weak reference
+    engine = get_execution_engine()
+    weak_ref = weakref.ref(engine)
+
+    # Verify it's cached
+    assert is_engine_cached()
+
+    # Delete strong reference
+    del engine
+    gc.collect()
+
+    # Weak reference should be dead
+    assert weak_ref() is None
+
+    # Cache should report no engine
+    assert not is_engine_cached()
+
+
+def test_clear_cache(clean_env):
+    """Test that clear_cache removes cached engine."""
+    engine = get_execution_engine()
+    assert is_engine_cached()
+
+    # Clear cache
+    clear_engine_cache()
+
+    # Cache should be empty
+    assert not is_engine_cached()
+
+    # Engine should still be alive (we have strong reference)
+    assert engine is not None
+
+
+def test_env_var_object_cache(clean_env):
+    """Test WAVE_ENABLE_OBJECT_CACHE environment variable."""
+    os.environ["WAVE_ENABLE_OBJECT_CACHE"] = "1"
+
+    engine = get_execution_engine()
+    assert engine is not None
+    # We can't directly test if object cache is enabled without
+    # accessing internal state, but we verify the engine was created
+
+
+def test_env_var_gdb_listener(clean_env):
+    """Test WAVE_ENABLE_GDB_LISTENER environment variable."""
+    os.environ["WAVE_ENABLE_GDB_LISTENER"] = "1"
+
+    engine = get_execution_engine()
+    assert engine is not None
+
+
+def test_env_var_perf_listener(clean_env):
+    """Test WAVE_ENABLE_PERF_LISTENER environment variable."""
+    os.environ["WAVE_ENABLE_PERF_LISTENER"] = "1"
+
+    engine = get_execution_engine()
+    assert engine is not None
+
+
+def test_all_env_vars(clean_env):
+    """Test all environment variables together."""
+    os.environ["WAVE_ENABLE_OBJECT_CACHE"] = "1"
+    os.environ["WAVE_ENABLE_GDB_LISTENER"] = "1"
+    os.environ["WAVE_ENABLE_PERF_LISTENER"] = "1"
+
+    engine = get_execution_engine()
+    assert engine is not None
+    assert is_engine_cached()
+
+
+def test_multiple_references(clean_env):
+    """Test that multiple references to the same engine work correctly."""
+    engine1 = get_execution_engine()
+    engine2 = get_execution_engine()
+    engine3 = get_execution_engine()
+
+    # All should be the same object
+    assert engine1 is engine2
+    assert engine2 is engine3
+
+
+def test_engine_interface(clean_env):
+    """Test that the engine provides the correct interface."""
+    engine = get_execution_engine()
+
+    # Check that required methods exist
+    assert hasattr(engine, "load_module")
+    assert hasattr(engine, "release_module")
+    assert hasattr(engine, "lookup")
+    assert hasattr(engine, "dump_to_object_file")
+
+    # Check that methods are callable
+    assert callable(engine.load_module)
+    assert callable(engine.release_module)
+    assert callable(engine.lookup)
+    assert callable(engine.dump_to_object_file)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/kernel/wave/wave_e2e_test.py

@@ -44,10 +44,11 @@
     param_bool,
     perf_test,
     require_cdna3,
-    require_e2e,
-    require_cdna_2_or_3_or_4,
     require_cdna4,
+    require_cdna_2_or_3_or_4,
+    require_e2e,
     require_rdna4,
+    require_water_and_ee,
 )
 from .common.shapes import get_test_shapes as get_common_test_shape
 
@@ -138,8 +139,13 @@ def test(
 @require_e2e
 @pytest.mark.parametrize("shape", get_test_shapes("test_copy"))
 @param_bool("use_buffer_ops", "buf_ops")
+@param_bool(
+    "use_water_pipeline",
+    "water",
+    values=[False, pytest.param(True, marks=require_water_and_ee)],
+)
 @check_leaks
-def test_copy(shape, use_buffer_ops, run_bench):
+def test_copy(shape, use_buffer_ops, run_bench, use_water_pipeline):
     M = tkl.sym.M
     N = tkl.sym.N
     ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE
@@ -183,6 +189,7 @@ def test(
         canonicalize=True,
         run_bench=run_bench,
         use_buffer_ops=use_buffer_ops,
+        use_water_pipeline=use_water_pipeline,
     )
     options = set_default_run_config(options)
     test = wave_compile(options, test)
@@ -194,7 +201,12 @@ def test(
 @require_e2e
 @pytest.mark.parametrize("shape", get_test_shapes("test_copy"))
 @param_bool("use_buffer_ops", "buf_ops")
-def test_dynamic_copy(shape, use_buffer_ops, run_bench):
+@param_bool(
+    "use_water_pipeline",
+    "water",
+    values=[False, pytest.param(True, marks=require_water_and_ee)],
+)
+def test_dynamic_copy(shape, use_buffer_ops, run_bench, use_water_pipeline):
     M = tkl.sym.M
     N = tkl.sym.N
     ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE
@@ -238,6 +250,7 @@ def test(
         canonicalize=True,
         run_bench=run_bench,
         use_buffer_ops=use_buffer_ops,
+        use_water_pipeline=use_water_pipeline,
     )
     options = set_default_run_config(options)
     test = wave_compile(options, test)

water/tools/water-opt/CMakeLists.txt

@@ -12,6 +12,11 @@ set(LIBS
 
         MLIRWaterTestTransforms
         MLIRWaterTestDialect
+
+        MLIRRegisterAllDialects
+        MLIRRegisterAllExtensions
+        MLIRRegisterAllPasses
+        MLIRToLLVMIRTranslationRegistration
 )
 
 add_llvm_executable(water-opt

water/tools/water-opt/water-opt.cpp

@@ -20,7 +20,9 @@
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllExtensions.h"
 #include "mlir/InitAllPasses.h"
+#include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Transforms/Passes.h"
 
@@ -40,20 +42,15 @@ int main(int argc, char **argv) {
   mlir::water::registerPasses();
   mlir::water::test::registerAllPasses();
   wave::registerPasses();
-  mlir::arith::registerArithIntRangeOptsPass();
-  mlir::registerCanonicalizerPass();
-  mlir::registerCSEPass();
-  mlir::registerLoopInvariantCodeMotionPass();
-  mlir::registerLowerAffinePass();
 
+  mlir::registerAllPasses();
   mlir::DialectRegistry registry;
-  registry.insert<mlir::affine::AffineDialect, mlir::amdgpu::AMDGPUDialect,
-                  mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect,
-                  mlir::func::FuncDialect, mlir::gpu::GPUDialect,
-                  mlir::LLVM::LLVMDialect, mlir::ROCDL::ROCDLDialect,
-                  mlir::memref::MemRefDialect, mlir::scf::SCFDialect,
-                  mlir::vector::VectorDialect, wave::WaveDialect>();
+  mlir::registerAllDialects(registry);
+  mlir::registerAllExtensions(registry);
 
+  mlir::registerAllGPUToLLVMIRTranslations(registry);
+
+  registry.insert<wave::WaveDialect>();
   mlir::water::test::registerWaterTestDialect(registry);
 
   return mlir::asMainReturnCode(