diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7f1e00f..040baf2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,40 +3,29 @@ jobs: test: defaults: run: - shell: bash -el {0} + shell: bash strategy: matrix: os: [ubuntu-latest] - python: ['3.11', '3.12', '3.13'] + environment: [test-py311, test-py312, test-py313] include: - os: macos-latest - python: '3.11' + environment: test-py311 - os: windows-latest - python: '3.11' + environment: test-py311 fail-fast: false runs-on: ${{ matrix.os }} steps: - name: Checkout Repo - uses: actions/checkout@v3 - - uses: actions/setup-python@v5 + uses: actions/checkout@v7 + - uses: prefix-dev/setup-pixi@v0.10.0 with: - python-version: ${{ matrix.python }} - - name: Install Poetry - uses: snok/install-poetry@v1 - - name: Configure Poetry - run: | - if [[ "${{ runner.os }}" == "Windows" ]]; then - poetry env use python - else - poetry env use python${{ matrix.python }} - fi - - name: Install package - run: | - poetry install --extras test + pixi-version: v0.72.2 + environments: ${{ matrix.environment }} - name: Run tests - run: | - poetry run pytest --junit-xml=test-${{ matrix.os }}-Python-${{ matrix.python }}.xml - - uses: codecov/codecov-action@v3 + # junit and coverage reports go where pytest.ini says (junit/, coverage.xml) + run: pixi run -e ${{ matrix.environment }} test + - uses: codecov/codecov-action@v7 on: # Trigger the workflow on push or pull request, diff --git a/.github/workflows/sysimage.yml b/.github/workflows/sysimage.yml new file mode 100644 index 0000000..fa246cb --- /dev/null +++ b/.github/workflows/sysimage.yml @@ -0,0 +1,53 @@ +name: Sysimage + +# Builds the Julia sysimage for each platform and publishes it to the GitHub +# release tagged `sysimage-`, where identifies the exact Julia +# environment. Users download it with `pixi run fetch-sysimage`. +on: + workflow_dispatch: + push: + branches: + - main + paths: + - src/galley_jl_python/juliapkg.json + - scripts/sysimage/** + - pyproject.toml + - pixi.lock + - .github/workflows/sysimage.yml + +permissions: + contents: write + +jobs: + build: + strategy: + matrix: + os: [ubuntu-latest, ubuntu-24.04-arm, macos-latest, windows-latest] + fail-fast: false + runs-on: ${{ matrix.os }} + # GitHub-hosted maximum; the generic multi-CPU build is slow + timeout-minutes: 360 + defaults: + run: + shell: bash + steps: + - name: Checkout Repo + uses: actions/checkout@v7 + - uses: prefix-dev/setup-pixi@v0.10.0 + with: + pixi-version: v0.72.2 + - name: Build sysimage + id: build + run: pixi run build-sysimage + - name: Publish to release + env: + GH_TOKEN: ${{ github.token }} + TAG: ${{ steps.build.outputs.tag }} + IMAGE: ${{ steps.build.outputs.path }} + run: | + # the matrix jobs race to create the release; losing the race is fine + gh release view "$TAG" > /dev/null 2>&1 \ + || gh release create "$TAG" --prerelease --title "$TAG" \ + --notes "Prebuilt Julia sysimages. Fetch with \`pixi run fetch-sysimage\`." \ + || gh release view "$TAG" > /dev/null + gh release upload "$TAG" "$IMAGE" --clobber diff --git a/.gitignore b/.gitignore index 2b2b89f..52ec6f2 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,6 @@ cython_debug/ .pixi/ *.egg-info pixi.lock + +# PackageCompiler environment used by scripts/sysimage/build_sysimage.jl +scripts/sysimage/.build-env/ diff --git a/README.md b/README.md index 9f01527..90e80e5 100644 --- a/README.md +++ b/README.md @@ -16,18 +16,35 @@ pip install galley-jl-python ## Contributing -### Packaging +### Development environment -Galley uses [poetry](https://python-poetry.org/) for packaging. +Galley uses [pixi](https://pixi.sh) for development. All configuration lives in +`pyproject.toml`: the package metadata and dependencies are under `[project]`, +and pixi's settings are under `[tool.pixi]`. -To install for development, clone the repository and run: +To set up, clone the repository and run: ```bash -poetry install --with test +pixi install # the default environment +pixi install -e test # adds the test dependencies (the `test` extra) ``` -to install the current project and dev dependencies. + +pixi installs the package in editable mode. Julia is not a pixi dependency: +[juliapkg](https://github.com/JuliaPy/pyjuliapkg) installs the pinned Julia +version and packages the first time `galley_jl_python` is imported. To trigger +that, and to fetch the sysimage (see below), run: +```bash +pixi run compile +``` + +Run any other command inside an environment with `pixi run`, for example +`pixi run -e test python`. + +The package is still built and published with Poetry (see +[Publishing](#publishing)). ### Working with a local copy of Finch.jl -The `develop.py ` script can be used to set up a local copy of Finch.jl for development. +The `develop.py` script can be used to set up a local copy of Finch.jl for +development. Run it with `pixi run python develop.py`. ``` Usage: @@ -38,6 +55,36 @@ Options: --path Path to the local copy of Finch.jl [default: ../Finch.jl]. ``` +### Julia sysimage + +Most of Galley's startup time is Julia compiling Finch itself. A prebuilt Julia +sysimage removes it: the first operations of a session drop from minutes to a +few seconds. + +- `pixi run fetch-sysimage` downloads the image for your platform into + `~/.cache/galley-jl-python/` (`GALLEY_JL_PYTHON_CACHE` overrides this). + `pixi run compile` and `pixi run test` run this step first. +- `pixi run build-sysimage` builds the image locally instead. This takes hours. +- `import galley_jl_python` loads a cached image automatically when it matches + the current Julia environment. Otherwise Julia starts without it. Set + `GALLEY_JL_PYTHON_SYSIMAGE=0` to turn this off. + +An image works only with the exact Julia and package versions it was built +from. For this reason `src/galley_jl_python/juliapkg.json` pins Julia and every +Julia package. Each image's name includes a hash of that environment. A local +Finch.jl from `develop.py` therefore runs without the image. + +To update the Julia dependencies: + +1. Loosen the pins you want to change. +2. Resolve with `pixi run compile`. +3. Re-pin with `python scripts/sysimage/pin_julia_deps.py`. + +Pushing the new pins to `main` runs the "Sysimage" GitHub Action. It builds +images for Linux, macOS and Windows and publishes them to a `sysimage-` +GitHub release, where `fetch-sysimage` finds them. The action can also be run +manually from the Actions tab. + ### Publishing The "Publish" GitHub Action is a manual workflow for publishing Python packages to PyPI using Poetry. It handles the version management based on the `pyproject.toml` file and automates tagging and creating GitHub releases. @@ -61,18 +108,41 @@ On successful execution, the action publishes the package to PyPI and tags the r ### Pre-commit hooks -To add pre-commit hooks, run: +The hooks run in their own `pre-commit` environment, which doesn't install +galley's dependencies or Julia. + ```bash -poetry run pre-commit install +pixi run pre-commit-install # run the hooks on every `git commit` +pixi run pre-commit # run every hook on every file now +pixi run pre-commit ruff # the same, skipping the listed hook ids ``` +To run a single hook, use `pixi shell -e pre-commit` and then +`pre-commit run -a`. + ### Testing -Finch uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the +Galley uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the tests: ```bash -poetry run pytest +pixi run test # one pytest-xdist worker per CPU +pixi run test 4 # or a fixed number of workers +``` + +This runs `compile` first, then two suites one after the other, each spread +over the workers: + +- `pixi run test-unit` runs the unit tests. +- `pixi run test-array-api` runs the Array API tests described below. + +Each worker is a separate Julia process that uses 1–2 GB of memory, so lower +the worker count on machines with little memory. + +To run a subset, call pytest in the test environment directly: + +```bash +pixi run -e test pytest tests/test_fused.py ``` Array API tests are included in `tests/test_array_api.py`. These tests invoke @@ -82,7 +152,7 @@ To forward `pytest` options to the nested `--array-api-pytest-args`): ```bash -poetry run pytest tests/test_array_api.py \ +pixi run -e test pytest tests/test_array_api.py \ --array-api="-k creation_functions" \ --array-api="-x" ``` diff --git a/array-api-skips.txt b/array-api-skips.txt index 797379d..29a70d4 100644 --- a/array-api-skips.txt +++ b/array-api-skips.txt @@ -6,6 +6,12 @@ array_api_tests/test_special_cases.py::test_unary[sign((x_i is -0 or x_i == +0)) array_api_tests/test_searching_functions.py::test_where # `test_solve` is not defined in Finch, hangs as xfail array_api_tests/test_linalg.py::test_solve +# Galley drops the array's shape for `x & False`, `x * 0` and other annihilating +# constants (https://github.com/finch-tensor/Finch.jl/issues/820), and binary +# functions don't accept a scalar as the first argument. Remove these once both +# are fixed. +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_bool[logical_and] +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_bitwise[bitwise_and] # test_signatures @@ -249,6 +255,7 @@ array_api_tests/test_indexing_functions.py::test_take # test_linalg # not implemented +array_api_tests/test_linalg.py::test_cross array_api_tests/test_linalg.py::test_matrix_transpose array_api_tests/test_linalg.py::test_vecdot array_api_tests/test_linalg.py::test_eigh diff --git a/pixi.toml b/pixi.toml deleted file mode 100644 index fcdcdfa..0000000 --- a/pixi.toml +++ /dev/null @@ -1,35 +0,0 @@ -[project] -authors = ["Hameer Abbasi <2190658+hameerabbasi@users.noreply.github.com>"] -channels = ["conda-forge"] -description = "Add a short description here" -name = "galley-jl-python" -platforms = ["osx-arm64"] -version = "0.2.12" - -[tasks] -compile = "python -c 'import galley_jl_python'" - -[dependencies] -python = ">=3.10,<3.13" -juliaup = ">=1.17.10,<2" - -[pypi-dependencies] -galley-jl-python = { path = ".", editable = true } -juliapkg = ">=0.1.16,<0.2" -juliacall = ">=0.9.24,<0.10" -numpy = ">=1.19" - -[feature.test.pypi-dependencies] -pytest = "*" -pytest-cov = "*" -sparse = ">=0.16,<0.17" -numba = ">=0.61" -scipy = "*" -numpy = "==2.*" -pytest-xdist = ">=3.6.1,<4" - -[feature.test.tasks] -test = { cmd = "pytest", depends-on = ["compile"] } - -[environments] -test = ["test"] diff --git a/pyproject.toml b/pyproject.toml index 915143b..ad0cddb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "numba>=0.61,<0.63.1", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", + "finch-tensor>=0.7.0,<0.8.0", ] [tool.poetry] @@ -20,19 +21,97 @@ requires-poetry = ">=2.0,<3.0" [project.optional-dependencies] test = [ "pytest>=7.4.4,<7.6", - "pre-commit>=3.6.0,<3.9", "pytest-cov>=4.1.0,<4.2", "sparse>=0.17.0,<0.17.1", "scipy>=1.7,<1.17", "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", + "pytest-xdist>=3.6.1,<4", ] [build-system] requires = ["poetry-core>=1.0.8"] build-backend = "poetry.core.masonry.api" +[tool.pixi.workspace] +channels = ["conda-forge"] +platforms = [ + { name = "linux-amd64", platform = "linux-64", glibc = "2.34" }, + { name = "linux-arm64", platform = "linux-aarch64", glibc = "2.34" }, + { name = "macos-arm64", platform = "osx-arm64", macos = "15.0" }, + { name = "win-amd64", platform = "win-64" }, +] + +[tool.pixi.dependencies] +python = ">=3.11,<3.14" + +[tool.pixi.pypi-dependencies] +galley-jl-python = { path = ".", editable = true } + +[tool.pixi.tasks] +# pixi has no post-install hook, so the entry-point tasks fetch the sysimage +compile = { cmd = "python -c 'import galley_jl_python'", depends-on = ["fetch-sysimage"] } +build-sysimage = "python scripts/sysimage/build.py" +fetch-sysimage = "python scripts/sysimage/fetch.py" + +# `pixi run test [njobs]` runs both suites, one after the other, each spread +# over `njobs` pytest-xdist workers (default: one per CPU). +[tool.pixi.feature.test.tasks.test] +args = [{ arg = "njobs", default = "auto" }] +# a no-op cmd, since `default-environment` is only allowed on tasks with one +cmd = "echo all test suites passed" +default-environment = "test" +depends-on = [ + { task = "test-unit", args = ["{{ njobs }}"] }, + { task = "test-array-api", args = ["{{ njobs }}"] }, +] + +[tool.pixi.feature.test.tasks.test-unit] +args = [{ arg = "njobs", default = "auto" }] +default-environment = "test" +cmd = "pytest -n {{ njobs }} --ignore=tests/test_array_api.py" +depends-on = ["compile"] + +# The Array API wrapper is a single test that runs array-api-tests in a nested +# pytest, so the workers go to the nested run. It skips coverage and writes its +# own junit file so it doesn't overwrite test-unit's reports. +[tool.pixi.feature.test.tasks.test-array-api] +args = [{ arg = "njobs", default = "auto" }] +default-environment = "test" +cmd = "pytest tests/test_array_api.py --no-cov --junitxml=junit/array-api-results.xml --array-api='-n {{ njobs }}'" +depends-on = ["compile"] + +# Hooks run in their own environment, without galley's dependencies or Julia. +[tool.pixi.feature.pre-commit.dependencies] +pre-commit = "*" + +[tool.pixi.feature.pre-commit.tasks] +pre-commit-install = "pre-commit install" + +# `pixi run pre-commit [skip]` runs every hook on every file; `skip` is a +# comma-separated list of hook ids to leave out. +[tool.pixi.feature.pre-commit.tasks.pre-commit] +args = [{ arg = "skip", default = "" }] +env = { SKIP = "{{ skip }}" } +cmd = "pre-commit run -a" + +[tool.pixi.feature.py311.dependencies] +python = "3.11.*" + +[tool.pixi.feature.py312.dependencies] +python = "3.12.*" + +[tool.pixi.feature.py313.dependencies] +python = "3.13.*" + +[tool.pixi.environments] +test = { features = ["test"], solve-group = "default" } +test-py311 = ["test", "py311"] +test-py312 = ["test", "py312"] +test-py313 = ["test", "py313"] +pre-commit = { features = ["pre-commit"], no-default-feature = true } + [tool.ruff] line-length = 88 @@ -61,4 +140,3 @@ section-order = [ [tool.mypy] ignore_missing_imports = true -exclude = ["tests/reference"] diff --git a/scripts/sysimage/_load.py b/scripts/sysimage/_load.py new file mode 100644 index 0000000..cd04e5b --- /dev/null +++ b/scripts/sysimage/_load.py @@ -0,0 +1,15 @@ +"""Load ``galley_jl_python._sysimage`` by path. + +Importing it through the package would run ``galley_jl_python/__init__.py``, +which starts Julia. +""" + +import importlib.util +from pathlib import Path + +_PATH = ( + Path(__file__).resolve().parents[2] / "src" / "galley_jl_python" / "_sysimage.py" +) +_spec = importlib.util.spec_from_file_location("_galley_sysimage", _PATH) +sysimage = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(sysimage) diff --git a/scripts/sysimage/build.py b/scripts/sysimage/build.py new file mode 100644 index 0000000..950a39c --- /dev/null +++ b/scripts/sysimage/build.py @@ -0,0 +1,79 @@ +"""Build a Julia sysimage that removes most of galley's first-call compile time. + +Traces the Julia methods compiled by ``workload.py``, compiles them into a +sysimage with PackageCompiler, and stores it in the per-user cache under a name +tied to the current Julia environment, where ``galley_jl_python`` finds it. +""" + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import juliapkg + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from _load import sysimage # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--statements", + type=Path, + help="reuse precompile statements from an earlier trace instead of tracing", + ) + args = parser.parse_args() + + project = juliapkg.project() + key = sysimage.env_key(project) + image = sysimage.image_path(key) + image.parent.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory() as tmp: + statements = args.statements or Path(tmp) / "precompile_statements.jl" + if args.statements is None: + print(f"[sysimage] tracing workload -> {statements}", flush=True) + env = { + k: v for k, v in os.environ.items() if k != "PYTHON_JULIACALL_SYSIMAGE" + } + subprocess.run( + [sys.executable, str(HERE / "workload.py")], + env={ + **env, + "GALLEY_JL_PYTHON_SYSIMAGE": "0", + "PYTHON_JULIACALL_TRACE_COMPILE": str(statements), + }, + check=True, + ) + + # build next to the target and rename, so an interrupted build never + # leaves a partial image where it would be loaded + partial = Path(tmp) / image.name + print(f"[sysimage] building {image.name} (this takes a while)", flush=True) + subprocess.run( + [ + juliapkg.executable(), + "--startup-file=no", + str(HERE / "build_sysimage.jl"), + project, + str(statements), + str(partial), + ], + check=True, + ) + shutil.move(partial, image) + + print(f"[sysimage] wrote {image}", flush=True) + if github_output := os.environ.get("GITHUB_OUTPUT"): + with open(github_output, "a") as f: + f.write(f"key={key}\npath={image}\ntag={sysimage.release_tag(key)}\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/build_sysimage.jl b/scripts/sysimage/build_sysimage.jl new file mode 100644 index 0000000..fcd8b59 --- /dev/null +++ b/scripts/sysimage/build_sysimage.jl @@ -0,0 +1,26 @@ +# Usage: julia build_sysimage.jl +# +# PackageCompiler lives in its own environment so it never enters the +# project that juliapkg manages. +using Pkg + +project, statements_file, output = ARGS + +build_env = joinpath(@__DIR__, ".build-env") +Pkg.activate(build_env; io=devnull) +if Base.find_package("PackageCompiler") === nothing + Pkg.add("PackageCompiler") +end +using PackageCompiler + +packages = ["Finch", "HDF5", "NPZ", "TensorMarket", "PythonCall"] + +create_sysimage( + packages; + project=project, + sysimage_path=output, + precompile_statements_file=statements_file, + # a generic target so the image runs on any x86_64 CPU, not just this one + cpu_target=PackageCompiler.default_app_cpu_target(), + include_transitive_dependencies=true, +) diff --git a/scripts/sysimage/fetch.py b/scripts/sysimage/fetch.py new file mode 100644 index 0000000..725569e --- /dev/null +++ b/scripts/sysimage/fetch.py @@ -0,0 +1,58 @@ +"""Download the prebuilt sysimage for this platform and Julia environment. + +CI publishes one image per platform to the GitHub release tagged +``sysimage-``, where the key identifies the exact Julia environment. +""" + +import os +import sys +import tempfile +import urllib.error +import urllib.request +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from _load import sysimage # noqa: E402 + + +def main() -> None: + key = sysimage.current_key() + image = sysimage.image_path(key) + if image.is_file(): + print(f"[sysimage] already present: {image}") + return + + url = f"{sysimage.RELEASES_URL}/{sysimage.release_tag(key)}/{image.name}" + print(f"[sysimage] downloading {url}", flush=True) + image.parent.mkdir(parents=True, exist_ok=True) + fd, partial = tempfile.mkstemp(dir=image.parent, suffix=".part") + try: + with os.fdopen(fd, "wb") as out, urllib.request.urlopen(url) as resp: + while chunk := resp.read(1 << 20): + out.write(chunk) + os.replace(partial, image) + # never fail: other tasks depend on this one, and galley works without an + # image, just with a slower first call + except urllib.error.HTTPError as e: + if e.code != 404: + print(f"[sysimage] download failed ({e}); continuing without it.") + return + print( + f"[sysimage] no prebuilt image for this platform and Julia environment " + f"({image.name}); continuing without it. Build one locally with " + "`pixi run build-sysimage`." + ) + return + except (urllib.error.URLError, OSError) as e: + print(f"[sysimage] download failed ({e}); continuing without it.") + return + finally: + if os.path.exists(partial): + os.remove(partial) + print(f"[sysimage] saved {image}") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/pin_julia_deps.py b/scripts/sysimage/pin_julia_deps.py new file mode 100644 index 0000000..3a0c25d --- /dev/null +++ b/scripts/sysimage/pin_julia_deps.py @@ -0,0 +1,66 @@ +"""Pin galley's Julia dependencies to the currently resolved environment. + +Rewrites ``src/galley_jl_python/juliapkg.json`` so that Julia and every +registered package in the resolved ``Manifest.toml`` (transitive dependencies +included) are pinned to exact versions. Everyone then resolves the same +manifest, which the prebuilt sysimage requires. + +Packages pinned to a git commit (``url`` and ``rev`` in juliapkg.json) are +left as they are; change their ``rev`` by hand. + +To update: loosen the pins you want to move, resolve (e.g. ``pixi run +compile``), run this script, then rebuild with ``pixi run build-sysimage``. +""" + +import argparse +import json +import sys +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +JULIAPKG_JSON = ROOT / "src" / "galley_jl_python" / "juliapkg.json" +# juliacall pins these itself (PythonCall to its own release, OpenSSL_jll to +# match Python's OpenSSL), so they must stay unpinned here +UNPINNED = {"OpenSSL_jll", "PythonCall"} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--manifest", + type=Path, + default=Path(sys.prefix) / "julia_env" / "Manifest.toml", + ) + args = parser.parse_args() + + manifest = tomllib.loads(args.manifest.read_text()) + config = json.loads(JULIAPKG_JSON.read_text()) + config["julia"] = f"={manifest['julia_version']}" + + git_pinned = { + name: spec for name, spec in config["packages"].items() if "rev" in spec + } + packages = {} + for name, (entry, *_) in sorted(manifest["deps"].items()): + if name in git_pinned: + packages[name] = git_pinned[name] + continue + # stdlibs have no tree hash and ship with Julia itself + if "git-tree-sha1" not in entry or name in UNPINNED: + continue + # drop build metadata (e.g. JLL "+0"): juliapkg can't parse it and Pkg + # compat ignores it anyway + version = entry["version"].split("+")[0] + packages[name] = {"uuid": entry["uuid"], "version": f"={version}"} + config["packages"] = packages + + JULIAPKG_JSON.write_text(json.dumps(config, indent=2, sort_keys=True) + "\n") + print( + f"pinned julia {config['julia']} and {len(packages)} packages " + f"({len(git_pinned)} to git commits: {', '.join(git_pinned) or 'none'})" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/workload.py b/scripts/sysimage/workload.py new file mode 100644 index 0000000..030167d --- /dev/null +++ b/scripts/sysimage/workload.py @@ -0,0 +1,62 @@ +"""Representative workload traced to build the Julia sysimage. + +Every Julia method compiled while this runs is recorded (via juliacall's +``trace_compile`` option) and baked into the sysimage by ``build.py``. Add +operations here to cover more of the library ahead of time. +""" + +import numpy as np + +import galley_jl_python as gl + + +def run_ops(A, B, C, v): + gl.matmul(A, B) + gl.sum(gl.matmul(gl.matmul(A, B), C)) + gl.add(A, B) + gl.multiply(A, B) + gl.subtract(A, B) + gl.sum(A, axis=0) + gl.sum(A, axis=1) + gl.max(A) + gl.tensordot(A, v, axes=1) + gl.permute_dims(A, (1, 0)) + A[0] + A[0, 0] + A[1:3, :] + A.todense() + + @gl.jit + def sum_abc(A, B, C): + return gl.sum(gl.matmul(gl.matmul(A, B), C)) + + @gl.jit + def chain(A, B, C): + return gl.add(gl.matmul(A, B), C) + + sum_abc(A, B, C) + chain(A, B, C) + + +def main(): + rng = np.random.default_rng(0) + for scheduler in (gl.GalleyScheduler(), gl.DefaultScheduler()): + gl.set_optimizer(scheduler) + for n in (50, 400): + A, B, C = ( + gl.random((n, n), density=0.02, random_state=s) for s in (1, 2, 3) + ) + v = gl.asarray(rng.random(n)) + run_ops(A, B, C, v) + for fmt in ("coo", "csr", "csc", "dense"): + for dtype in (np.float64, np.int64): + dense = (rng.random((60, 60)) < 0.05).astype(dtype) + X = gl.asarray(dense, format=fmt) + gl.matmul(X, X) + gl.sum(X) + gl.add(X, X) + gl.set_optimizer(gl.GalleyScheduler()) + + +if __name__ == "__main__": + main() diff --git a/src/galley_jl_python/__init__.py b/src/galley_jl_python/__init__.py index 5658656..0a19a2c 100644 --- a/src/galley_jl_python/__init__.py +++ b/src/galley_jl_python/__init__.py @@ -115,6 +115,7 @@ uint32, uint64, ) +from .finch_fused import jit from .io import ( read, write, @@ -292,6 +293,7 @@ "isfinite", "isinf", "isnan", + "jit", "lazy", "less", "less_equal", diff --git a/src/galley_jl_python/_sysimage.py b/src/galley_jl_python/_sysimage.py new file mode 100644 index 0000000..a96ffbe --- /dev/null +++ b/src/galley_jl_python/_sysimage.py @@ -0,0 +1,79 @@ +"""Locate a prebuilt Julia sysimage and point juliacall at it. + +This module must not import juliacall: it runs before Julia starts. + +A sysimage only works with the exact Julia version and package versions it was +built from, so each image is named after a hash of both (see ``env_key``) and +lives in a per-user cache directory. An image is used only when its name +matches the current environment; otherwise Julia starts normally. + +Get an image with ``pixi run fetch-sysimage`` (downloads the one CI built for +this environment) or ``pixi run build-sysimage`` (builds it locally). Set +``GALLEY_JL_PYTHON_SYSIMAGE=0`` to never load one, or +``PYTHON_JULIACALL_SYSIMAGE`` to use a specific image. +""" + +import hashlib +import json +import os +import platform +import sys +import tomllib +from pathlib import Path + +RELEASES_URL = "https://github.com/finch-tensor/galley-jl-python/releases/download" +_EXTENSIONS = {"linux": "so", "darwin": "dylib", "win32": "dll"} + + +def cache_dir() -> Path: + if custom := os.environ.get("GALLEY_JL_PYTHON_CACHE"): + return Path(custom) + if sys.platform == "win32": + base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local")) + elif sys.platform == "darwin": + base = Path.home() / "Library" / "Caches" + else: + base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + return base / "galley-jl-python" + + +def env_key(project: str) -> str: + """Hash the Julia version and resolved package set of a juliapkg project.""" + manifest = tomllib.loads((Path(project) / "Manifest.toml").read_text()) + resolved = sorted( + (name, entry.get("uuid"), entry.get("version"), entry.get("git-tree-sha1")) + for name, entries in manifest["deps"].items() + for entry in entries + ) + payload = json.dumps([manifest["julia_version"], resolved]) + return hashlib.sha256(payload.encode()).hexdigest()[:16] + + +def release_tag(key: str) -> str: + return f"sysimage-{key}" + + +def image_name(key: str) -> str: + ext = _EXTENSIONS.get(sys.platform, "so") + return f"galley-{sys.platform}-{platform.machine().lower()}-{key}.{ext}" + + +def image_path(key: str) -> Path: + return cache_dir() / image_name(key) + + +def current_key() -> str: + import juliapkg + + return env_key(juliapkg.project()) + + +def configure() -> None: + """Point juliacall at the cached sysimage for this environment, if any.""" + if os.environ.get("GALLEY_JL_PYTHON_SYSIMAGE") == "0": + return + if "PYTHON_JULIACALL_SYSIMAGE" in os.environ: + return + image = image_path(current_key()) + if image.is_file(): + os.environ["PYTHON_JULIACALL_SYSIMAGE"] = str(image) diff --git a/src/galley_jl_python/compiled.py b/src/galley_jl_python/compiled.py index af05973..83907ad 100644 --- a/src/galley_jl_python/compiled.py +++ b/src/galley_jl_python/compiled.py @@ -141,6 +141,9 @@ def set_optimizer(opt: AbstractScheduler) -> None: jl.Finch.set_scheduler_b(opt.get_julia_scheduler()) +set_optimizer(GalleyScheduler()) + + def lazy(tensor: Tensor) -> Tensor: from .tensor import Tensor diff --git a/src/galley_jl_python/finch_fused.py b/src/galley_jl_python/finch_fused.py new file mode 100644 index 0000000..2f64efc --- /dev/null +++ b/src/galley_jl_python/finch_fused.py @@ -0,0 +1,33 @@ +from collections.abc import Callable +from typing import Any + +from finch.finch_fused import jit as _finch_jit + +from .compiled import compute, lazy +from .tensor import Tensor + + +def _defer(value: Any) -> Any: + if isinstance(value, tuple): + return tuple(_defer(item) for item in value) + return lazy(value) if isinstance(value, Tensor) else value + + +def _compute(value: Any) -> Any: + if isinstance(value, tuple): + return tuple(_compute(item) for item in value) + return compute(value) if isinstance(value, Tensor) else value + + +def _maybedefer(values: tuple[Any, ...]) -> tuple[Any, ...]: + return tuple(_defer(value) for value in values) + + +def jit(f: Callable[..., Any], /, ctx: Any = None) -> Callable[..., Any]: + transformed = _finch_jit(f, ctx=ctx) + transformed.__globals__.update( + defer=_defer, + compute=_compute, + maybedefer=_maybedefer, + ) + return transformed diff --git a/src/galley_jl_python/julia.py b/src/galley_jl_python/julia.py index c8f2191..7ac7934 100644 --- a/src/galley_jl_python/julia.py +++ b/src/galley_jl_python/julia.py @@ -1,10 +1,17 @@ import juliapkg # noqa: I001, F401 +from . import _sysimage + +_sysimage.configure() + # To change the version of Finch used, see the documentation for pyjuliapkg here: https://github.com/JuliaPy/pyjuliapkg # Use pyjuliapkg to modify the `juliapkg.json` file in the root of this repo. # You can also run `develop.py` to quickly use a local copy of Finch.jl. # An example development json is found in `juliapkg_dev.json` -import juliacall as jc # noqa: F401 +# Julia and all packages are pinned so the prebuilt sysimage matches; update the +# pins with `scripts/sysimage/pin_julia_deps.py`. Finch is pinned to a git +# commit (`rev`), which that script leaves alone. +import juliacall as jc # noqa: E402, F401 from juliacall import Main as jl # noqa: E402, F401 jl.seval("using Finch") diff --git a/src/galley_jl_python/juliapkg.json b/src/galley_jl_python/juliapkg.json index ad999ee..84c5be3 100644 --- a/src/galley_jl_python/juliapkg.json +++ b/src/galley_jl_python/juliapkg.json @@ -1,21 +1,386 @@ { - "julia": "1.10 - 1.12", + "julia": "=1.12.2", "packages": { + "AbstractTrees": { + "uuid": "1520ce14-60c1-5f80-bbc7-55ef81b5835c", + "version": "=0.4.5" + }, + "AcceleratedKernels": { + "uuid": "6a4ca0a5-0e36-4168-a932-d9be78d558f1", + "version": "=0.4.3" + }, + "Accessors": { + "uuid": "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697", + "version": "=0.1.45" + }, + "Adapt": { + "uuid": "79e6a3ab-5dfb-504d-930d-738a2a938a0e", + "version": "=4.7.1" + }, + "AliasTables": { + "uuid": "66dad0bd-aa9a-41b7-9441-69ab47430ed8", + "version": "=1.1.3" + }, + "ArgCheck": { + "uuid": "dce04be8-c92d-5529-be00-80e4d2c0e197", + "version": "=2.5.0" + }, + "Atomix": { + "uuid": "a9b6321e-bd34-4604-b9c9-b65b8de01458", + "version": "=1.2.1" + }, + "AutoHashEquals": { + "uuid": "15f4f7f2-30c1-5605-9d31-71845cf9641f", + "version": "=2.2.0" + }, + "CommonSolve": { + "uuid": "38540f10-b2f7-11e9-35d8-d573e4eb0ff2", + "version": "=0.2.14" + }, + "Compat": { + "uuid": "34da2185-b29b-5c13-b0c7-acf172513d20", + "version": "=4.18.1" + }, + "CompositionsBase": { + "uuid": "a33af91c-f02d-484b-be07-31d278c5ca2b", + "version": "=0.1.2" + }, + "CondaPkg": { + "uuid": "992eb4ea-22a4-4c89-a5bb-47a3300528ab", + "version": "=0.2.36" + }, + "ConstructionBase": { + "uuid": "187b0558-2788-49d3-abe0-74a17ed4e7c9", + "version": "=1.6.0" + }, + "DataAPI": { + "uuid": "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a", + "version": "=1.16.0" + }, + "DataStructures": { + "uuid": "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8", + "version": "=0.19.6" + }, + "DataValueInterfaces": { + "uuid": "e2d170a0-9d28-54be-80f0-106bbe20a464", + "version": "=1.0.0" + }, + "Distributions": { + "uuid": "31c24e10-a181-5473-b8eb-7969acd0382f", + "version": "=0.25.131" + }, + "DocStringExtensions": { + "uuid": "ffbed154-4ef7-542d-bbb7-c09d3a79fcae", + "version": "=0.9.5" + }, + "FileIO": { + "uuid": "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549", + "version": "=1.20.0" + }, + "FillArrays": { + "uuid": "1a297f60-69ca-5386-bcde-b61e274b549b", + "version": "=1.17.1" + }, "Finch": { - "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", - "version": "1.2.10" + "rev": "368602490c2a106778afc9ca9d567a2201393012", + "url": "https://github.com/finch-tensor/Finch.jl", + "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" + }, + "GPUArraysCore": { + "uuid": "46192b85-c4d5-4398-a991-12ede77f4527", + "version": "=0.2.1" + }, + "Gamma": { + "uuid": "a0844989-3bd2-4988-8bea-c9407ab0941b", + "version": "=1.2.0" }, "HDF5": { "uuid": "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f", - "version": "0.17.2" + "version": "=0.17.4" + }, + "HDF5_jll": { + "uuid": "0234f1f7-429e-5d53-9886-15a909be8d59", + "version": "=2.2.2" + }, + "Hwloc_jll": { + "uuid": "e33a78d0-f292-5ffc-b300-72abe9b543c8", + "version": "=2.14.1" + }, + "HypergeometricFunctions": { + "uuid": "34004b35-14d8-5ef3-9330-4cdb6864b03a", + "version": "=0.3.30" + }, + "InverseFunctions": { + "uuid": "3587e190-3f89-42d0-90ee-14403ec27112", + "version": "=0.1.17" + }, + "IrrationalConstants": { + "uuid": "92d709cd-6900-40b7-9082-c6be49f344b6", + "version": "=0.2.6" + }, + "IterTools": { + "uuid": "c8e1da08-722c-5040-9ed9-7db0dc04731e", + "version": "=1.10.0" + }, + "IteratorInterfaceExtensions": { + "uuid": "82899510-4779-5014-852e-03e436cf321d", + "version": "=1.0.0" + }, + "JLLWrappers": { + "uuid": "692b3bcd-3c85-4b1f-b108-f13ce0eb3210", + "version": "=1.8.0" + }, + "JSON": { + "uuid": "682c06a0-de6a-54ab-a142-c8b1cf79cde6", + "version": "=1.9.0" + }, + "KernelAbstractions": { + "uuid": "63c18a36-062a-441e-b654-da1e3ab1ce7c", + "version": "=0.9.42" + }, + "Libiconv_jll": { + "uuid": "94ce4f54-9a6c-5748-9c1c-f9c7231a4531", + "version": "=1.18.0" + }, + "LogExpFunctions": { + "uuid": "2ab3a3ac-af41-5b50-aa03-7779005ae688", + "version": "=1.0.2" + }, + "MPIABI_jll": { + "uuid": "b5ada748-db0f-5fc0-8972-9331c762740c", + "version": "=1.0.0" + }, + "MPICH_jll": { + "uuid": "7cb0a576-ebde-5e09-9194-50597f1243b4", + "version": "=5.0.1" + }, + "MPIPreferences": { + "uuid": "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267", + "version": "=0.1.12" + }, + "MPItrampoline_jll": { + "uuid": "f1f71cc9-e9ae-5b93-9b94-4fe0e1ad3748", + "version": "=5.5.6" + }, + "MacroTools": { + "uuid": "1914dd2f-81c6-5fcd-8719-6d5c9610ff09", + "version": "=0.5.16" + }, + "MicroMamba": { + "uuid": "0b3b1443-0f03-428d-bdfb-f27f9c1191ea", + "version": "=0.1.15" + }, + "MicrosoftMPI_jll": { + "uuid": "9237b28f-5490-5468-be7b-bb81f5f5e6cf", + "version": "=10.1.4" + }, + "Missings": { + "uuid": "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28", + "version": "=1.2.0" }, "NPZ": { "uuid": "15e1cf62-19b3-5cfa-8e77-841668bca605", - "version": "0.4.3" + "version": "=0.4.3" + }, + "OpenMPI_jll": { + "uuid": "fe0851c0-eecd-5654-98d4-656369965a5c", + "version": "=5.0.12" + }, + "OpenSpecFun_jll": { + "uuid": "efe28fd5-8261-553b-a9e1-b2916fc3738e", + "version": "=0.5.6" + }, + "OrderedCollections": { + "uuid": "bac558e1-5e72-5ebc-8fee-abe8a469f55d", + "version": "=2.0.1" + }, + "PDMats": { + "uuid": "90014a1f-27ba-587c-ab20-58faa44d9150", + "version": "=0.11.41" + }, + "Parsers": { + "uuid": "69de0a69-1ddd-5017-9359-2bf0b02dc9f0", + "version": "=3.0.0" + }, + "Pidfile": { + "uuid": "fa939f87-e72e-5be4-a000-7fc836dbe307", + "version": "=1.3.0" + }, + "PrecompileTools": { + "uuid": "aea7be01-6a6a-4083-8856-8a6e6704d82a", + "version": "=1.3.4" + }, + "Preferences": { + "uuid": "21216c6a-2e73-6563-6e65-726566657250", + "version": "=1.6.0" + }, + "PtrArrays": { + "uuid": "43287f4e-b6f4-7ad1-bb20-aadabca52c3d", + "version": "=1.4.0" + }, + "QuadGK": { + "uuid": "1fd47b50-473d-5c70-9696-f719f8f3bcdc", + "version": "=2.11.3" + }, + "Reexport": { + "uuid": "189a3867-3050-52da-a836-e630ba90ab69", + "version": "=1.2.2" + }, + "Requires": { + "uuid": "ae029012-a4dd-5104-9daa-d747884805df", + "version": "=1.3.1" + }, + "RewriteTools": { + "uuid": "5969e224-3634-4c61-9f66-721b69e98b8a", + "version": "=0.2.4" + }, + "Rmath": { + "uuid": "79098fc4-a85e-5d69-aa6a-4863f24498fa", + "version": "=0.9.0" + }, + "Rmath_jll": { + "uuid": "f50d1b31-88e8-58de-be2c-1cc44531875f", + "version": "=0.5.2" + }, + "Roots": { + "uuid": "f2b01f46-fcfa-551c-844a-d8ac1e96c665", + "version": "=3.0.8" + }, + "Scratch": { + "uuid": "6c6a2e73-6563-6170-7368-637461726353", + "version": "=1.3.0" + }, + "SortingAlgorithms": { + "uuid": "a2af1166-a08f-5f64-846c-94a0d3cef48c", + "version": "=1.2.3" + }, + "SpecialFunctions": { + "uuid": "276daf66-3868-5448-9aa4-cd146d93841b", + "version": "=2.9.0" + }, + "StaticArrays": { + "uuid": "90137ffa-7385-5640-81b9-e52037218182", + "version": "=1.9.22" + }, + "StaticArraysCore": { + "uuid": "1e83bf80-4336-4d27-bf5d-d5a4f845583c", + "version": "=1.4.4" + }, + "Statistics": { + "uuid": "10745b16-79ce-11e8-11f9-7d13ad32a3b2", + "version": "=1.11.5" + }, + "StatsAPI": { + "uuid": "82ae8749-77ed-4fe6-ae5f-f523153014b0", + "version": "=1.8.0" + }, + "StatsBase": { + "uuid": "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91", + "version": "=0.34.13" + }, + "StatsFuns": { + "uuid": "4c63d2b9-4356-54db-8cca-17b64c39e42c", + "version": "=2.2.1" + }, + "StructUtils": { + "uuid": "ec057cc2-7a8d-4b58-b3b3-92acb9f63b42", + "version": "=2.9.2" + }, + "SyntaxInterface": { + "uuid": "b33eeca9-aacb-4496-a840-e75f1646a4fb", + "version": "=0.2.1" + }, + "TableTraits": { + "uuid": "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c", + "version": "=1.0.1" + }, + "Tables": { + "uuid": "bd369af6-aec1-5ad0-b16a-f7cc5008161c", + "version": "=1.14.0" }, "TensorMarket": { "uuid": "8b7d4fe7-0b45-4d0d-9dd8-5cc9b23b4b77", - "version": "0.2.0" + "version": "=0.2.2" + }, + "UnsafeAtomics": { + "uuid": "013be700-e6cd-48c3-b4a1-df204f14c38f", + "version": "=0.3.2" + }, + "UnsafePointers": { + "uuid": "e17b2a0c-0bdf-430a-bd0c-3a23cae4ff39", + "version": "=1.0.0" + }, + "XML2_jll": { + "uuid": "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a", + "version": "=2.15.4" + }, + "Xorg_libpciaccess_jll": { + "uuid": "a65dc6b1-eb27-53a1-bb3e-dea574b5389e", + "version": "=0.19.0" + }, + "ZipFile": { + "uuid": "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea", + "version": "=0.10.1" + }, + "aws_c_auth_jll": { + "uuid": "2b3700d1-4306-52e2-a478-c162f0c514be", + "version": "=0.9.6" + }, + "aws_c_cal_jll": { + "uuid": "70f11efc-bab2-57f1-b0f3-22aad4e67c4b", + "version": "=0.9.13" + }, + "aws_c_common_jll": { + "uuid": "73048d1d-b8c4-5092-a58d-866c5e8d1e50", + "version": "=0.12.6" + }, + "aws_c_compression_jll": { + "uuid": "73a04cd5-f3d7-5bac-9290-e8adb709f224", + "version": "=0.3.2" + }, + "aws_c_http_jll": { + "uuid": "3254fc65-9028-534d-aa9d-d76d128babc6", + "version": "=0.10.15" + }, + "aws_c_io_jll": { + "uuid": "13c41daa-f319-5298-b5eb-5754e0170d52", + "version": "=0.26.3" + }, + "aws_c_s3_jll": { + "uuid": "bd1f34fb-993f-5903-a121-aaf302eed6d4", + "version": "=0.11.5" + }, + "aws_c_sdkutils_jll": { + "uuid": "1282aa60-004d-510b-9f52-12498d409daa", + "version": "=0.2.4" + }, + "aws_checksums_jll": { + "uuid": "b2a88e68-78e7-5e94-8c20-c02986ec140e", + "version": "=0.2.10" + }, + "dlfcn_win32_jll": { + "uuid": "c4b69c83-5512-53e3-94e6-de98773c479f", + "version": "=1.4.2" + }, + "libaec_jll": { + "uuid": "477f73a3-ac25-53e9-8cc3-50b2fa2566f0", + "version": "=1.1.7" + }, + "micromamba_jll": { + "uuid": "f8abcde7-e9b7-5caa-b8af-a437887ae8e4", + "version": "=2.3.1" + }, + "mpif_jll": { + "uuid": "9aeb927a-4695-514f-a259-621a69f20ec0", + "version": "=1.0.0" + }, + "pixi_jll": { + "uuid": "4d7b5844-a134-5dcd-ac86-c8f19cd51bed", + "version": "=0.76.2" + }, + "s2n_tls_jll": { + "uuid": "cddc5d3d-934d-5d3a-9747-62fc12ea3f48", + "version": "=1.7.10" } } } diff --git a/src/galley_jl_python/tensor.py b/src/galley_jl_python/tensor.py index 06195c7..8f98335 100644 --- a/src/galley_jl_python/tensor.py +++ b/src/galley_jl_python/tensor.py @@ -25,6 +25,11 @@ ) from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix +# Builds a 0-d Finch tensor in Julia so the element keeps its Julia eltype. +_scalar_to_0d = jl.seval( + "(T, fv, x) -> Finch.Tensor(Finch.Element{convert(T, fv)}(T[convert(T, x)]))" +) + class SparseArray: """ @@ -279,7 +284,8 @@ def __getitem__(self, key): return Tensor(result) if jl.isa(result, jl.Finch.Tensor): return Tensor(jl.swizzle(result, *range(1, jl.ndims(result) + 1))) - return result + # the Array API requires a 0-d array, not a scalar, for integer-only keys + return Tensor(jl.swizzle(_scalar_to_0d(self.dtype, self.fill_value, result))) @property def dtype(self) -> DType: diff --git a/tests/conftest.py b/tests/conftest.py index 61800c9..14c22c7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,12 @@ import numpy as np +# Start Julia here, on the main thread, in every pytest process. Otherwise the +# pytest-xdist controller first imports galley when it rebuilds a +# PerformanceWarning sent by a worker, which starts Julia on execnet's receiver +# thread, and Julia's exit hook then deadlocks at interpreter shutdown. +import galley_jl_python # noqa: E402, F401 + def pytest_addoption(parser): parser.addoption( diff --git a/tests/test_fused.py b/tests/test_fused.py new file mode 100644 index 0000000..ac6a456 --- /dev/null +++ b/tests/test_fused.py @@ -0,0 +1,75 @@ +import numpy as np +from numpy.testing import assert_allclose + +import galley_jl_python as gl + + +def assert_galley_allclose(actual, expected): + assert isinstance(actual, gl.Tensor) + assert_allclose(actual.todense(), expected) + + +def test_jit_straight_line_uses_galley_tensors(): + @gl.jit + def opt_fn(A, B): + return gl.matmul(A, B) + + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[5, 6], [7, 8]])) + + assert_galley_allclose(opt_fn(A, B), np.matmul(A.todense(), B.todense())) + + +def test_jit_return_expr_uses_galley_tensors(): + @gl.jit + def opt_fn(A, B): + return gl.matmul(A, B), gl.add(A, B) + + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[5, 6], [7, 8]])) + + actual_matmul, actual_add = opt_fn(A, B) + assert_galley_allclose(actual_matmul, np.matmul(A.todense(), B.todense())) + assert_galley_allclose(actual_add, A.todense() + B.todense()) + + +def test_jit_scalar_loop_uses_galley_tensors(): + @gl.jit + def opt_fn(A, n): + B = A + for _i in range(n): + B = gl.add(B, A) + return B + + A = gl.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) + + assert_galley_allclose(opt_fn(A, 3), 4 * A.todense()) + + +def test_jit_if_branch_uses_galley_tensors(): + @gl.jit + def opt_fn(A, B, use_matmul): + if use_matmul: + return gl.matmul(A, B) + return gl.add(A, B) + + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[1, 0], [0, 1]])) + + assert_galley_allclose(opt_fn(A, B, True), np.matmul(A.todense(), B.todense())) + assert_galley_allclose(opt_fn(A, B, False), A.todense() + B.todense()) + + +def test_jit_while_loop_uses_galley_tensors(): + @gl.jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = gl.add(C, B) + n = n - 1 + return C + + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[1, 0], [0, 1]])) + + assert_galley_allclose(opt_fn(A, B, 3), A.todense() + 3 * B.todense())