add jax as a test of upstream bindings

Author: makslevental

.github/workflows/test.yml

@@ -49,3 +49,40 @@ jobs:
         shell: bash
         run: |
           pytest tests
+
+
+  test-against-jax-bindings:
+
+    runs-on: ${{ matrix.os }}
+    continue-on-error: true
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-22.04, macos-11, windows-2022 ]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v2
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
+      - name: Install and configure
+        shell: bash
+        run: |
+          pip install .[jax-test] \
+            -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest
+          if [ ${{ matrix.os }} == 'windows-2022' ]; then
+            # configure-mlir-python-utils.exe -y jaxlib.mlir
+            pushd /tmp && python -m mlir_utils._configuration -y jaxlib.mlir && popd
+          else
+            configure-mlir-python-utils -y jaxlib.mlir
+          fi
+
+      - name: Test
+        shell: bash
+        run: |
+          pytest tests/test_smoke.py

mlir_utils/_configuration/generate_trampolines.py

@@ -80,7 +80,7 @@ def generate_op_trampoline(op_class):
         for k, d in zip(args.kwonlyargs, args.kw_defaults)
     ]
 
-    fun_name = op_class.OPERATION_NAME.split(".")[-1]
+    fun_name = op_class.OPERATION_NAME.split(".")[-1].replace("-", "_")
     if keyword.iskeyword(fun_name):
         fun_name = fun_name + "_"
     op_class_name = op_class.__name__

mlir_utils/dialects/.gitignore

@@ -3,18 +3,21 @@
 /bufferization.py
 /builtin.py
 /cf.py
+/chlo.py
 /complex.py
 /func.py
 /gpu.py
 /linalg.py
 /math.py
 /memref.py
+/mhlo.py
 /ml_program.py
 /pdl.py
 /quant.py
 /scf.py
 /shape.py
 /sparse_tensor.py
+/stablehlo.py
 /tensor.py
 /torch.py
 /tosa.py

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
 
 [project.optional-dependencies]
 torch-mlir-test = ["pytest", "torch-mlir-core", "mlir-native-tools"]
+jax-test = ["pytest", "jax[cpu]", "mlir-native-tools"]
 mlir-test = ["pytest", "mlir-python-bindings", "mlir-native-tools"]
 mlir = ["mlir-python-bindings"]
 

tests/test_smoke.py

@@ -11,6 +11,7 @@
 
 # noinspection PyUnresolvedReferences
 from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
+from util import skip_jax_not_installed, skip_torch_mlir_not_installed
 
 # needed since the fix isn't defined here nor conftest.py
 pytest.mark.usefixtures("ctx")
@@ -26,6 +27,8 @@ def test_smoke(ctx: MLIRContext):
     filecheck(correct, ctx.module)
 
 
+# skip if jax *is* installed because jax doesn't generate almost any of the upstream dialects
+@pytest.mark.skipif(lambda: not skip_jax_not_installed(), reason="jax installed")
 def test_dialect_trampolines_smoke():
     # noinspection PyUnresolvedReferences
     from mlir_utils.dialects import (
@@ -51,15 +54,18 @@ def test_dialect_trampolines_smoke():
     )
 
 
-def skip_torch_mlir_not_installed():
-    try:
-        from torch_mlir.dialects import torch
-
-        # don't skip
-        return False
-    except ImportError:
-        # skip
-        return True
+@pytest.mark.skipif(skip_jax_not_installed(), reason="jax not installed")
+def test_dialect_trampolines_smoke():
+    # noinspection PyUnresolvedReferences
+    from mlir_utils.dialects import (
+        builtin,
+        chlo,
+        func,
+        mhlo,
+        ml_program,
+        sparse_tensor,
+        stablehlo,
+    )
 
 
 @pytest.mark.skipif(skip_torch_mlir_not_installed(), reason="torch_mlir not installed")
@@ -74,3 +80,26 @@ def test_torch_dialect_trampolines_smoke():
     )
     # noinspection PyUnresolvedReferences
     from mlir_utils.dialects import torch
+
+
+@pytest.mark.skipif(skip_jax_not_installed(), reason="jax not installed")
+def test_jax_trampolines_smoke():
+    for mod in ["chlo", "mhlo", "stablehlo"]:
+        try:
+            modu = __import__(f"mlir_utils.dialects.{mod}", fromlist=["*"])
+            os.remove(modu.__file__)
+        except ModuleNotFoundError:
+            pass
+        generate_trampolines(
+            f"jaxlib.mlir.dialects.{mod}", Path(mlir_utils.dialects.__path__[0]), mod
+        )
+    # noinspection PyUnresolvedReferences
+    from mlir_utils.dialects import (
+        builtin,
+        chlo,
+        func,
+        mhlo,
+        ml_program,
+        sparse_tensor,
+        stablehlo,
+    )

tests/util.py

@@ -0,0 +1,22 @@
+def skip_torch_mlir_not_installed():
+    try:
+        from torch_mlir.dialects import torch
+
+        # don't skip
+        return False
+
+    except ImportError:
+        # skip
+        return True
+
+
+def skip_jax_not_installed():
+    try:
+        from jaxlib import mlir
+
+        # don't skip
+        return False
+
+    except ImportError:
+        # skip
+        return True