Skip to content

Commit 8d6bed1

Browse files
authored
[Python] Basic dependency management with pyproject and uv (#6)
Sets us up to do a uv sync to set up the current Python environment with the correct dependencies. Note that this is just a facilitating mechanism for preparing the python environment; users are still free to set up the environment via any means they like. By default the only dependency installed is the Python bindings (from llvm/eudsl) -- that is, enabling import mlir in Python. This is the common dependency for generating and running pipelines and for ingress. Optionally one can pass (a variant of) --extra ingress-torch to uv sync to also get the torch and torch-mlir dependencies -- that is, enabling import torch and import torch_mlir in Python. Per the included docs, it is recommended to first set up a fresh virtualenv, e.g. via uv venv. The pyproject.toml comes with uv-specific sections so as to enable specifying from which "indices", i.e. Python package repositories, each specific package is expected to come from. This file pins particular nightlies/versions from these package repositories.
1 parent 20c7e72 commit 8d6bed1

File tree

5 files changed

+83
-73
lines changed

5 files changed

+83
-73
lines changed

README.md

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,22 @@ The planned work is:
4949
* **Schedule**: MLIR schedules to combine into pipelines and combination strategy on branch [TBD].
5050
* **Runtime**: Dependencies, tools and environments to run on various hardware on branch [TBD].
5151
* **Build System**: CMake / Bazel magic to check dependencies, track repositories, run CI on branch [TBD].
52+
53+
## Getting up and running
54+
55+
For the time being, `lighthouse` depends on just the Python bindings for [`mlir`](https://github.com/llvm/eudsl/releases).
56+
To install this dependency, obtain the [`uv`](https://docs.astral.sh/uv/getting-started/installation/#pypi) Python package manager and run the following in the root of the project:
57+
```
58+
$ uv venv # Create a .venv virtualenv
59+
$ uv sync # Install the `mlir-python-bindings` dependency into the virtualenv
60+
$ uv sync --extra ingress-torch-cpu # Optionally install the dependencies for torch ingress
61+
```
62+
63+
<details>
64+
<summary>
65+
A note on vendor-specific `torch` versions.
66+
</summary>
67+
For vendor-specific versions of `torch` use the targets `ingress-torch-nvidia`, `ingress-torch-rocm` or `ingress-torch-xpu` for Nvidia, AMD, and Intel-enabled versions, respectively.
68+
</details>
69+
70+
To run the Python programs in this repo, either enter the virtual environment (`$ source .venv/bin/activate`) and execute a program _or_ execute each of the programs through `uv` (i.e. `$ uv run $EXE`), which will automatically run them inside the virtualenv.

ingress/Torch-MLIR/install-virtualenv.sh

Lines changed: 0 additions & 58 deletions
This file was deleted.

ingress/mlir-gen/generate-linalg-3layer-mlp.sh

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
#!/usr/bin/env bash
22

3-
source mlir-gen-venv/bin/activate
4-
53
LAYERS=1024,2048,4096,512
64

75
mkdir -p cache

ingress/mlir-gen/install-virtualenv.sh

Lines changed: 0 additions & 13 deletions
This file was deleted.

pyproject.toml

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
[project]
2+
name = "lighthouse"
3+
version = "0.1.0a1"
4+
requires-python = ">=3.10,<3.13" # Bounds are due to torch-mlir's packaging
5+
dependencies = [
6+
"mlir-python-bindings==20251011+3af43e303",
7+
]
8+
9+
[project.optional-dependencies]
10+
# Additional "targets" which pull in optional dependencies -- use `uv sync --extra TARGET`
11+
ingress_torch_cpu = [
12+
"torch==v2.8.0+cpu",
13+
"torch-mlir==20251003.589"
14+
]
15+
ingress_torch_nvidia = [
16+
"torch==2.8.0", # Nvidia-enabled version of torch
17+
"torch-mlir==20251003.589"
18+
]
19+
ingress_torch_rocm = [
20+
"torch==2.8.0+rocm6.4", # AMD-enabled version of torch
21+
"torch-mlir==20251003.589",
22+
"pytorch_triton_rocm" # Transitive dependency listed explicitly so that we can state which package repository it is supposed to come from
23+
]
24+
ingress_torch_xpu = [
25+
"torch==2.8.0+xpu", # Intel-enabled version of torch
26+
"torch-mlir==20251003.589",
27+
"pytorch_triton_xpu" # Transitive dependency listed explicitly so that we can state which package repository it is supposed to come from
28+
]
29+
30+
[tool.uv]
31+
# Declare that the following "targets" are mutually exclusive of one another
32+
conflicts = [
33+
[
34+
{extra = "ingress_torch_cpu" },
35+
{extra = "ingress_torch_nvidia" },
36+
{extra = "ingress_torch_rocm" },
37+
{extra = "ingress_torch_xpu" }
38+
]
39+
]
40+
41+
[tool.uv.sources]
42+
# Bind packages to particular package repositories
43+
mlir_python_bindings = { index = "eudsl" }
44+
torch = { index = "pytorch" }
45+
pytorch_triton_xpu = { index = "pytorch" }
46+
pytorch_triton_rocm = { index = "pytorch" }
47+
torch_mlir = { index = "torch_mlir" }
48+
49+
# The following are the different non-Pypi package repositories we depend on
50+
[[tool.uv.index]]
51+
name = "eudsl"
52+
url = "https://llvm.github.io/eudsl"
53+
explicit = true
54+
format = "flat"
55+
56+
[[tool.uv.index]]
57+
name = "pytorch"
58+
url = "https://download.pytorch.org/whl"
59+
explicit = true
60+
61+
[[tool.uv.index]]
62+
name = "torch_mlir"
63+
url = "https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels"
64+
explicit = true

0 commit comments

Comments
 (0)