From c5c0317871c7fe1cb5d3de3c24340a30df44c40c Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 11:27:58 +0000 Subject: [PATCH 01/16] WIP: Improve install for cpu, ipu, development, and continuous integration use cases --- .github/workflows/cli.yml | 29 ++++++++++++++ .github/workflows/notebooks.yaml | 4 +- pyproject.toml | 13 ------- pyscf_ipu/electron_repulsion/direct.py | 6 ++- pyscf_ipu/nanoDFT/nanoDFT.py | 4 +- requirements.txt | 1 + setup.py | 54 ++++++++++++++++++++++++++ 7 files changed, 92 insertions(+), 19 deletions(-) create mode 100644 .github/workflows/cli.yml delete mode 100644 pyproject.toml create mode 100644 setup.py diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml new file mode 100644 index 0000000..d6f2586 --- /dev/null +++ b/.github/workflows/cli.yml @@ -0,0 +1,29 @@ +name: nanoDFT CLI + +on: + pull_request: + push: + branches: [main] + +jobs: + nanoDFT-cli: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v3 + + - name: Install default requirements + run: | + pip install -U pip + pip install -e . + + - name: Log installed environment + run: | + python3 -m pip freeze + + - name: Test nanoDFT CLI on CPU + run: | + nanoDFT + + + \ No newline at end of file diff --git a/.github/workflows/notebooks.yaml b/.github/workflows/notebooks.yaml index 6871a20..1877211 100644 --- a/.github/workflows/notebooks.yaml +++ b/.github/workflows/notebooks.yaml @@ -21,9 +21,7 @@ jobs: - name: Install requirements run: | pip install -U pip - pip install -r requirements_ipu.txt - pip install -r requirements_test.txt - pip install -e . + pip install -e ".[test,ipu]" - name: Log installed environment run: | diff --git a/pyproject.toml b/pyproject.toml deleted file mode 100644 index ae3cf9d..0000000 --- a/pyproject.toml +++ /dev/null @@ -1,13 +0,0 @@ -[project] -name = "pyscf_ipu" -description = "PySCF on the Graphcore IPU" -version = "0.1" -authors = [ - { name = "Alex Mathiasen", email = "alexm@graphcore.ai" } -] - -[tool.setuptools] -packages = [ - "pyscf_ipu" -] - diff --git a/pyscf_ipu/electron_repulsion/direct.py b/pyscf_ipu/electron_repulsion/direct.py index fc27513..35afc25 100644 --- a/pyscf_ipu/electron_repulsion/direct.py +++ b/pyscf_ipu/electron_repulsion/direct.py @@ -39,8 +39,10 @@ def get_atom_string(atoms, locs): str += "%s %4f %4f %4f; "%((atom,) + tuple(loc) ) return atom_string, str - -NUM_TILES = jax.devices("ipu")[0].num_tiles +try: + NUM_TILES = jax.devices("ipu")[0].num_tiles +except: + NUM_TILES = 1472 @partial(jax.jit, backend="ipu", static_argnums=(3,)) def single(input_floats, input_ints, input_ijkl, tiles): diff --git a/pyscf_ipu/nanoDFT/nanoDFT.py b/pyscf_ipu/nanoDFT/nanoDFT.py index fd63643..39fa0e3 100644 --- a/pyscf_ipu/nanoDFT/nanoDFT.py +++ b/pyscf_ipu/nanoDFT/nanoDFT.py @@ -493,7 +493,7 @@ def nanoDFT_options( jax.config.update('jax_enable_x64', not float32) return args, mol_str -if __name__ == "__main__": +def main(): # Limit PySCF threads to mitigate problem with NUMA nodes. import os os.environ['OMP_NUM_THREADS'] = "16" @@ -561,3 +561,5 @@ def nanoDFT_options( writer.append_data(imageio.v2.imread("_tmp/tmp.jpg")) writer.close() +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt index fdad7c3..58850a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -27,3 +27,4 @@ setuptools < 60.0 jax==0.3.16 jaxlib==0.3.15 jaxtyping==0.2.8 +chex diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..f238bb1 --- /dev/null +++ b/setup.py @@ -0,0 +1,54 @@ +# Copyright (c) 2023 Graphcore Ltd. All rights reserved. +from pathlib import Path + +from setuptools import setup + +__version__ = "0.0.1" + + +def read_requirements(file): + pwd = Path(__file__).parent.resolve() + txt = (pwd / file).read_text(encoding="utf-8").split("\n") + + def remove_comments(line: str): + return len(line) > 0 and not line.startswith(("#", "-")) + + return list(filter(remove_comments, txt)) + + +install_requires = read_requirements("requirements.txt") +ipu_requires = read_requirements("requirements_ipu.txt") +test_requires = read_requirements("requirements_test.txt") + +# url = "https://github.com/graphcore-research/jax-experimental/releases/latest/download/" +# ipu_requires = [ +# f"jaxlib @ {url}jaxlib-0.3.15-cp38-none-manylinux2014_x86_64.whl", +# f"jax @ {url}jax-0.3.16-py3-none-any.whl", +# ] + +setup( + name="pyscf-ipu", + version=__version__, + description="PySCF on IPU", + long_description="file: README.md", + long_description_content_type="text/markdown", + license="Apache License 2.0", + author="Graphcore Research", + author_email="contact@graphcore.ai", + url="https://github.com/graphcore-research/pyscf-ipu", + project_urls={ + "Code": "https://github.com/graphcore-research/pyscf-ipu", + }, + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Topic :: Scientific/Engineering", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + ], + install_requires=install_requires, + extras_require={"test": test_requires, "ipu": ipu_requires}, + python_requires=">=3.8", + packages=["pyscf_ipu"], + entry_points={"console_scripts": ["nanoDFT=pyscf_ipu.nanoDFT:main"]}, +) From 7939db7a93d417fe87520590156b9e90db7529f4 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 11:31:18 +0000 Subject: [PATCH 02/16] remove jaxlib from cpu req --- requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 58850a9..209fbe2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,6 +25,5 @@ py3Dmol setuptools < 60.0 jax==0.3.16 -jaxlib==0.3.15 jaxtyping==0.2.8 chex From 4490092e48f94b81caa7c811e71b9ba8e65583af Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 11:38:35 +0000 Subject: [PATCH 03/16] fix python version to 3.8 in cli workflow --- .github/workflows/cli.yml | 2 ++ requirements.txt | 1 + 2 files changed, 3 insertions(+) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index d6f2586..c96f910 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -11,6 +11,8 @@ jobs: steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v3 + with: + python-version: "3.8.10" - name: Install default requirements run: | diff --git a/requirements.txt b/requirements.txt index 209fbe2..58850a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,5 +25,6 @@ py3Dmol setuptools < 60.0 jax==0.3.16 +jaxlib==0.3.15 jaxtyping==0.2.8 chex From 448b45f7675e69a5be96ac8e1fa211c7139393fb Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 11:44:00 +0000 Subject: [PATCH 04/16] fix ubuntu version --- .github/workflows/cli.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index c96f910..331151f 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -7,7 +7,7 @@ on: jobs: nanoDFT-cli: - runs-on: ubuntu-latest + runs-on: ubuntu-20.04 steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v3 From ff879bdaca169b7b587cabdf2911aab759c417cc Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 12:19:59 +0000 Subject: [PATCH 05/16] rework requirments --- .github/workflows/cli.yml | 2 +- requirements.txt => requirements_core.txt | 5 ++--- requirements_cpu.txt | 9 +++++++++ requirements_ipu.txt | 3 ++- requirements_test.txt | 3 ++- setup.py | 5 +++-- 6 files changed, 19 insertions(+), 8 deletions(-) rename requirements.txt => requirements_core.txt (79%) create mode 100644 requirements_cpu.txt diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 331151f..38dc3ae 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -17,7 +17,7 @@ jobs: - name: Install default requirements run: | pip install -U pip - pip install -e . + pip install -e ".[cpu] " - name: Log installed environment run: | diff --git a/requirements.txt b/requirements_core.txt similarity index 79% rename from requirements.txt rename to requirements_core.txt index 58850a9..17e5ef5 100644 --- a/requirements.txt +++ b/requirements_core.txt @@ -1,6 +1,7 @@ -# Runtime dependencies for pyscf-ipu +# Core dependencies for pyscf-ipu # # See also: +# requirements_cpu.txt for cpu backend configuration # requirements_ipu.txt for ipu backend configuration # requirements_test.txt for test-only dependencies numpy @@ -24,7 +25,5 @@ py3Dmol # silence warnings about setuptools + numpy setuptools < 60.0 -jax==0.3.16 -jaxlib==0.3.15 jaxtyping==0.2.8 chex diff --git a/requirements_cpu.txt b/requirements_cpu.txt new file mode 100644 index 0000000..6dd46ad --- /dev/null +++ b/requirements_cpu.txt @@ -0,0 +1,9 @@ +# Runtime dependencies for pyscf-ipu with CPU backend +# +# See also: +# requirements_core.txt for core configuration +# requirements_ipu.txt for ipu backend configuration +# requirements_test.txt for test-only dependencies + +jax==0.3.16 +jaxlib==0.3.15 diff --git a/requirements_ipu.txt b/requirements_ipu.txt index fddd17a..09d262b 100644 --- a/requirements_ipu.txt +++ b/requirements_ipu.txt @@ -1,7 +1,8 @@ # Runtime dependencies for pyscf-ipu with IPU backend # # See also: -# requirements.txt for core runtime configuration +# requirements_core.txt for core runtime configuration +# requirements_cpu.txt for cpu backend configuration # requirements_test.txt for test-only dependencies -r requirements.txt diff --git a/requirements_test.txt b/requirements_test.txt index c3320fb..9d07eb3 100644 --- a/requirements_test.txt +++ b/requirements_test.txt @@ -1,7 +1,8 @@ # Test dependencies for pyscf-ipu # # See also: -# requirements.txt for core runtime configuration +# requirements_core.txt for core runtime configuration +# requirements_cpu.txt for cpu backend configuration # requirements_ipu.txt for ipu backend configuration pytest nbmake diff --git a/setup.py b/setup.py index f238bb1..2b905f3 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,8 @@ def remove_comments(line: str): return list(filter(remove_comments, txt)) -install_requires = read_requirements("requirements.txt") +install_requires = read_requirements("requirements_core.txt") +cpu_requires = read_requirements("requirements_cpu.txt") ipu_requires = read_requirements("requirements_ipu.txt") test_requires = read_requirements("requirements_test.txt") @@ -47,7 +48,7 @@ def remove_comments(line: str): "Programming Language :: Python :: 3", ], install_requires=install_requires, - extras_require={"test": test_requires, "ipu": ipu_requires}, + extras_require={"cpu": cpu_requires, "ipu": ipu_requires, "test": test_requires}, python_requires=">=3.8", packages=["pyscf_ipu"], entry_points={"console_scripts": ["nanoDFT=pyscf_ipu.nanoDFT:main"]}, From 5b56f0c707446cf656992b8c2b8cf907c9c1c1b1 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 12:23:39 +0000 Subject: [PATCH 06/16] fix whitespace --- .github/workflows/cli.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 38dc3ae..80e4692 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -17,7 +17,7 @@ jobs: - name: Install default requirements run: | pip install -U pip - pip install -e ".[cpu] " + pip install -e ".[cpu]" - name: Log installed environment run: | From 50fe4dc4e8c7e738ed813dddd5a6eb308a47b731 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 11 Sep 2023 19:54:57 +0000 Subject: [PATCH 07/16] use urls for ipu req --- requirements_ipu.txt | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/requirements_ipu.txt b/requirements_ipu.txt index 09d262b..8b63733 100644 --- a/requirements_ipu.txt +++ b/requirements_ipu.txt @@ -4,10 +4,6 @@ # requirements_core.txt for core runtime configuration # requirements_cpu.txt for cpu backend configuration # requirements_test.txt for test-only dependencies --r requirements.txt - -# Experimental JAX for IPU --f https://graphcore-research.github.io/jax-experimental/wheels.html -jax==0.3.16+ipu -jaxlib==0.3.15+ipu.sdk320 +jax @ https://github.com/graphcore-research/jax-experimental/releases/download/jax-v0.3.16-ipu-beta3-sdk3/jax-0.3.16%2Bipu-py3-none-any.whl +jaxlib @ https://github.com/graphcore-research/jax-experimental/releases/download/jax-v0.3.16-ipu-beta3-sdk3/jaxlib-0.3.15%2Bipu.sdk320-cp38-none-manylinux2014_x86_64.whl tessellate-ipu@git+https://github.com/graphcore-research/tessellate-ipu.git@main From 78e2e98ab5b564179c28450341e2c517dd3c3473 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Tue, 12 Sep 2023 06:55:08 +0000 Subject: [PATCH 08/16] remove comment strings --- setup.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/setup.py b/setup.py index 2b905f3..a7f6c4b 100644 --- a/setup.py +++ b/setup.py @@ -21,12 +21,6 @@ def remove_comments(line: str): ipu_requires = read_requirements("requirements_ipu.txt") test_requires = read_requirements("requirements_test.txt") -# url = "https://github.com/graphcore-research/jax-experimental/releases/latest/download/" -# ipu_requires = [ -# f"jaxlib @ {url}jaxlib-0.3.15-cp38-none-manylinux2014_x86_64.whl", -# f"jax @ {url}jax-0.3.16-py3-none-any.whl", -# ] - setup( name="pyscf-ipu", version=__version__, From cdcfbbc06e224309058be34563150bfb002ae7b5 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Tue, 12 Sep 2023 08:11:44 +0000 Subject: [PATCH 09/16] readme: adding cli badge and draft install instructions --- README.md | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 5ca8090..4a7f7f4 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,7 @@ :red_circle: :warning: **Experimental and non-official Graphcore product** :warning: :red_circle: [![notebook-tests](https://github.com/graphcore-research/pyscf-ipu/actions/workflows/notebooks.yaml/badge.svg)](https://github.com/graphcore-research/pyscf-ipu/actions/workflows/notebooks.yaml) +[![nanoDFT CLI](https://github.com/graphcore-research/pyscf-ipu/actions/workflows/cli.yaml/badge.svg)](https://github.com/graphcore-research/pyscf-ipu/actions/workflows/cli.yaml) [**Installation guide**](#installation) | [**Example DFT Computations**](#example-dft-computations) @@ -47,26 +48,20 @@ Additional notebooks in [notebooks](notebooks) demonstrate other aspects of the PySCF on IPU requires Python 3.8, [JAX IPU experimental](https://github.com/graphcore-research/jax-experimental), [TessellateIPU library](https://github.com/graphcore-research/tessellate-ipu) and [Graphcore Poplar SDK 3.2](https://www.graphcore.ai/downloads). -To run this package on a standard CPU machine (laptop or server), -install the base Python requirements: +We recommend upgrading `pip` to the latest stable release to prepare your enviroment. ```bash -pip install -r requirements.txt +pip install -U pip ``` -On IPU machines, please additionally use the IPU requirements file: +This project is currently under active development. +For CPU simulations, we recommend installing `pyscf-ipu` from latest `main` branch as: ```bash -pip install -U pip -pip install -r requirements_ipu.txt +pip install pyscf-ipu[cpu] @ git+https://github.com/graphcore-research/pyscf-ipu ``` -This will configure Graphcore research experimental JAX support in your python environment. - -We recommend upgrading `pip` to the latest stable release when using the IPU -requirements. This may be an optional step depending on the overall configuration of -your python environment. -And finally, make our sub-packages available: +and on IPU equipped machines: ```bash -pip install -e . +pip install pyscf-ipu[ipu] @ git+https://github.com/graphcore-research/pyscf-ipu ``` ## Example DFT Computations From 212c85025b24bc2584dc21e2c9881c6f6234a7d8 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Tue, 12 Sep 2023 08:15:16 +0000 Subject: [PATCH 10/16] can't have spaces around @ --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 4a7f7f4..f7d55de 100644 --- a/README.md +++ b/README.md @@ -56,12 +56,12 @@ pip install -U pip This project is currently under active development. For CPU simulations, we recommend installing `pyscf-ipu` from latest `main` branch as: ```bash -pip install pyscf-ipu[cpu] @ git+https://github.com/graphcore-research/pyscf-ipu +pip install pyscf-ipu[cpu]@git+https://github.com/graphcore-research/pyscf-ipu ``` and on IPU equipped machines: ```bash -pip install pyscf-ipu[ipu] @ git+https://github.com/graphcore-research/pyscf-ipu +pip install pyscf-ipu[ipu]@git+https://github.com/graphcore-research/pyscf-ipu ``` ## Example DFT Computations From 62f3985596f5fcb9d3d7cbfd8cb0c03a6fe499d9 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Tue, 12 Sep 2023 12:12:02 +0000 Subject: [PATCH 11/16] update notebooks to use setup.py --- notebooks/DFT-dataset-generation.ipynb | 12 +++++++++--- notebooks/nanoDFT-demo.ipynb | 3 +-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/notebooks/DFT-dataset-generation.ipynb b/notebooks/DFT-dataset-generation.ipynb index a29a1ee..bfd8f5e 100644 --- a/notebooks/DFT-dataset-generation.ipynb +++ b/notebooks/DFT-dataset-generation.ipynb @@ -34,6 +34,14 @@ "print(\"Working directory:\", os.getcwd())" ] }, + { + "cell_type": "markdown", + "id": "e5faf299", + "metadata": {}, + "source": [ + "Install `pyscf-ipu`:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -41,9 +49,7 @@ "metadata": {}, "outputs": [], "source": [ - "# PySCF IPU dependencies \n", - "%pip install -r requirements.txt\n", - "%pip install -r requirements_ipu.txt" + "%pip install -e \"..[ipu]\"" ] }, { diff --git a/notebooks/nanoDFT-demo.ipynb b/notebooks/nanoDFT-demo.ipynb index ab4938e..30c2d61 100644 --- a/notebooks/nanoDFT-demo.ipynb +++ b/notebooks/nanoDFT-demo.ipynb @@ -27,8 +27,7 @@ "metadata": {}, "outputs": [], "source": [ - "%pip install -r ../requirements_ipu.txt\n", - "%pip install -e .." + "%pip install -e \"..[ipu]\"" ] }, { From b10df05c20e5302756683db671ffcf6ebf2bfb51 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Fri, 15 Sep 2023 13:35:49 +0000 Subject: [PATCH 12/16] fix ext --- .github/workflows/{cli.yml => cli.yaml} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename .github/workflows/{cli.yml => cli.yaml} (100%) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yaml similarity index 100% rename from .github/workflows/cli.yml rename to .github/workflows/cli.yaml From b641e1cb4419bd616aee1f89e5d21bf0adc0a2db Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Fri, 15 Sep 2023 13:47:06 +0000 Subject: [PATCH 13/16] default to cpu install --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index a7f6c4b..2ff36da 100644 --- a/setup.py +++ b/setup.py @@ -41,8 +41,8 @@ def remove_comments(line: str): "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3", ], - install_requires=install_requires, - extras_require={"cpu": cpu_requires, "ipu": ipu_requires, "test": test_requires}, + install_requires=install_requires + cpu_requires, + extras_require={"ipu": ipu_requires, "test": test_requires}, python_requires=">=3.8", packages=["pyscf_ipu"], entry_points={"console_scripts": ["nanoDFT=pyscf_ipu.nanoDFT:main"]}, From 362aab71c7040d3fcb560306b52f6d7dee32228b Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Fri, 15 Sep 2023 13:47:28 +0000 Subject: [PATCH 14/16] update install doc --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index f7d55de..e8ae7a4 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ pip install -U pip This project is currently under active development. For CPU simulations, we recommend installing `pyscf-ipu` from latest `main` branch as: ```bash -pip install pyscf-ipu[cpu]@git+https://github.com/graphcore-research/pyscf-ipu +pip install pyscf-ipu@git+https://github.com/graphcore-research/pyscf-ipu ``` and on IPU equipped machines: From 3bda71ed3639586f83c1eb3f3db59207752d9a7a Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 18 Sep 2023 08:07:05 +0000 Subject: [PATCH 15/16] wrap num tiles lookup in function scope --- pyscf_ipu/electron_repulsion/direct.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/pyscf_ipu/electron_repulsion/direct.py b/pyscf_ipu/electron_repulsion/direct.py index 35afc25..e23231e 100644 --- a/pyscf_ipu/electron_repulsion/direct.py +++ b/pyscf_ipu/electron_repulsion/direct.py @@ -39,10 +39,8 @@ def get_atom_string(atoms, locs): str += "%s %4f %4f %4f; "%((atom,) + tuple(loc) ) return atom_string, str -try: - NUM_TILES = jax.devices("ipu")[0].num_tiles -except: - NUM_TILES = 1472 +def num_ipu_tiles() -> int: + return jax.devices("ipu")[0].num_tiles @partial(jax.jit, backend="ipu", static_argnums=(3,)) def single(input_floats, input_ints, input_ijkl, tiles): @@ -89,7 +87,7 @@ def ipu_direct_mult_v0(__out3, dm, indices, do_lists, N, num_tiles, indxs_inv, i #print(threads) num_threads = threads - num_tiles = NUM_TILES-1 + num_tiles = num_ipu_tiles()-1 tiles = tuple((np.arange(num_tiles*num_threads)%num_tiles+1).tolist()) dm = tile_put_replicated( dm, tiles=tiles) @@ -178,7 +176,7 @@ def ipu_direct_mult_v1(__out3, dm, indices, do_lists, N, num_tiles, indxs_inv, i #print(threads) num_threads = threads - num_tiles = NUM_TILES-1 + num_tiles = num_ipu_tiles()-1 tiles = tuple((np.arange(num_tiles*num_threads)%num_tiles+1).tolist()) @@ -313,7 +311,7 @@ def ipu_direct_mult_v2(__out3, dm, indices, do_lists, N, num_tiles, indxs_inv, i #print(threads) num_threads = threads - num_tiles = NUM_TILES-1 + num_tiles = num_ipu_tiles()-1 tiles = tuple((np.arange(num_tiles*num_threads)%num_tiles+1).tolist()) @@ -534,7 +532,7 @@ def integrals_v0(input_floats, input_ints, input_ijkl, shapes, sizes, counts, in np.random.seed(42) # - num_tiles = NUM_TILES-1 + num_tiles = num_ipu_tiles()-1 num_calls = len(indxs_inv) #print(num_calls) @@ -609,7 +607,7 @@ def integrals_v1(input_floats, input_ints, input_ijkl, shapes, sizes, counts, in start, stop = 0, 0 np.random.seed(42) - num_tiles = NUM_TILES-1 + num_tiles = num_ipu_tiles()-1 num_calls = len(indxs_inv) #print(num_calls) @@ -1420,7 +1418,7 @@ def compute_integrals(input_floats, input_ints, input_ijkl, num_calls, num_tiles out3 = np.zeros(out.shape) else: num_threads = int(args.threads) - ipu_num_tiles = NUM_TILES + ipu_num_tiles = num_ipu_tiles() if num_calls >= ipu_num_tiles*num_threads: # If enough calls allocate all threads and all tiles. tiles = [i for i in range(1, ipu_num_tiles) for _ in range(num_threads)] From 80643d8d7fc2e3a79aedc38a6020cf70ee1937b0 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Mon, 18 Sep 2023 08:09:51 +0000 Subject: [PATCH 16/16] removed cpu arg --- .github/workflows/cli.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cli.yaml b/.github/workflows/cli.yaml index 80e4692..0d26ccd 100644 --- a/.github/workflows/cli.yaml +++ b/.github/workflows/cli.yaml @@ -17,7 +17,7 @@ jobs: - name: Install default requirements run: | pip install -U pip - pip install -e ".[cpu]" + pip install -e "." - name: Log installed environment run: |