Make wheel usable from Python and register dialects (#5)

These dialects appear in Wave output and passes are useful to simplify
IR before processing or printing.

Author: ftynse

README.md

@@ -59,7 +59,7 @@ A Python wheel containing the `water-opt` binary can be produced using the
 following commands:
 
 ```sh
-cd build_tools
+cd build_tools/wheel
 WATER_MLIR_DIR=$BUILD_DIR/lib/cmake/mlir python -m pip wheel .
 ```
 

build_tools/pyproject.toml renamed to build_tools/wheel/pyproject.toml

@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from setuptools import setup, Extension
+from setuptools import setup, Extension, find_packages
 import subprocess
 import os
 import shutil
@@ -43,8 +43,10 @@ def prepare_installation():
         "-DBUILD_SHARED_LIBS=Off",
         f"-DCMAKE_BUILD_TYPE={build_type}",
     ]
-    source_dir = os.path.dirname(SETUPPY_DIR)
-    os.makedirs(BINARY_DIR, exist_ok=True)
+    source_dir = os.path.dirname(os.path.dirname(SETUPPY_DIR))
+    if os.path.exists(BINARY_DIR):
+        shutil.rmtree(BINARY_DIR)
+    os.makedirs(BINARY_DIR)
     subprocess.check_call(["cmake", source_dir, *cmake_args], cwd=BINARY_DIR)
     subprocess.check_call(
         ["cmake", "--build", ".", "--target", "water-opt"], cwd=BINARY_DIR
@@ -55,12 +57,13 @@ class CMakeBuildPy(_build_py):
     """Pretends to be building but actually just copies pre-built binaries."""
 
     def run(self):
-        target_dir = os.path.abspath(self.build_lib)
+        target_dir = os.path.join(os.path.abspath(self.build_lib), "water_mlir")
         print(f"Building in target dir: {target_dir}", file=sys.stderr)
         if os.path.exists(target_dir):
             shutil.rmtree(target_dir)
         os.makedirs(target_dir)
         shutil.copy(os.path.join(BINARY_DIR, "bin", "water-opt"), target_dir)
+        shutil.copy(os.path.join(SETUPPY_DIR, "tools", "binaries.py"), target_dir)
 
 
 # This is needed to create at least some binary understood by Python.
@@ -107,7 +110,8 @@ def run(self):
 setup(
     name="water_mlir",
     version="0.1.0",
-    ext_modules=[CMakeExtension("water_mlir")],
+    packages=find_packages(include=["tools"]),
+    ext_modules=[CMakeExtension("water_mlir_ext")],
     cmdclass={
         "build": CustomBuild,
         "built_ext": NoopBuildExtension,

build_tools/setup.py renamed to build_tools/wheel/setup.py

@@ -0,0 +1,5 @@
+# Copyright 2025 The Water 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

build_tools/wheel/tools/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2025 The Water 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 os
+
+
+def find_binary(name: str) -> str:
+    this_path = os.path.dirname(__file__)
+    tool_path = os.path.join(this_path, name)
+    assert os.path.isfile(tool_path) and os.access(
+        tool_path, os.X_OK
+    ), f"Could not find the {name} executable in package."
+    return tool_path

build_tools/wheel/tools/binaries.py

@@ -6,12 +6,15 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Conversion/Passes.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 
+#include "mlir/Transforms/Passes.h"
 #include "water/Transforms/Passes.h"
 
 // Forward-declare test passes so we don't have a dependency on the test
@@ -24,9 +27,17 @@ int main(int argc, char **argv) {
   mlir::water::registerPasses();
   mlir::water::test::registerAllPasses();
 
+  mlir::arith::registerArithIntRangeOptsPass();
+  mlir::registerCanonicalizerPass();
+  mlir::registerCSEPass();
+  mlir::registerLoopInvariantCodeMotionPass();
+  mlir::registerLowerAffinePass();
+
   mlir::DialectRegistry registry;
-  registry.insert<mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect,
-                  mlir::func::FuncDialect, mlir::memref::MemRefDialect,
+  registry.insert<mlir::affine::AffineDialect, mlir::amdgpu::AMDGPUDialect,
+                  mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect,
+                  mlir::func::FuncDialect, mlir::gpu::GPUDialect,
+                  mlir::LLVM::LLVMDialect, mlir::memref::MemRefDialect,
                   mlir::scf::SCFDialect, mlir::vector::VectorDialect>();
 
   return mlir::asMainReturnCode(