diff --git a/.cirrus.yml b/.cirrus.yml deleted file mode 100644 index fc68ca1b0..000000000 --- a/.cirrus.yml +++ /dev/null @@ -1,25 +0,0 @@ -task: - name: Tests (linux) - container: - image: ubuntu:latest - install_script: - - apt-get update - - apt-get -y install python3 python3-venv python3-dev g++ - test_script: - - bash ./pytest.sh - - -task: - name: Tests (macOS) - macos_instance: - image: ghcr.io/cirruslabs/macos-ventura-xcode:latest - - script: - - brew install libomp - - brew install python@3.11 - - $(brew --prefix python@3.11)/bin/python3.11 -m venv .test_venv - - source .test_venv/bin/activate - - pip install --upgrade pip - - sh ./pytest.sh - - diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 69934f854..27bb56b8b 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -1,12 +1,19 @@ -name: Lint +name: Format Check -on: [push, pull_request] +on: + push: + pull_request: jobs: - lint: + format: + name: Format (Black) runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - - uses: psf/black@stable + - name: Checkout + uses: actions/checkout@v4 + + - name: Black format check + uses: psf/black@stable + with: + options: "--check ." diff --git a/.github/workflows/cuda_test.yml b/.github/workflows/cuda_test.yml deleted file mode 100644 index 1d727b2f1..000000000 --- a/.github/workflows/cuda_test.yml +++ /dev/null @@ -1,15 +0,0 @@ -# Launch test on a cuda machine -name: Cuda test - -on: - push - -jobs: - test: - runs-on: self-hosted - steps: - - name: Clone Repository - uses: actions/checkout@v2 - - name: Test with pytest - run: | - sh ./pytest.sh diff --git a/.github/workflows/linux_cpu_test.yml b/.github/workflows/linux_cpu_test.yml new file mode 100644 index 000000000..51821b78b --- /dev/null +++ b/.github/workflows/linux_cpu_test.yml @@ -0,0 +1,21 @@ +name: Linux CPU Test + +on: + push: + pull_request: + +jobs: + tests-linux: + name: Tests (Linux CPU) + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + + - name: Run tests + run: bash ./pytest.sh \ No newline at end of file diff --git a/.github/workflows/linux_cuda_test.yml b/.github/workflows/linux_cuda_test.yml new file mode 100644 index 000000000..9cef38911 --- /dev/null +++ b/.github/workflows/linux_cuda_test.yml @@ -0,0 +1,16 @@ +name: Linux CUDA Test + +on: + push: + pull_request: + +jobs: + tests-linux-cuda: + name: Tests (Linux CUDA) + runs-on: self-hosted + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Run tests + run: bash ./pytest.sh diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml new file mode 100644 index 000000000..75d159032 --- /dev/null +++ b/.github/workflows/macOs_test.yml @@ -0,0 +1,46 @@ +name: macOS Test + +on: + push: + pull_request: + +jobs: + tests-macos-latest: + name: Tests (macOS latest) + runs-on: macos-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + + - name: Install dependencies + run: brew install libomp + + - name: Run tests + run: bash ./pytest.sh + + tests-macos-intel: + name: Tests (macOS Intel) + runs-on: macos-26-intel + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Pin NumPy for Intel CI + run: | + echo "numpy<2" > /tmp/pip-constraints.txt + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install dependencies + run: brew install libomp + + - name: Run tests + run: bash ./pytest.sh --pip-constraint=/tmp/pip-constraints.txt -v 2 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..79294ee1c --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,14 @@ +repos: + - repo: https://github.com/psf/black + rev: 24.10.0 + hooks: + - id: black + + - repo: local + hooks: + - id: pytest + name: pytest + entry: ./pytest.sh + language: system + pass_filenames: false + always_run: true diff --git a/Jenkinsfile b/Jenkinsfile deleted file mode 100644 index 01c1a4661..000000000 --- a/Jenkinsfile +++ /dev/null @@ -1,25 +0,0 @@ -pipeline { - agent none - stages { - stage('Test PyKeOps') { - parallel { - - stage('Test Cuda') { - agent { - label 'cuda' - } - steps { - echo 'Clean KeOps Cache...' - sh 'rm -rf $HOME/.cache/keops*' - echo 'Testing..' - sh '''#!/bin/bash - srun -n 1 -c 16 --mem=8G --gpus=1 --gres-flags=enforce-binding -J keops_ci pytest.sh -v - ''' - } - } - - } - } - - } -} diff --git a/benchmarks/PyTorch_GPU.py b/benchmarks/PyTorch_GPU.py index bfecaa956..1836f7f9a 100644 --- a/benchmarks/PyTorch_GPU.py +++ b/benchmarks/PyTorch_GPU.py @@ -12,7 +12,6 @@ import numpy as np from time import time - nits = 100 Ns, D = [10000, 100000, 1000000], 3 diff --git a/benchmarks/TVM.py b/benchmarks/TVM.py index 837021116..ec4be25cd 100644 --- a/benchmarks/TVM.py +++ b/benchmarks/TVM.py @@ -41,7 +41,6 @@ import numpy as np from time import time - # Global declarations of environment. tgt_host = "llvm" tgt = "cuda" diff --git a/doc/conf.py b/doc/conf.py index 1c12a72cd..365c591e3 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -108,6 +108,7 @@ def find_source(): # Add patterns # 'filename_pattern': r'../pykeops/pykeops/tutorials/*', "ignore_pattern": r"__init__\.py|benchmark_utils\.py|dataset_utils\.py", + "parallel": int(os.environ.get("SPHINX_GALLERY_JOBS", "1")), } # Generate the API documentation when building diff --git a/doc/python/api/common/Config.rst b/doc/python/api/common/Config.rst new file mode 100644 index 000000000..51aa10c2f --- /dev/null +++ b/doc/python/api/common/Config.rst @@ -0,0 +1,19 @@ +Configuration +------------- + +This section documents lower-level KeOps configuration helpers that are useful +when troubleshooting PyKeOps installations. + +.. rubric:: Summary + +.. currentmodule:: keopscore.config.Cuda + +.. autosummary:: + :nosignatures: + + CudaConfig + +.. rubric:: Syntax + +.. autoclass:: CudaConfig + :members: diff --git a/doc/python/api/common/index.rst b/doc/python/api/common/index.rst index 263f096e2..ed97c421f 100644 --- a/doc/python/api/common/index.rst +++ b/doc/python/api/common/index.rst @@ -6,5 +6,6 @@ Common Python API .. toctree:: GenericLazyTensor + Config Utils diff --git a/doc/python/installation.rst b/doc/python/installation.rst index 599a7c591..b1f3d7efa 100644 --- a/doc/python/installation.rst +++ b/doc/python/installation.rst @@ -1,69 +1,156 @@ -Python install -############## +Python installation +################### -PyKeOps is a **Python 3 wrapper** around the low-level KeOpsCore library which is written in **C++/CUDA**. -It provides functions that can be used in any **NumPy** or **PyTorch** script. +The ``pykeops`` Python module provides the NumPy and PyTorch bindings for +KeOps. It relies on the ``keopscore`` Python module, the KeOps +metaprogramming engine, to generate and compile the C++/CUDA routines that +evaluate symbolic kernel formulas on the fly. + + +.. _`part.PyKeOpsRequirements`: Requirements ============ -- **Python** (>= 3.8) with the **numpy** package. -- A C++ compiler compatible with ``std=c++11``: **g++** version >=7 or **clang++** version >=8. -- The **Cuda** toolkit: version >=10.0 is recommended. -- **PyTorch** (optional): version >= 1.5. +Required dependencies: + +- **Python** (>= 3.8) with the **NumPy** package. +- A C++ compiler such as ``gcc`` (on Linux) or ``clang`` (on macOS). + +Optional, but highly recommended: + +- An NVIDIA GPU with the **NVIDIA drivers** and the **CUDA** toolkit. +- **PyTorch** (version >= 2 is recommended). +- The **OpenMP** libraries and headers for CPU-only systems. Using pip (recommended) ======================= -1. Just in case: in a terminal, check the **consistency** of the outputs of the commands ``which python``, ``python --version``, ``which pip`` and ``pip --version``. +1. In a terminal, check that ``python`` and ``pip`` point to the same + environment with: -2. In a terminal, type: + .. prompt:: bash $ - .. prompt:: bash $ + which python + python --version + which pip + pip --version - pip install pykeops + You can also create a fresh Python virtual environment: - Note that compiled shared objects (``.so`` files on Unix, ``.dylib`` on macOS) will be stored in the folder ``~/.cache/keops/``, where ``~`` is the path to your home folder. If you want to change this default location, define the environment variable ``KEOPS_CACHE_FOLDER`` to another folder prior to importing pykeops. + .. prompt:: bash $ + + python -m venv keops_venv + source keops_venv/bin/activate + +2. **Install PyKeOps:** + + .. prompt:: bash $ + + pip install pykeops + + Compiled shared objects (``.so`` files on Unix, ``.dylib`` files on + macOS) are stored in ``~/.cache/keops``, where ``~`` is your + home folder and ```` is the installed ``pykeops`` version. To + change this location, define the ``KEOPS_CACHE_FOLDER`` environment + variable before importing ``pykeops``. + +3. Test your installation: + + .. prompt:: bash $ + + python -c "import pykeops; pykeops.test_numpy_bindings()" + +More details are in the :ref:`dedicated section `. -3. Test your installation, as described in the :ref:`next section `. On Google Colab =============== -Google provides free virtual machines where KeOps runs -out-of-the-box. -In a new `Colab notebook `_, typing: +Google provides free virtual machines, running on Ubuntu Linux, where KeOps +runs out of the box. In a new +`Colab notebook `_, run: -.. prompt:: bash $ +.. prompt:: python >>> !pip install pykeops > install.log + import pykeops + pykeops.test_numpy_bindings() should allow you to get a working version of KeOps in less than twenty seconds. +.. _`part.CondaConfig`: + +Using a Conda/Miniconda/Mamba environment +========================================= + +Conda environments can be a convenient way to get a working configuration +without the root permissions that may be needed to install system dependencies +such as the CUDA toolkit or OpenMP. They are not fully isolated from every other +installation mechanism, though. If things do not behave as expected, check +whether older packages are already installed somewhere else on your system. + +The dependencies for PyKeOps are listed :ref:`above `, +but the exact packages needed depend on your system. For instance, the following +commands should create a working environment on Ubuntu 24.04 equipped with +a GPU: + +.. prompt:: bash $ + + conda create --name keops_env python=3.14 + conda activate keops_env + + conda install libgomp + conda install nvidia::cuda-toolkit + pip install pykeops + + python -c "import pykeops; pykeops.test_numpy_bindings()" + + + +On macOS +======== + +We recommend installing the OpenMP libraries with Homebrew: + +.. prompt:: bash $ + + brew install libomp + pip install pykeops + +You should now be able to test your installation: + +.. prompt:: bash $ + + python -c "import pykeops; pykeops.test_numpy_bindings()" + +More help can be found in the :ref:`dedicated section `. + Using Docker or Singularity ============================ -We provide a reference -`Dockerfile `_ -and publish full containers on our -`DockerHub channel `_ -using the -`docker-images.sh `_ script. -These environments contain a full installation of CUDA, NumPy, PyTorch, R, KeOps and GeomLoss. -Their PYTHONPATH are configured to ensure that git installations of KeOps or GeomLoss -mounted in ``/opt/keops`` or ``/opt/geomloss`` take precedence over the -pre-installed Pip versions. - -As an example, here are the steps that we follow to render this website on the +We provide a reference +`Dockerfile `_ and +publish full containers on our +`DockerHub channel `_ +using the `docker-images.sh `_ +script. These environments contain full installations of CUDA, NumPy, PyTorch, +R, KeOps (for Python and R) and GeomLoss. + +The container's ``PYTHONPATH`` environment variable is configured so that Git +installations of KeOps or GeomLoss mounted in ``/opt/keops`` or +``/opt/geomloss`` take precedence over the pre-installed pip versions. + +As an example, here are the steps that we follow to render this website on the `Jean Zay `_ scientific cluster: -.. code-block:: bash +.. prompt:: bash $ # First, clone the latest release of the KeOps repository in ~/code/keops: - mkdir ~/code - cd ~/code + mkdir ~/code + cd ~/code git clone git@github.com:getkeops/keops.git # Load singularity in our environment: @@ -74,10 +161,11 @@ As an example, here are the steps that we follow to render this website on the cd ~/scratch/containers # Download the Docker image and store it as an immutable Singularity Image File: - # N.B.: Our image is pretty heavy (~7 Gb), so it is safer to create + # N.B.: Our image is fairly large (~7 GB), so it is safer to create # cache folders on the hard drive instead of relying on the RAM-only tmpfs: - # N.B.: This step may take 15mn to 60mn, so you may prefer to execute it on your - # local computer and then copy the resulting file `keops-full.sif` to the cluster. + # N.B.: This step may take 15 to 60 minutes, so you may prefer to execute it on + # your local computer and then copy the resulting file `keops-full.sif` + # to the cluster. # Alternatively, on the Jean Zay cluster, you may use the `prepost` partition # to have access to both a large RAM and an internet connection. mkdir cache @@ -88,15 +176,15 @@ As an example, here are the steps that we follow to render this website on the # At this point, on the Jean Zay cluster, you should use a command like: # idrcontmgr cp keops-full.sif - # to add our new environment to the cluster's container registry as explained here: + # to add our new environment to the cluster's container registry as explained here: # http://www.idris.fr/jean-zay/cpu/jean-zay-utilisation-singularity.html - - # Then, create a separate home folder for this image. This is to ensure - # that we won't see any conflict between different versions of the KeOps binaries, + + # Then, create a separate home folder for this image. This helps ensure + # that we do not encounter conflicts between different versions of the KeOps binaries, # stored in the ~/.cache folder of the virtual machine: mkdir -p ~/containers/singularity_homes/keops-full - # Ask the slurm scheduler to render our documentation. + # Ask the Slurm scheduler to render our documentation. sbatch keops-doc.batch @@ -106,16 +194,16 @@ Where ``keops-doc.batch`` is an executable file that contains: #!/bin/bash - #SBATCH -A dvd@a100 # Use a A100 GPU - dvd@v100 is also available - #SBATCH -C a100 + #SBATCH -A dvd@a100 # Use an A100 GPU - dvd@v100 is also available + #SBATCH -C a100 #SBATCH --partition=gpu_p5 #SBATCH --job-name=keops_doc # create a short name for your job #SBATCH --mail-type=ALL # Mail events (NONE, BEGIN, END, FAIL, ALL) - #SBATCH --mail-user=your.name@inria.fr # Where to send mail + #SBATCH --mail-user=your.name@inria.fr # Where to send mail #SBATCH --nodes=1 # node count #SBATCH --ntasks=1 # total number of tasks across all nodes #SBATCH --cpus-per-task=8 # cpu-cores per task (>1 if multi-threaded tasks) - #SBATCH --gres=gpu:1 # GPU nodes are only available in gpu partition + #SBATCH --gres=gpu:1 # GPU nodes are only available in gpu partition #SBATCH --time=03:00:00 # total run time limit (HH:MM:SS) #SBATCH --output=logs/keops_doc.out # output file name #SBATCH --error=logs/keops_doc.err # error file name @@ -171,7 +259,7 @@ And ``keops-doc.sh`` is an executable file that contains: # Render the website: make html - # Re-render the doc to remove compilation messages: + # Re-render the documentation to remove compilation messages: make clean make html @@ -179,53 +267,51 @@ And ``keops-doc.sh`` is an executable file that contains: -From source using git +From source using Git ===================== - -The simplest way of installing a specific version -of KeOps is to use `some advanced pip syntax `_: - +The simplest way to install a specific version of KeOps is to use pip's +`Git URL syntax `_: .. prompt:: bash $ - pip install git+https://github.com/getkeops/keops.git@main#subdirectory=keopscore - pip install git+https://github.com/getkeops/keops.git@main#subdirectory=pykeops + pip install git+https://github.com/getkeops/keops.git@main#subdirectory=keopscore + pip install git+https://github.com/getkeops/keops.git@main#subdirectory=pykeops Alternatively, you may: -1. Clone the KeOps repo at a location of your choice (denoted here as ``/path/to``): +1. Clone the KeOps repository at a location of your choice: - .. prompt:: bash $ + .. prompt:: bash $ - git clone --recursive https://github.com/getkeops/keops.git /path/to/libkeops + git clone https://github.com/getkeops/keops.git /path/to/keops_cloned_repo - Note that compiled **.so** routines will be stored in the folder ``/path/to/libkeops/pykeops/build``: this directory must have **write permission**. +2. Install the Python packages in editable mode: + .. prompt:: bash $ -2. Install via pip in editable mode as follows : - - .. prompt:: bash $ + pip install -e /path/to/keops_cloned_repo/keopscore -e /path/to/keops_cloned_repo/pykeops - pip install -e /path/to/libkeops/keopscore -e /path/to/libkeops/pykeops + If you prefer not to install the packages, you can add + ``/path/to/keops_cloned_repo/keopscore`` and + ``/path/to/keops_cloned_repo/pykeops`` to your Python path. To do this once + permanently, add the paths to your ``~/.bashrc``: - + Otherwise you may add the directories ``/path/to/libkeops/keopscore`` and ``/path/to/libkeops/pykeops`` to your python path. This can be done once and for all, by adding the path to to your ``~/.bashrc``. In a terminal, type: - - .. prompt:: bash $ + .. prompt:: bash $ + + echo "export PYTHONPATH=$PYTHONPATH:/path/to/keops_cloned_repo/keopscore:/path/to/keops_cloned_repo/pykeops" >> ~/.bashrc - echo "export PYTHONPATH=$PYTHONPATH:/path/to/libkeops/keopscore:/path/to/libkeops/pykeops" >> ~/.bashrc + Alternatively, add these lines at the beginning of your Python scripts: - + Alternatively, you may add the following line to the beginning of your python scripts: - - .. code-block:: python + .. code-block:: python - import os.path - import sys - sys.path.append('/path/to/libkeops/keopscore') - sys.path.append('/path/to/libkeops/pykeops') + import sys -3. Test your installation, as described in the :ref:`next section. ` + sys.path.append("/path/to/keops_cloned_repo/keopscore") + sys.path.append("/path/to/keops_cloned_repo/pykeops") + +3. Test your installation, as described in the :ref:`next section `. .. _`part.checkPython`: @@ -233,37 +319,40 @@ Alternatively, you may: Testing your installation ========================= -You can use the following test functions to compile and run simple KeOps formulas. If the compilation fails, it returns the full log. +You can use the following test functions to compile and run simple KeOps +formulas. If compilation fails, they return the full log. -1. In a python terminal, run :func:`pykeops.test_numpy_bindings `. +1. In a Python terminal, run + :func:`pykeops.test_numpy_bindings `. - .. code-block:: python + .. prompt:: python >>> - import pykeops - pykeops.test_numpy_bindings() # perform the compilation - - should return: + import pykeops + assert pykeops.test_numpy_bindings() # perform the compilation - .. code-block:: text + It should print: - pyKeOps with numpy bindings is working! + .. code-block:: text -2. If you use PyTorch, run :func:`pykeops.test_torch_bindings `. + pyKeOps with numpy bindings is working! - .. code-block:: python +2. If you use PyTorch, run + :func:`pykeops.test_torch_bindings `. - import pykeops - pykeops.test_torch_bindings() # perform the compilation - - should return: + .. prompt:: python >>> + + import pykeops + assert pykeops.test_torch_bindings() # perform the compilation + + It should print: - .. code-block:: text + .. code-block:: text - pyKeOps with torch bindings is working! + pyKeOps with torch bindings is working! -Please note that running ``pytest -v`` in a copy of our git repository will also -let you perform an in-depth test of the entire KeOps codebase. +Running ``pytest -v`` in a copy of our Git repository will also let you perform +an in-depth test of the entire KeOps codebase. Troubleshooting @@ -272,24 +361,137 @@ Troubleshooting KeOps health check ------------------ -To get an overview of your KeOps installation (along with any related issues), including relevant paths, environments, compilation flags, and more, it’s recommended to run the :func:`pykeops.check_health ` function. Simply type the following in a Python shell: +To get an overview of your KeOps installation, including relevant paths, +environments, compilation flags and possible issues, we recommend running the +:func:`pykeops.check_health ` function: -.. code-block:: python +.. prompt:: python >>> import pykeops + pykeops.clean_pykeops() pykeops.check_health() +You can inspect the paths found by KeOps to check whether the external +libraries are properly detected. Compilation issues ------------------ -First of all, make sure that you are using a C++ compiler which is compatible with the **C++11 revision**. Otherwise, compilation of formulas may fail in unexpected ways. Depending on your system, you can: +KeOps compiles small code fragments to compute kernel operations on a device +(CPU or GPU). Most installation issues occur during this compilation step. They +are usually caused by misconfigured environments, overlapping package +installations or non-standard installation paths. The common failure modes have +evolved over the years, and the +`KeOps issue tracker `_ is a good +up-to-date starting point. + +We detail some common issues below, along with generic recommendations. + + +CUDA toolkit detection +....................... + +Detecting the CUDA toolkit is not always straightforward because several CUDA +versions may be present on the same system. For instance, PyTorch installations +may bring a partial CUDA toolkit installation into the active Python environment (see +detail below). + +KeOps does not ship its own CUDA toolkit. It tries to detect a working toolkit +with :class:`keopscore.config.CudaConfig ` in +the following order: + +1. *Environment variables:* ``CUDA_PATH``, ``CUDA_HOME``, ``CUDA_ROOT`` and + ``CUDA_TOOLKIT_ROOT_DIR`` (in this order). + +2. *Conda installation:* if you use Conda, as described + :ref:`above `. + +3. *System installation from your distribution* **(recommended)**: this is the + best way to ensure that the CUDA toolkit and driver versions match and that + all paths are set consistently. + + On common Linux distributions, system CUDA packages can be installed for instance with: + + .. prompt:: bash $ + + # Debian/Ubuntu, with distribution packages: + sudo apt install nvidia-cuda-dev nvidia-cuda-toolkit + + # Debian/Ubuntu, with NVIDIA CUDA repositories enabled: + sudo apt install cuda-dev nvidia-cuda-toolkit + + # Arch Linux and derivatives: + yay -S cuda + +4. *`NVIDIA PyPI packages `_*: needed package could be installed + with the ``pip install pykeops[cuda]`` or ``pip install pykeops[cu12]`` or or ``pip install pykeops[cu13]`` recipes. Beware, many cuda install can co-exist in the same virtual environment. + + +We illustrate here how the CUDA toolkit could be manually selected. On Arch linux, the following command: + +.. prompt:: bash $ + python -m venv keops_venv + source keops_venv/bin/activate + + pip install pykeops + PYKEOPS_VERBOSE=0 python -c "import pykeops; pykeops.config.cuda.print_all()" | head -n 9 + +yields the detection of a system-wide CUDA installation under ``/opt/cuda``: + +.. code-block:: text + + ============================================================ + CUDA Support + ============================================================ + CUDA Support: Enabled ✅ + Number of GPUs Detected: 4 + CUDA Version: 13.2 + Libcuda Path: /usr/lib64/libcuda.so.1 + Libnvrtc Path: /opt/cuda/lib64/libnvrtc.so.13 + Libcudart Path: /opt/cuda/lib64/libcudart.so.13 + +Now, forcing KeOps to use the NVIDIA PyPI packages installed in the active Python virtual + environment may be done by running: + +.. prompt:: bash $ + python -m venv keops_venv_cuda_pip + source keops_venv_cuda_pip/bin/activate + + pip install pykeops[cu13] + CUDA_PATH="$VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13" \ + PYKEOPS_VERBOSE=0 python -c "import pykeops; pykeops.config.cuda.print_all()" | head -n 9 + +This gives: + +.. code-block:: text + + ============================================================ + CUDA Support + ============================================================ + CUDA Support: Enabled ✅ + Number of GPUs Detected: 4 + CUDA Version: 13.0 + Libcuda Path: /usr/lib64/libcuda.so.1 + Libnvrtc Path: $VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13/lib/libnvrtc.so.13 + Libcudart Path: $VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13/lib/libcudart.so.13 + + +Compiler +........ -1. Install a compiler **system-wide**: for instance, on Debian-based Linux distributions, you can install g++ with apt and then use `update-alternatives `_ to choose a suitable compiler as default. Don't forget to pick compatible versions for both **gcc** and **g++**. +Recent Linux and macOS distributions usually provide suitable compilers. If +compilation fails, make sure that you are using a C++ compiler compatible with +the **C++11 revision** (``std=c++11`` flag supported). Otherwise, formula compilation + may fail in unexpected ways. -2. Install a compiler **locally**: if you are using a conda environment, you can install a new instance of gcc and g++ by following the `documentation of conda `_. +1. Install a compiler **system-wide**: for instance, on Debian-based Linux + distributions, you can install g++ with apt and then use + `update-alternatives `_ + to choose a suitable compiler as default. -3. If you have a conda environment with CUDA toolkit and pyKeOps, the compiling test with ``pykeops.test_numpy_bindings()``will fail unless you also have a system-wide CUDA toolkit installation, due to missing ``cuda.h`` file. See `_, question "How can I compile CUDA (host or device) codes in my environment?" +2. Install a compiler **locally**: if you are using a conda environment, you can + install a new instance of gcc and g++ by following the + `Conda documentation `_. @@ -298,37 +500,44 @@ First of all, make sure that you are using a C++ compiler which is compatible wi Cache directory --------------- -If you experience problems with compilation, it may be a good idea to **flush the build folder** that KeOps uses as a cache for already-compiled formulas. To do this, just type: +If you experience compilation problems, it may be a good idea to **flush the +build folder** that KeOps uses as a cache for already-compiled formulas. To do +this, type: -.. code-block:: python +.. prompt:: python >>> import pykeops pykeops.clean_pykeops() -You can change the build folder by using the ``set_build_folder()`` function: +You can change the build folder with the ``set_build_folder()`` function: -.. code-block:: python +.. prompt:: python >>> import pykeops print(pykeops.get_build_folder()) # display current build_folder - pykeops.set_build_folder("/my/new/location") # change the build folder + pykeops.set_build_folder("/tmp/keops_cache_new_location") # change the build folder print(pykeops.get_build_folder()) # display new build_folder -Note that the command ``set_build_folder()`` without any argument will reset the location to the default one (``~/.keops/build`` on unix-like systems) +Calling ``set_build_folder()`` without any argument resets the location to the +default one (``~/.cache/keops`` on Unix-like systems). Verbosity level --------------- -You can deactivate all messages and warnings by setting the environment variable `PYKEOPS_VERBOSE` to 0. In a terminal, type: +The KeOps verbosity level is an integer equal to 0 (silent), 1 (default) or 2 +(verbose). You can deactivate all messages and warnings by setting the +``PYKEOPS_VERBOSE`` environment variable to 0. In a terminal, type: .. prompt:: bash $ - export PYKEOPS_VERBOSE=0 - python my_script_calling_pykeops.py + PYKEOPS_VERBOSE=0 python my_script_calling_pykeops.py -Alternatively, you can disable verbose compilation from your python script using the function ``pykeops.set_verbose()``. In a python shell, type: +Alternatively, you can disable verbose compilation from your Python script with +``pykeops.set_verbose()``: -.. code-block:: python +.. prompt:: python >>> import pykeops - pykeops.set_verbose(False) + pykeops.set_verbose(0) # no output + pykeops.set_verbose(1) # default verbosity level + pykeops.set_verbose(2) # maximum verbosity level diff --git a/doc/tools/nb_to_doc.py b/doc/tools/nb_to_doc.py index 7c594e837..0b1fe76e3 100755 --- a/doc/tools/nb_to_doc.py +++ b/doc/tools/nb_to_doc.py @@ -3,6 +3,7 @@ Convert empty IPython notebook to a sphinx doc page. """ + import sys from subprocess import check_call as sh diff --git a/keops_version b/keops_version index bb576dbde..7451fe42d 100644 --- a/keops_version +++ b/keops_version @@ -1 +1 @@ -2.3 +2.4b diff --git a/keopscore/keopscore/__init__.py b/keopscore/keopscore/__init__.py index e7beb752e..741b7158d 100644 --- a/keopscore/keopscore/__init__.py +++ b/keopscore/keopscore/__init__.py @@ -1,47 +1,23 @@ import os -from os import path -########################################################### -# Verbosity level -verbose = True -if os.getenv("KEOPS_VERBOSE") == "0": - verbose = False - -here = path.abspath(path.dirname(__file__)) +# Version +here = os.path.abspath(os.path.dirname(__file__)) with open(os.path.join(here, "keops_version"), encoding="utf-8") as v: __version__ = v.read().rstrip() -from keopscore.config import * -from keopscore.utils.code_gen_utils import clean_keops, check_health -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK - -set_build_folder = config.set_different_build_folder - -# flags for debugging : -# prints information about atomic operations during code building -debug_ops = False -# adds C++ code for printing all input and output values -# for all atomic operations during computations -debug_ops_at_exec = False - -# flag for automatic factorization : apply automatic factorization for all formulas before reduction. -auto_factorize = False - -# Initialize CUDA libraries if CUDA is used -if cuda_config.get_use_cuda(): - # Initialize CUDA libraries if necessary - cuda_config._cuda_libraries_available() - from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile - from keopscore.binders.nvrtc.Gpu_link_compile import jit_compile_dll - - if not os.path.exists(jit_compile_dll()): - Gpu_link_compile.compile_jit_compile_dll() +# Config +import keopscore.config +# Create build folder if it does not exist and set it as the default build folder +try: + keopscore.config.path.set_build_folder(reset_all=False) +except Exception as e: + from .utils.messages import KeOps_Warning -# Retrieve the current build folder -build_folder = config.get_build_folder() + KeOps_Warning( + f"An error occurred while setting up keopcore: {e}. Use keopscore.config.check_health() to get details on the current configuration.", + level=1, + ) -# Retrieve details about the current CUDA configuration -show_gpu_config = cuda_config -show_cuda_status = cuda_config.get_use_cuda() -cuda_block_size = cuda_config.get_cuda_block_size() +# expose to the user +set_build_folder = keopscore.config.path.set_build_folder diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index aae5d70ea..ab79001f2 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -1,10 +1,16 @@ import os -from keopscore.config import config + +import keopscore.config from keopscore.utils.code_gen_utils import get_hash_name -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Message +from keopscore.utils.messages import KeOps_Error, KeOps_Message -cpp_flags = config.get_cpp_flags() -get_build_folder = config.get_build_folder +cpp_flag = keopscore.config.cxx.get_compile_options() +cpp_flag += keopscore.config.cxx.get_linking_options() +cpp_flag += keopscore.config.openmp.get_compile_options() +cpp_flag += keopscore.config.openmp.get_include_options() +cpp_flag += keopscore.config.openmp.get_linking_options() +cpp_flag += keopscore.config.cuda.get_include_options() +cpp_flag += keopscore.config.cuda.get_linking_options() class LinkCompile: @@ -31,17 +37,17 @@ def __init__(self): self.use_half, self.use_fast_math, self.device_id, - cpp_flags, + cpp_flag, ) # info_file is the name of the file that will contain some meta-information required by the bindings, e.g. 7b9a611f7e.nfo self.info_file = os.path.join( - get_build_folder(), self.gencode_filename + ".nfo" + keopscore.config.path.get_build_folder(), self.gencode_filename + ".nfo" ) # gencode_file is the name of the source file to be created and then compiled, e.g. 7b9a611f7e.cpp or 7b9a611f7e.cu self.gencode_file = os.path.join( - get_build_folder(), + keopscore.config.path.get_build_folder(), self.gencode_filename + "." + self.source_code_extension, ) @@ -51,23 +57,21 @@ def save_info(self): # dimy (sum of dimensions of j-indexed vectors) f = open(self.info_file, "w") f.write( - f"red_formula={self.red_formula_string}\ndim={self.dim}\ntagI={self.tagI}\ndimy={self.dimy}" + f"red_formula={self.red_formula_string}\ndim={self.dim}\ntagI={self.tagI}\ndimy={self.dimy}\n\n" + + f"dtype={self.dtype}\nsum_scheme={self.sum_scheme_string} in {self.dtypeacc}\n" + + f"tagHostDevice={self.tagHostDevice}\ntagCpuGpu={self.tagCpuGpu}\ntag1D2D={self.tag1D2D}\n\ndevice_id={self.device_id}\n" ) f.close() def read_info(self): # read info_file to retreive dim, tagI, dimy f = open(self.info_file, "r") - string = f.read() + f.readline() # skip line 0 + tmp_dim = f.readline().rstrip("\n").split("=") + tmp_tag = f.readline().rstrip("\n").split("=") + tmp_dimy = f.readline().rstrip("\n").split("=") f.close() - tmp = string.split("\n") - if len(tmp) != 4: - KeOps_Error("Incorrect info file") - tmp_dim, tmp_tag, tmp_dimy = ( - tmp[1].split("="), - tmp[2].split("="), - tmp[3].split("="), - ) + if ( len(tmp_dim) != 2 or tmp_dim[0] != "dim" @@ -77,9 +81,9 @@ def read_info(self): or tmp_dimy[0] != "dimy" ): KeOps_Error("Incorrect info file") - self.dim = eval(tmp_dim[1]) - self.tagI = eval(tmp_tag[1]) - self.dimy = eval(tmp_dimy[1]) + self.dim = int(tmp_dim[1]) + self.tagI = int(tmp_tag[1]) + self.dimy = int(tmp_dimy[1]) def write_code(self): # write the generated code in the source file ; this is used as a subfunction of compile_code @@ -99,15 +103,15 @@ def get_dll_and_params(self): "Generating code for " + self.red_formula.__str__() + " ... ", flush=True, end="", + level=2, ) self.generate_code() self.save_info() - KeOps_Message("OK", use_tag=False, flush=True) + KeOps_Message("OK", use_tag=False, flush=True, level=2) else: self.read_info() return dict( tag=self.gencode_filename, - source_file=self.true_dllname, low_level_code_file=self.low_level_code_file, tagI=self.tagI, use_half=self.use_half, diff --git a/keopscore/keopscore/binders/cpp/Cpu_link_compile.py b/keopscore/keopscore/binders/cpp/Cpu_link_compile.py index 1328f94e8..17305196b 100644 --- a/keopscore/keopscore/binders/cpp/Cpu_link_compile.py +++ b/keopscore/keopscore/binders/cpp/Cpu_link_compile.py @@ -1,4 +1,22 @@ +import os +import sysconfig + +import keopscore.config from keopscore.binders.LinkCompile import LinkCompile +from keopscore.utils.messages import KeOps_Error, KeOps_Message +from keopscore.utils.system_utils import KeOps_OS_Run + +cpp_dtype = { + "float": "float", + "double": "double", +} + +cpu_runtime_src = os.path.join( + keopscore.config.path.get_base_dir_path(), + "binders", + "cpp", + "keops_cpu_runtime.cpp", +) class Cpu_link_compile(LinkCompile): @@ -6,20 +24,78 @@ class Cpu_link_compile(LinkCompile): def __init__(self): LinkCompile.__init__(self) - # these are used for command line compiling mode - self.low_level_code_file = "".encode("utf-8") + self.dllname = os.path.join( + keopscore.config.path.get_build_folder(), + self.gencode_filename + "_cpp" + sysconfig.get_config_var("SHLIB_SUFFIX"), + ) + self.low_level_code_file = self.dllname - # actual dll to be called - self.true_dllname = self.gencode_file # file to check for existence to detect compilation is needed - self.file_to_check = self.gencode_file + self.file_to_check = self.dllname def generate_code(self): # method to generate the code and compile it # generate the code and save it in self.code, by calling get_code method from CpuReduc class : self.get_code() + self.code += self.get_cpu_dll_entrypoint_code() # write the code in the source file self.write_code() + self.compile_dll() # retreive some parameters that will be saved into info_file. self.tagI = self.red_formula.tagI self.dim = self.red_formula.dim + + def compile_dll(self): + compile_command = ( + f"{keopscore.config.cxx.get_cxx_compiler()} " + f"{keopscore.config.cxx.get_compile_options()} " + f"{keopscore.config.openmp.get_compile_options()} " + f"{keopscore.config.openmp.get_include_options()} " + f"{keopscore.config.path.get_include_options()} " + f"{self.gencode_file} " + f"{cpu_runtime_src} " + f"{keopscore.config.cxx.get_linking_options()} " + f"{keopscore.config.openmp.get_linking_options()} " + f"-o {self.dllname}" + ) + KeOps_Message( + "Compiling cpp formula " + self.gencode_filename + " module", + flush=True, + end="", + level=1, + ) + KeOps_Message( + " in cache folder " + keopscore.config.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + KeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) + + out = KeOps_OS_Run(compile_command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling cpp formula {self.gencode_filename}. See compiler output above." + ) + else: + KeOps_Message("OK", use_tag=False, flush=True, level=1) + + def get_cpu_dll_entrypoint_code(self): + if self.dtype not in cpp_dtype: + KeOps_Error( + "The cpp backend only supports float and double formulas, got " + + self.dtype + + "." + ) + + scalar_type = cpp_dtype[self.dtype] + + return f""" + +#include "binders/cpp/keops_cpu_runtime.h" + +extern "C" KEOPS_CPU_EXPORT int keops_cpu_launch(const KeOpsCpuArgs *runtime_args) {{ + return keops_cpu_launch_impl< {scalar_type} >( + runtime_args, &launch_keops_cpu_{self.gencode_filename}< {scalar_type} >); +}} +""" diff --git a/keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp new file mode 100644 index 000000000..ae7385813 --- /dev/null +++ b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp @@ -0,0 +1,66 @@ +#include "binders/cpp/keops_cpu_runtime.h" + +namespace { + +template std::vector keops_cpu_make_vector(const T *data, int size) { + if (data == nullptr || size <= 0) { + return std::vector(); + } + return std::vector(data, data + size); +} + +} // namespace + +template +int keops_cpu_launch_impl(const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher) { + if (runtime_args == nullptr) { + return -1; + } + + std::vector indsi_v = + keops_cpu_make_vector(runtime_args->indsi, runtime_args->n_indsi); + std::vector indsj_v = + keops_cpu_make_vector(runtime_args->indsj, runtime_args->n_indsj); + std::vector indsp_v = + keops_cpu_make_vector(runtime_args->indsp, runtime_args->n_indsp); + + std::vector dimsx_v = + keops_cpu_make_vector(runtime_args->dimsx, runtime_args->n_dimsx); + std::vector dimsy_v = + keops_cpu_make_vector(runtime_args->dimsy, runtime_args->n_dimsy); + std::vector dimsp_v = + keops_cpu_make_vector(runtime_args->dimsp, runtime_args->n_dimsp); + std::vector shapeout_v = + keops_cpu_make_vector(runtime_args->shapeout, runtime_args->n_shapeout); + + std::vector arg_v(runtime_args->nargs); + for (int i = 0; i < runtime_args->nargs; ++i) { + arg_v[i] = reinterpret_cast(runtime_args->arg[i]); + } + + std::vector> argshape_v(runtime_args->nargs); + for (int i = 0; i < runtime_args->nargs; ++i) { + const signed long int *shape = runtime_args->argshape[i]; + int shape_size = runtime_args->argshape_sizes[i]; + if (shape != nullptr && shape_size > 0) { + argshape_v[i].assign(shape, shape + shape_size); + } + } + + return formula_launcher( + runtime_args->dimY, runtime_args->nx, runtime_args->ny, runtime_args->tagI, + runtime_args->tagZero, runtime_args->use_half, runtime_args->dimred, + runtime_args->use_chunk_mode, indsi_v, indsj_v, indsp_v, + runtime_args->dimout, dimsx_v, dimsy_v, dimsp_v, runtime_args->ranges, + shapeout_v, reinterpret_cast(runtime_args->out), arg_v.data(), + argshape_v); +} + +template int keops_cpu_launch_impl( + const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher); + +template int keops_cpu_launch_impl( + const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher); diff --git a/keopscore/keopscore/binders/cpp/keops_cpu_runtime.h b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.h new file mode 100644 index 000000000..aec14e7a0 --- /dev/null +++ b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.h @@ -0,0 +1,65 @@ +#pragma once + +#if defined(_WIN32) +#define KEOPS_CPU_EXPORT __declspec(dllexport) +#define KEOPS_CPU_INTERNAL +#else +#define KEOPS_CPU_EXPORT __attribute__((visibility("default"))) +#define KEOPS_CPU_INTERNAL __attribute__((visibility("hidden"))) +#endif + +#include + +struct KeOpsCpuArgs { + signed long int dimY; + signed long int nx; + signed long int ny; + int tagI; + int tagZero; + int use_half; + signed long int dimred; + int use_chunk_mode; + + int n_indsi; + const int *indsi; + int n_indsj; + const int *indsj; + int n_indsp; + const int *indsp; + + signed long int dimout; + int n_dimsx; + const signed long int *dimsx; + int n_dimsy; + const signed long int *dimsy; + int n_dimsp; + const signed long int *dimsp; + + signed long int **ranges; + + int n_shapeout; + const signed long int *shapeout; + void *out; + + int nargs; + void **arg; + const signed long int **argshape; + const int *argshape_sizes; +}; + +typedef int (*KeOpsCpuLaunchFn)(const KeOpsCpuArgs *); + +template +using KeOpsCpuFormulaLaunchFn = int (*)( + signed long int dimY, signed long int nx, signed long int ny, int tagI, + int tagZero, int use_half, signed long int dimred, int use_chunk_mode, + std::vector indsi, std::vector indsj, std::vector indsp, + signed long int dimout, std::vector dimsx, + std::vector dimsy, std::vector dimsp, + signed long int **ranges, std::vector shapeout, TYPE *out, + TYPE **arg, std::vector> argshape); + +template +KEOPS_CPU_INTERNAL int +keops_cpu_launch_impl(const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher); diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index ccd6eb00d..bcbd72f42 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -1,53 +1,25 @@ +import ctypes import os -from ctypes import create_string_buffer, CDLL, c_int -from os import RTLD_LAZY -import sysconfig -from os.path import join -import keopscore +import keopscore.config from keopscore.binders.LinkCompile import LinkCompile -from keopscore.config import * - -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Message, KeOps_OS_Run -from keopscore.utils.gpu_utils import custom_cuda_include_fp16_path - -cuda_version = cuda_config.get_cuda_version() -jit_binary = config.get_jit_binary() -cxx_compiler = config.get_cxx_compiler() -nvrtc_flags = cuda_config.get_nvrtc_flags() -cuda_available = cuda_config.get_use_cuda() -build_folder = config.get_build_folder() -get_gpu_props = cuda_config.get_gpu_props() - -nvrtc_include = " -I" + config.get_bindings_source_dir() -cuda_include_path = cuda_config.get_cuda_include_path() -if cuda_include_path: - nvrtc_include += " -I" + cuda_include_path - -jit_source_file = join( - config.get_base_dir_path(), "binders", "nvrtc", "keops_nvrtc.cpp" -) - -jit_compile_src = os.path.join( - os.path.abspath(os.path.dirname(__file__)), "nvrtc_jit.cpp" -) - - -def jit_compile_dll(): - return os.path.join( - build_folder, - "nvrtc_jit" + sysconfig.get_config_var("SHLIB_SUFFIX"), - ) +from keopscore.utils.messages import KeOps_Error, KeOps_Message +from keopscore.utils.system_utils import KeOps_OS_Run class Gpu_link_compile(LinkCompile): + """ + This class is used to compile the main dll that will be used to JIT compile the cuda code for each formula. + """ + source_code_extension = "cu" - low_level_code_prefix = "cubin_" if cuda_version >= 11010 else "ptx_" - ngpu, gpu_props_compile_flags = get_gpu_props def __init__(self): # checking that the system has a Gpu : - if not (cuda_available and Gpu_link_compile.ngpu): + if not ( + keopscore.config.cuda.get_use_cuda() + and keopscore.config.cuda.get_n_visible_devices() + ): KeOps_Error( "Trying to compile cuda code... but we detected that the system has no properly configured cuda lib." ) @@ -55,15 +27,28 @@ def __init__(self): LinkCompile.__init__(self) # these are used for JIT compiling mode # low_level_code_file is filename of low level code (PTX for Cuda) or binary (CUBIN for Cuda) - # generated by the JIT compiler, e.g. ptx_7b9a611f7e + # generated by the JIT compiler, e.g. 7b9a611f7e_nvrtc.{ptx,cubin} self.low_level_code_file = os.path.join( - build_folder, - self.low_level_code_prefix + self.gencode_filename, + keopscore.config.path.get_build_folder(), + self.gencode_filename + + "_nvrtc." + + keopscore.config.cuda.get_ir_type().lower(), ).encode("utf-8") - self.my_c_dll = CDLL(jit_compile_dll(), mode=RTLD_LAZY) - # actual dll to be called is the jit binary, TODO: check if this is relevent - self.true_dllname = jit_binary + # Load nvrtc_jit dll, that will be used to compile the cuda code in jit mode + self.my_c_dll = ctypes.CDLL( + keopscore.config.keops_jit_compile_name(type="target"), mode=os.RTLD_LAZY + ) + self.my_c_dll.Compile.argtypes = [ + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_char_p, + ] + self.my_c_dll.Compile.restype = ctypes.c_int + # file to check for existence to detect compilation is needed self.file_to_check = self.low_level_code_file @@ -71,23 +56,25 @@ def generate_code(self): # method to generate the code and compile it # generate the code and save it in self.code, by calling get_code method from GpuReduc class : self.get_code() - # write the code in the source file + # write the code in the source .cu file self.write_code() # we execute the main dll, passing the code as argument, and the name of the low level code file to save the assembly instructions + cuda_include_paths = "\n".join( + path for path in keopscore.config.cuda.get_cuda_include_path() if path + ) + res = self.my_c_dll.Compile( - create_string_buffer(self.low_level_code_file), - create_string_buffer(self.code.encode("utf-8")), - c_int(self.use_half), - c_int(self.use_fast_math), - c_int(self.device_id), - create_string_buffer( - (custom_cuda_include_fp16_path() + os.path.sep).encode("utf-8") - ), + ctypes.create_string_buffer(self.low_level_code_file), + ctypes.create_string_buffer(self.code.encode("utf-8")), + ctypes.c_int(self.use_half), + ctypes.c_int(self.use_fast_math), + ctypes.c_int(self.device_id), + ctypes.create_string_buffer(cuda_include_paths.encode("utf-8")), ) - if res != 0: + if res != keopscore.config.cuda.CUDA_SUCCESS: KeOps_Error( - f"Error when compiling formula (error in nvrtcCompileProgram, nvrtcResult={res})" + f"Error when compiling {self.low_level_code_file} formula (error in nvrtcCompileProgram, nvrtcResult={res})" ) # retreive some parameters that will be saved into info_file. self.tagI = self.red_formula.tagI @@ -95,29 +82,51 @@ def generate_code(self): @staticmethod def get_compile_command( - sourcename=jit_source_file, dllname=jit_binary, extra_flags="" + sourcename, + dllname, + extra_flags="", ): # This is about the main KeOps binary (dll) that will be used to JIT compile all formulas. # If the dll is not present, it compiles it from source, except if check_compile is False. - target_tag = ( - "CUBIN" if Gpu_link_compile.low_level_code_prefix == "cubin_" else "PTX" + return ( + f"{keopscore.config.cxx.get_cxx_compiler()} " + f"{keopscore.config.cuda.get_preprocessing_options()} " + f"{keopscore.config.cxx.get_compile_options()} " + f"{keopscore.config.path.get_include_options()} " + f"{keopscore.config.cuda.get_include_options()} " + f" {extra_flags} " + f"{sourcename} " + f"{keopscore.config.cxx.get_linking_options()} " + f"{keopscore.config.cuda.get_linking_options()} " + f"-o {dllname}" ) - nvrtcGetTARGET = "nvrtcGet" + target_tag - nvrtcGetTARGETSize = nvrtcGetTARGET + "Size" - arch_tag = ( - '\\"sm\\"' - if Gpu_link_compile.low_level_code_prefix == "cubin_" - else '\\"compute\\"' - ) - target_type_define = f"-DnvrtcGetTARGET={nvrtcGetTARGET} -DnvrtcGetTARGETSize={nvrtcGetTARGETSize} -DARCHTAG={arch_tag}" - return f"{cxx_compiler} {nvrtc_flags} {extra_flags} {target_type_define} {nvrtc_include} {Gpu_link_compile.gpu_props_compile_flags} {sourcename} -o {dllname}" @staticmethod - def compile_jit_compile_dll(): - KeOps_Message("Compiling cuda jit compiler engine ... ", flush=True, end="") + def compile_jit_compile_dll(force_recompile=False): + + if ( + os.path.exists(keopscore.config.keops_jit_compile_name(type="target")) + and not force_recompile + ): + return + + KeOps_Message("Compiling cuda jit compiler engine", flush=True, end="", level=1) + KeOps_Message( + " in cache folder " + keopscore.config.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + KeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) command = Gpu_link_compile.get_compile_command( - sourcename=jit_compile_src, dllname=jit_compile_dll() + keopscore.config.keops_jit_compile_name(type="src"), + keopscore.config.keops_jit_compile_name(type="target"), ) - KeOps_OS_Run(command) - KeOps_Message("OK", use_tag=False, flush=True) + out = KeOps_OS_Run(command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling cuda/nvrtc binder {keopscore.config.keops_jit_compile_name(type='target')}. See compiler output above." + ) + else: + KeOps_Message("OK", use_tag=False, flush=True, level=1) diff --git a/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp b/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp index bad02d6d7..bf87be432 100644 --- a/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp +++ b/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp @@ -1,28 +1,30 @@ // nvcc -shared -Xcompiler -fPIC -lnvrtc -lcuda keops_nvrtc.cu -o keops_nvrtc.so // g++ --verbose -L/opt/cuda/lib64 -L/opt/cuda/targets/x86_64-linux/lib/ // -I/opt/cuda/targets/x86_64-linux/include/ -I../../include -shared -fPIC -// -lcuda -lnvrtc -fpermissive -DMAXIDGPU=0 -DMAXTHREADSPERBLOCK0=1024 +// -lcuda -lnvrtc -DMAXIDGPU=0 -DMAXTHREADSPERBLOCK0=1024 // -DSHAREDMEMPERBLOCK0=49152 -DnvrtcGetTARGET=nvrtcGetCUBIN // -DnvrtcGetTARGETSize=nvrtcGetCUBINSize -DARCHTAG=\"sm\" keops_nvrtc.cpp -o // keops_nvrtc.so g++ -std=c++11 -shared -fPIC -O3 -fpermissive -L /usr/lib -L // /opt/cuda/lib64 -lcuda -lnvrtc -DnvrtcGetTARGET=nvrtcGetCUBIN -// -DnvrtcGetTARGETSize=nvrtcGetCUBINSize -DARCHTAG=\"sm\" -// -I/home/bcharlier/projets/keops/keops/keops/include -I/opt/cuda/include +// -DnvrtcGetTARGETSize=nvrtcGetCUBINSize -DARCHTAG=\"sm\" -I/opt/cuda/include // -I/usr/include/python3.10/ -DMAXIDGPU=0 -DMAXTHREADSPERBLOCK0=1024 // -DSHAREDMEMPERBLOCK0=49152 -// /home/bcharlier/projets/keops/keops/keops/binders/nvrtc/keops_nvrtc.cpp -o +// ./keopscore/binders/nvrtc/keops_nvrtc.cpp -o // keops_nvrtc.cpython-310-x86_64-linux-gnu.so -#include + #include #include #include -#include #include #include #include #include -// #include + + +#include +#include +#include #define C_CONTIGUOUS 1 #define USE_HALF 0 @@ -33,7 +35,6 @@ #include "include/utils_pe.h" #include "include/CudaSizes.h" -#include // ------------------------------------------------------------------------ // DevicePointer wrapper diff --git a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp index 439481110..9b219e787 100644 --- a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp +++ b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp @@ -13,16 +13,17 @@ // /home/bcharlier/projets/keops/keops/keops/binders/nvrtc/keops_nvrtc.cpp -o // keops_nvrtc.cpython-310-x86_64-linux-gnu.so -#include #include #include -#include #include #include #include +#include #include #include -// #include + +#include +#include #define C_CONTIGUOUS 1 #define USE_HALF 0 @@ -33,26 +34,24 @@ #include "include/utils_pe.h" #include "include/CudaSizes.h" -#include extern "C" int Compile(const char *target_file_name, const char *cu_code, int use_half, int use_fast_math, int device_id, - const char *cuda_include_path) { + const char *cuda_include_paths) { nvrtcProgram prog; - int numHeaders; - const char *header_names[1]; - const char *header_sources[1]; - - if (use_half) { - numHeaders = 1; - std::ostringstream header_path; - header_path << cuda_include_path << "cuda_fp16.h"; - header_names[0] = "cuda_fp16.h"; - header_sources[0] = read_text_file(header_path.str().c_str()); - } else { - numHeaders = 0; + std::vector compile_options_storage; + std::vector compile_options; + + if (cuda_include_paths != nullptr && cuda_include_paths[0] != '\0') { + std::istringstream include_paths_stream(cuda_include_paths); + std::string include_path; + while (std::getline(include_paths_stream, include_path)) { + if (!include_path.empty()) { + compile_options_storage.emplace_back("--include-path=" + include_path); + } + } } // Get device id from Driver API @@ -72,47 +71,58 @@ extern "C" int Compile(const char *target_file_name, const char *cu_code, arch_flag << "-arch=" << ARCHTAG << "_" << deviceProp_major << deviceProp_minor; - char *arch_flag_char = new char[arch_flag.str().length()]; - arch_flag_char = strdup(arch_flag.str().c_str()); + compile_options_storage.push_back(arch_flag.str()); + if (use_fast_math) { + compile_options_storage.emplace_back("-use_fast_math"); + } + compile_options.reserve(compile_options_storage.size()); + for (const std::string &option : compile_options_storage) { + compile_options.push_back(option.c_str()); + } NVRTC_SAFE_CALL(nvrtcCreateProgram(&prog, // prog cu_code, // buffer NULL, // name - numHeaders, // numHeaders - header_sources, // headers - header_names // includeNames + 0, // numHeaders + NULL, // headers + NULL // includeNames )); - nvrtcResult compileResult; - if (use_fast_math) { - const char *opts[] = {arch_flag_char, "-use_fast_math"}; - compileResult = nvrtcCompileProgram(prog, // prog - 2, // numOptions - opts); // options - } else { - const char *opts[] = {arch_flag_char}; - compileResult = nvrtcCompileProgram(prog, // prog - 1, // numOptions - opts); // options - } + nvrtcResult compileResult = + nvrtcCompileProgram(prog, // prog + compile_options.size(), // numOptions + compile_options.data()); // options // following "if" block is when there is a mismatch between // the device compute capability and the cuda libs versions : typically // when the device is more recent than the lib, the -arch flag may fail to // compile. if (compileResult == NVRTC_ERROR_INVALID_OPTION) { - const char *new_opts[] = {"-use_fast_math"}; - compileResult = nvrtcCompileProgram(prog, // prog - 1, // numOptions - new_opts); // options + std::vector fallback_options; + fallback_options.reserve(compile_options_storage.size()); + for (const std::string &option : compile_options_storage) { + if (option.rfind("-arch=", 0) != 0) { + fallback_options.push_back(option.c_str()); + } + } + compileResult = nvrtcCompileProgram(prog, // prog + fallback_options.size(), // numOptions + fallback_options.data()); // options } if (compileResult != NVRTC_SUCCESS) { + size_t logSize = 0; + nvrtcGetProgramLogSize(prog, &logSize); + if (logSize > 1) { + std::vector log(logSize); + nvrtcGetProgramLog(prog, log.data()); + std::cerr << "[KeOps] NVRTC compile log:\n" << log.data() << std::endl; + } + std::cerr << "[KeOps] nvrtcCompileProgram failed: " + << nvrtcGetErrorString(compileResult) << std::endl; return compileResult; } - delete[] arch_flag_char; - // Obtain PTX or CUBIN from the program. size_t targetSize; NVRTC_SAFE_CALL(nvrtcGetTARGETSize(prog, &targetSize)); diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py new file mode 100644 index 000000000..d9716a7f2 --- /dev/null +++ b/keopscore/keopscore/config/Cuda.py @@ -0,0 +1,754 @@ +import ctypes +import glob +import os + +from keopscore.utils.messages import print_envs +from keopscore.utils.messages import enabled_dict, not_found_str +from keopscore.utils.messages import KeOps_Error, KeOps_Warning +from keopscore.utils.path_utils import ( + _first_matching_file, + _ordered_search_roots, + _path_candidates, +) +from keopscore.utils.system_utils import _find_library_by_names + + +class CudaConfig: + """ + CUDA detection and configuration. + + The main state is stored in the ``_libcuda_info``, ``_libnvrtc_info`` + and ``_headers_*_info`` dictionaries. These dictionaries are filled with the paths + to the corresponding libraries and headers, when found, and with their + ctypes handles when successfully loaded. Detection is performed by + ``_cuda_libraries_available``, which is called by ``set_use_cuda``. The + remaining configuration, including the CUDA version, include options and + preprocessing options, is set from the detection results. + """ + + # CUDA constants + CUDA_SUCCESS = 0 + CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 + CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 + CUDA_BLOCK_SIZE = 192 + + # Cuda detection variables + _use_cuda = False + _cuda_version = -1 + _ir_type = "" # "PTX" or "CUBIN" + _cuda_include_path = [""] # str or list of str + + _visible_devices = "" + _n_visible_devices = -1 + _MaxThreadsPerBlock = [] # list of int + _SharedMemPerBlock = [] # list of int + _cuda_block_size = -1 + + _preprocessing_options = "" + _include_options = "" + _linking_options = "" + + # ------------------------ # + # Search location # + # ------------------------ # + + # default order of precedence for searching CUDA libraries and headers + order_precedence = ("env_vars", "conda", "pip", "system") + + # environment variables that may point to CUDA installations + cuda_env_vars = [ + "CUDA_PATH", + "CUDA_HOME", + "CUDA_ROOT", + "CUDA_TOOLKIT_ROOT_DIR", + "CUDA_VISIBLE_DEVICES", + ] + + # pip installation suffixes to look for when searching for CUDA libraries and headers + pip_suffixes = { + "cu12": [ + os.path.join("nvidia", "cuda_runtime"), + os.path.join("nvidia", "cuda_nvrtc"), + os.path.join("nvidia", "cuda_nvcc"), + os.path.join("nvidia", "cuda_cccl"), + ], + "cu13": [ + os.path.join("nvidia", "cu13"), + ], + } + + # standard system suffixes to look for when searching for CUDA libraries and headers + system_prefixes = ( + os.path.join(os.path.sep, "usr", "local", "cuda"), + os.path.join(os.path.sep, "usr", "local"), + os.path.join(os.path.sep, "opt", "cuda"), + os.path.join(os.path.sep, "usr"), + os.path.join(os.path.sep, "lib"), + ) + + library_suffixes = ( + "lib64", + "lib", + os.path.join("lib", "x86_64-linux-gnu"), + "", + ) + + include_suffixes = ( + "include", + os.path.join("targets", "x86_64-linux", "include"), + os.path.join("targets", "sbsa-linux", "include"), + os.path.join("targets", "aarch64-linux", "include"), + ) + + # ------------------------- # + # Library info dicts # + # ------------------------- # + + _libcuda_info = { + "name": "cuda", + "lib_basename_candidate": ["libcuda.so.*"], + "library": "", # to be filled later + "ctype_handle": None, # to be filled later + } + _libnvrtc_info = { + "name": "nvrtc", + "lib_basename_candidate": ["libnvrtc.so.*"], + "library": "", # to be filled later + } + _libnvrtc_builtins_info = { + "name": "nvrtc-builtins", + "lib_basename_candidate": [ + "libnvrtc-builtins.so*", + "libnvrtc-builtins.alt.so*", + ], + "library": "", # to be filled later + } + _headers_nvrtc_info = { + "header_basename": "nvrtc.h", + "header": "", # to be filled later + } + _headers_cuda_info = { + "header_basename": "cuda.h", + "header": "", # to be filled later + } + _headers_crt_info = { + "header_basename": os.path.join("crt", "device_functions.h"), + "header": "", # to be filled later + } + _headers_nv_info = { + "header_basename": os.path.join("nv", "target"), + "header": "", # to be filled later + } + _headers_fp16_info = { + "header_basename": "cuda_fp16.h", + "header": "", # to be filled later + } + + def __init__(self, platform): + + self.platform = platform + + self.set_visible_devices() + + super().__init__() + + self.set_use_cuda() + + # If cuda is enabled, then we finalize the rest of the config + if self.get_use_cuda(): + self.set_ir_type() + self.set_cuda_include_path() + self.set_linking_options() + self.set_cuda_block_size() + self.set_preprocessing_options() + self.set_include_options() + + def find_library_path(self, lib_dict_info, where_to_search): + """ + Locate a cuda library using an explicit ordered search. + + Arguments: + lib_dict_info (dict): A dictionary containing at least the keys 'name' and 'lib_basename_candidate' for the library to find. This allows the function to be used for finding libcuda or nvrtc by passing the appropriate info dict. + + where_to_search (dict): A dictionary containing the search locations, with keys corresponding to the search types (e.g., "env_vars", "conda", "system", "pip") and values containing the relevant paths or environment variable names. + + Returns: + str: The absolute path to the library file if found, or None if not found. + """ + result = lib_dict_info.copy() + + candidate_roots = _ordered_search_roots( + **where_to_search, + order=self.order_precedence, + ) + + # First try to find the library file using the candidate roots and library suffixes + + result["library"] = _first_matching_file( + _path_candidates(candidate_roots, self.library_suffixes), + result["lib_basename_candidate"], + ) + if result["library"] is None: + result["library"] = _find_library_by_names(result["name"]) + return result + + def find_header_path(self, header_dict_info, where_to_search): + """ + Locate a header file using an explicit ordered search. + + Arguments: + header_basename (str): The basename of the header to find. + + where_to_search (dict): A dictionary containing the search locations, with keys corresponding to the search types (e.g., "env_vars", "conda", "system", "pip") and values containing the relevant paths or environment variable names. + + Returns: + str: The absolute path to the header file if found, or None if not found. + """ + + candidate_roots = _ordered_search_roots( + **where_to_search, + order=self.order_precedence, + ) + + header_dict_info["header"] = _first_matching_file( + _path_candidates(candidate_roots, self.include_suffixes), + header_dict_info["header_basename"], + ) + + if not header_dict_info.get("header"): + return ( + False, + f"{header_dict_info['header_basename']} not found. Make sure the CUDA headers are installed and accessible.", + ) + + return True, "" + + def _find_and_load_libcuda(self, where_to_search): + """ + Locate, load, and initialize the CUDA driver library. + Get the version of the CUDA driver and the number of visible GPUs, as well as some of their attributes. + Update the config state with this information. + """ + self._libcuda_info = self.find_library_path(self._libcuda_info, where_to_search) + libcuda_path = self._libcuda_info["library"] + if not libcuda_path: + return ( + False, + "libcuda not found. Make sure the CUDA driver is installed and accessible.", + ) + + try: + libcuda = ctypes.CDLL(libcuda_path, mode=ctypes.RTLD_GLOBAL) + except OSError as e: + return ( + False, + f"Failed to load library '{libcuda_path}': {e}", + ) + + if libcuda.cuInit(0) != self.CUDA_SUCCESS: + return ( + False, + "libcuda was detected, but driver API could not be initialized (flushing KeOps caches and/or rebooting the system may help).", + ) + + cuda_version = ctypes.c_int() + if libcuda.cuDriverGetVersion(ctypes.byref(cuda_version)) != self.CUDA_SUCCESS: + return ( + False, + "libcuda was detected and initialized, but failed to query CUDA driver API version.", + ) + self._cuda_version = int(cuda_version.value) + + # If we successfully loaded libcuda and initialized it, store the handle in the config for potential future use + self._libcuda_info["ctype_handle"] = libcuda + + nGpus = ctypes.c_int() + if ( + libcuda.cuDeviceGetCount(ctypes.byref(nGpus)) != self.CUDA_SUCCESS + or nGpus.value == 0 + ): + return ( + False, + "libcuda was detected and driver API was initialized, but no working GPU found.", + ) + + self._MaxThreadsPerBlock = [0] * nGpus.value + self._SharedMemPerBlock = [0] * nGpus.value + + for d in range(nGpus.value): + self._MaxThreadsPerBlock[d], self._SharedMemPerBlock[d], err_msg = ( + self._get_device_attributes(libcuda, d) + ) + if err_msg: + return ( + False, + f"libcuda was detected and driver API was initialized, but " + + err_msg, + ) + + self._n_visible_devices = nGpus.value + + return True, "" + + def _find_libnvrtc(self, where_to_search): + """Locate the NVRTC runtime compilation library.""" + self._libnvrtc_info = self.find_library_path( + self._libnvrtc_info, where_to_search + ) + libnvrtc_path = self._libnvrtc_info["library"] + if not libnvrtc_path: + return ( + False, + "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible.", + ) + + return True, "" + + def _find_libnvrtc_builtins(self, folder_to_search): + """ + Locate NVRTC builtins from the same folder as libnvrtc. + + This check is non-blocking: warnings are emitted on failure, and CUDA + detection continues. We intentionally do not preload this library with + ctypes; runtime resolution is handled through linker rpath flags. + """ + + self._libnvrtc_builtins_info = self.find_library_path( + self._libnvrtc_builtins_info, {"system": (folder_to_search,)} + ) + libnvrtc_builtins_path = self._libnvrtc_builtins_info["library"] + if not libnvrtc_builtins_path: + return ( + True, + f"NVRTC builtins library not found in {folder_to_search}. This may cause runtime compilation to fail in environments with multiple CUDA toolkit versions installed.", + ) + + return True, "" + + def _cuda_libraries_available(self): + """ + Check if libcuda (driver) and nvrtc (toolkit) libraries are available. + This is where ```_cuda_include_path``` is set. + + Returns: + True if all three are loadable, False otherwise. + This is also where we handle one single warning if needed. + """ + + # If CUDA_VISIBLE_DEVICES is explicitly set to empty, no GPU is requested. + if os.getenv("CUDA_VISIBLE_DEVICES") == "": + KeOps_Warning( + "CUDA_VISIBLE_DEVICES is set to empty, no GPU will be used. Switching to CPU only." + ) + return False + + where_to_search = { + "env_vars": self.cuda_env_vars, + "conda": "CONDA_PREFIX", + "system": self.system_prefixes, + "pip": (), + } + + # Libcuda (driver) loaded globally so it is available to KeOps shared objects. + # This is usually provided by the NVIDIA Driver (install system-wide) + success_cuda, err_cuda = self._find_and_load_libcuda(where_to_search) + if not success_cuda: + KeOps_Warning(f"{err_cuda}. Switching to CPU only.") + return False + + # Restrict the pip search to the suffixes corresponding to the detected CUDA version, if any. + where_to_search["pip"] = _path_candidates( + self.platform.get_python_package_roots(), + self.pip_suffixes.get(f"cu{self.get_cuda_version(out_type='major')}", ()), + ) + + # libnvrtc as well, since it's required for the runtime compilation of CUDA code. + # This is usually provided by the CUDA Toolkit (install system-wide or in conda/pip). + success_nvrtc, err_nvrtc = self._find_libnvrtc(where_to_search) + if not success_nvrtc: + KeOps_Warning(f"{err_nvrtc}. Switching to CPU only.") + return False + + # Locate NVRTC builtins in the same toolkit folder. + # This helps set runtime search paths in environments where + # multiple CUDA toolkit versions are installed. + _, warning_nvrtc = self._find_libnvrtc_builtins( + os.path.dirname(self._libnvrtc_info["library"]) + ) + if warning_nvrtc: + KeOps_Warning(warning_nvrtc, level=2) + + # Finally, check that we can find the CUDA toolkit headers + for header_dict_info in ( + self._headers_cuda_info, + self._headers_nvrtc_info, + self._headers_crt_info, + self._headers_nv_info, + self._headers_fp16_info, + ): + success, err = self.find_header_path(header_dict_info, where_to_search) + if not success: + KeOps_Warning(f"{err}. Switching to CPU only.") + return False + + return True + + # CUDA Support + def set_use_cuda(self): + """Determine and set whether to use CUDA.""" + self._use_cuda = self._cuda_libraries_available() + + def get_use_cuda(self): + return self._use_cuda + + def is_available(self): + return self.get_use_cuda() + + def print_use_cuda(self): + print(f"CUDA Support: {enabled_dict[self.get_use_cuda() or False]}") + + # CUDA Block Size + def set_cuda_block_size(self): + """Sets default cuda block size.""" + self._cuda_block_size = self.CUDA_BLOCK_SIZE + + def get_cuda_block_size(self): + return self._cuda_block_size + + def print_cuda_block_size(self): + print(f"CUDA Block Size: {self.get_cuda_block_size()}") + + # Visibles GPUs devices + def set_visible_devices(self): + """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" + cuda_visible = os.getenv("CUDA_VISIBLE_DEVICES") + + if cuda_visible is not None: + self._visible_devices = ( + cuda_visible.replace(",", "_") if cuda_visible else "empty" + ) + + def get_visible_devices(self): + """Get the specific GPUs.""" + return self._visible_devices + + # Number of GPUs + def set_n_visible_devices(self): + """Set the number of GPUs detected. This is done in _find_and_load_libcuda.""" + pass + + def get_n_visible_devices(self): + """Get the number of GPUs detected.""" + return self._n_visible_devices + + def print_n_visible_devices(self): + """Print the number of GPUs detected.""" + print(f"Number of GPUs Detected: {self.get_n_visible_devices()}") + + # Libcuda folder + def set_libcuda_folder(self): + """ + Is set in _cuda_libraries_available. + """ + pass + + def get_libcuda_folder(self): + return self._libcuda_info["library"] and os.path.dirname( + self._libcuda_info["library"] + ) + + # Libcuda path + def set_libcuda_path(self): + """ + Is set in _cuda_libraries_available. + """ + pass + + def get_libcuda_path(self): + return self._libcuda_info["library"] + + def print_libcuda_path(self): + print(f"Libcuda Path: {self.get_libcuda_path() or not_found_str}") + + # Libnvrtc folder + def set_libnvrtc_folder(self): + """ + Return nothing if not using cuda + self.libnvrtc_folder is already set in _cuda_libraries_available. + """ + pass + + def get_libnvrtc_folder(self): + return self._libnvrtc_info["library"] and os.path.dirname( + self._libnvrtc_info["library"] + ) + + # Libnvrtc path + def set_libnvrtc_path(self): + """ + Return nothing if not using cuda + self.libnvrtc_path is already set in _cuda_libraries_available. + """ + pass + + def get_libnvrtc_path(self): + return self._libnvrtc_info["library"] + + def print_libnvrtc_path(self): + print(f"Libnvrtc Path: {self.get_libnvrtc_path() or not_found_str}") + + # CUDA Version + def set_cuda_version(self): + """Set in _find_and_load_libcuda.""" + pass + + def get_cuda_version(self, out_type="single_value"): + + major = self._cuda_version // 1000 + minor = (self._cuda_version % 1000) // 10 + + if out_type == "major,minor": + return major, minor + if out_type == "major": + return major + elif out_type == "string": + return f"{major}.{minor}" + else: + return self._cuda_version + + def print_cuda_version(self): + str = f"CUDA Version: {self.get_cuda_version(out_type='string') if self.get_cuda_version() else not_found_str}" + print(str) + + # CUDA Include Path + def set_cuda_include_path(self): + """Set the CUDA include path by searching for cuda.h and nvrtc.h.""" + include_dirs = [ + os.path.dirname(header) + for header in ( + self._libnvrtc_info.get("header"), + self._headers_cuda_info.get("header"), + self._headers_fp16_info.get("header"), + ) + if header + ] + # Remove the crt in the header + include_dirs.extend( + [ + os.path.realpath(os.path.join(os.path.dirname(header), "..")) + for header in ( + self._headers_crt_info.get("header"), + self._headers_nv_info.get("header"), + ) + ] + ) + self._cuda_include_path = list(set(include_dirs)) if include_dirs else ("",) + + def get_cuda_include_path(self): + """ + Attempt to find CUDA headers (cuda.h, nvrtc.h) using an explicit + ordered search over environment variables and standard locations. + """ + return self._cuda_include_path + + def print_cuda_include_path(self): + print( + f"CUDA Include Path: {':'.join(self.get_cuda_include_path()) or not_found_str}" + ) + + # IR type + def set_ir_type(self): + """Set the IR type to be used for nvrtc compilation based on the CUDA version.""" + if self.get_cuda_version() >= 11010: + self._ir_type = "CUBIN" + else: + self._ir_type = "PTX" + + def get_ir_type(self): + return self._ir_type + + # NVRTC include options + def set_include_options(self): + self._include_options += "".join( + f" -I{p}" for p in list(set(self.get_cuda_include_path())) + ) + + def get_include_options(self): + return self._include_options.strip() + + def print_include_options(self): + print(f"CUDA Include Options: {self.get_include_options() or not_found_str}") + + # NVRTC preprocessins options + def set_preprocessing_options(self): + """Set GPU preprocessing compile flags based on detected GPU properties.""" + if self.get_n_visible_devices() == 0: + return + + self._preprocessing_options = f"-DMAXIDGPU={self.get_n_visible_devices() - 1}" + + for d in range(self.get_n_visible_devices()): + self.add_to_preprocessing_options( + f"-DMAXTHREADSPERBLOCK{d}={self.get_MaxThreadsPerBlock()[d]}" + ) + self.add_to_preprocessing_options( + f"-DSHAREDMEMPERBLOCK{d}={self.get_SharedMemPerBlock()[d]}" + ) + + target_tag = "CUBIN" if self.get_ir_type() == "CUBIN" else "PTX" + nvrtcGetTARGET = "nvrtcGet" + target_tag + self.add_to_preprocessing_options(f"-DnvrtcGetTARGET={nvrtcGetTARGET}") + + nvrtcGetTARGETSize = nvrtcGetTARGET + "Size" + self.add_to_preprocessing_options(f"-DnvrtcGetTARGETSize={nvrtcGetTARGETSize}") + + arch_tag = '\\"sm\\"' if self.get_ir_type() == "CUBIN" else '\\"compute\\"' + self.add_to_preprocessing_options(f"-DARCHTAG={arch_tag}") + + def add_to_preprocessing_options(self, flags): + self._preprocessing_options += " " + flags + + def get_preprocessing_options(self): + """Get GPU preprocessing compile flags based on detected GPU properties.""" + return self._preprocessing_options + + def print_preprocessing_options(self): + print( + f"GPU Preprocessing Options: {self.get_preprocessing_options() or not_found_str}" + ) + + # NVRTC linking options + def set_linking_options(self): + """Set the Linking option for nvrt/cuda entry point compilation.""" + + link_options = [] + for lib_info in [self._libcuda_info, self._libnvrtc_info]: + link_options.append( + lib_info["library"] if lib_info["library"] else f"-l{lib_info['name']}" + ) + + # Ensure runtime loader can resolve CUDA toolkit side dependencies + # (e.g. nvrtc-builtins) without requiring ctypes preloading. + rpath_dirs = [] + for lib_info in [ + self._libcuda_info, + self._libnvrtc_info, + self._libnvrtc_builtins_info, + ]: + lib_path = lib_info.get("library") + if lib_path: + lib_dir = os.path.dirname(lib_path) + if lib_dir and lib_dir not in rpath_dirs: + rpath_dirs.append(lib_dir) + + for lib_dir in rpath_dirs: + link_options.append(f"-Wl,-rpath,{lib_dir}") + + self._linking_options = " ".join(link_options) + + def get_linking_options(self): + """Get the Linking option for nvrt/cuda entry point compilation.""" + return self._linking_options + + def print_linking_options(self): + """Print the Linking option for nvrt/cuda entry point compilation.""" + print(f"Linking Options: {self.get_linking_options()}") + + # Helper functions for CUDA attribute detection + def _get_device_attributes(self, libcuda, device_index): + """Return CUDA attributes needed to build compile flags for one device.""" + + device = ctypes.c_int() + if libcuda.cuDeviceGet(ctypes.byref(device), device_index) != self.CUDA_SUCCESS: + return ( + 0, + 0, + f"failed to get device handle for device index {device_index}.", + ) + + output = ctypes.c_int() + if ( + libcuda.cuDeviceGetAttribute( + ctypes.byref(output), + self.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, + device, + ) + != self.CUDA_SUCCESS + ): + return ( + 0, + 0, + f"failed to get max threads per block for device index {device_index}.", + ) + max_threads_per_block = output.value + + if ( + libcuda.cuDeviceGetAttribute( + ctypes.byref(output), + self.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, + device, + ) + != self.CUDA_SUCCESS + ): + return ( + 0, + 0, + f"failed to get shared memory per block for device index {device_index}.", + ) + shared_mem_per_block = output.value + + return max_threads_per_block, shared_mem_per_block, "" + + def set_MaxThreadsPerBlock(self): + pass + + def get_MaxThreadsPerBlock(self): + return self._MaxThreadsPerBlock + + def set_SharedMemPerBlock(self): + pass + + def get_SharedMemPerBlock(self): + return self._SharedMemPerBlock + + def print_all(self): + """Print all CUDA-related configuration""" + + # CUDA Support + print("=" * 60) + print("CUDA Support") + print("=" * 60) + + self.print_use_cuda() + if self.get_use_cuda(): + self.print_n_visible_devices() + self.print_cuda_version() + self.print_libcuda_path() + self.print_libnvrtc_path() + self.print_cuda_include_path() + + self.print_preprocessing_options() + self.print_include_options() + self.print_linking_options() + + # Print relevant environment variables. + print_envs(self.cuda_env_vars) + + +if __name__ == "__main__": + from keopscore.config.Platform import PlatformConfig + from keopscore.config.CxxCompiler import CxxCompilerConfig + + # from keopscore.config.OpenMP import OpenMPConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() + + # openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + # openmp_info.print_all() + + cuda_info = CudaConfig(platform_info) + cuda_info.print_all() diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py new file mode 100644 index 000000000..bbf4f4602 --- /dev/null +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -0,0 +1,233 @@ +import os +import re +import shutil + +from keopscore.utils.messages import print_envs +from keopscore.utils.messages import not_found_str +from keopscore.utils.messages import KeOps_Error +from keopscore.utils.system_utils import KeOps_OS_Run + + +class CxxCompilerConfig: + """ + Common platform and C++ compiler configuration. + """ + + _cxx_compiler = "" + _cxx_env_flags = "" + _compile_options = "" + _linking_options = "" + + _disable_pragma_unrolls = True + _use_clang_on_macos = False + + cxx_envs = ["CXX", "CXXFLAGS"] + + def __init__(self, platform): + + # Platform instance get info on system + self.platform = platform + + self.set_cxx_compiler() + self.set_use_clang_on_macos() + self.set_cxx_env_flags() + self.set_compile_options() + self.set_linking_options() + self.set_disable_pragma_unrolls() + + # C++ Compiler Detection + def set_cxx_compiler(self): + """Set the C++ compiler.""" + self._cxx_compiler = self.detect_cxx_compiler() + + def get_cxx_compiler(self): + """Get the C++ compiler.""" + return self._cxx_compiler + + def print_cxx_compiler(self): + """Print the C++ compiler information.""" + print(f"C++ Compiler Path: {self.get_cxx_compiler() or not_found_str}") + + def detect_cxx_compiler(self): + """Return the best available C++ compiler for the current platform.""" + preferred_compilers = [] + + env_cxx = os.getenv("CXX") + if env_cxx: + preferred_compilers.append(env_cxx) + + # Common alias for the default C++ compiler + preferred_compilers.append("c++") + + if self.platform.get_platform() == "Darwin": + # On macOS, prioritize clang++ over g++ + preferred_compilers.extend(("clang++", "g++")) + else: + preferred_compilers.append(["g++", "clang++"]) + + # Return the first available compiler from the preferred list. + for compiler in preferred_compilers: + cxx_compiler_path = shutil.which(compiler) + if cxx_compiler_path: + return cxx_compiler_path + else: + KeOps_Error( + "No C++ compiler found. Define CXX environment variable or install g++." + ) + + return None + + def get_cxx_compiler_version(self): + """Detect if using Apple Clang.""" + compiler_info = ( + KeOps_OS_Run(f"{self.get_cxx_compiler()} --version") + .stdout.decode("utf-8") + .splitlines()[0] + .strip() + ) + return compiler_info + + def print_cxx_compiler_version(self): + """Print the C++ compiler version information.""" + print( + f"C++ Compiler Version: {self.get_cxx_compiler_version() or not_found_str}" + ) + + def set_use_clang_on_macos(self): + """Detect if using clang on macOS.""" + self._use_clang_on_macos = ( + "clang" in self.get_cxx_compiler_version() + and self.platform.get_platform() == "Darwin" + if self.get_cxx_compiler_version() + else False + ) + + def get_use_clang_on_macos(self): + return self._use_clang_on_macos + + # Disable Pragma Unrolls + def set_disable_pragma_unrolls(self): + """Set the flag for disabling pragma unrolls.""" + self._disable_pragma_unrolls = True + + def get_disable_pragma_unrolls(self): + """Get the flag for disabling pragma unrolls.""" + return self._disable_pragma_unrolls + + def print_disable_pragma_unrolls(self): + """Print the flag for disabling pragma unrolls.""" + status = "no (default)" if self.get_disable_pragma_unrolls() else "yes" + print(f"Pragma unroll: {status}") + + # Compile Options + def set_compile_options(self): + """Set the compile options.""" + + self._compile_options = self.get_cxx_env_flags() + + # compilation / behavior + self.add_to_compile_option("-std=c++11 -O3 -flto=auto -fpermissive") + # code model + self.add_to_compile_option("-fPIC") + + # architecture (macOS ARM) + if self.platform.get_platform() == "Darwin" and self.platform.get_machine() in [ + "arm64", + "arm64e", + ]: + self.add_to_compile_option("-arch arm64") + + def add_to_compile_option(self, flags): + self._compile_options += " " + flags + + def get_compile_options(self): + """Get the compile options.""" + return self._compile_options + + def print_compile_options(self): + """Print the compile options.""" + print(f"Compile Options: {self.get_compile_options()}") + + # Linking options + def set_linking_options(self): + """Set the linking options.""" + # linking / output + self.add_to_linking_options("-shared") + + # linker behavior (macOS) + if self.platform.get_platform() == "Darwin" and self.get_use_clang_on_macos(): + self.add_to_linking_options("-undefined dynamic_lookup") + + # architecture (macOS ARM) + if self.platform.get_platform() == "Darwin" and self.platform.get_machine() in [ + "arm64", + "arm64e", + ]: + self.add_to_linking_options("-arch arm64") + + def add_to_linking_options(self, flags): + self._linking_options += " " + flags + + def get_linking_options(self): + """Get the linking options.""" + return self._linking_options + + def print_linking_options(self): + print(f"Linking Options: {self.get_linking_options()}") + + def get_dynamic_loader_linking_options(self): + """Return extra linker options needed for dlopen/dlsym support.""" + if self.platform.get_platform() == "Linux": + return "-ldl" + # macOS: dlopen/dlsym are in standard C library, no extra flag needed + return "" + + # C++ Environment Flags. + def set_cxx_env_flags(self): + """Recover CXXFLAGS and sanitize incompatible flags when needed.""" + cxxflags = os.getenv("CXXFLAGS", "") + + # Apple clang does not support some GCC tuning flags that may leak in + # from external environments (e.g. conda, torch build stacks). + if ( + cxxflags + and self.platform.get_platform() == "Darwin" + and self.get_use_clang_on_macos() + ): + cxxflags = re.sub( + r"(?:^|\s)-(?:march|mtune|mfpmath)(?:=\S+|\s+\S+)?", + "", + cxxflags, + ).strip() + + self._cxx_env_flags = cxxflags + + def get_cxx_env_flags(self): + """Get the C++ environment flags.""" + return self._cxx_env_flags + + # Comprehensive Platform and Compiler Information + def print_all(self): + """Print all platform and C++ compiler information.""" + + print("=" * 60) + print("C/C++ Compiler Information") + print("=" * 60) + + self.print_cxx_compiler() + self.print_cxx_compiler_version() + self.print_compile_options() + self.print_linking_options() + + # Print relevant environment variables. + print_envs(self.cxx_envs) + + +if __name__ == "__main__": + from keopscore.config.Platform import PlatformConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() diff --git a/keopscore/keopscore/config/Debug.py b/keopscore/keopscore/config/Debug.py new file mode 100644 index 000000000..be50213e7 --- /dev/null +++ b/keopscore/keopscore/config/Debug.py @@ -0,0 +1,49 @@ +import os + +from keopscore.utils.messages import KeOps_Warning + + +class DebugConfig: + + # prints information about atomic operations during code building + _debug_ops = False + # adds C++ code for printing all input and output values for all atomic operations during computations + _debug_ops_at_exec = False + + def __init__(self): + env_val = os.getenv("KEOPS_VERBOSE") + if env_val is None: + self._verbose = 1 + else: + val = int(env_val) + if val in (0, 1, 2): + self._verbose = val + else: + KeOps_Warning( + f"Invalid KEOPS_VERBOSE value: {env_val}. Verbose level must be 0, 1 or 2. Defaulting to 1." + ) + self._verbose = 1 + + def set_debug_ops(self, debug_ops): + self._debug_ops = debug_ops + + def get_debug_ops(self): + return self._debug_ops + + def set_debug_ops_at_exec(self, debug_ops_at_exec): + self._debug_ops_at_exec = debug_ops_at_exec + + def get_debug_ops_at_exec(self): + return self._debug_ops_at_exec + + def get_verbose(self): + return self._verbose + + def set_verbose(self, val): + if val in (0, 1, 2): + self._verbose = val + else: + KeOps_Warning( + f"Invalid verbose value: {val}. Verbose level must be 0, 1 or 2. Keeping previous value: {self._verbose}." + ) + return diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py new file mode 100644 index 000000000..52bd26009 --- /dev/null +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -0,0 +1,241 @@ +import os +import sys +import sysconfig + +import keopscore +from keopscore.utils.path_utils import ensure_directory +from keopscore.utils.messages import not_found_str, print_envs + + +class KeOpsPathConfig: + """ + Class to manage the path to the KeOps library. + """ + + _base_dir_path = "" + + _keops_cache_folder = "" + _build_folder = "" + + _default_build_folder_name = "" + _default_build_path = "" + + _include_options = "" + + path_env_vars = ("KEOPS_CACHE_FOLDER",) + + def __init__(self, platform, cuda): + + self.platform = platform + self.cuda = cuda + + # Initialize common configuration settings + self.set_base_dir_path() + self.set_keops_cache_folder() + + self.set_default_build_folder_name() + self.set_default_build_path() # should be done at the end.. + self.set_include_options() + + # Base Directory Path + def set_base_dir_path(self): + """Set the base directory path.""" + self._base_dir_path = os.path.abspath( + os.path.join(os.path.dirname(os.path.realpath(__file__)), "..") + ) + + def get_base_dir_path(self): + """Get the base directory path.""" + return self._base_dir_path + + def print_base_dir_path(self): + """Print the base directory path.""" + print(f"Base Directory Path: {self.get_base_dir_path() or not_found_str}") + + # KeOps Cache Folder + def set_keops_cache_folder(self): + """Set the KeOps cache folder.""" + + if os.getenv("KEOPS_CACHE_FOLDER"): + cache_folder = os.getenv("KEOPS_CACHE_FOLDER") + else: # fallback to default cache folder in user home directory + version = getattr(keopscore, "__version__", "") + cache_folder = os.path.join( + os.path.expanduser("~"), ".cache", f"keops{version}" + ) + + os.makedirs(cache_folder, exist_ok=True) + self._keops_cache_folder = cache_folder + + def get_keops_cache_folder(self): + """Get the KeOps cache folder.""" + return self._keops_cache_folder + + def print_keops_cache_folder(self): + """Print the KeOps cache folder.""" + print(f"KeOps Cache Folder: {self.get_keops_cache_folder() or not_found_str}") + + # Build Folder Management + def set_build_folder( + self, path=None, read_save_file=False, write_save_file=True, reset_all=True + ): + """ + Set or update the build folder path for KeOps. + + Parameters: + - path: The new build folder path. If None, it will be determined based on saved settings or defaults. + - read_save_file: If True, read the build folder path from a save file if path is not provided. + - write_save_file: If True, write the new build folder path to the save file. + """ + + # If path is not given, we either read the save file or use the default build path + save_file = os.path.join( + self.get_keops_cache_folder(), "build_folder_location.txt" + ) + if not path: + if read_save_file and os.path.isfile(save_file): + with open(save_file, "r") as f: + path = f.read() + else: + path = self.get_default_build_path() + + # Remove the old build path from sys.path if it's there. + old_build_folder = self.get_build_folder() + if old_build_folder and old_build_folder in sys.path: + sys.path.remove(old_build_folder) + + # Create the folder if not yet done and add the new build path to sys.path + ensure_directory(path, add_to_syspath=True) + # Update _build_folder to the new path + self._build_folder = path + + # Saving the location of the build path in a file + if write_save_file: + with open(save_file, "w") as f: + f.write(path) + + # Reset all cached formulas if needed + if reset_all: + # Reset cached formulas + keopscore.get_keops_dll.get_keops_dll.reset( + new_save_folder=self._build_folder + ) + + # Handle CUDA-specific recompilation if CUDA is used + if self.cuda.get_use_cuda(): + # Recompile the nvrtc binder if needed + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile + + Gpu_link_compile.compile_jit_compile_dll(force_recompile=reset_all) + + def get_build_folder(self): + return self._build_folder + + # Default Build Path + def set_default_build_folder_name(self): + """Set the default build folder name.""" + name_parts = [ + "_".join(self.platform.get_uname()[:3]), + f"python{self.platform.get_python_version()}", + ] + if self.cuda.get_use_cuda(): + name_parts.append(f"CUDA{self.cuda.get_cuda_version()}") + + visible_devices = self.cuda.get_visible_devices() + if visible_devices: + name_parts.append(f"VISIBLE_DEVICES{visible_devices}") + + self._default_build_folder_name = "_".join(name_parts) + + def get_default_build_folder_name(self): + """Return the platform-specific default build folder suffix.""" + return self._default_build_folder_name + + def print_default_build_folder_name(self): + """Print the default build folder name.""" + print(f"Default Build Folder Name: {self.get_default_build_folder_name()}") + + def set_default_build_path(self): + """Set the default build path.""" + self._default_build_path = ensure_directory( + os.path.join( + self.get_keops_cache_folder(), self.get_default_build_folder_name() + ), + add_to_syspath=True, + ) + + def get_default_build_path(self): + """Get the default build path.""" + return self._default_build_path + + def print_default_build_path(self): + """Print the default build path.""" + print(f"Default Build Path: {self.get_default_build_path() or not_found_str}") + + # include options management + def set_include_options(self): + """Set the include options for compilation.""" + self.add_to_include_option(f" -I{self.get_base_dir_path()}") + + def add_to_include_option(self, option): + """Add an include option for compilation.""" + self._include_options += f" {option}" + + def get_include_options(self): + """Get the include options for compilation.""" + return self._include_options + + def print_include_options(self): + """Print the include options for compilation.""" + print(f"Include Options: {self.get_include_options() or not_found_str}") + + # helpers + def get_python_extension_path(self, basename, suffix="EXT_SUFFIX"): + return os.path.join( + self.get_build_folder(), basename + sysconfig.get_config_var(suffix) + ) + + def target_needs_update(self, target, source): + return not os.path.exists(target) or os.path.getmtime( + source + ) > os.path.getmtime(target) + + # Comprehensive Path Information + def print_all(self): + """Print all path-related information.""" + + # Base Configuration + print("=" * 60) + print(f"KeOps Paths Configuration") + print("=" * 60) + + self.print_base_dir_path() + self.print_keops_cache_folder() + self.print_default_build_folder_name() + self.print_default_build_path() + self.print_include_options() + + # Print relevant environment variables. + print_envs(self.path_env_vars) + + +if __name__ == "__main__": + from keopscore.config.Platform import PlatformConfig + from keopscore.config.CxxCompiler import CxxCompilerConfig + from keopscore.config.OpenMP import OpenMPConfig + from keopscore.config.Cuda import CudaConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() + + openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + openmp_info.print_all() + + cuda_info = CudaConfig(platform_info) + cuda_info.print_all() + + keops_info = KeOpsPathConfig(platform_info, cuda_info) + keops_info.print_all() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py new file mode 100644 index 000000000..dbc14235f --- /dev/null +++ b/keopscore/keopscore/config/OpenMP.py @@ -0,0 +1,394 @@ +import importlib.util +import os +import tempfile + +from keopscore.utils.messages import KeOps_Message, print_envs +from keopscore.utils.messages import enabled_dict, not_found_str +from keopscore.utils.messages import KeOps_Warning +from keopscore.utils.path_utils import ( + _first_matching_file, + _ordered_search_roots, + _path_candidates, +) +from keopscore.utils.system_utils import ( + _find_library_by_names, + get_include_file_abspath, + KeOps_OS_Run, +) + + +class OpenMPConfig: + """ + Class for OpenMP detection and configuration. + """ + + _use_OpenMP = False + + _libomp_folder = "" + _libomp_include_path = "" + + _compile_options = "" + _include_options = "" + _linking_options = "" + + openmp_env_vars = ( + "OMP_PATH", + "LIBOMP_PATH", + "OpenMP_ROOT", + "OpenMP_ROOT_DIR", + ) + + system_prefixes = [ + # self.get_brew_prefix() added later on, + os.path.join(os.path.sep, "opt", "homebrew"), + os.path.join(os.path.sep, "usr", "local"), + os.path.join(os.path.sep, "usr"), + ] + + _library_suffixes = ( + "lib", + "lib64", + os.path.join("opt", "libomp", "lib"), + ) + + _include_suffixes = ( + "include", + os.path.join("opt", "libomp", "include"), + ) + + _omp_info = { + "name": ["omp", "gomp"], + "lib_basename_candidate": [ + "libomp.dylib", + "libomp.so*", + "libgomp.dylib", + "libgomp.so*", + ], + "header_basename": "omp.h", + "library": "", # to be filled later + "header": "", # not needed + "ctype_handle": None, # to be filled later + } + + def __init__(self, platform, cxx): + + self.platform = platform + self.cxx = cxx + + # On MacOS: Add brew prefix to OpenMP search paths if it exists + if self.platform.get_brew_prefix(): + self.system_prefixes.insert( + 0, + self.platform.get_brew_prefix(), + ) + + # Detect OpenMP once and set related configuration variables. + omp_available = self._omp_is_available() + if omp_available: + self.set_libomp_include_path() + self.set_libomp_folder() + self.set_compile_options() + self.set_include_options() + self.set_linking_options() + + # Check if the compiler supports OpenMP. + self.set_use_OpenMP(omp_available) + + def find_install_path(self, lib_dict_info): + result = lib_dict_info.copy() + + # First try to find OpenMP library using standard names via ctypes /cxx compiler. + header_path = get_include_file_abspath( + lib_dict_info["header_basename"], self.cxx.get_cxx_compiler() + ) + + #### + KeOps_Message("OpenMP header search using standard names:", level=2) + KeOps_Message( + f" Trying header name: {lib_dict_info['header_basename']}", level=2 + ) + KeOps_Message(f" Found header path: {header_path or not_found_str}", level=2) + #### + + if not header_path: + KeOps_Message( + " Standard-name library probing skipped: omp.h was not found.", + level=2, + ) + else: + result["header"] = header_path + + for name in lib_dict_info["name"]: + library_path = _find_library_by_names((name,)) + + #### + KeOps_Message("OpenMP library search using standard names:", level=2) + KeOps_Message(f" Trying library name: {name}", level=2) + KeOps_Message( + f" Found library path: {library_path or not_found_str}", + level=2, + ) + #### + + if not library_path: + continue + + if header_path: + result["library"] = library_path + return result + + # If that fails, search for OpenMP headers and libraries in common locations. + candidate_roots = _ordered_search_roots( + env_vars=self.openmp_env_vars, + conda="CONDA_PREFIX", + system=self.system_prefixes, + ) + + # libraries + result["library"] = _first_matching_file( + _path_candidates(candidate_roots, self._library_suffixes), + lib_dict_info["lib_basename_candidate"], + ) + + #### + KeOps_Message("OpenMP library search in common locations:", level=2) + KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) + KeOps_Message(f" Library search suffixes: {self._library_suffixes}", level=2) + KeOps_Message( + f" Found library path: {result['library'] or not_found_str}", level=2 + ) + #### + + # headers + result["header"] = _first_matching_file( + _path_candidates(candidate_roots, self._include_suffixes), + lib_dict_info["header_basename"], + ) + + #### + KeOps_Message("OpenMP header search in common locations:", level=2) + KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) + KeOps_Message(f" Header search suffixes: {self._include_suffixes}", level=2) + KeOps_Message( + f" Found header path: {result['header'] or not_found_str}", level=2 + ) + #### + + return result + + def _omp_is_available(self): + self._omp_info = self.find_install_path(self._omp_info) + library_path = self._omp_info.get("library") + header_path = self._omp_info.get("header") + liomp_path = bool( + library_path + and header_path + and os.path.exists(library_path) + and os.path.exists(header_path) + ) + if not liomp_path: + KeOps_Warning( + "libomp not found. Set OMP_PATH, LIBOMP_PATH, or OpenMP_ROOT if it is installed in a non-standard location." + ) + return False + return True + + # OpenMP support + def set_use_OpenMP(self, omp_available=None): + if omp_available is None: + omp_available = self._omp_is_available() + self._use_OpenMP = omp_available and self.check_compiler_for_openmp() + if not self._use_OpenMP: + self._compile_options = "" + self._include_options = "" + self._linking_options = "" + + def get_use_OpenMP(self): + """Boolean to determine if OpenMP is available *and* can be used through cxx compiler""" + return self._use_OpenMP + + def print_use_OpenMP(self): + print(f"OpenMP Support: {enabled_dict[self.get_use_OpenMP() or False]}") + + # OpenMP library path + def get_libomp_path(self): + """try to locate OpenMP libraries""" + return self._omp_info["library"] + + def print_libomp_path(self): + print(f"OpenMP Library Path: {self._omp_info['library'] or not_found_str}") + + def set_libomp_folder(self): + """Set the OpenMP library directory (containing .so or .dylib files).""" + # This is set in set_openmplib_path if the library is found, otherwise it remains None. + self._libomp_folder = self._omp_info["library"] and os.path.dirname( + self._omp_info["library"] + ) + + def get_libomp_folder(self): + """Get the OpenMP library directory (containing .so or .dylib files).""" + return self._libomp_folder + + # OpenMP header path + def set_libomp_include_path(self): + """Set the OpenMP include directory (containing headers).""" + self._libomp_include_path = self._omp_info["header"] + + def get_libomp_include_path(self): + """Get the OpenMP include directory (containing headers).""" + return self._libomp_include_path + + def print_libomp_include_path(self): + print(f"OpenMP Header Path: {self.get_libomp_include_path() or not_found_str}") + + def get_openmp_include_dir(self): + """Get the OpenMP include directory (containing headers).""" + return self._omp_info["header"] and os.path.dirname(self._omp_info["header"]) + + # Helper functions + def check_compiler_for_openmp(self): + """Attempt to compile a simple OpenMP program to check if the compiler supports OpenMP.""" + + if not self.cxx.get_cxx_compiler(): + KeOps_Warning("No C++ compiler available to check for OpenMP support.") + return False + + test_program = """ + #include + int main() { + #pragma omp parallel + {} + return 0; + } + """ + with tempfile.NamedTemporaryFile("w", suffix=".cpp", delete=False) as f: + f.write(test_program) + test_file = f.name + + # Try to compile with compiler's OpenMP options directly + compile_command_str = ( + f"{self.cxx.get_cxx_compiler()} " + f"{self.cxx.get_compile_options()} " + f"{self.get_compile_options()} " + f"{test_file} " + f"{self.get_include_options()} " + f"{self.get_linking_options()} " + f"-o {test_file}.out" + ) + + out = KeOps_OS_Run(compile_command_str, print_warning=True) + if out.returncode == 0: + os.remove(test_file) + os.remove(test_file + ".out") + return True + else: + os.remove(test_file) + if os.path.exists(test_file + ".out"): + os.remove(test_file + ".out") + KeOps_Warning( + f"{self.cxx.get_cxx_compiler()} does not support OpenMP. OpenMP support will be disabled." + ) + return False + + # C++ Compiler Options + def set_compile_options(self): + # Apple clang does not support -fopenmp directly; -Xpreprocessor is required. + if self.cxx.get_use_clang_on_macos(): + self._compile_options += " -Xpreprocessor" + + self._compile_options += " -fopenmp" + + def get_compile_options(self): + return self._compile_options.strip() + + def print_compile_options(self): + print(f"Compile Options: {self.get_compile_options()}") + + # C++ Compiler Include Options + def set_include_options(self): + self._include_options = ( + f"-I{self.get_openmp_include_dir()}" + if self.get_openmp_include_dir() + else "" + ) + + def get_include_options(self): + return self._include_options + + def print_include_options(self): + print(f"Include Options: {self.get_include_options()}") + + # C++ linking Options + def set_linking_options(self): + link_flags = [] + libomp_path = self.get_libomp_path() + + libomp_folder = libomp_path and os.path.dirname(libomp_path) + + if libomp_folder: + link_flags.append(f"-L{libomp_folder}") + + # Ensure runtime loader can find libomp on macOS non-system paths. + if ( + self.platform.get_platform() == "Darwin" + and libomp_folder + and not importlib.util.find_spec("torch") + ): + # Force-link OpenMP runtime to avoid unresolved symbols at dlopen time. + # N.B. (Joan) We added this because when libomp was installed via homebrew, we had these unresolved symbols at runtime. + # However I cannot reproduce the issue now, so it may be unnecessary. + + lib_basename = os.path.basename(libomp_path) + if lib_basename.startswith("lib"): + lib_name = lib_basename[3:].split(".")[0] + if lib_name: + link_flags.append(f"-l{lib_name}") + + if ( + "homebrew" in libomp_folder + ): # only when libomp is installed via homebrew, because when it is installed via conda, it is useless and produces a "duplicate rpath" warning. + link_flags.append(f"-Wl,-rpath,{libomp_folder}") + + self._linking_options = " ".join(link_flags) + + def get_linking_options(self): + return self._linking_options.rstrip() + + def print_linking_options(self): + print(f"Linking Options: {self.get_linking_options()}") + + # OpenMP configuration printing + def print_all(self): + """ + Print all OpenMP-related configuration and system health status. + """ + + print("=" * 60) + print(f"OpenMP Configuration") + print("=" * 60) + + self.print_use_OpenMP() + self.print_libomp_path() + self.print_libomp_include_path() + + if self.get_use_OpenMP(): + self.print_compile_options() + self.print_include_options() + self.print_linking_options() + + # Print relevant environment variables. + print_envs(self.openmp_env_vars) + + +if __name__ == "__main__": + from keopscore.config.Platform import PlatformConfig + from keopscore.config.CxxCompiler import CxxCompilerConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_info = CxxCompilerConfig(platform_info) + cxx_info.print_all() + + openmp_info = OpenMPConfig(platform_info, cxx_info) + openmp_info.print_all() diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index 238b44042..aba21036b 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -1,91 +1,222 @@ +import os import platform +import site import sys -import os +import sysconfig + +from keopscore.utils.path_utils import _unique_paths +from keopscore.utils.messages import not_found_str, print_envs +from keopscore.utils.system_utils import KeOps_OS_Run -class DetectPlatform: +class PlatformConfig: """ Class for detecting the operating system, Python version, and environment type. """ + _os = "" + _platform = "" + _machine = "" + _uname = "" + _python_version = "" + _python_executable = "" + _python_package_roots = [] + _env_type = "" + _brew_prefix = "" + + platform_envs = [ + "PYTHONPATH", + "PATH", + "VIRTUAL_ENV", + "CONDA_DEFAULT_ENV", + "CONDA_PREFIX", + ] + def __init__(self): - self.os = None - self.python_version = None - self.env_type = None self.set_os() + self.set_platform() + self.set_machine() + self.set_uname() self.set_python_version() + self.set_python_executable() + self.set_python_package_roots() self.set_env_type() + self.set_brew_prefix() + # OS Detection (Distribution Name and Version) def set_os(self): """Set the operating system.""" + self._os = self.detect_os() + + def get_os(self): + return self._os + + def print_os(self): + print(f"Operating System: {self.get_os()}") + + @staticmethod + def detect_os(): + """Return a human-readable operating system description.""" if platform.system() == "Linux": try: with open("/etc/os-release") as f: info = dict(line.strip().split("=", 1) for line in f if "=" in line) name = info.get("NAME", "Linux").strip('"') version = info.get("VERSION_ID", "").strip('"') - self.os = f"{platform.system()} {name} {version}" + return f"{platform.system()} {name} {version}" except FileNotFoundError: - self.os = "Linux (distribution info not found)" - else: - self.os = platform.system() + " " + platform.version() + return "Linux (distribution info not found)" - def get_os(self): - return self.os + return platform.system() + " " + platform.version() - def print_os(self): - print(f"Operating System: {self.os}") + # Platform detection (Darwin, Windows, Linux, etc.) + def set_platform(self): + self._platform = platform.system() + def get_platform(self): + return self._platform + + # Machine architecture detection (x86_64, arm64, etc.) + def set_machine(self): + self._machine = platform.machine() + + def get_machine(self): + return self._machine + + def print_machine(self): + print(f"Machine Architecture: {self.get_machine()}") + + # uname detection + def set_uname(self): + self._uname = platform.uname() + + def get_uname(self): + return self._uname + + # Python Version Detection def set_python_version(self): """Set the Python version.""" - self.python_version = platform.python_version() + self._python_version = platform.python_version() def get_python_version(self): - return self.python_version + return self._python_version def print_python_version(self): - print(f"Python Version: {self.python_version}") + print(f"Python Version: {self.get_python_version()}") + + # Python Executable Detection + def set_python_executable(self): + """Set the Python executable path.""" + self._python_executable = sys.executable + + def get_python_executable(self): + return self._python_executable + + def print_python_executable(self): + print(f"Python Executable: {self.get_python_executable()}") + + # Python package roots detection + def set_python_package_roots(self): + """Set the Python package roots where pip-installed wheels may live.""" + self._python_package_roots = self.detect_python_package_roots() + + def get_python_package_roots(self): + return self._python_package_roots + def print_python_package_roots(self): + roots = self.get_python_package_roots() + rendered_roots = ":".join(str(path) for path in roots) + print(f"Python Package Roots: {rendered_roots or not_found_str}") + + @staticmethod + def detect_python_package_roots(): + """Return Python package roots where pip-installed wheels may live.""" + package_roots = [] + getters = [ + lambda: getattr(site, "getsitepackages", lambda: [])(), + lambda: [site.getusersitepackages()], + lambda: [sysconfig.get_path("purelib")], + lambda: [sysconfig.get_path("platlib")], + lambda: sys.path, + ] + for getter in getters: + try: + # Some site helpers are unavailable in embedded or non-standard Python builds. + package_roots.extend(getter() or []) + except Exception: + continue + return _unique_paths(package_roots) + + # Environment Type Detection def set_env_type(self): """Set the environment type (conda, virtualenv, or system).""" - if "CONDA_DEFAULT_ENV" in os.environ: - self.env_type = f"conda ({os.environ['CONDA_DEFAULT_ENV']})" - elif hasattr(sys, "real_prefix") or ( - hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix - ): - self.env_type = "virtualenv" - else: - self.env_type = "system" + self._env_type = self.detect_env_type() def get_env_type(self): - return self.env_type + return self._env_type def print_env_type(self): - print(f"Environment Type: {self.env_type}") + print(f"Environment Type: {self.get_env_type()} {sys.prefix}", end="") + if self.get_env_type() != "system": + print(f" (base at {sys.base_prefix})", end="") + print() + + # Brew package system + def set_brew_prefix(self): + """Get Homebrew prefix path using KeOps_OS_Run""" + if self.get_platform() != "Darwin": + return + + out = KeOps_OS_Run(f"brew --prefix", print_warning=False) + self._brew_prefix = ( + out.stdout.decode("utf-8").strip() if out.returncode == 0 else None + ) + + def get_brew_prefix(self): + """Get Homebrew prefix path using KeOps_OS_Run""" + return self._brew_prefix + + def print_brew_prefix(self): + if self.get_platform() == "Darwin": + brew_prefix = self.get_brew_prefix() + if brew_prefix: + print(f"Homebrew Prefix: {brew_prefix}") + else: + print("Homebrew not found or not installed.") + @staticmethod + def detect_env_type(): + """Return whether Python runs in conda, virtualenv, or the system env.""" + if "CONDA_DEFAULT_ENV" in os.environ and os.environ["CONDA_DEFAULT_ENV"]: + return f"conda ({os.environ['CONDA_DEFAULT_ENV']})" + if hasattr(sys, "real_prefix") or ( + hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix + ): + return "virtualenv" + return "system" + + # Comprehensive Platform Information def print_all(self): """ - Print all platform-related information, including relevant environment variables. + Print all platform-related information. """ - print("\nPlatform Information") + + print("=" * 60) + print("Platform Information") print("=" * 60) + self.print_os() self.print_python_version() self.print_env_type() + self.print_python_executable() + self.print_python_package_roots() - # Print relevant environment variables - print("\nRelevant Environment Variables") - print("-" * 60) - env_vars = [ - "CONDA_DEFAULT_ENV", - "VIRTUAL_ENV", - "PATH", - "PYTHONPATH", - "CXX", - ] - for var in env_vars: - value = os.environ.get(var) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") + self.print_brew_prefix() + + # Print relevant environment variables. + print_envs(self.platform_envs) + + +if __name__ == "__main__": + platform_info = PlatformConfig() + platform_info.print_all() diff --git a/keopscore/keopscore/config/ReductionTuning.py b/keopscore/keopscore/config/ReductionTuning.py new file mode 100644 index 000000000..895bde41b --- /dev/null +++ b/keopscore/keopscore/config/ReductionTuning.py @@ -0,0 +1,93 @@ +class ReductionTuningConfig: + """ + Configuration and state management for chunked computation schemes. + Special computation scheme for dim>100 + """ + + # Constants for chunking strategy + _dimchunk = 64 + _dim_treshold_chunk = 146 + _specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] + + _enable_chunks = True + _dimfinalchunk = 64 + _enable_final_chunk = True + _mult_var_highdim = False + + # Automatic factorization formulas + _auto_factorize = False + + def __init__(self): + pass + + def get_dimchunk(self): + """Get the current chunk dimension.""" + return self._dimchunk + + def get_dim_treshold_chunk(self): + """Get the current dimension threshold for chunking.""" + return self._dim_treshold_chunk + + def get_specdims_use_chunk(self): + """Get the current list of specific dimensions for which chunking is used.""" + return self._specdims_use_chunk + + def get_enable_chunks(self): + """Get the current enable_chunk state.""" + return self._enable_chunks + + def set_enable_chunks(self, val): + """Set enable_chunk state from int (1=True, 0=False, -1=no change).""" + if val == 1: + self._enable_chunks = True + elif val == 0: + self._enable_chunks = False + # val == -1 means keep previous value + + def get_dimfinalchunk(self): + """Get the current final chunk dimension.""" + return self._dimfinalchunk + + def set_dimfinalchunk(self, val): + """Set the final chunk dimension.""" + self._dimfinalchunk = val + + def set_enable_finalchunk(self, val): + """Set enable_final_chunk state from int (1=True, 0=False, -1=no change).""" + if val == 1: + self._enable_final_chunk = True + elif val == 0: + self._enable_final_chunk = False + # val == -1 means keep previous value + + def get_enable_final_chunk(self): + """Get the current enable_final_chunk state.""" + return self._enable_final_chunk + + def set_mult_var_highdim(self, val): + """Set mult_var_highdim state from int (1=True, 0=False, -1=no change).""" + if val == 1: + self._mult_var_highdim = True + elif val == 0: + self._mult_var_highdim = False + # val == -1 means keep previous value + + def get_mult_var_highdim(self): + """Get the current mult_var_highdim state.""" + return self._mult_var_highdim + + def set_auto_factorize(self, val): + """Set automatic factorization state.""" + self._auto_factorize = bool(val) + + def get_auto_factorize(self): + """Get automatic factorization state.""" + return self._auto_factorize + + def use_final_chunks(self, red_formula): + """Determine if final chunks mode should be used for this formula.""" + return ( + self.get_enable_final_chunk() + and self.get_mult_var_highdim() + and red_formula.dim > self.get_dim_treshold_chunk() + ) diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index fc31a0b9f..20a456230 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -1,47 +1,89 @@ +import os +import shutil + # Import the configuration classes -from .base_config import Config -from .cuda import CUDAConfig -from .openmp import OpenMPConfig -from .Platform import DetectPlatform +from .Cuda import CudaConfig +from .CxxCompiler import CxxCompilerConfig +from .Debug import DebugConfig +from .KeOpsPath import KeOpsPathConfig +from .OpenMP import OpenMPConfig +from .Platform import PlatformConfig +from .ReductionTuning import ReductionTuningConfig +from keopscore.utils.messages import KeOps_Error, KeOps_Message -# Instantiate the configurations -config = Config() -platform_detector = DetectPlatform() -cuda_config = CUDAConfig() -openmp_config = OpenMPConfig() +# Instantiate the configurations once at import time to preserve the existing API. +debug = DebugConfig() +platform = PlatformConfig() +cxx = CxxCompilerConfig(platform) +openmp = OpenMPConfig(platform, cxx) +cuda = CudaConfig(platform) +path = KeOpsPathConfig(platform, cuda) +reduction = ReductionTuningConfig() -__all__ = [ - "config", - "platform_detector", - "cuda_config", - "openmp_config", - "get_config", - "get_platform_config", - "get_cuda_config", - "get_openmp_config", -] -# Lazy initializers -_instances = {} +def keops_jit_compile_name(type="src"): + basename = "nvrtc_jit" + if type == "src": + return os.path.join( + path.get_base_dir_path(), "binders", "nvrtc", basename + ".cpp" + ) + + return path.get_python_extension_path(basename, suffix="SHLIB_SUFFIX") -def get_instance(key, factory): - if key not in _instances: - _instances[key] = factory - return _instances[key] +def check_health(infos="all"): + """ + Check the health of the specified configuration. + """ + if infos == "all" or infos == "platform": + platform.print_all() + if infos == "all" or infos == "cxx": + cxx.print_all() + if infos == "all" or infos == "openmp": + openmp.print_all() + if infos == "all" or infos == "cuda": + cuda.print_all() + if infos == "all" or infos == "path": + path.print_all() -def get_config(): - return get_instance("config", config) +def clean_keops(recompile_jit_binary=True, verbose=True): + build_path = path.get_build_folder() + default_build_path = path.get_default_build_path() -def get_cuda_config(): - return get_instance("cuda_config", cuda_config) + if os.path.abspath(build_path) != os.path.abspath(default_build_path): + KeOps_Error( + f"Your build folder is set to {build_path}, which is not the default build folder. For safety reasons, the clean_keops function will not delete files in this folder. If you want to clean this folder, please do it manually." + ) + if build_path and os.path.isdir(build_path): + jit_binary = keops_jit_compile_name(type="target") + # TODO: Add a safety check to prevent accidental deletion of important directories. + for entry in os.scandir(build_path): + if recompile_jit_binary or os.path.abspath(entry.path) != os.path.abspath( + jit_binary + ): + if entry.is_dir(follow_symlinks=False): + shutil.rmtree(entry.path) + else: + os.remove(entry.path) -def get_openmp_config(): - return get_instance("openmp_config", openmp_config) + if verbose: + KeOps_Message(f"{build_path} has been cleaned.") + # Re-initialize the build folder after cleaning. + path.set_build_folder(reset_all=True) -def get_platform_config(): - return get_instance("platform_detector", platform_config) + +__all__ = [ + "platform", + "cxx", + "openmp", + "cuda", + "path", + "debug", + "reduction", + "check_health", + "clean_keops", +] diff --git a/keopscore/keopscore/config/base_config.py b/keopscore/keopscore/config/base_config.py deleted file mode 100644 index 3519bc5d2..000000000 --- a/keopscore/keopscore/config/base_config.py +++ /dev/null @@ -1,440 +0,0 @@ -import os -from os.path import join -import platform -import sys -import shutil -from pathlib import Path -import keopscore -from keopscore.utils.misc_utils import KeOps_Warning, KeOps_OS_Run, get_brew_prefix -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK - - -class Config: - """ - Base configuration class for the KeOps library. - This class contains common attributes and methods shared by other configuration classes. - """ - - # Common attributes - base_dir_path = None - bindings_source_dir = None - keops_cache_folder = None - default_build_folder_name = None - default_build_path = None - jit_binary = None - cxx_compiler = None - cpp_env_flags = None - compile_options = None - cpp_flags = None - disable_pragma_unrolls = None - os = None - _build_folder = None - - def __init__(self): - - # Initialize common configuration settings - self.set_base_dir_path() - self.set_bindings_source_dir() - self.set_keops_cache_folder() - self.set_default_build_folder_name() - self.set_default_build_path() - self.set_jit_binary() - self.set_cxx_compiler() - self.set_cpp_env_flags() - self.set_compile_options() - self.set_cpp_flags() - self.set_disable_pragma_unrolls() - self.set_os() - - # Setters, getters, and print methods for common attributes - - def set_base_dir_path(self): - """Set the base directory path.""" - self.base_dir_path = os.path.abspath( - join(os.path.dirname(os.path.realpath(__file__)), "..") - ) - - def get_base_dir_path(self): - """Get the base directory path.""" - return self.base_dir_path - - def print_base_dir_path(self): - """Print the base directory path.""" - print(f"Base Directory Path: {self.base_dir_path}") - - def set_bindings_source_dir(self): - """Set the bindings source directory.""" - self.bindings_source_dir = self.get_base_dir_path() - - def get_bindings_source_dir(self): - """Get the bindings source directory.""" - return self.bindings_source_dir - - def print_bindings_source_dir(self): - """Print the bindings source directory.""" - print(f"Bindings Source Directory: {self.bindings_source_dir}") - - def set_keops_cache_folder(self): - """Set the KeOps cache folder.""" - self.keops_cache_folder = os.getenv("KEOPS_CACHE_FOLDER") - if self.keops_cache_folder is None: - self.keops_cache_folder = join( - os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" - ) - # Ensure the cache folder exists - os.makedirs(self.keops_cache_folder, exist_ok=True) - - def get_keops_cache_folder(self): - """Get the KeOps cache folder.""" - return self.keops_cache_folder - - def print_keops_cache_folder(self): - """Print the KeOps cache folder.""" - print(f"KeOps Cache Folder: {self.keops_cache_folder}") - - def set_default_build_folder_name(self): - """Set the default build folder name.""" - uname = platform.uname() - self.default_build_folder_name = ( - "_".join(uname[:3]) + f"_p{sys.version.split(' ')[0]}" - ) - - def get_default_build_folder_name(self): - """Get the default build folder name.""" - return self.default_build_folder_name - - def print_default_build_folder_name(self): - """Print the default build folder name.""" - print(f"Default Build Folder Name: {self.default_build_folder_name}") - - def set_default_build_path(self): - """Set the default build path.""" - self.default_build_path = join( - self.keops_cache_folder, self.default_build_folder_name - ) - # Ensure the build path exists - os.makedirs(self.default_build_path, exist_ok=True) - # Add the build path to sys.path - if self.default_build_path not in sys.path: - sys.path.append(self.default_build_path) - # Initialize _build_path - self._build_folder = self.default_build_path - - def get_default_build_path(self): - """Get the default build path.""" - return self.default_build_path - - def print_default_build_path(self): - """Print the default build path.""" - print(f"Default Build Path: {self.default_build_path}") - - def set_jit_binary(self): - """Set the path to the JIT binary.""" - self.jit_binary = join(self.default_build_path) - - def get_jit_binary(self): - """Get the path to the JIT binary.""" - return self.jit_binary - - def print_jit_binary(self): - """Print the path to the JIT binary.""" - print(f"JIT Binary Path: {self.jit_binary}") - - def set_os(self): - """Set the operating system.""" - if platform.system() == "Linux": - try: - with open("/etc/os-release") as f: - info = dict(line.strip().split("=", 1) for line in f if "=" in line) - name = info.get("NAME", "Linux").strip('"') - version = info.get("VERSION_ID", "").strip('"') - self.os = f"{platform.system()} {name} {version}" - except FileNotFoundError: - self.os = "Linux (distribution info not found)" - else: - self.os = platform.system() + " " + platform.version() - - def get_os(self): - return self.os - - def print_os(self): - print(f"Operating System: {self.os}") - - def set_cxx_compiler(self): - """Set the C++ compiler.""" - env_cxx = os.getenv("CXX") - if env_cxx and shutil.which(env_cxx): - self.cxx_compiler = env_cxx - else: - # On macOS, try clang++ first, as it's the canonical C++ compiler driver for Apple Clang - if platform.system() == "Darwin": - if shutil.which("clang++"): - self.cxx_compiler = "clang++" - elif shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No suitable C++ compiler (clang++ or g++) found on macOS." - ) - else: - # On Linux or other systems, fall back to g++ if available - if shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No C++ compiler found. Define CXX environment variable or install g++." - ) - - def get_cxx_compiler(self): - """Get the C++ compiler.""" - return self.cxx_compiler - - def print_cxx_compiler(self): - """Print the C++ compiler.""" - print(f"C++ Compiler: {self.cxx_compiler}") - - def set_cpp_env_flags(self): - """Set the C++ environment flags.""" - self.cpp_env_flags = os.getenv("CXXFLAGS") if "CXXFLAGS" in os.environ else "" - - def get_cpp_env_flags(self): - """Get the C++ environment flags.""" - return self.cpp_env_flags - - def print_cpp_env_flags(self): - """Print the C++ environment flags.""" - print(f"C++ Environment Flags (CXXFLAGS): {self.cpp_env_flags}") - - def set_compile_options(self): - """Set the compile options.""" - self.compile_options = " -shared -fPIC -O3 -std=c++11" - - def get_compile_options(self): - """Get the compile options.""" - return self.compile_options - - def print_compile_options(self): - """Print the compile options.""" - print(f"Compile Options: {self.compile_options}") - - def get_use_Apple_clang(self): - """Detect if using Apple Clang.""" - is_apple_clang = False - if platform.system() == "Darwin": - compiler_info = KeOps_OS_Run(f"c++ --version").stdout.decode("utf-8") - # Check if 'Apple clang' appears in the output - is_apple_clang = "Apple clang" in compiler_info - return is_apple_clang - - def set_cpp_flags(self): - """Set the C++ compiler flags.""" - self.cpp_flags = f"{self.cpp_env_flags} {self.compile_options}" - self.cpp_flags += f" -I{self.base_dir_path}/include" - self.cpp_flags += f" -I{self.bindings_source_dir}" - - # Add OpenMP flags based on compiler - if self.get_use_Apple_clang(): - # For Apple Clang, you need to specify OpenMP library location - brew_prefix = get_brew_prefix() - if brew_prefix is not None: - self.cpp_flags += f" -Xpreprocessor -fopenmp" - self.cpp_flags += f" -I{brew_prefix}/opt/libomp/include" - self.cpp_flags += f" -L{brew_prefix}/opt/libomp/lib" - else: - # For GCC and other compilers - self.cpp_flags += " -fopenmp" - - # Specific check for Apple Silicon chips - if platform.system() == "Darwin" and platform.machine() in ["arm64", "arm64e"]: - self.cpp_flags += " -arch arm64" - - if platform.system() == "Darwin": - self.cpp_flags += " -undefined dynamic_lookup" - self.cpp_flags += " -flto" - - self.cpp_flags += " -flto=auto" - - def get_cpp_flags(self): - """Get the C++ compiler flags.""" - return self.cpp_flags - - def print_cpp_flags(self): - """Print the C++ compiler flags.""" - print(f"C++ Compiler Flags: {self.cpp_flags}") - - def set_disable_pragma_unrolls(self): - """Set the flag for disabling pragma unrolls.""" - self.disable_pragma_unrolls = True - - def get_disable_pragma_unrolls(self): - """Get the flag for disabling pragma unrolls.""" - return self.disable_pragma_unrolls - - def print_disable_pragma_unrolls(self): - """Print the flag for disabling pragma unrolls.""" - status = "Enabled" if self.disable_pragma_unrolls else "Disabled" - print(f"Disable Pragma Unrolls: {status}") - - def set_different_build_folder( - self, path=None, read_save_file=False, write_save_file=True, reset_all=True - ): - """ - Set or update the build folder path for KeOps. - - Parameters: - - path: The new build folder path. If None, it will be determined based on saved settings or defaults. - - read_save_file: If True, read the build folder path from a save file if path is not provided. - - write_save_file: If True, write the new build folder path to the save file. - - reset_all: If True, reset all cached formulas and recompile necessary components. - """ - # If path is not given, we either read the save file or use the default build path - save_file = join(self.keops_cache_folder, "build_folder_location.txt") - if not path: - if read_save_file and os.path.isfile(save_file): - with open(save_file, "r") as f: - path = f.read() - else: - path = self.default_build_path - - # Create the folder if not yet done - os.makedirs(path, exist_ok=True) - - # Remove the old build path from sys.path if it's there - if self._build_folder and self._build_folder in sys.path: - sys.path.remove(self._build_folder) - # Update _build_folder to the new path - self._build_folder = path - # Add the new build path to sys.path - if self._build_folder not in sys.path: - sys.path.append(self._build_folder) - - # Saving the location of the build path in a file - if write_save_file: - with open(save_file, "w") as f: - f.write(path) - - # Reset all cached formulas if needed - if reset_all: - # Reset cached formulas - keopscore.get_keops_dll.get_keops_dll.reset( - new_save_folder=self._build_folder - ) - # Handle CUDA-specific recompilation if CUDA is used - if hasattr(self, "_use_cuda") and self._use_cuda: - from keopscore.binders.nvrtc.Gpu_link_compile import ( - Gpu_link_compile, - jit_compile_dll, - ) - - if not os.path.exists(jit_compile_dll()): - Gpu_link_compile.compile_jit_compile_dll() - - def get_build_folder(self): - return self._build_folder - - # Environment variables printing method - def print_environment_variables(self): - """Print relevant environment variables.""" - print("\nRelevant Environment Variables:") - env_vars = [ - "KEOPS_CACHE_FOLDER", - "CXX", - "CXXFLAGS", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") - - def print_all(self): - """ - Print all base configuration - """ - - # Base Configuration - print(f"\nBase Configuration") - print("-" * 60) - - # Base Directory Path - base_dir_path = self.get_base_dir_path() - base_dir_status = ( - CHECK_MARK - if base_dir_path and os.path.exists(base_dir_path) - else CROSS_MARK - ) - print(f"Base Directory Path: {base_dir_path or 'Not Found'} {base_dir_status}") - - # Bindings Source Directory - bindings_source_dir = self.get_bindings_source_dir() - bindings_source_dir_status = ( - CHECK_MARK - if bindings_source_dir and os.path.exists(bindings_source_dir) - else CROSS_MARK - ) - print( - f"Bindings Source Directory: {bindings_source_dir or 'Not Found'} {bindings_source_dir_status}" - ) - - # KeOps Cache Folder - keops_cache_folder = self.get_keops_cache_folder() - keops_cache_folder_status = ( - CHECK_MARK - if keops_cache_folder and os.path.exists(keops_cache_folder) - else CROSS_MARK - ) - print( - f"KeOps Cache Folder: {keops_cache_folder or 'Not Found'} {keops_cache_folder_status}" - ) - - # Default Build Folder Name - default_build_folder_name = self.get_default_build_folder_name() - print(f"Default Build Folder Name: {default_build_folder_name}") - - # Default Build Path - default_build_path = self.get_default_build_path() - default_build_path_status = ( - CHECK_MARK - if default_build_path and os.path.exists(default_build_path) - else CROSS_MARK - ) - print( - f"Default Build Path: {default_build_path or 'Not Found'} {default_build_path_status}" - ) - - # JIT Binary Path - jit_binary = self.get_jit_binary() - jit_binary_status = ( - CHECK_MARK if jit_binary and os.path.exists(jit_binary) else CROSS_MARK - ) - print(f"JIT Binary Path: {jit_binary or 'Not Found'} {jit_binary_status}") - - # Disable Pragma Unrolls - disable_pragma_unrolls = self.get_disable_pragma_unrolls() - disable_status = CHECK_MARK if disable_pragma_unrolls else CROSS_MARK - status_text = "Enabled" if disable_pragma_unrolls else "Disabled" - print(f"Disable Pragma Unrolls: {status_text} {disable_status}") - - # Compile Options - compile_options = self.get_compile_options() - print(f"Compile Options: {compile_options}") - - # C++ Compiler Flags - cpp_flags = self.get_cpp_flags() - print(f"C++ Compiler Flags: {cpp_flags}") - - # Print relevant environment variables. - print("\nRelevant Environment Variables:") - env_vars = [ - "KEOPS_CACHE_FOLDER", - "CXXFLAGS", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") diff --git a/keopscore/keopscore/config/chunks.py b/keopscore/keopscore/config/chunks.py deleted file mode 100644 index 4efb79c97..000000000 --- a/keopscore/keopscore/config/chunks.py +++ /dev/null @@ -1,65 +0,0 @@ -# special computation scheme for dim>100 - -enable_chunk = True - - -def get_enable_chunk(): - global enable_chunk - return enable_chunk - - -def set_enable_chunk(val): - global enable_chunk - if val == 1: - enable_chunk = True - elif val == 0: - enable_chunk = False - - -dimchunk = 64 -dim_treshold_chunk = 146 -specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] - -# special mode for formula of the type sum_j k(x_i,y_j)*b_j with high dimensional b_j -enable_final_chunk = True - - -def set_enable_finalchunk(val): - global enable_final_chunk - if val == 1: - enable_final_chunk = True - elif val == 0: - enable_final_chunk = False - - -dimfinalchunk = 64 - - -def get_dimfinalchunk(): - global dimfinalchunk - return dimfinalchunk - - -def set_dimfinalchunk(val): - global dimfinalchunk - dimfinalchunk = val - - -def use_final_chunks(red_formula): - global enable_final_chunk - global mult_var_highdim - global dim_treshold_chunk - return ( - enable_final_chunk and mult_var_highdim and red_formula.dim > dim_treshold_chunk - ) - - -mult_var_highdim = False - - -def set_mult_var_highdim(val): - global mult_var_highdim - if val == 1: - mult_var_highdim = True - elif val == 0: - mult_var_highdim = False diff --git a/keopscore/keopscore/config/cuda.py b/keopscore/keopscore/config/cuda.py deleted file mode 100644 index 320f700e9..000000000 --- a/keopscore/keopscore/config/cuda.py +++ /dev/null @@ -1,518 +0,0 @@ -import os -import ctypes -from ctypes.util import find_library -from ctypes import ( - c_int, - c_void_p, - c_char_p, - CDLL, - byref, - cast, - POINTER, - Structure, - RTLD_GLOBAL, -) -from pathlib import Path -import shutil -from os.path import join -import platform -import subprocess -import sys -import keopscore -from keopscore.utils.misc_utils import KeOps_Warning -from keopscore.utils.misc_utils import KeOps_OS_Run -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK, get_include_file_abspath - - -class CUDAConfig: - """ - Class for CUDA detection and configuration. - """ - - # CUDA constants - CUDA_SUCCESS = 0 - CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 - CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 - - # Cuda attributes - libcuda_folder = None - libnvrtc_folder = None - cuda_include_path = None - nvrtc_flags = None - cuda_version = None - n_gpus = 0 - gpu_compile_flags = "" - cuda_message = "" - specific_gpus = None - cuda_block_size = None - - def __init__(self): - self.set_keops_cache_folder() - self.set_default_build_folder_name() - self.set_specific_gpus() - self.set_build_folder() - self.set_cxx_compiler() - self.set_use_cuda() - # If cuda is enabled, then we finalize the rest of the config - if self._use_cuda: - self.set_libcuda_folder() - self.set_libnvrtc_folder() - self.set_nvrtc_flags() - self.set_cuda_block_size() - - def _try_load_library(self, lib_name): - """ - Attempt to locate and load libraties. - Returns: - success (bool): True if the library was found and loaded. - abspath (str): Absolute path to the loaded library if success==True, else "". - error_msg (str): Contains error details if success==False, else "". - """ - # Find library - found_path = find_library(lib_name) - if not found_path: - return (False, "", f"Library '{lib_name}' not found on this system.") - - # Try to load it - try: - lib_handle = CDLL(found_path, mode=RTLD_GLOBAL) - except OSError as e: - return (False, "", f"Failed to load library '{lib_name}': {e}") - - class LINKMAP(Structure): - _fields_ = [("l_addr", c_void_p), ("l_name", c_char_p)] - - try: - # Attempt to load libdl to use dlinfo - libdl_path = find_library("dl") - if not libdl_path: - # If we can't find libdl, we can't do dlinfo; fallback - return (True, found_path, "") - - libdl = CDLL(libdl_path) - dlinfo = libdl.dlinfo - dlinfo.argtypes = (c_void_p, c_int, c_void_p) - dlinfo.restype = c_int - - lmptr = c_void_p() - # RTLD_DI_LINKMAP = 2 - result = dlinfo(lib_handle._handle, 2, byref(lmptr)) - if result != 0: - # dlinfo call failed, fallback - return (True, found_path, "") - - abspath_bytes = cast(lmptr, POINTER(LINKMAP)).contents.l_name - abspath_str = abspath_bytes.decode("utf-8") - if abspath_str: - return (True, abspath_str, "") - else: - return (True, found_path, "") - except Exception as err: - # If anything goes wrong, fallback to found_path - return (True, found_path, "") - - def _cuda_libraries_available(self): - """ - Check if both cuda and nvrtc libraries are available. - Returns: - True if both cuda and nvrtc are loadable, False otherwise. - This is also where we handle one single warning if needed. - """ - - # This step loads "libcuda.so (driver) and libnvrtc (cuda tool kit) **Globaly** to - # make cuda avalaible to keops shared objects - success_cuda, cuda_path, err_cuda = self._try_load_library("cuda") - success_nvrtc, nvrtc_path, err_nvrtc = self._try_load_library("nvrtc") - - if not success_cuda or not success_nvrtc: - self.cuda_message = "CUDA libraries not found or could not be loaded; Switching to CPU only." - KeOps_Warning(self.cuda_message) - - return False - - # If both succeeded, store their folder paths - self.libcuda_folder = os.path.dirname(cuda_path) - self.libnvrtc_folder = os.path.dirname(nvrtc_path) - return True - - def set_use_cuda(self): - """Determine and set whether to use CUDA.""" - self._use_cuda = True - if not self._cuda_libraries_available(): - self._use_cuda = False - - self.get_cuda_version() - self.get_cuda_include_path() - self.get_gpu_props() - if self.n_gpus == 0 and self._use_cuda: - self._use_cuda = False - self.cuda_message = "CUDA libraries detected, but no GPUs found on this system; Switching to CPU only." - KeOps_Warning(self.cuda_message) - - def get_use_cuda(self): - return self._use_cuda - - def print_use_cuda(self): - status = "Enabled ✅" if self._use_cuda else "Disabled ❌" - print(f"CUDA Support: {status}") - - def set_cuda_block_size(self, cuda_block_size=192): - """Sets default cuda block size.""" - self.cuda_block_size = cuda_block_size - - def get_cuda_block_size(self): - return self.cuda_block_size - - def print_cuda_block_size(self): - print(f"CUDA Block Size: {self.cuda_block_size}") - - def set_specific_gpus(self): - """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" - self.specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES") - if self.specific_gpus: - # Modify the build folder name to include GPU specifics - gpu_suffix = self.specific_gpus.replace(",", "_") - self.default_build_folder_name += f"_CUDA_VISIBLE_DEVICES_{gpu_suffix}" - - def get_specific_gpus(self): - """Get the specific GPUs.""" - return self.specific_gpus - - def print_specific_gpus(self): - """Print the specific GPUs.""" - if self.specific_gpus: - print(f"Specific GPUs (CUDA_VISIBLE_DEVICES): {self.specific_gpus}") - else: - print("Specific GPUs (CUDA_VISIBLE_DEVICES): Not Set") - - def set_cxx_compiler(self): - """Set the C++ compiler.""" - env_cxx = os.getenv("CXX") - if env_cxx and shutil.which(env_cxx): - self.cxx_compiler = env_cxx - elif shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No C++ compiler found. You need to either define the CXX environment variable pointing to a valid compiler, or ensure that 'g++' is installed and in your PATH." - ) - - def set_keops_cache_folder(self): - """Set the KeOps cache folder.""" - self.keops_cache_folder = os.getenv("KEOPS_CACHE_FOLDER") - if self.keops_cache_folder is None: - self.keops_cache_folder = join( - os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" - ) - # Ensure the cache folder exists - os.makedirs(self.keops_cache_folder, exist_ok=True) - - def set_default_build_folder_name(self): - """Set the default build folder name.""" - uname = platform.uname() - self.default_build_folder_name = ( - "_".join(uname[:3]) + f"_p{sys.version.split(' ')[0]}" - ) - - def set_build_folder(self): - self.build_folder = join( - self.keops_cache_folder, self.default_build_folder_name - ) - - def get_build_folder(self): - return self.build_folder - - def set_libcuda_folder(self): - """ - Return nothing if not using cuda - self.libcuda_folder is already set in _cuda_libraries_available. - """ - if not self._use_cuda: - return - - def get_libcuda_folder(self): - return self.libcuda_folder - - def set_libnvrtc_folder(self): - """ - Return nothing if not using cuda - self.libnvrtc_folder is already set in _cuda_libraries_available. - """ - if not self._use_cuda: - return - - def get_libnvrtc_folder(self): - return self.libnvrtc_folder - - def get_cuda_version(self, out_type="single_value"): - if not self._use_cuda: - self.cuda_version = None - return None - try: - libcudart_path = find_library("cudart") - if not libcudart_path: - self.cuda_version = None - return None - - libcudart = ctypes.CDLL(libcudart_path) - cuda_version = ctypes.c_int() - libcudart.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) - cuda_version_value = int(cuda_version.value) - - if out_type == "single_value": - self.cuda_version = cuda_version_value - return cuda_version_value - - major = cuda_version_value // 1000 - minor = (cuda_version_value % 1000) // 10 - - if out_type == "major,minor": - return major, minor - elif out_type == "string": - return f"{major}.{minor}" - except Exception as e: - KeOps_Warning(f"Could not determine CUDA version: {e}") - self.cuda_version = None - return None - - def get_cuda_include_path(self): - """ - Attempt to find CUDA headers (cuda.h, nvrtc.h) in standard - places or environment variables. - """ - if not self._use_cuda: - self.cuda_include_path = None - return None - - # Check the CUDA_PATH and CUDA_HOME environment variables - for env_var in ["CUDA_PATH", "CUDA_HOME"]: - path = os.getenv(env_var) - if path: - include_path = Path(path) / "include" - if (include_path / "cuda.h").is_file() and ( - include_path / "nvrtc.h" - ).is_file(): - self.cuda_include_path = str(include_path) - return self.cuda_include_path - - # Check if CUDA is installed via conda - conda_prefix = os.getenv("CONDA_PREFIX") - if conda_prefix: - include_path = Path(conda_prefix) / "include" - if (include_path / "cuda.h").is_file() and ( - include_path / "nvrtc.h" - ).is_file(): - self.cuda_include_path = str(include_path) - return self.cuda_include_path - - # Check standard locations - cuda_version_str = self.get_cuda_version(out_type="string") - possible_paths = [ - Path("/usr/local/cuda"), - Path(f"/usr/local/cuda-{cuda_version_str}") if cuda_version_str else None, - Path("/opt/cuda"), - ] - # Filter out Nones (if cuda_version_str is None) - possible_paths = [p for p in possible_paths if p is not None] - - for base_path in possible_paths: - include_path = base_path / "include" - if (include_path / "cuda.h").is_file() and ( - include_path / "nvrtc.h" - ).is_file(): - self.cuda_include_path = str(include_path) - return self.cuda_include_path - - # If not found in any known location, try the compiler approach: - cuda_h_path = self.get_include_file_abspath("cuda.h") - nvrtc_h_path = self.get_include_file_abspath("nvrtc.h") - if cuda_h_path and nvrtc_h_path: - if os.path.dirname(cuda_h_path) == os.path.dirname(nvrtc_h_path): - self.cuda_include_path = os.path.dirname(cuda_h_path) - return self.cuda_include_path - - # If still not found, issue a warning - KeOps_Warning( - "CUDA include path not found. Please set the CUDA_PATH or CUDA_HOME environment variable." - ) - self.cuda_include_path = None - return self.cuda_include_path - - def get_include_file_abspath(self, filename): - return get_include_file_abspath(filename, self.cxx_compiler) - - def set_nvrtc_flags(self): - """Set the NVRTC flags for CUDA compilation.""" - # Ensure that compile_options is set (inherited from ConfigNew) - compile_options = " -shared -fPIC -O3 -std=c++11" - - # Ensure that libcuda_folder and libnvrtc_folder are set - libcuda_folder = self.libcuda_folder - libnvrtc_folder = self.libnvrtc_folder - - # Set the NVRTC flags - self.nvrtc_flags = ( - compile_options - + f" -fpermissive -L{libcuda_folder} -L{libnvrtc_folder} -lcuda -lnvrtc" - ) - - def get_nvrtc_flags(self): - """Get the NVRTC flags for CUDA compilation.""" - return self.nvrtc_flags - - def print_nvrtc_flags(self): - """Print the NVRTC flags for CUDA compilation.""" - print(f"NVRTC Flags: {self.nvrtc_flags}") - - def get_gpu_props(self): - """ - Getting GPU properties and related attributes. - """ - if not self._use_cuda: - # Already determined that CUDA is unavailable - self.n_gpus = 0 - self.gpu_compile_flags = "" - return (self.n_gpus, self.gpu_compile_flags) - - # Attempt to load the CUDA driver library - success, libcuda_path, err_msg = self._try_load_library("cuda") - if not success: - # Something is off at driver level => revert to CPU - KeOps_Warning( - "cuda library not fully accessible. " - + err_msg - + " Switching to CPU only." - ) - self.n_gpus = 0 - self.gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) - - # We have a handle, let's proceed - libcuda = ctypes.CDLL(libcuda_path) - result = libcuda.cuInit(0) - if result != self.CUDA_SUCCESS: - KeOps_Warning( - "CUDA was detected, but driver API could not be initialized. Switching to CPU only." - ) - self.n_gpus = 0 - self.gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) - - # Get GPU count - nGpus = ctypes.c_int() - result = libcuda.cuDeviceGetCount(ctypes.byref(nGpus)) - if result != self.CUDA_SUCCESS: - KeOps_Warning( - "CUDA was detected and driver API was initialized, but no working GPU found. " - "Switching to CPU only." - ) - self.n_gpus = 0 - self.gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) - - self.n_gpus = nGpus.value - if self.n_gpus == 0: - self.gpu_compile_flags = "" - return (self.n_gpus, self.gpu_compile_flags) - - # Query each GPU for properties - MaxThreadsPerBlock = [0] * self.n_gpus - SharedMemPerBlock = [0] * self.n_gpus - - def safe_call(dev_idx, result_code): - if result_code != self.CUDA_SUCCESS: - KeOps_Warning( - f"Error detecting properties for GPU device {dev_idx}. " - "Switching to CPU only." - ) - return False - return True - - for d in range(self.n_gpus): - device = ctypes.c_int() - if not safe_call(d, libcuda.cuDeviceGet(ctypes.byref(device), d)): - self.n_gpus = 0 - self.gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) - - output = ctypes.c_int() - if not safe_call( - d, - libcuda.cuDeviceGetAttribute( - byref(output), - self.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, - device, - ), - ): - self.n_gpus = 0 - self.gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) - MaxThreadsPerBlock[d] = output.value - - if not safe_call( - d, - libcuda.cuDeviceGetAttribute( - byref(output), - self.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, - device, - ), - ): - self.n_gpus = 0 - self.gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) - SharedMemPerBlock[d] = output.value - - # Build compile flags string - self.gpu_compile_flags = f"-DMAXIDGPU={self.n_gpus - 1} " - for d in range(self.n_gpus): - self.gpu_compile_flags += ( - f"-DMAXTHREADSPERBLOCK{d}={MaxThreadsPerBlock[d]} " - ) - self.gpu_compile_flags += f"-DSHAREDMEMPERBLOCK{d}={SharedMemPerBlock[d]} " - - return self.n_gpus, self.gpu_compile_flags - - def print_all(self): - """ - Print all CUDA-related configuration and system health status. - """ - - # CUDA Support - cuda_status = CHECK_MARK if self.get_use_cuda() else CROSS_MARK - print(f"\nCUDA Support") - print("-" * 60) - self.print_use_cuda() - if self.get_use_cuda(): - print(f"Libcuda Path: {self.libcuda_folder}") - print(f"Libnvrtc Path: {self.libnvrtc_folder}") - print(f"CUDA Version: {self.cuda_version}") - print(f"Number of GPUs: {self.n_gpus}") - print(f"GPU Compile Flags: {self.gpu_compile_flags}") - # CUDA Include Path - cuda_include_path = self.cuda_include_path - print(f"CUDA Include Path: {cuda_include_path or 'Not Found'}") - - # Print relevant environment variables. - print("\nRelevant Environment Variables:") - env_vars = [ - "CUDA_VISIBLE_DEVICES", - "CUDA_PATH", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") - - -if __name__ == "__main__": - cudastuff = CUDAConfig() - cudastuff.print_all() diff --git a/keopscore/keopscore/config/openmp.py b/keopscore/keopscore/config/openmp.py deleted file mode 100644 index 66ab72317..000000000 --- a/keopscore/keopscore/config/openmp.py +++ /dev/null @@ -1,156 +0,0 @@ -import os -import shutil -import tempfile -import subprocess -import platform -from ctypes.util import find_library - -from keopscore.utils.misc_utils import KeOps_Warning, KeOps_OS_Run, get_brew_prefix -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK - - -class OpenMPConfig: - """ - Class for OpenMP detection and configuration. - """ - - def __init__(self): - self._use_OpenMP = None - self.openmp_lib_path = None - self.os = platform.system() - self.set_cxx_compiler() - self.set_use_OpenMP() - - def set_cxx_compiler(self): - """Set the C++ compiler.""" - env_cxx = os.getenv("CXX") - if env_cxx and shutil.which(env_cxx): - self.cxx_compiler = env_cxx - elif shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No C++ compiler found. You need to either define the CXX environment variable pointing to a valid compiler, or ensure that 'g++' is installed and in your PATH." - ) - - def set_use_OpenMP(self): - """Determine and set whether to use OpenMP.""" - compiler_supports_openmp = self.check_compiler_for_openmp() - openmp_libs_available = self.check_openmp_libraries() - self._use_OpenMP = compiler_supports_openmp or openmp_libs_available - if not self._use_OpenMP: - KeOps_Warning("OpenMP support is not available. Disabling OpenMP.") - - def get_use_OpenMP(self): - return self._use_OpenMP - - def print_use_OpenMP(self): - status = "Enabled ✅" if self._use_OpenMP else "Disabled ❌" - print(f"OpenMP Support: {status}") - - def check_compiler_for_openmp(self): - if not self.cxx_compiler: - KeOps_Warning("No C++ compiler available to check for OpenMP support.") - return False - - test_program = """ - #include - int main() { - #pragma omp parallel - {} - return 0; - } - """ - with tempfile.NamedTemporaryFile("w", suffix=".cpp", delete=False) as f: - f.write(test_program) - test_file = f.name - - compile_command = [ - self.cxx_compiler, - test_file, - "-fopenmp", - "-o", - test_file + ".out", - ] - try: - # Warning : subprocess is used below to compile the test program (using subprocess.check_output to capture stderr) - subprocess.check_output(compile_command, stderr=subprocess.STDOUT) - os.remove(test_file) - os.remove(test_file + ".out") - return True - except subprocess.CalledProcessError: - os.remove(test_file) - return False - - def check_openmp_libraries(self): - if self.os.startswith("Linux"): - openmp_lib = find_library("gomp") - if not openmp_lib: - KeOps_Warning("OpenMP library 'libgomp' not found.") - return False - else: - self.openmp_lib_path = openmp_lib - return True - # Specific check for M1/M2/M3 apple Silicon chips - elif self.os.startswith("Darwin") and platform.machine() in ["arm64", "arm64e"]: - brew_prefix = get_brew_prefix() - if brew_prefix is not None: - openmp_path = f"{brew_prefix}/opt/libomp/lib/libomp.dylib" - openmp_lib = openmp_path if os.path.exists(openmp_path) else None - else: - openmp_lib = None - if openmp_lib is None: - KeOps_Warning( - "OpenMP library not found, it must be downloaded through Homebrew for apple Silicon chips" - ) - return False - else: - self.openmp_lib_path = openmp_lib - return True - elif self.os.startswith("Darwin"): - openmp_lib = find_library("omp") - if not openmp_lib: - KeOps_Warning("OpenMP library 'libomp' not found.") - return False - else: - self.openmp_lib_path = openmp_lib - return True - else: - self.openmp_lib_path = None - return False - - def print_all(self): - """ - Print all OpenMP-related configuration and system health status. - """ - # OpenMP Support - openmp_status = CHECK_MARK if self.get_use_OpenMP() else CROSS_MARK - print(f"\nOpenMP Support") - print("-" * 60) - self.print_use_OpenMP() - if self.get_use_OpenMP(): - openmp_lib_path = self.openmp_lib_path or "Not Found" - print(f"OpenMP Library Path: {openmp_lib_path}") - # Compiler path - compiler_path = ( - shutil.which(self.cxx_compiler) if self.cxx_compiler else None - ) - print(f"C++ Compiler: {self.cxx_compiler}") - if not compiler_path: - print( - f"Compiler '{self.cxx_compiler}' not found on the system.{CROSS_MARK}" - ) - else: - print(f"OpenMP support is disabled or not available.{CROSS_MARK}") - # Print relevant environment variables. - print("\nRelevant Environment Variables:") - env_vars = [ - "OMP_PATH", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") diff --git a/keopscore/keopscore/formulas/Chunkable_Op.py b/keopscore/keopscore/formulas/Chunkable_Op.py index 783edfba1..23cfed7ef 100644 --- a/keopscore/keopscore/formulas/Chunkable_Op.py +++ b/keopscore/keopscore/formulas/Chunkable_Op.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Var import Var -from keopscore.config.chunks import enable_chunk, dim_treshold_chunk, specdims_use_chunk +from keopscore.config import reduction class Chunkable_Op(Operation): @@ -19,9 +19,13 @@ def notchunked_vars(self, cat): @property def use_chunk(self): - test = enable_chunk & all(child.is_chunkable for child in self.children) + test = reduction.get_enable_chunks() & all( + child.is_chunkable for child in self.children + ) child = self.children[0] - subtest = (child.dim >= dim_treshold_chunk) | (child.dim in specdims_use_chunk) + subtest = (child.dim >= reduction.get_dim_treshold_chunk()) | ( + child.dim in reduction.get_specdims_use_chunk() + ) test &= subtest return test diff --git a/keopscore/keopscore/formulas/GetReduction.py b/keopscore/keopscore/formulas/GetReduction.py index 60afe3e5c..cd3658326 100644 --- a/keopscore/keopscore/formulas/GetReduction.py +++ b/keopscore/keopscore/formulas/GetReduction.py @@ -1,16 +1,9 @@ import ast, inspect -import keopscore -import keopscore.formulas -from keopscore.utils.misc_utils import KeOps_Print +import keopscore.config +from keopscore.utils.messages import KeOps_Print from keopscore.utils.code_gen_utils import get_hash_name -from keopscore.formulas.reductions import * -from keopscore.formulas.maths import * -from keopscore.formulas.complex import * -from keopscore.formulas.variables import * -from keopscore.formulas.autodiff import * -from keopscore.formulas.LinearOperators import * -from keopscore.formulas.factorization import * +from keopscore.formulas import * class GetReduction: @@ -18,7 +11,9 @@ class GetReduction: def __new__(self, red_formula_string, aliases=[]): string_id_hash = get_hash_name( - red_formula_string, aliases, keopscore.auto_factorize + red_formula_string, + aliases, + keopscore.config.reduction.get_auto_factorize(), ) if string_id_hash in GetReduction.library: return GetReduction.library[string_id_hash] @@ -31,7 +26,7 @@ def __new__(self, red_formula_string, aliases=[]): varname, var = alias.split("=") aliases_dict[varname] = eval(var) reduction = eval(red_formula_string, globals(), aliases_dict) - if keopscore.auto_factorize: + if keopscore.config.reduction.get_auto_factorize(): formula = reduction.children[0] new_formula = AutoFactorize(formula) reduction.children[0] = new_formula diff --git a/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py b/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py index 9411ca848..ff1ab0544 100644 --- a/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py @@ -2,7 +2,7 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl @@ -21,7 +21,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// ADJOINT OPERATOR //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py b/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py index bd667da26..df453d94e 100644 --- a/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py @@ -1,6 +1,6 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl @@ -15,7 +15,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// LINEAR OPERATORS //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py b/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py index 548077a41..55f8bcd2e 100644 --- a/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py @@ -2,7 +2,7 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl @@ -17,7 +17,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// SUM_LIN OPERATOR //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py b/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py index d5ab12ac4..2eff1ffce 100644 --- a/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py @@ -3,7 +3,7 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl @@ -18,7 +18,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// TRACE OPERATOR //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/__init__.py b/keopscore/keopscore/formulas/LinearOperators/__init__.py index 8e2515c48..b68ec01ea 100644 --- a/keopscore/keopscore/formulas/LinearOperators/__init__.py +++ b/keopscore/keopscore/formulas/LinearOperators/__init__.py @@ -1,3 +1,11 @@ -from .TraceOperator import TraceOperator from .AdjointOperator import AdjointOperator from .SumLinOperator import SumLinOperator +from .TraceOperator import TraceOperator + +_exports = [ + TraceOperator, + AdjointOperator, + SumLinOperator, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/Operation.py b/keopscore/keopscore/formulas/Operation.py index 62019af5c..663485761 100644 --- a/keopscore/keopscore/formulas/Operation.py +++ b/keopscore/keopscore/formulas/Operation.py @@ -1,7 +1,7 @@ from keopscore.utils.code_gen_utils import new_c_varname, c_array from keopscore.utils.Tree import Tree import keopscore -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Print +from keopscore.utils.messages import KeOps_Error, KeOps_Print ################### ## Base class @@ -71,14 +71,14 @@ def __call__(self, out, table): from keopscore.formulas.variables.Var import Var string = f"\n{{\n// Starting code block for {self.__repr__()}.\n\n" - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Building code block for {self.__repr__()}") KeOps_Print("out=", out) KeOps_Print("dim of out : ", out.dim) KeOps_Print("table=", table) for v in table: KeOps_Print(f"dim of {v} : ", v.dim) - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): string += f'printf("\\n\\nComputing {self.__repr__()} :\\n");\n' args = [] # Evaluation of the child operations @@ -104,12 +104,12 @@ def __call__(self, out, table): string += self.Op(out, table, *args) # some debugging helper : - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): for arg in args: string += arg.c_print string += out.c_print string += f'printf("\\n\\n");\n' - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Finished building code block for {self.__repr__()}") string += f"\n\n// Finished code block for {self.__repr__()}.\n}}\n\n" diff --git a/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py b/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py index 07207d331..b221d5ce5 100644 --- a/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py +++ b/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import ComplexVectApply from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class VectorizedComplexScalarOp(Operation): diff --git a/keopscore/keopscore/formulas/VectorizedScalarOp.py b/keopscore/keopscore/formulas/VectorizedScalarOp.py index efee4d041..0750a9d47 100644 --- a/keopscore/keopscore/formulas/VectorizedScalarOp.py +++ b/keopscore/keopscore/formulas/VectorizedScalarOp.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import VectApply from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class VectorizedScalarOp(Operation): diff --git a/keopscore/keopscore/formulas/autodiff/Divergence.py b/keopscore/keopscore/formulas/autodiff/Divergence.py index 18c69affc..f9bcadb34 100644 --- a/keopscore/keopscore/formulas/autodiff/Divergence.py +++ b/keopscore/keopscore/formulas/autodiff/Divergence.py @@ -2,7 +2,7 @@ from keopscore.formulas.LinearOperators import TraceOperator from keopscore.utils.code_gen_utils import GetInds from keopscore.formulas.variables.IntCst import IntCst -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////// # /// DIVERGENCE OPERATOR : Divergence< F, V, U > //// diff --git a/keopscore/keopscore/formulas/autodiff/Grad.py b/keopscore/keopscore/formulas/autodiff/Grad.py index ae8280058..2603b8f19 100644 --- a/keopscore/keopscore/formulas/autodiff/Grad.py +++ b/keopscore/keopscore/formulas/autodiff/Grad.py @@ -1,6 +1,6 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////// # /// GRADIENT OPERATOR : Grad< F, V, Gradin > //// diff --git a/keopscore/keopscore/formulas/autodiff/Laplacian.py b/keopscore/keopscore/formulas/autodiff/Laplacian.py index cf217c048..752748c46 100644 --- a/keopscore/keopscore/formulas/autodiff/Laplacian.py +++ b/keopscore/keopscore/formulas/autodiff/Laplacian.py @@ -3,7 +3,7 @@ from keopscore.formulas.autodiff.Divergence import Divergence from keopscore.utils.code_gen_utils import GetInds from keopscore.formulas.variables.IntCst import IntCst -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////// # /// LAPLACIAN OPERATOR : Laplacian< F, V, U > //// diff --git a/keopscore/keopscore/formulas/autodiff/__init__.py b/keopscore/keopscore/formulas/autodiff/__init__.py index f29bf31cf..f5685f9a1 100644 --- a/keopscore/keopscore/formulas/autodiff/__init__.py +++ b/keopscore/keopscore/formulas/autodiff/__init__.py @@ -1,6 +1,17 @@ -from .Grad import Grad -from .Grad_WithSavedForward import Grad_WithSavedForward from .Diff import Diff from .Diff_WithSavedForward import Diff_WithSavedForward -from .Laplacian import Laplacian from .Divergence import Divergence +from .Grad import Grad +from .Grad_WithSavedForward import Grad_WithSavedForward +from .Laplacian import Laplacian + +_exports = [ + Grad, + Grad_WithSavedForward, + Diff, + Diff_WithSavedForward, + Laplacian, + Divergence, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/complex/ComplexAbs.py b/keopscore/keopscore/formulas/complex/ComplexAbs.py index 8f0c1e41c..758050cbb 100644 --- a/keopscore/keopscore/formulas/complex/ComplexAbs.py +++ b/keopscore/keopscore/formulas/complex/ComplexAbs.py @@ -1,7 +1,6 @@ from keopscore.formulas.complex.ComplexSquareAbs import ComplexSquareAbs from keopscore.formulas.maths.Sqrt import Sqrt - # ///////////////////////////////////////////////////////////////////////// # //// ComplexAbs //// # ///////////////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/complex/ComplexExp.py b/keopscore/keopscore/formulas/complex/ComplexExp.py index c14d6c1b5..8f3a19340 100644 --- a/keopscore/keopscore/formulas/complex/ComplexExp.py +++ b/keopscore/keopscore/formulas/complex/ComplexExp.py @@ -14,7 +14,6 @@ from keopscore.formulas.maths.Cos import Cos from keopscore.formulas.maths.Sin import Sin - # ///////////////////////////////////////////////////////////////////////// # //// ComplexExp //// # ///////////////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/complex/ComplexExp1j.py b/keopscore/keopscore/formulas/complex/ComplexExp1j.py index 395a71ae8..68ec03823 100644 --- a/keopscore/keopscore/formulas/complex/ComplexExp1j.py +++ b/keopscore/keopscore/formulas/complex/ComplexExp1j.py @@ -5,7 +5,7 @@ from keopscore.formulas.complex.Imag2Complex import Imag2Complex from keopscore.formulas.maths.Cos import Cos from keopscore.formulas.maths.Sin import Sin -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexExp1j //// diff --git a/keopscore/keopscore/formulas/complex/ComplexImag.py b/keopscore/keopscore/formulas/complex/ComplexImag.py index 8c1e3a8e2..3309a29ee 100644 --- a/keopscore/keopscore/formulas/complex/ComplexImag.py +++ b/keopscore/keopscore/formulas/complex/ComplexImag.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_for_loop from keopscore.formulas.complex.Imag2Complex import Imag2Complex -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexImag //// diff --git a/keopscore/keopscore/formulas/complex/ComplexReal.py b/keopscore/keopscore/formulas/complex/ComplexReal.py index d2ffbb1db..c08e73d14 100644 --- a/keopscore/keopscore/formulas/complex/ComplexReal.py +++ b/keopscore/keopscore/formulas/complex/ComplexReal.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_for_loop from keopscore.formulas.complex.Real2Complex import Real2Complex -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexReal //// diff --git a/keopscore/keopscore/formulas/complex/ComplexRealScal.py b/keopscore/keopscore/formulas/complex/ComplexRealScal.py index e283ef00f..dd2b33384 100644 --- a/keopscore/keopscore/formulas/complex/ComplexRealScal.py +++ b/keopscore/keopscore/formulas/complex/ComplexRealScal.py @@ -2,7 +2,7 @@ from keopscore.utils.code_gen_utils import c_for_loop from keopscore.formulas.complex.Real2Complex import Real2Complex from keopscore.formulas.complex.ComplexMult import ComplexMult -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexRealScal //// diff --git a/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py b/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py index a95a5626f..184448a77 100644 --- a/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py +++ b/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py @@ -3,7 +3,7 @@ from keopscore.formulas.complex.ComplexReal import ComplexReal from keopscore.formulas.complex.ComplexMult import ComplexMult from keopscore.formulas.complex.Conj import Conj -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexSquareAbs //// diff --git a/keopscore/keopscore/formulas/complex/ComplexSum.py b/keopscore/keopscore/formulas/complex/ComplexSum.py index c244b952c..0fa3c8ab4 100644 --- a/keopscore/keopscore/formulas/complex/ComplexSum.py +++ b/keopscore/keopscore/formulas/complex/ComplexSum.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_zero_float, c_for_loop -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexSum //// diff --git a/keopscore/keopscore/formulas/complex/ComplexSumT.py b/keopscore/keopscore/formulas/complex/ComplexSumT.py index 634f6b283..d2ec2797b 100644 --- a/keopscore/keopscore/formulas/complex/ComplexSumT.py +++ b/keopscore/keopscore/formulas/complex/ComplexSumT.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_for_loop -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// adjoint of ComplexSum //// diff --git a/keopscore/keopscore/formulas/complex/__init__.py b/keopscore/keopscore/formulas/complex/__init__.py index 0008a4202..6f4add17f 100644 --- a/keopscore/keopscore/formulas/complex/__init__.py +++ b/keopscore/keopscore/formulas/complex/__init__.py @@ -15,3 +15,25 @@ from .Conj import Conj from .Imag2Complex import Imag2Complex from .Real2Complex import Real2Complex + +_exports = [ + ComplexAbs, + ComplexAdd, + ComplexAngle, + ComplexDivide, + ComplexExp, + ComplexExp1j, + ComplexImag, + ComplexMult, + ComplexReal, + ComplexRealScal, + ComplexSquareAbs, + ComplexSubtract, + ComplexSum, + ComplexSumT, + Conj, + Imag2Complex, + Real2Complex, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/factorization/Factorize.py b/keopscore/keopscore/formulas/factorization/Factorize.py index cf9f4907c..e92d4d147 100644 --- a/keopscore/keopscore/formulas/factorization/Factorize.py +++ b/keopscore/keopscore/formulas/factorization/Factorize.py @@ -4,7 +4,7 @@ from keopscore.formulas.Operation import Operation import keopscore from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print class Factorize_Impl(Operation): @@ -32,14 +32,14 @@ def __call__(self, out, table): from keopscore.formulas.variables.Var import Var string = f"\n{{\n// Starting code block for {self.__repr__()}.\n\n" - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Building code block for {self.__repr__()}") KeOps_Print("out=", out) KeOps_Print("dim of out : ", out.dim) KeOps_Print("table=", table) for v in table: KeOps_Print(f"dim of {v} : ", v.dim) - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): string += f'printf("\\n\\nComputing {self.__repr__()} :\\n");\n' f, g = self.children @@ -68,7 +68,7 @@ def __call__(self, out, table): # Evaluation of f string += newf(out, table) - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Finished building code block for {self.__repr__()}") string += f"\n\n// Finished code block for {self.__repr__()}.\n}}\n\n" diff --git a/keopscore/keopscore/formulas/factorization/__init__.py b/keopscore/keopscore/formulas/factorization/__init__.py index 970407e18..8668dab65 100644 --- a/keopscore/keopscore/formulas/factorization/__init__.py +++ b/keopscore/keopscore/formulas/factorization/__init__.py @@ -1 +1,8 @@ from .Factorize import Factorize, AutoFactorize + +_exports = [ + Factorize, + AutoFactorize, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/maths/Add.py b/keopscore/keopscore/formulas/maths/Add.py index ecb42f534..7be4aafd8 100644 --- a/keopscore/keopscore/formulas/maths/Add.py +++ b/keopscore/keopscore/formulas/maths/Add.py @@ -7,7 +7,6 @@ from keopscore.formulas.variables.RatCst import RatCst, RatCst_Impl from keopscore.formulas.variables.Zero import Zero - ########################## ###### Add ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/ArgMax.py b/keopscore/keopscore/formulas/maths/ArgMax.py index bf42c1237..dc10f9f56 100644 --- a/keopscore/keopscore/formulas/maths/ArgMax.py +++ b/keopscore/keopscore/formulas/maths/ArgMax.py @@ -7,7 +7,7 @@ value, c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ArgMax ##### diff --git a/keopscore/keopscore/formulas/maths/ArgMin.py b/keopscore/keopscore/formulas/maths/ArgMin.py index d8fcba980..0a798ab3a 100644 --- a/keopscore/keopscore/formulas/maths/ArgMin.py +++ b/keopscore/keopscore/formulas/maths/ArgMin.py @@ -7,7 +7,7 @@ value, c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ArgMin ##### diff --git a/keopscore/keopscore/formulas/maths/Atan2.py b/keopscore/keopscore/formulas/maths/Atan2.py index c3cb4fb8e..b9f3d7086 100644 --- a/keopscore/keopscore/formulas/maths/Atan2.py +++ b/keopscore/keopscore/formulas/maths/Atan2.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.utils.math_functions import keops_atan2 - # ////////////////////////////////////////////////////////////// # //// ATAN2 : Atan2< F, G > //// # ////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/maths/BSpline.py b/keopscore/keopscore/formulas/maths/BSpline.py index 44f231cb9..929344598 100644 --- a/keopscore/keopscore/formulas/maths/BSpline.py +++ b/keopscore/keopscore/formulas/maths/BSpline.py @@ -2,7 +2,7 @@ from keopscore.formulas.maths.Extract import Extract from keopscore.utils.code_gen_utils import c_variable, c_for_loop, c_zero_float from keopscore.utils.code_gen_utils import c_array, VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// BSPLINE VECTOR : BSPLINE //// diff --git a/keopscore/keopscore/formulas/maths/Concat.py b/keopscore/keopscore/formulas/maths/Concat.py index b10f55886..4dafd03ad 100644 --- a/keopscore/keopscore/formulas/maths/Concat.py +++ b/keopscore/keopscore/formulas/maths/Concat.py @@ -1,8 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.maths.Extract import Extract from keopscore.utils.code_gen_utils import VectCopy -from keopscore.utils.misc_utils import KeOps_Error - +from keopscore.utils.messages import KeOps_Error ############################ ###### Concat ##### diff --git a/keopscore/keopscore/formulas/maths/Divide.py b/keopscore/keopscore/formulas/maths/Divide.py index cd62f304e..9b61f0de8 100644 --- a/keopscore/keopscore/formulas/maths/Divide.py +++ b/keopscore/keopscore/formulas/maths/Divide.py @@ -6,7 +6,7 @@ from keopscore.formulas.variables.Zero import Zero from keopscore.formulas.variables.IntCst import IntCst, IntCst_Impl from keopscore.formulas.variables.RatCst import RatCst, RatCst_Impl -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ###### Divide ##### diff --git a/keopscore/keopscore/formulas/maths/Elem.py b/keopscore/keopscore/formulas/maths/Elem.py index e802e2dfb..d061b4048 100644 --- a/keopscore/keopscore/formulas/maths/Elem.py +++ b/keopscore/keopscore/formulas/maths/Elem.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ELEMENT EXTRACTION : Elem(f,m) (aka get_item) ##### diff --git a/keopscore/keopscore/formulas/maths/ElemT.py b/keopscore/keopscore/formulas/maths/ElemT.py index 652ad4c79..aaa1fa88e 100644 --- a/keopscore/keopscore/formulas/maths/ElemT.py +++ b/keopscore/keopscore/formulas/maths/ElemT.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import value, c_zero_float, c_for_loop -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ELEMENT "INJECTION" : ElemT(f,n,m) diff --git a/keopscore/keopscore/formulas/maths/Exp.py b/keopscore/keopscore/formulas/maths/Exp.py index b52514cbd..a6b1a7be3 100644 --- a/keopscore/keopscore/formulas/maths/Exp.py +++ b/keopscore/keopscore/formulas/maths/Exp.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.utils.math_functions import keops_exp - ########################## ###### Exp ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Extract.py b/keopscore/keopscore/formulas/maths/Extract.py index 62bc63761..9d2cdd8ae 100644 --- a/keopscore/keopscore/formulas/maths/Extract.py +++ b/keopscore/keopscore/formulas/maths/Extract.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_array, VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// VECTOR EXTRACTION : Extract //// diff --git a/keopscore/keopscore/formulas/maths/ExtractT.py b/keopscore/keopscore/formulas/maths/ExtractT.py index 9263a2b8c..11f89a1b1 100644 --- a/keopscore/keopscore/formulas/maths/ExtractT.py +++ b/keopscore/keopscore/formulas/maths/ExtractT.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_zero_float, VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// VECTOR "INJECTION" : ExtractT //// diff --git a/keopscore/keopscore/formulas/maths/GradMatrix.py b/keopscore/keopscore/formulas/maths/GradMatrix.py index a6c8495ef..215b5474e 100644 --- a/keopscore/keopscore/formulas/maths/GradMatrix.py +++ b/keopscore/keopscore/formulas/maths/GradMatrix.py @@ -4,7 +4,6 @@ from keopscore.formulas.variables.Var import Var from keopscore.utils.code_gen_utils import GetInds - # ////////////////////////////////////////////////////////////////////////////////////////////// # //// Standard basis of R^DIM : < (1,0,0,...) , (0,1,0,...) , ... , (0,...,0,1) > //// # ////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/maths/Inv.py b/keopscore/keopscore/formulas/maths/Inv.py index a0355ffec..0e3905017 100644 --- a/keopscore/keopscore/formulas/maths/Inv.py +++ b/keopscore/keopscore/formulas/maths/Inv.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.utils.math_functions import keops_rcp - ########################## ###### INVERSE : Inv ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/MatVecMult.py b/keopscore/keopscore/formulas/maths/MatVecMult.py index 0369384de..927400f9b 100644 --- a/keopscore/keopscore/formulas/maths/MatVecMult.py +++ b/keopscore/keopscore/formulas/maths/MatVecMult.py @@ -4,7 +4,7 @@ c_for_loop, c_zero_float, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// Matrix-vector product A x b //// diff --git a/keopscore/keopscore/formulas/maths/Max.py b/keopscore/keopscore/formulas/maths/Max.py index cdbf13f3e..d780b33d4 100644 --- a/keopscore/keopscore/formulas/maths/Max.py +++ b/keopscore/keopscore/formulas/maths/Max.py @@ -2,7 +2,7 @@ from keopscore.formulas.maths.ArgMax import ArgMax from keopscore.formulas.maths.OneHot import OneHot from keopscore.utils.code_gen_utils import c_for_loop, c_if, value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### Max ##### diff --git a/keopscore/keopscore/formulas/maths/Min.py b/keopscore/keopscore/formulas/maths/Min.py index 6995c2278..1956b80d1 100644 --- a/keopscore/keopscore/formulas/maths/Min.py +++ b/keopscore/keopscore/formulas/maths/Min.py @@ -2,7 +2,7 @@ from keopscore.formulas.maths.ArgMin import ArgMin from keopscore.formulas.maths.OneHot import OneHot from keopscore.utils.code_gen_utils import c_for_loop, c_if, value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### Min ##### diff --git a/keopscore/keopscore/formulas/maths/Minus.py b/keopscore/keopscore/formulas/maths/Minus.py index f374d1867..5a59c193b 100644 --- a/keopscore/keopscore/formulas/maths/Minus.py +++ b/keopscore/keopscore/formulas/maths/Minus.py @@ -4,7 +4,6 @@ from keopscore.formulas.variables.IntCst import IntCst_Impl, IntCst from keopscore.formulas.variables.RatCst import RatCst_Impl, RatCst - ########################## ###### Minus ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Mult.py b/keopscore/keopscore/formulas/maths/Mult.py index c8f09c6a5..870f8160d 100644 --- a/keopscore/keopscore/formulas/maths/Mult.py +++ b/keopscore/keopscore/formulas/maths/Mult.py @@ -8,7 +8,7 @@ from keopscore.formulas.variables.IntCst import IntCst_Impl, IntCst from keopscore.formulas.variables.RatCst import RatCst_Impl, RatCst from keopscore.formulas.maths.SumT import SumT, SumT_Impl -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ###### Mult ##### diff --git a/keopscore/keopscore/formulas/maths/OneHot.py b/keopscore/keopscore/formulas/maths/OneHot.py index 8cf93fb7c..7261ae005 100644 --- a/keopscore/keopscore/formulas/maths/OneHot.py +++ b/keopscore/keopscore/formulas/maths/OneHot.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_zero_float -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// ONE-HOT REPRESENTATION : OneHot //// diff --git a/keopscore/keopscore/formulas/maths/Rsqrt.py b/keopscore/keopscore/formulas/maths/Rsqrt.py index 62454ccc6..f8f930831 100644 --- a/keopscore/keopscore/formulas/maths/Rsqrt.py +++ b/keopscore/keopscore/formulas/maths/Rsqrt.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.formulas.maths.IntInv import IntInv - ########################## ###### Rsqrt ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Scalprod.py b/keopscore/keopscore/formulas/maths/Scalprod.py index 4da68f8b2..5d72d4fd6 100644 --- a/keopscore/keopscore/formulas/maths/Scalprod.py +++ b/keopscore/keopscore/formulas/maths/Scalprod.py @@ -5,7 +5,7 @@ VectApply, ) from keopscore.utils.math_functions import keops_fma -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ##### Scalprod #### diff --git a/keopscore/keopscore/formulas/maths/Sign.py b/keopscore/keopscore/formulas/maths/Sign.py index 760e38b10..122faea48 100644 --- a/keopscore/keopscore/formulas/maths/Sign.py +++ b/keopscore/keopscore/formulas/maths/Sign.py @@ -2,7 +2,6 @@ from keopscore.formulas.variables.Zero import Zero from keopscore.utils.math_functions import keops_sign - ########################## ###### Sign ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py b/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py index 6aff05757..9387d9f2f 100644 --- a/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py +++ b/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py @@ -3,7 +3,7 @@ #################################################################### from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, diff --git a/keopscore/keopscore/formulas/maths/Sqrt.py b/keopscore/keopscore/formulas/maths/Sqrt.py index ee00b1b51..12696cfde 100644 --- a/keopscore/keopscore/formulas/maths/Sqrt.py +++ b/keopscore/keopscore/formulas/maths/Sqrt.py @@ -3,7 +3,6 @@ from keopscore.formulas.maths.Rsqrt import Rsqrt from keopscore.utils.math_functions import keops_sqrt - ########################## ###### Sqrt ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Square.py b/keopscore/keopscore/formulas/maths/Square.py index 808bc8698..1ca9ced5b 100644 --- a/keopscore/keopscore/formulas/maths/Square.py +++ b/keopscore/keopscore/formulas/maths/Square.py @@ -2,7 +2,6 @@ from keopscore.formulas.variables.Zero import Zero - ########################## ###### Square ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Sum.py b/keopscore/keopscore/formulas/maths/Sum.py index f0483ae6a..45632508a 100644 --- a/keopscore/keopscore/formulas/maths/Sum.py +++ b/keopscore/keopscore/formulas/maths/Sum.py @@ -3,7 +3,6 @@ from keopscore.utils.code_gen_utils import c_zero_float, VectApply from keopscore.formulas.maths.Square import Square_Impl - ########################## ###### Sum ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/SumT.py b/keopscore/keopscore/formulas/maths/SumT.py index 786c44536..795b18b32 100644 --- a/keopscore/keopscore/formulas/maths/SumT.py +++ b/keopscore/keopscore/formulas/maths/SumT.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Zero import Zero from keopscore.utils.code_gen_utils import value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ###### SumT ##### diff --git a/keopscore/keopscore/formulas/maths/TensorProd.py b/keopscore/keopscore/formulas/maths/TensorProd.py index f03001e86..2320a18ff 100644 --- a/keopscore/keopscore/formulas/maths/TensorProd.py +++ b/keopscore/keopscore/formulas/maths/TensorProd.py @@ -3,8 +3,7 @@ c_variable, c_for_loop, ) -from keopscore.utils.misc_utils import KeOps_Error - +from keopscore.utils.messages import KeOps_Error #################################### ###### Tensor product ##### diff --git a/keopscore/keopscore/formulas/maths/VecMatMult.py b/keopscore/keopscore/formulas/maths/VecMatMult.py index 93b400529..369f0baa3 100644 --- a/keopscore/keopscore/formulas/maths/VecMatMult.py +++ b/keopscore/keopscore/formulas/maths/VecMatMult.py @@ -4,7 +4,7 @@ c_for_loop, c_zero_float, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// Vector-matrix product b x A //// diff --git a/keopscore/keopscore/formulas/maths/__init__.py b/keopscore/keopscore/formulas/maths/__init__.py index c036915f3..ce39f4a04 100644 --- a/keopscore/keopscore/formulas/maths/__init__.py +++ b/keopscore/keopscore/formulas/maths/__init__.py @@ -66,72 +66,74 @@ from .WeightedSqNorm import WeightedSqNorm from .XLogX import XLogX -__all__ = [ - "Abs", - "Acos", - "Add", - "ArgMax", - "ArgMin", - "Asin", - "Atan", - "Atan2", - "BSpline", - "Clamp", - "ClampInt", - "Concat", - "Cos", - "DiffClampInt", - "Divide", - "Elem", - "ElemT", - "Equal", - "Exp", - "Extract", - "ExtractT", - "GradMatrix", - "Floor", - "IfElse", - "IntInv", - "Inv", - "Kron", - "LessOrEqual", - "LessThan", - "Log", - "MatVecMult", - "Max", - "Min", - "Minus", - "Mod", - "Mult", - "Norm2", - "Normalize", - "NotEqual", - "OneHot", - "Pow", - "Powf", - "ReLU", - "Round", - "Rsqrt", - "Scalprod", - "Sign", - "Sin", - "SinXDivX", - "SoftDTW_SqDist", - "SqDist", - "SqNorm2", - "SqNormDiag", - "SqNormIso", - "Sqrt", - "Square", - "Step", - "Subtract", - "Sum", - "SumT", - "SymSqNorm", - "TensorDot", - "TensorProd", - "VecMatMult", - "WeightedSqDist", - "WeightedSqNorm", - "XLogX", +_exports = [ + Abs, + Acos, + Add, + ArgMax, + ArgMin, + Asin, + Atan, + Atan2, + BSpline, + Clamp, + ClampInt, + Concat, + Cos, + DiffClampInt, + Divide, + Elem, + ElemT, + Equal, + Exp, + Extract, + ExtractT, + Floor, + GradMatrix, + IfElse, + IntInv, + Inv, + Kron, + LessOrEqual, + LessThan, + Log, + MatVecMult, + Max, + Min, + Minus, + Mod, + Mult, + Norm2, + Normalize, + NotEqual, + OneHot, + Pow, + Powf, + ReLU, + Round, + Rsqrt, + Scalprod, + Sign, + Sin, + SinXDivX, + SoftDTW_SqDist, + SqDist, + SqNorm2, + SqNormDiag, + SqNormIso, + Sqrt, + Square, + Step, + Subtract, + Sum, + SumT, + SymSqNorm, + TensorDot, + TensorProd, + VecMatMult, + WeightedSqDist, + WeightedSqNorm, + XLogX, ] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py b/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py index cb19df7a1..a4d5d38d8 100644 --- a/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py @@ -10,7 +10,7 @@ use_pragma_unroll, ) from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class KMin_ArgKMin_Reduction(Reduction): diff --git a/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py b/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py index 6963e2195..530489a64 100644 --- a/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py +++ b/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py @@ -5,7 +5,7 @@ c_if, ) from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import c_variable diff --git a/keopscore/keopscore/formulas/reductions/Max_Reduction.py b/keopscore/keopscore/formulas/reductions/Max_Reduction.py index ca0837bd8..91a439ad2 100644 --- a/keopscore/keopscore/formulas/reductions/Max_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/Max_Reduction.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import neg_infinity, c_if from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class Max_Reduction(Reduction): diff --git a/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py b/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py index 364e19e4f..c3760c446 100644 --- a/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py @@ -10,7 +10,7 @@ c_for_loop, ) from keopscore.utils.math_functions import keops_exp -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class Max_SumShiftExpWeight_Reduction(Reduction): @@ -120,4 +120,5 @@ def Diff(self, v, diffin, MS): ) -Max_SumShiftExp_Reduction = Max_SumShiftExpWeight_Reduction +def Max_SumShiftExp_Reduction(*args, **kwargs): + return Max_SumShiftExpWeight_Reduction(*args, **kwargs) diff --git a/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py b/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py index aad023907..f5e66858a 100644 --- a/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py +++ b/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py @@ -5,7 +5,7 @@ c_if, ) from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import c_variable diff --git a/keopscore/keopscore/formulas/reductions/Min_Reduction.py b/keopscore/keopscore/formulas/reductions/Min_Reduction.py index 33e014f6d..f8b3000aa 100644 --- a/keopscore/keopscore/formulas/reductions/Min_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/Min_Reduction.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import infinity, c_if from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class Min_Reduction(Reduction): diff --git a/keopscore/keopscore/formulas/reductions/__init__.py b/keopscore/keopscore/formulas/reductions/__init__.py index 956363d88..5ad75016f 100644 --- a/keopscore/keopscore/formulas/reductions/__init__.py +++ b/keopscore/keopscore/formulas/reductions/__init__.py @@ -14,5 +14,28 @@ from .Min_Reduction import Min_Reduction from .Sum_Reduction import Sum_Reduction from .Zero_Reduction import Zero_Reduction +from .sum_schemes import block_sum, kahan_scheme, direct_sum, make_sum_scheme -from .sum_schemes import * +_exports = [ + Reduction, + ArgKMin_Reduction, + ArgMax_Reduction, + ArgMin_Reduction, + KMin_ArgKMin_Reduction, + KMin_Reduction, + Max_ArgMax_Reduction, + Max_Reduction, + Max_SumShiftExpWeight_Reduction, + Max_SumShiftExp_Reduction, + Min_ArgMin_Reduction, + Min_Reduction, + Sum_Reduction, + Zero_Reduction, + block_sum, + kahan_scheme, + direct_sum, + make_sum_scheme, +] + + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/reductions/sum_schemes.py b/keopscore/keopscore/formulas/reductions/sum_schemes.py index 437581a94..97fba9e46 100644 --- a/keopscore/keopscore/formulas/reductions/sum_schemes.py +++ b/keopscore/keopscore/formulas/reductions/sum_schemes.py @@ -5,6 +5,12 @@ c_variable, ) +""" +This module defines the different schemes for performing the reduction of a formula. + +The schemes must be added in the _SUM_SCHEME_CLASSES dictionary, and must be subclasses of Sum_Scheme, which defines the interface for the different schemes. +""" + class Sum_Scheme: def __init__(self, red_formula, dtype, dimred=None): @@ -84,3 +90,15 @@ def initialize_temporary_accumulator_first_init(self): def accumulate_result(self, acc, fout, j, hack=False): return self.red_formula.KahanScheme(acc, fout, self.tmp_acc) + + +_SUM_SCHEME_CLASSES = { + "direct_sum": direct_sum, + "block_sum": block_sum, + "kahan_scheme": kahan_scheme, +} + + +def make_sum_scheme(sum_scheme_string, red_formula, dtype, dimred=None): + sum_scheme_class = _SUM_SCHEME_CLASSES.get(sum_scheme_string, None) + return sum_scheme_class(red_formula, dtype, dimred=dimred) diff --git a/keopscore/keopscore/formulas/variables/Var.py b/keopscore/keopscore/formulas/variables/Var.py index ae299df4e..2d009ba67 100644 --- a/keopscore/keopscore/formulas/variables/Var.py +++ b/keopscore/keopscore/formulas/variables/Var.py @@ -1,7 +1,6 @@ from keopscore.utils.code_gen_utils import VectCopy from keopscore.formulas.Operation import Operation - ####################### ## Var operation ####################### diff --git a/keopscore/keopscore/formulas/variables/__init__.py b/keopscore/keopscore/formulas/variables/__init__.py index cc9af2fb4..2e2bc23ba 100644 --- a/keopscore/keopscore/formulas/variables/__init__.py +++ b/keopscore/keopscore/formulas/variables/__init__.py @@ -2,3 +2,12 @@ from .RatCst import RatCst from .Var import Var from .Zero import Zero + +_exports = [ + IntCst, + RatCst, + Var, + Zero, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/get_keops_dll.py b/keopscore/keopscore/get_keops_dll.py index 6d655609a..595f147bf 100644 --- a/keopscore/keopscore/get_keops_dll.py +++ b/keopscore/keopscore/get_keops_dll.py @@ -19,9 +19,6 @@ It returns - tag : string, hash code used as id for the input formula and parameters - - source_file : string, either : - - in CPU mode : name of the source file to be compiled - - in GPU mode : name of the main dll to be called - low_level_code_file : string, either : - in CPU mode : the empty string "" - in GPU mode : name of the low level code (ptx) or binary file (cubin) to be passed to the main dll. @@ -52,24 +49,13 @@ import inspect import sys -import keopscore -from keopscore.config import * - import keopscore.mapreduce -from keopscore import cuda_block_size -from keopscore.config.chunks import ( - get_enable_chunk, - set_enable_chunk, - dimchunk, - set_enable_finalchunk, - use_final_chunks, - set_mult_var_highdim, -) +from keopscore.config import path, cuda, debug, reduction from keopscore.formulas import Zero_Reduction, Sum_Reduction from keopscore.formulas.GetReduction import GetReduction from keopscore.formulas.variables.Zero import Zero from keopscore.utils.Cache import Cache -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Print +from keopscore.utils.messages import KeOps_Error, KeOps_Print # Get every classes in mapreduce map_reduce = dict(inspect.getmembers(keopscore.mapreduce, inspect.isclass)) @@ -87,19 +73,19 @@ def get_keops_dll_impl( # detecting the need for special chunked computation modes : use_chunk_mode = 0 if "Gpu" in map_reduce_id: - if not cuda_config.get_use_cuda(): + if not cuda.get_use_cuda(): KeOps_Error( "You selected a Gpu reduce scheme but KeOps is in Cpu only mode." ) - set_enable_chunk(enable_chunks) - set_enable_finalchunk(enable_finalchunks) - set_mult_var_highdim(mul_var_highdim) + reduction.set_enable_chunks(enable_chunks) + reduction.set_enable_finalchunk(enable_finalchunks) + reduction.set_mult_var_highdim(mul_var_highdim) red_formula = GetReduction(red_formula_string, aliases) - if use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D": + if reduction.use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D": use_chunk_mode = 2 map_reduce_id += "_finalchunks" - elif get_enable_chunk() and map_reduce_id != "GpuReduc2D": - if len(red_formula.formula.chunked_formulas(dimchunk)) == 1: + elif reduction.get_enable_chunks() and map_reduce_id != "GpuReduc2D": + if len(red_formula.formula.chunked_formulas(reduction.get_dimchunk())) == 1: from keopscore.mapreduce.Chunk_Mode_Constants import ( Chunk_Mode_Constants, ) @@ -115,7 +101,7 @@ def get_keops_dll_impl( rf = map_reduce_obj.red_formula - if keopscore.debug_ops: + if debug.get_debug_ops(): KeOps_Print("In get_keops_dll, formula is :", rf) KeOps_Print("formula.__repr__() is : ", rf.__repr__()) rf.make_dot() @@ -138,13 +124,12 @@ def get_keops_dll_impl( return ( res["tag"], - res["source_file"], res["low_level_code_file"], res["tagI"], tagZero, res["use_half"], res["use_fast_math"], - cuda_block_size, + cuda.get_cuda_block_size(), use_chunk_mode, tag1D2D, res["dimred"], @@ -162,7 +147,7 @@ def get_keops_dll_impl( get_keops_dll = Cache( get_keops_dll_impl, use_cache_file=True, - save_folder=config.get_build_folder(), + save_folder=path.get_build_folder(), ) diff --git a/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py b/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py index 59e05e301..1f2f8d6ec 100644 --- a/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py +++ b/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py @@ -1,4 +1,4 @@ -from keopscore.config.chunks import dimchunk +from keopscore.config import reduction from keopscore.utils.code_gen_utils import GetDims, GetInds, Var_loader from keopscore.formulas.variables.Var import Var @@ -15,10 +15,10 @@ def __init__(self, red_formula): formula = red_formula.formula self.dimfout = formula.dim # dimension of output variable of inner function - chunked_formula = formula.chunked_formulas(dimchunk)[0] + chunked_formula = formula.chunked_formulas(reduction.get_dimchunk())[0] self.dim_org = chunked_formula["dim_org"] - self.nchunks = 1 + (self.dim_org - 1) // dimchunk - self.dimlastchunk = self.dim_org - (self.nchunks - 1) * dimchunk + self.nchunks = 1 + (self.dim_org - 1) // reduction.get_dimchunk() + self.dimlastchunk = self.dim_org - (self.nchunks - 1) * reduction.get_dimchunk() self.nminargs = varloader.nminargs self.fun_chunked = chunked_formula["formula"] self.dimout_chunk = self.fun_chunked.dim diff --git a/keopscore/keopscore/mapreduce/MapReduce.py b/keopscore/keopscore/mapreduce/MapReduce.py index efd9fc80c..70e48df65 100644 --- a/keopscore/keopscore/mapreduce/MapReduce.py +++ b/keopscore/keopscore/mapreduce/MapReduce.py @@ -1,6 +1,13 @@ -from keopscore.formulas.reductions import * +from keopscore.formulas.reductions import make_sum_scheme from keopscore.formulas.GetReduction import GetReduction -from keopscore.utils.code_gen_utils import Var_loader, new_c_varname, pointer, c_include +from keopscore.utils.code_gen_utils import ( + Var_loader, + new_c_varname, + pointer, + c_include, + c_variable, + c_array, +) class MapReduce: @@ -56,7 +63,7 @@ def get_code(self): dtype = self.dtype dtypeacc = self.dtypeacc nargs = self.nargs - self.sum_scheme = eval(self.sum_scheme_string)(red_formula, dtype) + self.sum_scheme = make_sum_scheme(self.sum_scheme_string, red_formula, dtype) self.i = i = c_variable("signed long int", "i") self.j = j = c_variable("signed long int", "j") diff --git a/keopscore/keopscore/mapreduce/__init__.py b/keopscore/keopscore/mapreduce/__init__.py index 683989a1a..517b05fc5 100644 --- a/keopscore/keopscore/mapreduce/__init__.py +++ b/keopscore/keopscore/mapreduce/__init__.py @@ -1,7 +1,6 @@ -from .cpu import * -from ..config import get_cuda_config +from keopscore.config import cuda -cuda_config = get_cuda_config() +from .cpu import * -if cuda_config._use_cuda: +if cuda.get_use_cuda(): from .gpu import * diff --git a/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py b/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py index de1333ac2..e66103e4b 100644 --- a/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py +++ b/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py @@ -1,12 +1,10 @@ -import keopscore +import keopscore.config from keopscore.binders.cpp.Cpu_link_compile import Cpu_link_compile from keopscore.mapreduce.MapReduce import MapReduce from keopscore.utils.code_gen_utils import ( c_include, c_zero_float, ) -import keopscore -from keopscore.config import * class CpuAssignZero(MapReduce, Cpu_link_compile): @@ -26,9 +24,9 @@ def get_code(self): args = self.args headers = ["stdlib.h"] - if keopscore.openmp_config.get_use_OpenMP(): + if keopscore.config.openmp.get_use_OpenMP(): headers.append("omp.h") - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): headers.append("iostream") self.headers += c_include(*headers) diff --git a/keopscore/keopscore/mapreduce/cpu/CpuReduc.py b/keopscore/keopscore/mapreduce/cpu/CpuReduc.py index 29e8fb6b1..3b702e720 100644 --- a/keopscore/keopscore/mapreduce/cpu/CpuReduc.py +++ b/keopscore/keopscore/mapreduce/cpu/CpuReduc.py @@ -1,10 +1,8 @@ -import keopscore +import keopscore.config from keopscore.binders.cpp.Cpu_link_compile import Cpu_link_compile from keopscore.mapreduce.cpu.CpuAssignZero import CpuAssignZero from keopscore.mapreduce.MapReduce import MapReduce from keopscore.utils.code_gen_utils import c_include -import keopscore -from keopscore.config import * class CpuReduc(MapReduce, Cpu_link_compile): @@ -34,9 +32,9 @@ def get_code(self): sum_scheme = self.sum_scheme headers = ["cmath", "stdlib.h"] - if keopscore.openmp_config.get_use_OpenMP(): + if keopscore.config.openmp.get_use_OpenMP(): headers.append("omp.h") - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): headers.append("iostream") self.headers += c_include(*headers) diff --git a/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py b/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py index 374c9382d..c4dfdade4 100644 --- a/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py +++ b/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py @@ -1,4 +1,4 @@ -import keopscore +import keopscore.config from keopscore.binders.cpp.Cpu_link_compile import Cpu_link_compile from keopscore.mapreduce.cpu.CpuAssignZero import CpuAssignZero @@ -8,8 +8,6 @@ c_array, c_include, ) -import keopscore -from keopscore.config import * class CpuReduc_ranges(MapReduce, Cpu_link_compile): @@ -65,9 +63,9 @@ def get_code(self): jmstarty = c_variable("int", "j-start_y") headers = ["cmath", "stdlib.h"] - if keopscore.openmp_config.get_use_OpenMP: + if keopscore.config.openmp.get_use_OpenMP(): headers.append("omp.h") - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): headers.append("iostream") self.headers += c_include(*headers) diff --git a/keopscore/keopscore/mapreduce/cpu/__init__.py b/keopscore/keopscore/mapreduce/cpu/__init__.py index 49e1c5474..27dffc967 100644 --- a/keopscore/keopscore/mapreduce/cpu/__init__.py +++ b/keopscore/keopscore/mapreduce/cpu/__init__.py @@ -1,3 +1,11 @@ -from .CpuReduc_ranges import CpuReduc_ranges -from .CpuReduc import CpuReduc from .CpuAssignZero import CpuAssignZero +from .CpuReduc import CpuReduc +from .CpuReduc_ranges import CpuReduc_ranges + +_exports = [ + CpuAssignZero, + CpuReduc, + CpuReduc_ranges, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index d664a3fa3..a394c06b6 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -1,9 +1,9 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimchunk from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.config import cuda, reduction +from keopscore.formulas.reductions import make_sum_scheme +from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -12,8 +12,9 @@ table, table4, use_pragma_unroll, + c_variable, + c_array, ) -from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants def do_chunk_sub( @@ -53,7 +54,7 @@ def do_chunk_sub( ) load_chunks_routine_i = load_vars_chunks( indsi_chunked, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -63,7 +64,7 @@ def do_chunk_sub( ) load_chunks_routine_j = load_vars_chunks( indsj_chunked, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -73,7 +74,7 @@ def do_chunk_sub( ) load_chunks_routine_p = load_vars_chunks( indsp_chunked, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -129,7 +130,9 @@ def __init__(self, *args): self.chk = Chunk_Mode_Constants(self.red_formula) self.dimy = self.chk.dimy self.blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) def get_code(self): @@ -153,7 +156,9 @@ def get_code(self): chk = self.chk param_loc = c_array(dtype, chk.dimp, "param_loc") acc = c_array(dtypeacc, chk.dimred, "acc") - sum_scheme = eval(self.sum_scheme_string)(red_formula, dtype, dimred=chk.dimred) + sum_scheme = make_sum_scheme( + self.sum_scheme_string, red_formula, dtype, dimred=chk.dimred + ) xi = c_array(dtype, chk.dimx, "xi") fout_chunk = c_array( dtype, self.blocksize_chunks * chk.dimout_chunk, "fout_chunk" @@ -176,7 +181,7 @@ def get_code(self): dtype, red_formula, chk.fun_chunked, - dimchunk, + reduction.get_dimchunk(), chk.dimsx, chk.dimsy, chk.dimsp, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index 9e7ad6c7f..0d77b2db6 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -1,10 +1,9 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimfinalchunk from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile +from keopscore.config import cuda, reduction from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction -from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -12,8 +11,11 @@ pointer, Var_loader, use_pragma_unroll, + c_variable, + c_zero_float, + c_array, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error def do_finalchunk_sub( @@ -33,10 +35,12 @@ def do_finalchunk_sub( out, ): dimout = varfinal.dim - yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {dimfinalchunk})") + yjloc = c_variable( + pointer(dtype), f"({yj.id} + threadIdx.x * {reduction.get_dimfinalchunk()})" + ) load_chunks_routine_j = load_vars_chunks( [varfinal.ind], - dimfinalchunk, + reduction.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -51,7 +55,7 @@ def do_finalchunk_sub( {load_chunks_routine_j} }} __syncthreads(); - for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {ny.id} - {jstart.id}); jrel++, yjrel += {dimfinalchunk}) {{ + for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {ny.id} - {jstart.id}); jrel++, yjrel += {reduction.get_dimfinalchunk()}) {{ if ({i.id} < {nx.id}) {{ // we compute only if needed {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) {{ @@ -63,7 +67,7 @@ def do_finalchunk_sub( if ({i.id} < {nx.id}) {{ {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) - {out.id}[i*{dimout}+{chunk.id}*{dimfinalchunk}+k] += {acc.id}[k]; + {out.id}[i*{dimout}+{chunk.id}*{reduction.get_dimfinalchunk()}+k] += {acc.id}[k]; }} __syncthreads(); """ @@ -98,8 +102,8 @@ def get_code(self): ) formula = fun_internal.formula varfinal = self.red_formula.formula.children[1 - ind_fun_internal] - nchunks = 1 + (varfinal.dim - 1) // dimfinalchunk - dimlastfinalchunk = varfinal.dim - (nchunks - 1) * dimfinalchunk + nchunks = 1 + (varfinal.dim - 1) // reduction.get_dimfinalchunk() + dimlastfinalchunk = varfinal.dim - (nchunks - 1) * reduction.get_dimfinalchunk() varloader = Var_loader(fun_internal) dimsx = varloader.dimsx dimsy = varloader.dimsy @@ -116,9 +120,11 @@ def get_code(self): KeOps_Error("dimfout should be 1") sum_scheme = self.sum_scheme - self.dimy = max(dimfinalchunk, dimy) + self.dimy = max(reduction.get_dimfinalchunk(), dimy) blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) if not isinstance(sum_scheme, block_sum): @@ -126,7 +132,7 @@ def get_code(self): param_loc = c_array(dtype, dimp, "param_loc") fout = c_array(dtype, dimfout * blocksize_chunks, "fout") xi = c_array(dtype, dimx, "xi") - acc = c_array(dtypeacc, dimfinalchunk, "acc") + acc = c_array(dtypeacc, reduction.get_dimfinalchunk(), "acc") yjloc = c_array(dtype, dimy, f"(yj + threadIdx.x * {dimy})") foutjrel = c_array(dtype, dimfout, f"({fout.id}+jrel*{dimfout})") yjrel = c_array(dtype, dimy, "yjrel") @@ -137,7 +143,7 @@ def get_code(self): chunk_sub_routine = do_finalchunk_sub( dtype, varfinal, - dimfinalchunk, + reduction.get_dimfinalchunk(), acc, i, j, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index ca3dcb254..41bcdffd0 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -1,9 +1,10 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimchunk +from keopscore.config import cuda, reduction + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.formulas.reductions import make_sum_scheme +from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -14,8 +15,9 @@ table4, Var_loader, use_pragma_unroll, + c_array, + c_variable, ) -from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants def do_chunk_sub_ranges( @@ -61,7 +63,7 @@ def do_chunk_sub_ranges( load_chunks_routine_i = load_vars_chunks( indsi_chunked, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -79,7 +81,7 @@ def do_chunk_sub_ranges( load_chunks_routine_i_batches = load_vars_chunks_offsets( indsi_chunked, indsi_global, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -91,7 +93,7 @@ def do_chunk_sub_ranges( load_chunks_routine_j = load_vars_chunks( indsj_chunked, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -103,7 +105,7 @@ def do_chunk_sub_ranges( load_chunks_routine_j_batches = load_vars_chunks_offsets( indsj_chunked, indsj_global, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -115,7 +117,7 @@ def do_chunk_sub_ranges( load_chunks_routine_p = load_vars_chunks( indsp_chunked, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -126,7 +128,7 @@ def do_chunk_sub_ranges( load_chunks_routine_p_batches = load_vars_chunks_offsets( indsp_chunked, indsp_global, - dimchunk, + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -194,7 +196,9 @@ def __init__(self, *args): self.chk = Chunk_Mode_Constants(self.red_formula) self.dimy = self.chk.dimy self.blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) def get_code(self): @@ -245,7 +249,9 @@ def get_code(self): chk = self.chk param_loc = c_array(dtype, chk.dimp, "param_loc") acc = c_array(dtypeacc, chk.dimred, "acc") - sum_scheme = eval(self.sum_scheme_string)(red_formula, dtype, dimred=chk.dimred) + sum_scheme = make_sum_scheme( + self.sum_scheme_string, red_formula, dtype, dimred=chk.dimred + ) xi = c_array(dtype, chk.dimx, "xi") fout_chunk = c_array( dtype, self.blocksize_chunks * chk.dimout_chunk, "fout_chunk" @@ -275,7 +281,7 @@ def get_code(self): dtype, red_formula, chk.fun_chunked, - dimchunk, + reduction.get_dimchunk(), chk.dimsx, chk.dimsy, chk.dimsp, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index f968ecd9f..e727ba3d6 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -1,10 +1,10 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimfinalchunk +from keopscore.config import cuda, reduction + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction -from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -13,8 +13,11 @@ pointer, Var_loader, use_pragma_unroll, + c_array, + c_zero_float, + c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error def do_finalchunk_sub_ranges( @@ -38,11 +41,13 @@ def do_finalchunk_sub_ranges( out, ): dimout = varfinal.dim - yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {dimfinalchunk})") + yjloc = c_variable( + pointer(dtype), f"({yj.id} + threadIdx.x * {reduction.get_dimfinalchunk()})" + ) indsj_global = Var_loader(fun_global).indsj load_chunks_routine_j = load_vars_chunks( [varfinal.ind], - dimfinalchunk, + reduction.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -53,7 +58,7 @@ def do_finalchunk_sub_ranges( load_chunks_routine_j_ranges = load_vars_chunks_offsets( [varfinal.ind], indsj_global, - dimfinalchunk, + reduction.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -73,7 +78,7 @@ def do_finalchunk_sub_ranges( }} }} __syncthreads(); - for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {end_y.id} - {jstart.id}); jrel++, yjrel += {dimfinalchunk}) {{ + for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {end_y.id} - {jstart.id}); jrel++, yjrel += {reduction.get_dimfinalchunk()}) {{ if ({i.id} < {end_x.id}) {{ // we compute only if needed {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) {{ @@ -85,7 +90,7 @@ def do_finalchunk_sub_ranges( if ({i.id} < {end_x.id}) {{ {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) - {out.id}[i*{dimout}+{chunk.id}*{dimfinalchunk}+k] += {acc.id}[k]; + {out.id}[i*{dimout}+{chunk.id}*{reduction.get_dimfinalchunk()}+k] += {acc.id}[k]; }} __syncthreads(); """ @@ -121,8 +126,8 @@ def get_code(self): ) formula = fun_internal.formula varfinal = self.red_formula.formula.children[1 - ind_fun_internal] - nchunks = 1 + (varfinal.dim - 1) // dimfinalchunk - dimlastfinalchunk = varfinal.dim - (nchunks - 1) * dimfinalchunk + nchunks = 1 + (varfinal.dim - 1) // reduction.get_dimfinalchunk() + dimlastfinalchunk = varfinal.dim - (nchunks - 1) * reduction.get_dimfinalchunk() varloader = Var_loader(fun_internal) dimsx = varloader.dimsx dimsy = varloader.dimsy @@ -139,9 +144,11 @@ def get_code(self): KeOps_Error("dimfout should be 1") sum_scheme = self.sum_scheme - self.dimy = max(dimfinalchunk, dimy) + self.dimy = max(reduction.get_dimfinalchunk(), dimy) blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) if not isinstance(sum_scheme, block_sum): @@ -149,7 +156,7 @@ def get_code(self): param_loc = c_array(dtype, dimp, "param_loc") fout = c_array(dtype, dimfout * blocksize_chunks, "fout") xi = c_array(dtype, dimx, "xi") - acc = c_array(dtypeacc, dimfinalchunk, "acc") + acc = c_array(dtypeacc, reduction.get_dimfinalchunk(), "acc") yjloc = c_array(dtype, dimy, f"(yj + threadIdx.x * {dimy})") foutjrel = c_array(dtype, dimfout, f"({fout.id}+jrel*{dimfout})") yjrel = c_array(dtype, dimy, "yjrel") @@ -199,7 +206,7 @@ def get_code(self): dtype, fun_global, varfinal, - dimfinalchunk, + reduction.get_dimfinalchunk(), acc, i, j, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py index 6abc89fb9..552fa85fc 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py @@ -12,7 +12,7 @@ use_pragma_unroll, c_zero_float, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class GpuReduc2D(MapReduce, Gpu_link_compile): diff --git a/keopscore/keopscore/mapreduce/gpu/__init__.py b/keopscore/keopscore/mapreduce/gpu/__init__.py index cb53e5618..caa8a929e 100644 --- a/keopscore/keopscore/mapreduce/gpu/__init__.py +++ b/keopscore/keopscore/mapreduce/gpu/__init__.py @@ -1,10 +1,23 @@ from .GpuAssignZero import GpuAssignZero + from .GpuReduc1D import GpuReduc1D from .GpuReduc1D_chunks import GpuReduc1D_chunks from .GpuReduc1D_finalchunks import GpuReduc1D_finalchunks from .GpuReduc1D_ranges import GpuReduc1D_ranges from .GpuReduc1D_ranges_chunks import GpuReduc1D_ranges_chunks -from .GpuReduc1D_ranges_finalchunks import ( - GpuReduc1D_ranges_finalchunks, -) +from .GpuReduc1D_ranges_finalchunks import GpuReduc1D_ranges_finalchunks + from .GpuReduc2D import GpuReduc2D + +_exports = [ + GpuAssignZero, + GpuReduc1D, + GpuReduc1D_chunks, + GpuReduc1D_finalchunks, + GpuReduc1D_ranges, + GpuReduc1D_ranges_chunks, + GpuReduc1D_ranges_finalchunks, + GpuReduc2D, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/sandbox/Conv2D.py b/keopscore/keopscore/sandbox/Conv2D.py index 8b7cfd8f8..3fb2c776e 100644 --- a/keopscore/keopscore/sandbox/Conv2D.py +++ b/keopscore/keopscore/sandbox/Conv2D.py @@ -12,7 +12,7 @@ test_grad = False test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/chunks.py b/keopscore/keopscore/sandbox/chunks.py index 8b5024054..a8cadcc72 100644 --- a/keopscore/keopscore/sandbox/chunks.py +++ b/keopscore/keopscore/sandbox/chunks.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/chunks_ranges.py b/keopscore/keopscore/sandbox/chunks_ranges.py index 83cc28c61..419575726 100644 --- a/keopscore/keopscore/sandbox/chunks_ranges.py +++ b/keopscore/keopscore/sandbox/chunks_ranges.py @@ -16,7 +16,7 @@ pykeops.config.gpu_available = 0 -device_id = "cuda:0" if pykeops.config.gpu_available else "cpu" +device_id = "cuda" if pykeops.config.gpu_available else "cpu" do_warmup = False x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/do_clean_keops.py b/keopscore/keopscore/sandbox/do_clean_keops.py index 4fe078e82..70b357ac8 100644 --- a/keopscore/keopscore/sandbox/do_clean_keops.py +++ b/keopscore/keopscore/sandbox/do_clean_keops.py @@ -1,3 +1,3 @@ -from keopscore.utils.code_gen_utils import clean_keops +from keopscore.config import clean_keops clean_keops() diff --git a/keopscore/keopscore/sandbox/do_test_formula.py b/keopscore/keopscore/sandbox/do_test_formula.py index 85833e601..21f359de0 100644 --- a/keopscore/keopscore/sandbox/do_test_formula.py +++ b/keopscore/keopscore/sandbox/do_test_formula.py @@ -2,13 +2,17 @@ from keopscore.utils.TestFormula import TestFormula if __name__ == "__main__": - if len(sys.argv) == 1: - print( - """ + + if len(sys.argv) == 0: + # Add a default formula + res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") + + elif len(sys.argv) == 1: + print(""" you should provide a formula to test, e.g. 'python do_test_formula.py "Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)"' - """ - ) + """) + else: if len(sys.argv) == 2: res = TestFormula(sys.argv[1]) diff --git a/keopscore/keopscore/sandbox/finalchunks.py b/keopscore/keopscore/sandbox/finalchunks.py index 88125b7ee..0db91a099 100644 --- a/keopscore/keopscore/sandbox/finalchunks.py +++ b/keopscore/keopscore/sandbox/finalchunks.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/finalchunks_ranges.py b/keopscore/keopscore/sandbox/finalchunks_ranges.py index a5301f2c8..eaa8f8347 100644 --- a/keopscore/keopscore/sandbox/finalchunks_ranges.py +++ b/keopscore/keopscore/sandbox/finalchunks_ranges.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/formula.py b/keopscore/keopscore/sandbox/formula.py index 03e93c095..2b26ea517 100644 --- a/keopscore/keopscore/sandbox/formula.py +++ b/keopscore/keopscore/sandbox/formula.py @@ -1,10 +1,9 @@ # testing some formulas with keopscore +import keopscore.config from keopscore.formulas import * -import keopscore - -keopscore.auto_factorize = True +keopscore.config.reduction.set_auto_factorize(True) print("********************************") print("test 1") diff --git a/keopscore/keopscore/sandbox/genred_auto_factorize.py b/keopscore/keopscore/sandbox/genred_auto_factorize.py index c2ab468f8..48396ed51 100644 --- a/keopscore/keopscore/sandbox/genred_auto_factorize.py +++ b/keopscore/keopscore/sandbox/genred_auto_factorize.py @@ -8,7 +8,7 @@ dtype = torch.float32 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(M, D, device=device_id, dtype=dtype) / math.sqrt(D) y = torch.rand(N, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/genred_float16.py b/keopscore/keopscore/sandbox/genred_float16.py index f43c093b6..38c117d1c 100644 --- a/keopscore/keopscore/sandbox/genred_float16.py +++ b/keopscore/keopscore/sandbox/genred_float16.py @@ -6,17 +6,12 @@ import torch from pykeops.torch import Genred -( - M, - N, -) = ( - 5, - 5, -) +M = 5 +N = 5 dtype = torch.float16 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.zeros(M, 1, device=device_id, dtype=dtype) b = torch.ones(N, 1, device=device_id, dtype=dtype) diff --git a/keopscore/keopscore/sandbox/genred_gaussian.py b/keopscore/keopscore/sandbox/genred_gaussian.py index 51b7cfae1..bb9d36184 100644 --- a/keopscore/keopscore/sandbox/genred_gaussian.py +++ b/keopscore/keopscore/sandbox/genred_gaussian.py @@ -10,7 +10,7 @@ dtype = torch.float64 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/genred_inplace.py b/keopscore/keopscore/sandbox/genred_inplace.py index 149bd1439..f2e123de0 100644 --- a/keopscore/keopscore/sandbox/genred_inplace.py +++ b/keopscore/keopscore/sandbox/genred_inplace.py @@ -10,7 +10,7 @@ dtype = torch.float64 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/laplacian.py b/keopscore/keopscore/sandbox/laplacian.py index 4404eb9d3..a3be6eaf2 100644 --- a/keopscore/keopscore/sandbox/laplacian.py +++ b/keopscore/keopscore/sandbox/laplacian.py @@ -3,8 +3,9 @@ from time import time from keopscore.formulas import * -# import keopscore -# keopscore.debug_ops = True +import keopscore + +keopscore.config.debug.set_debug_ops(True) def GaussLapKernel(sigma, D): diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian.py b/keopscore/keopscore/sandbox/lazytensor_gaussian.py index 501abf387..5649fd678 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian.py @@ -16,7 +16,7 @@ pykeops.config.gpu_available = 0 -device_id = "cuda:0" if pykeops.config.gpu_available else "cpu" +device_id = "cuda" if pykeops.config.gpu_available else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py b/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py index e0a3aec8e..e6e8edbc1 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py @@ -7,7 +7,7 @@ D = 3 M, N = 2000, 3000 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float32 x = torch.rand(B1, 1, B3, M, 1, D, dtype=dtype, device=device_id) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py b/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py index 5defa71c7..814980cb7 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(3, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) y = torch.rand(1, 1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py b/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py index 1dd479273..1a1bd0a45 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_grad.py b/keopscore/keopscore/sandbox/lazytensor_grad.py index 39763fd14..798fac664 100644 --- a/keopscore/keopscore/sandbox/lazytensor_grad.py +++ b/keopscore/keopscore/sandbox/lazytensor_grad.py @@ -8,7 +8,7 @@ import keopscore -keopscore.auto_factorize = False +keopscore.config.reduction.set_auto_factorize(False) M, N, D, DV = ( (100000, 100000, 3, 1) if torch.cuda.is_available() else (10000, 10000, 3, 1) @@ -19,7 +19,7 @@ test_grad = True test_grad2 = True test_grad3 = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(M, 1, D, requires_grad=test_grad, device=device_id, dtype=dtype) y = torch.rand(1, N, 1, device=device_id, dtype=dtype) diff --git a/keopscore/keopscore/sandbox/lazytensor_sandbox.py b/keopscore/keopscore/sandbox/lazytensor_sandbox.py index eeb431916..f4eb0f485 100644 --- a/keopscore/keopscore/sandbox/lazytensor_sandbox.py +++ b/keopscore/keopscore/sandbox/lazytensor_sandbox.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_tensordot.py b/keopscore/keopscore/sandbox/lazytensor_tensordot.py index 9a9d44a6f..7a3ec76cd 100644 --- a/keopscore/keopscore/sandbox/lazytensor_tensordot.py +++ b/keopscore/keopscore/sandbox/lazytensor_tensordot.py @@ -11,14 +11,13 @@ from pykeops.torch import LazyTensor - M, N = 2, 10 ####################################################################################################################### # Matrix multiplication as a special case of Tensordot # ---------------------------------------------------- # -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True a = torch.randn(4 * 7, requires_grad=True, device=device_id, dtype=torch.float64) diff --git a/keopscore/keopscore/test/test_op.py b/keopscore/keopscore/test/test_op.py index 34512a12d..27f5b8536 100644 --- a/keopscore/keopscore/test/test_op.py +++ b/keopscore/keopscore/test/test_op.py @@ -8,7 +8,8 @@ import keopscore import keopscore.formulas -from keopscore.utils.misc_utils import KeOps_Error +import pykeops.config as pykeopsconfig +from keopscore.utils.messages import KeOps_Error from pykeops.torch import Genred # fix seed for reproducibility @@ -63,8 +64,9 @@ def perform_test(op_str, tol=1e-4, dtype="float32", verbose=True): M = 300 N = 500 - # Choose the storage place for our data : CPU (host) or GPU (device) memory. - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # Keep tensor placement aligned with the KeOps backend selected at runtime. + use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available + device = torch.device("cuda" if use_cuda else "cpu") if dtype == "float32": torchtype = torch.float32 diff --git a/keopscore/keopscore/test/test_openmp.py b/keopscore/keopscore/test/test_openmp.py new file mode 100644 index 000000000..63c50c70e --- /dev/null +++ b/keopscore/keopscore/test/test_openmp.py @@ -0,0 +1,26 @@ +import os +import warnings + +import keopscore.config +import pytest + + +def test_openmp_detection_non_critical(): + """OpenMP is optional: warn when unavailable, validate when available.""" + openmp = keopscore.config.openmp + assert isinstance(openmp.get_use_OpenMP(), bool) + + if not openmp.get_use_OpenMP(): + with pytest.warns(UserWarning, match="OpenMP is not available"): + warnings.warn( + "OpenMP is not available on this platform/configuration.", + UserWarning, + ) + return + + lib_path = openmp.get_libomp_path() + header_path = openmp.get_libomp_include_path() + + assert lib_path and os.path.exists(lib_path) + assert header_path and os.path.exists(header_path) + assert "-fopenmp" in openmp.get_compile_options() diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index 0a76ed597..87ae47f36 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -1,14 +1,21 @@ import os import pickle + import keopscore -from keopscore.config import * +import keopscore.config # global configuration parameter to be added for the lookup : # N.B we turn this into a function because the parameters need to be read dynamically. env_param = ( - lambda: keopscore.config.get_cpp_flags() - + " auto_factorize=" - + str(keopscore.auto_factorize) + lambda: keopscore.config.cxx.get_compile_options() + + keopscore.config.cxx.get_linking_options() + + keopscore.config.openmp.get_compile_options() + + keopscore.config.openmp.get_include_options() + + keopscore.config.openmp.get_linking_options() + + f" auto_factorize={keopscore.config.reduction.get_auto_factorize()}" + + keopscore.config.cuda.get_preprocessing_options() + + keopscore.config.cuda.get_include_options() + + keopscore.config.cuda.get_linking_options() ) @@ -17,6 +24,7 @@ def __init__(self, fun, use_cache_file=False, save_folder="."): self.fun = fun self.library = {} self.use_cache_file = use_cache_file + self.save_folder = save_folder if use_cache_file: self.cache_file = os.path.join(save_folder, fun.__name__ + "_cache.pkl") if os.path.isfile(self.cache_file) and os.path.getsize(self.cache_file) > 0: @@ -37,6 +45,10 @@ def reset(self, new_save_folder=None): self.library = {} if new_save_folder: self.save_folder = new_save_folder + if self.use_cache_file: + self.cache_file = os.path.join( + self.save_folder, self.fun.__name__ + "_cache.pkl" + ) def save_cache(self): f = open(self.cache_file, "wb") @@ -69,6 +81,7 @@ def __init__(self, cls, use_cache_file=False, save_folder="."): self.cls = cls self.library = {} self.use_cache_file = use_cache_file + self.save_folder = save_folder if self.use_cache_file: self.cache_file = os.path.join(save_folder, cls.__name__ + "_cache.pkl") if os.path.isfile(self.cache_file): @@ -82,7 +95,7 @@ def __init__(self, cls, use_cache_file=False, save_folder="."): atexit.register(self.save_cache) def __call__(self, *args): - str_id = "".join(list(str(arg) for arg in args)) + str(env_param) + str_id = "".join(list(str(arg) for arg in args)) + str(env_param()) if not str_id in self.library: if self.use_cache_file: if str_id in self.library_params: @@ -102,6 +115,10 @@ def reset(self, new_save_folder=None): self.library_params = {} if new_save_folder: self.save_folder = new_save_folder + if self.use_cache_file: + self.cache_file = os.path.join( + self.save_folder, self.cls.__name__ + "_cache.pkl" + ) def save_cache(self): f = open(self.cache_file, "wb") diff --git a/keopscore/keopscore/utils/TestFormula.py b/keopscore/keopscore/utils/TestFormula.py index a3dd253fc..bbff20bb4 100644 --- a/keopscore/keopscore/utils/TestFormula.py +++ b/keopscore/keopscore/utils/TestFormula.py @@ -5,7 +5,7 @@ from pykeops.torch import Genred from keopscore.formulas import * import types -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print def TestFormula(formula, tol=1e-4, dtype="float32", test_grad=False, randseed=None): @@ -80,3 +80,7 @@ def TestFormula(formula, tol=1e-4, dtype="float32", test_grad=False, randseed=No ) return c, g return c + + +if __name__ == "__main__": + res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") diff --git a/keopscore/keopscore/utils/TestOperation.py b/keopscore/keopscore/utils/TestOperation.py index 56c8ee355..b8ce7c7e6 100644 --- a/keopscore/keopscore/utils/TestOperation.py +++ b/keopscore/keopscore/utils/TestOperation.py @@ -5,7 +5,7 @@ from pykeops.torch import Genred from keopscore.formulas import * import types -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print def TestOperation(op_str, tol=1e-4, dtype="float32", test_grad=True): diff --git a/keopscore/keopscore/utils/Tree.py b/keopscore/keopscore/utils/Tree.py index 9bd3f84f7..ffc608d78 100644 --- a/keopscore/keopscore/utils/Tree.py +++ b/keopscore/keopscore/utils/Tree.py @@ -1,4 +1,4 @@ -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print class Tree: diff --git a/keopscore/keopscore/utils/code_gen_utils.py b/keopscore/keopscore/utils/code_gen_utils.py index 10647c239..c75b09573 100644 --- a/keopscore/keopscore/utils/code_gen_utils.py +++ b/keopscore/keopscore/utils/code_gen_utils.py @@ -1,11 +1,7 @@ -import os from hashlib import sha256 -import keopscore -from keopscore.config import * -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Message - -disable_pragma_unrolls = config.get_disable_pragma_unrolls() +import keopscore.config +from keopscore.utils.messages import KeOps_Error def get_hash_name(*args): @@ -216,7 +212,7 @@ def __getitem__(self, other): def use_pragma_unroll(n=64): - if disable_pragma_unrolls: + if keopscore.config.cxx.get_disable_pragma_unrolls(): return "\n" else: if n is None: @@ -826,52 +822,3 @@ def varseq_to_array(vars, vars_ptr_name): string += f""" {vars_ptr_name}[{i}] = {vars[i].id}; """ return string - - -def clean_keops(recompile_jit_binary=True, verbose=True): - - build_path = config.get_build_folder() - use_cuda = keopscore.cuda_config.get_use_cuda() - if use_cuda: - jit_binary = config.get_jit_binary() - else: - jit_binary = None - for f in os.scandir(build_path): - if recompile_jit_binary or f.path != jit_binary: - os.remove(f.path) - if verbose: - KeOps_Message(f"{build_path} has been cleaned.") - from keopscore.get_keops_dll import get_keops_dll - - get_keops_dll.reset() - if use_cuda and recompile_jit_binary: - from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile - - Gpu_link_compile.compile_jit_compile_dll() - - -def check_health(config_type="all"): - """ - Check the health of the specified configuration. - - Parameters: - config_type (str): The configuration to check. Options are: - 'cuda', 'openmp', 'platform', 'base', 'all'. - Default is 'all'. - """ - if config_type == "cuda": - cuda_config.print_all() - elif config_type == "openmp": - openmp_config.print_all() - elif config_type == "platform": - platform_detector.print_all() - elif config_type == ("base"): - config.print_all() - elif config_type == "all": - platform_detector.print_all() - config.print_all() - cuda_config.print_all() - openmp_config.print_all() - else: - print(f"Unknown configuration type: '{config_type}'") - print("Please specify one of: 'cuda', 'openmp', 'platform', 'base', 'all'") diff --git a/keopscore/keopscore/utils/file_utils.py b/keopscore/keopscore/utils/file_utils.py new file mode 100644 index 000000000..478dce8d6 --- /dev/null +++ b/keopscore/keopscore/utils/file_utils.py @@ -0,0 +1,14 @@ +import os +import re + + +def file_to_string(file_path): + """Read a text file and return its content as a string.""" + with open(file_path, "r", encoding="utf-8") as file: + return file.read() + + +def string_to_file(string, file_path): + """Write a string to a text file.""" + with open(file_path, "w", encoding="utf-8") as file: + file.write(string) diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py deleted file mode 100644 index df5b1debc..000000000 --- a/keopscore/keopscore/utils/gpu_utils.py +++ /dev/null @@ -1,244 +0,0 @@ -import ctypes -from ctypes.util import find_library - - -from keopscore.utils.misc_utils import ( - KeOps_Error, - KeOps_Warning, - find_library_abspath, - KeOps_OS_Run, - get_include_file_abspath, -) - -import keopscore -from keopscore.config import * - -import os -from os.path import join - -# Some constants taken from cuda.h -CUDA_SUCCESS = 0 -CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 -CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 - -# warning : libcuda.so is shipped with nvidia drivers (usually system-wide) -# but nvrtc is shipped with cuda-toolkit-dev (may be user installed) -libcuda_folder = os.path.dirname(find_library_abspath("cuda")) -libnvrtc_folder = os.path.dirname(find_library_abspath("nvrtc")) - - -def get_cuda_include_path(): - # auto detect location of cuda headers - - # First we look at CUDA_PATH env variable if it is set - path = os.getenv("CUDA_PATH") - if path: - path = join(path, "include") - if os.path.isfile(join(path, "cuda.h")) and os.path.isfile( - join(path, "nvrtc.h") - ): - return path - - # if not successfull, we try a few standard locations: - cuda_paths_to_try_start = [] - # if user has installed cuda toolkit via conda in the current env, - # we will find it via CONDA_PREFIX environment variable - path_conda = os.getenv("CONDA_PREFIX") - if path_conda is not None: - cuda_paths_to_try_start.append(path_conda) - - cuda_version = get_cuda_version(out_type="string") - cuda_paths_to_try_start += [ - join(os.path.sep, "opt", "cuda"), - join(os.path.sep, "usr", "local", "cuda"), - join(os.path.sep, "usr", "local", f"cuda-{cuda_version}"), - ] - - cuda_paths_to_try_end = [ - "include", - join("targets", "x86_64-linux", "include"), - ] - for path_start in cuda_paths_to_try_start: - for path_end in cuda_paths_to_try_end: - path = join(path_start, path_end) - if os.path.isfile(join(path, "cuda.h")) and os.path.isfile( - join(path, "nvrtc.h") - ): - return path - - # if not successfull, we try to infer location from the libs - cuda_include_path = None - for libpath in libcuda_folder, libnvrtc_folder: - for libtag in "lib", "lib64": - libtag = os.path.sep + libtag + os.path.sep - if libtag in libpath: - includetag = os.path.sep + "include" + os.path.sep - includepath = libpath.replace(libtag, includetag) - if os.path.isfile(join(includepath, "cuda.h")) and os.path.isfile( - join(includepath, "nvrtc.h") - ): - return includepath - - # last try, testing if by any chance the header is already in the default - # include path of gcc - path_cudah = get_include_file_abspath("cuda.h", config.cxx_compiler()) - if path_cudah: - path = os.path.dirname(path_cudah) - if os.path.isfile(join(path, "nvrtc.h")): - return path - - # finally nothing found, so we display a warning asking the user to do something - KeOps_Warning( - """ - The location of Cuda header files cuda.h and nvrtc.h could not be detected on your system. - You must determine their location and then define the environment variable CUDA_PATH, - either before launching Python or using os.environ before importing keops. For example - if these files are in /vol/cuda/10.2.89-cudnn7.6.4.38/include you can do : - import os - os.environ['CUDA_PATH'] = '/vol/cuda/10.2.89-cudnn7.6.4.38' - """ - ) - - -def orig_cuda_include_fp16_path(): - """ - We look for float 16 cuda headers cuda_fp16.h and cuda_fp16.hpp - based on cuda_path locations and return their directory - """ - cuda_include_path = cuda_config.get_cuda_include_path() - - if cuda_include_path: - return cuda_include_path - cuda_fp16_h_abspath = cuda_config.get_include_file_abspath("cuda_fp16.h") - cuda_fp16_hpp_abspath = cuda_config.get_include_file_abspath("cuda_fp16.hpp") - if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: - path = os.path.dirname(cuda_fp16_h_abspath) - if path != os.path.dirname(cuda_fp16_hpp_abspath): - KeOps_Error("cuda_fp16.h and cuda_fp16.hpp are not in the same folder !") - return path - else: - KeOps_Error("cuda_fp16.h and cuda_fp16.hpp were not found") - - -def custom_cuda_include_fp16_path(): - """ - Here we will create (if not done already) a custom cuda_fp16.h file - to be included in nvrtc code compilation, and put it in the keops - build folder. - We need to create this custom cuda_fp16.h header because the original - cuda_fp16.h includes other cuda headers, and for some unknown reason, - providing all the recursively required headers to the nvrtc compiler - does not work. Hence we produce a packed stand-alone version of cuda_fp16.h - by replacing all #include statements by the corresponding headers contents. - """ - from keopscore.utils.misc_utils import pack_header - - build_folder = config.get_build_folder() - fp16_header = "cuda_fp16.h" - fp16_header_path = join(build_folder, fp16_header) - if not os.path.isfile(fp16_header_path): - pack_header(fp16_header, orig_cuda_include_fp16_path(), build_folder) - return build_folder - - -def get_cuda_version(out_type="single_value"): - cuda = ctypes.CDLL(find_library("cudart")) - cuda_version = ctypes.c_int() - cuda.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) - cuda_version = int(cuda_version.value) - if out_type == "single_value": - return cuda_version - cuda_version_major = cuda_version // 1000 - cuda_version_minor = (cuda_version - (1000 * cuda_version_major)) // 10 - if out_type == "major,minor": - return cuda_version_major, cuda_version_minor - elif out_type == "string": - return f"{cuda_version_major}.{cuda_version_minor}" - - -def get_gpu_props(): - """ - Return number of GPU by reading libcuda. - Here we assume the system has cuda support (more precisely that libcuda can be loaded) - Adapted from https://gist.github.com/f0k/0d6431e3faa60bffc788f8b4daa029b1 - credit: Jan Schlüter - """ - cuda = ctypes.CDLL(find_library("cuda")) - - nGpus = ctypes.c_int() - error_str = ctypes.c_char_p() - - result = cuda.cuInit(0) - if result != CUDA_SUCCESS: - # cuda.cuGetErrorString(result, ctypes.byref(error_str)) - # KeOps_Warning("cuInit failed with error code %d: %s" % (result, error_str.value.decode())) - KeOps_Warning( - "cuda was detected, but driver API could not be initialized. Switching to cpu only." - ) - return 0, "" - - result = cuda.cuDeviceGetCount(ctypes.byref(nGpus)) - if result != CUDA_SUCCESS: - # cuda.cuGetErrorString(result, ctypes.byref(error_str)) - # KeOps_Warning("cuDeviceGetCount failed with error code %d: %s" % (result, error_str.value.decode())) - KeOps_Warning( - "cuda was detected, driver API has been initialized, but no working GPU has been found. Switching to cpu only." - ) - return 0, "" - - nGpus = nGpus.value - - def safe_call(d, result): - test = result == CUDA_SUCCESS - if not test: - KeOps_Warning( - f""" - cuda was detected, driver API has been initialized, - but there was an error for detecting properties of GPU device nr {d}. - Switching to cpu only. - """ - ) - return test - - test = True - MaxThreadsPerBlock = [0] * (nGpus) - SharedMemPerBlock = [0] * (nGpus) - for d in range(nGpus): - # getting handle to cuda device - device = ctypes.c_int() - result &= safe_call(d, cuda.cuDeviceGet(ctypes.byref(device), ctypes.c_int(d))) - - # getting MaxThreadsPerBlock info for device - output = ctypes.c_int() - result &= safe_call( - d, - cuda.cuDeviceGetAttribute( - ctypes.byref(output), - ctypes.c_int(CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK), - device, - ), - ) - MaxThreadsPerBlock[d] = output.value - - # getting SharedMemPerBlock info for device - result &= safe_call( - d, - cuda.cuDeviceGetAttribute( - ctypes.byref(output), - ctypes.c_int(CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK), - device, - ), - ) - SharedMemPerBlock[d] = output.value - - # Building compile flags in the form "-D..." options for further compilations - # (N.B. the purpose is to avoid the device query at runtime because it would slow down computations) - string_flags = f"-DMAXIDGPU={nGpus-1} " - for d in range(nGpus): - string_flags += f"-DMAXTHREADSPERBLOCK{d}={MaxThreadsPerBlock[d]} " - string_flags += f"-DSHAREDMEMPERBLOCK{d}={SharedMemPerBlock[d]} " - - if test: - return nGpus, string_flags - else: - return 0, 0, "" diff --git a/keopscore/keopscore/utils/math_functions.py b/keopscore/keopscore/utils/math_functions.py index e883d5072..48de80ba6 100644 --- a/keopscore/keopscore/utils/math_functions.py +++ b/keopscore/keopscore/utils/math_functions.py @@ -3,7 +3,7 @@ new_c_varname, c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error import keopscore from keopscore.config import * @@ -160,7 +160,7 @@ def call(*args): ) keops_round = math_function( cpu_code=lambda x, d: ( - f"round({x})" if eval(d) == 0 else f"(round({x}*{10**eval(d)})/{10**eval(d)})" + f"round({x})" if int(d) == 0 else f"(round({x}*{10**int(d)})/{10**int(d)})" ), gpu_half2_code="NA", ) diff --git a/keopscore/keopscore/utils/messages.py b/keopscore/keopscore/utils/messages.py new file mode 100644 index 000000000..8dffc64c7 --- /dev/null +++ b/keopscore/keopscore/utils/messages.py @@ -0,0 +1,64 @@ +import os +import sys + +CROSS_MARK = "❌" +CHECK_MARK = "✅" + +not_found_str = f"Not Found. {CROSS_MARK}" +enabled_dict = {True: f"Enabled {CHECK_MARK}", False: f"Disabled {CROSS_MARK}"} + + +def _keops_verbose(): + config = sys.modules.get("keopscore.config") + debug = getattr(config, "debug", None) + if debug is not None: + return debug.get_verbose() + else: + message = "[KeOps] Warning : Could not access to verbosity level. Defaulting to verbose level 1." + print(message) + + return 1 # default verbose level if config or debug is not available + + +def KeOps_Print(*messages, force_print=False, level=1, **kwargs): + if _keops_verbose() >= level or force_print: + print(*messages, **kwargs) + + +def KeOps_Message(message, use_tag=True, level=1, **kwargs): + if _keops_verbose() >= level: + tag = "[KeOps] " if use_tag else "" + message = tag + message + print(message, **kwargs) + + +def KeOps_Warning(message, newline=False, level=1, **kwargs): + if _keops_verbose() >= level: + message = ("\n" if newline else "") + "[KeOps] Warning: " + message + print(message, **kwargs) + + +def KeOps_Error(message, show_line_number=True): + message = "[KeOps] Error: " + message + if show_line_number: + from inspect import currentframe, getframeinfo + + frameinfo = getframeinfo(currentframe().f_back) + message += f" (error at line {frameinfo.lineno} in file {frameinfo.filename})" + raise ValueError(message) + + +def print_envs(env_vars): + """Print the values of specified environment variables.""" + + if not env_vars: + return + + print("\nRelevant Environment Variables") + print("-" * 60) + for var in env_vars: + value = os.environ.get(var) + if value: + print(f"{var} = {value}") + else: + print(f"{var} is not set") diff --git a/keopscore/keopscore/utils/misc_utils.py b/keopscore/keopscore/utils/misc_utils.py deleted file mode 100644 index 887127dba..000000000 --- a/keopscore/keopscore/utils/misc_utils.py +++ /dev/null @@ -1,166 +0,0 @@ -####################################################################### -# . Warnings, Errors, etc. -####################################################################### - -import keopscore -from os.path import join -import re -import platform - - -def KeOps_Print(message, force_print=False, **kwargs): - if keopscore.verbose or force_print: - print(message, **kwargs) - - -def KeOps_Message(message, use_tag=True, **kwargs): - if keopscore.verbose: - tag = "[KeOps] " if use_tag else "" - message = tag + message - print(message, **kwargs) - - -def KeOps_Warning(message, newline=False): - if keopscore.verbose: - message = ("\n" if newline else "") + "[KeOps] Warning : " + message - print(message) - - -def KeOps_Error(message, show_line_number=True): - message = "[KeOps] Error : " + message - if show_line_number: - from inspect import currentframe, getframeinfo - - frameinfo = getframeinfo(currentframe().f_back) - message += f" (error at line {frameinfo.lineno} in file {frameinfo.filename})" - raise ValueError(message) - - -def KeOps_OS_Run(command, print_warning=True): - import sys - - python_version = sys.version_info - if python_version >= (3, 7): - import subprocess - - out = subprocess.run(command, shell=True, capture_output=True) - if out.stderr != b"" and print_warning: - KeOps_Warning("There were warnings or errors :", newline=True) - KeOps_Print(out.stderr.decode("utf-8")) - elif python_version >= (3, 5): - import subprocess - - out = subprocess.run( - command, - shell=True, - ) - else: - KeOps_Error("Python version >= 3.5 required.", newline=True) - return out - - -def get_include_file_abspath(filename, compiler): - out = KeOps_OS_Run( - f'echo "#include <{filename}>" | {compiler} -M -E -x c++ - | head -n 2' - ) - strings = out.stdout.decode("utf8").split() - abspath = None - for s in strings: - if filename in s: - abspath = s - return abspath - - -def get_brew_prefix(): - """Get Homebrew prefix path using KeOps_OS_Run""" - if platform.system() != "Darwin": - return None - out = KeOps_OS_Run(f"brew --prefix", print_warning=False) - if out.stderr != b"": - return None - return out.stdout.decode("utf-8").strip() - - -def find_library_abspath(lib): - """ - wrapper around ctypes find_library that returns the full path - of the library. - Warning : it also opens the shared library ! - Adapted from - https://stackoverflow.com/questions/35682600/get-absolute-path-of-shared-library-in-python/35683698 - """ - from ctypes import c_int, c_void_p, c_char_p, CDLL, byref, cast, POINTER, Structure - from ctypes.util import find_library - - # linkmap structure, we only need the second entry - class LINKMAP(Structure): - _fields_ = [("l_addr", c_void_p), ("l_name", c_char_p)] - - res = find_library(lib) - if res is None: - return "" - - lib = CDLL(res) - libdl = CDLL(find_library("dl")) - - dlinfo = libdl.dlinfo - dlinfo.argtypes = c_void_p, c_int, c_void_p - dlinfo.restype = c_int - - # gets typecasted later, I dont know how to create a ctypes struct pointer instance - lmptr = c_void_p() - - # 2 equals RTLD_DI_LINKMAP, pass pointer by reference - dlinfo(lib._handle, 2, byref(lmptr)) - - # typecast to a linkmap pointer and retrieve the name. - abspath = cast(lmptr, POINTER(LINKMAP)).contents.l_name - - return abspath.decode("utf-8") - - -def file_to_string(file_path): - """ - reads text file and returns its content as a string - """ - f = open(file_path, "r", encoding="utf-8") - out = f.read() - f.close() - return out - - -def string_to_file(string, file_path): - """ - writes string to file - """ - f = open(file_path, "w", encoding="utf-8") - out = f.write(string) - f.close() - - -def pack_header(filename, origin_folder, target_folder): - """ - Given a C/C++ header file "filename", located in folder "origin_folder", - produces a stand-alone version of the header by recursiveley - replacing all #include "xxx" statements by the content of the corresponding - file. The resulting file is put into "target_folder". - """ - code = file_to_string(join(origin_folder, filename)) - used_headers = [] - while True: - match = re.search('#include *"([^"]*)"', code) - if match is None: - break - header = match.groups()[0] - if header in used_headers: - code_to_insert = "" - else: - code_to_insert = file_to_string(join(origin_folder, header)) - used_headers.append(header) - code = code[: match.start()] + code_to_insert + code[match.end() :] - string_to_file(code, join(target_folder, filename)) - - -# Factoring the check/cross marks used in the config classes prints -CHECK_MARK = "✅" -CROSS_MARK = "❌" diff --git a/keopscore/keopscore/utils/path_utils.py b/keopscore/keopscore/utils/path_utils.py new file mode 100644 index 000000000..2b12baa05 --- /dev/null +++ b/keopscore/keopscore/utils/path_utils.py @@ -0,0 +1,118 @@ +import os +import sys +from pathlib import Path + + +def _unique_paths(paths): + """Return existing paths in first-seen order with duplicates and None values removed.""" + unique = [] + seen = set() + for path in paths: + if path is None: + continue + path = Path(path) + if not path.exists() and not path.is_dir(): + continue + key = str(path) + if key not in seen: + seen.add(key) + unique.append(path) + return unique + + +def _env_roots(env_vars): + """Return existing environment variable values as normalized Path objects.""" + return _unique_paths(os.getenv(env_var) for env_var in env_vars) + + +def _path_candidates(roots, suffixes): + """Expand root directories with relative suffixes while preserving order.""" + candidates = [] + for root in _unique_paths(roots): + for suffix in suffixes: + candidates.append(root / suffix if suffix else root) + return _unique_paths(candidates) + + +def _ordered_search_roots(env_vars=(), pip=(), conda=None, system=(), order=None): + """Return normalized search roots in a caller-defined category order. + + Args: + env_vars (tuple[str] | list[str]): Environment variable names whose values + should be considered as root directories. + pip (tuple[str] | list[str]): Relative suffixes appended to Python + package roots discovered from the current interpreter. + conda (str | None): Name of an environment variable that points to a + Conda root directory. + system (tuple[str] | list[str]): Fallback root directories to append. + order (tuple[str] | list[str] | str | None): Ordered categories to apply. + Allowed items are ``env_vars``, ``pip``, ``conda``, and + ``system``. A comma-separated string is also accepted. + + Returns: + list[pathlib.Path]: Existing candidates, deduplicated in first-seen order. + """ + + if order is None: + order = ("env_vars", "conda", "system", "pip") + else: + isinstance(order, str) and (order := tuple(order.split(","))) + for item in order: + if item not in ( + "env_vars", + "pip", + "conda", + "system", + ): + raise ValueError(f"Invalid order item: {item}") + + roots = [] + + for item in order: + if item == "env_vars" and env_vars: + roots.extend(_env_roots(env_vars)) + + if item == "pip" and pip: + roots.extend(_unique_paths(pip)) + + if item == "conda" and conda: + roots.extend(_env_roots((conda,))) + + if item == "system" and system: + roots.extend(_unique_paths(system)) + + return _unique_paths(roots) + + +def _first_matching_file(directories, patterns): + """Return the first file matching glob patterns inside ordered directories.""" + if not directories or not patterns: + return None + + if isinstance(patterns, str): + patterns = (patterns,) + patterns = [pattern for pattern in patterns if pattern] + if not patterns: + return None + + for directory in _unique_paths(directories): + try: + directory = Path(directory) + if not directory.is_dir(): + continue + for pattern in patterns: + for match in sorted(directory.glob(pattern)): + if match.is_file(): + return str(match) + except Exception: + continue + + return None + + +def ensure_directory(path, add_to_syspath=False): + """Create a directory and optionally register it in sys.path.""" + os.makedirs(path, exist_ok=True) + if add_to_syspath and path not in sys.path: + sys.path.append(path) + return path diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py new file mode 100644 index 000000000..a9c082123 --- /dev/null +++ b/keopscore/keopscore/utils/system_utils.py @@ -0,0 +1,91 @@ +import os +import subprocess +from ctypes import c_int, c_void_p, c_char_p, CDLL, byref, cast, POINTER, Structure +from ctypes.util import find_library + +from keopscore.utils.messages import KeOps_Print, KeOps_Warning + + +def KeOps_OS_Run(command, print_warning=True): + out = subprocess.run(command, shell=True, capture_output=True) + if out.stderr != b"" and print_warning: + KeOps_Warning( + "There were warnings or errors while executing: " + command, newline=True + ) + KeOps_Print(out.stderr.decode("utf-8")) + return out + + +def get_include_file_abspath(filename, compiler): + """Return the full path of the header filename using compiler.""" + + cmd = f'echo "#include <{filename}>" | {compiler} -M -E -x c++ -' + out = KeOps_OS_Run(cmd, print_warning=False) + + text = out.stdout.decode("utf8") + text = text.replace("\\\n", " ") + + try: + deps = text.split(":", 1)[1] + except IndexError: + return None + + for path in deps.split(): + if os.path.basename(path) == filename: + return path + + return None + + +def find_library_abspath(lib): + """ + Wrapper around ctypes find_library that returns the full path of the library. + Warning: it also opens the shared library. + """ + + class LINKMAP(Structure): + _fields_ = [("l_addr", c_void_p), ("l_name", c_char_p)] + + res = find_library(lib) + if res is None: + return "" + + if os.path.isabs(res): + return res + + lib = CDLL(res) + libdl = CDLL(find_library("dl")) + + try: + dlinfo = libdl.dlinfo + except AttributeError: + dlinfo = None + + if dlinfo is not None: + try: + dlinfo.argtypes = c_void_p, c_int, c_void_p + dlinfo.restype = c_int + + lmptr = c_void_p() + dlinfo(lib._handle, 2, byref(lmptr)) + + abspath = cast(lmptr, POINTER(LINKMAP)).contents.l_name + if abspath: + return abspath.decode("utf-8") + except Exception: + pass + + return "" + + +def _find_library_by_names(library_names): + """Return the first library path resolved by ctypes for known library names.""" + if isinstance(library_names, str): + library_names = (library_names,) + + for library_name in library_names: + library_path = find_library(library_name) + if library_path: + return find_library_abspath(library_name) + + return None diff --git a/keopscore/setup.py b/keopscore/setup.py index a0854b4a2..5c8fccdc6 100644 --- a/keopscore/setup.py +++ b/keopscore/setup.py @@ -1,7 +1,7 @@ # Always prefer setuptools over distutils # To use a consistent encoding -from codecs import open import os +from codecs import open from os import path from setuptools import setup @@ -26,18 +26,20 @@ "Bug Reports": "https://github.com/getkeops/keops/issues", "Source": "https://github.com/getkeops/keops", }, - author="B. Charlier, J. Feydy, J. Glaunes", - author_email="benjamin.charlier@umontpellier.fr, jean.feydy@gmail.com, alexis.glaunes@parisdescartes.fr", + author="B. Charlier, J. Feydy, J. Glaunès", + author_email="benjamin.charlier@inrae.fr, jean.feydy@inria.com, alexis.glaunes@parisdescartes.fr", python_requires=">=3.8", classifiers=[ + "Topic :: Scientific/Engineering", "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", - "Topic :: Scientific/Engineering", - "License :: OSI Approved :: MIT License", + "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Operating System :: MacOS :: MacOS X", "Programming Language :: C++", "Programming Language :: Python :: 3 :: Only", + "Environment :: GPU :: NVIDIA CUDA", + "License :: OSI Approved :: MIT License", ], keywords="kernels gpu autodiff", packages=[ @@ -65,7 +67,8 @@ "readme.md", "licence.txt", "keops_version", - "config/libiomp5.dylib", + "binders/cpp/keops_cpu_runtime.cpp", + "binders/cpp/keops_cpu_runtime.h", "binders/nvrtc/keops_nvrtc.cpp", "binders/nvrtc/nvrtc_jit.cpp", "include/CudaSizes.h", diff --git a/pybuild.sh b/pybuild.sh index f49b3f406..4674cfb94 100755 --- a/pybuild.sh +++ b/pybuild.sh @@ -1,4 +1,4 @@ -#!/bin/bash +#!/usr/bin/env bash ################################################################################ # Instructions for Creating a New Release @@ -6,7 +6,7 @@ # 0) Generate a Twine API token and configure your `.pypirc` file # Ensure both TestPyPI and PyPI are properly set up. # -# 1) Update the version number in the file: ./keops_version and Changes log: ./CHANGELOG.md +# 1) Update the version number in the file: ./keops_version and Changes log: ./CHANGELOG.md # # 2) Build the packages using the build script: # sh ./pybuild.sh @@ -20,157 +20,129 @@ # twine upload ./build/dist/pykeops-XXXXX.tar.gz --repository testpypi # pip install -i https://test.pypi.org/simple/ pykeops # -# ⚠️ Note: TestPyPI may have dependency resolution issues. -# If problems occur, install pykeops from PyPI, uninstall it, -# then reinstall pykeops from TestPyPI. +# Note: TestPyPI may have dependency resolution issues. +# If problems occur, install pykeops from PyPI, uninstall it, +# then reinstall pykeops from TestPyPI. # -# ⚠️ Note: Do not forget to remove the install from TestPyPI... +# Note: Do not forget to remove the install from TestPyPI. # # 5) Once validated, upload to the official PyPI: # twine upload ./build/dist/keopscore-XXXXX.tar.gz # twine upload ./build/dist/pykeops-XXXXX.tar.gz ################################################################################ +set -euo pipefail +readonly PYTHON_BIN="python3" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +readonly BUILD_VENV="${SCRIPT_DIR}/.build_venv" +readonly BUILD_REQUIREMENTS=(pip build pyclean) +readonly VERSION="$(<"${SCRIPT_DIR}/keops_version")" -# exit in case of any errors -set -e - -################################################################################ -# help # -################################################################################ -function print_help() { - # Display Help - echo "Build script for keopscore/pykeops packages." - echo - echo "Usage: $0 [option...]" - echo - echo " -h Print the help" - echo " -l Build in local mode (without hard-coded keopscore version requirement in pykeops)" - echo " -v Verbose mode" - echo - echo "Note: by default, the keopscore version requirement is hard-coded in pykeops." - echo - exit 1 -} - -################################################################################ -# utils # -################################################################################ - -# log with verbosity management -function logging() { - if [[ ${PYBUILD_VERBOSE} == 1 ]]; then - echo -e $1 - fi -} - -################################################################################ -# process script options # -################################################################################ - -# default options LOCAL_PYBUILD=0 PYBUILD_VERBOSE=0 -# Get the options -while getopts 'hlv' option; do - case $option in - h) # display Help - print_help - ;; - l) # local build (no hard-coded keopscore version requirements) - LOCAL_PYBUILD=1 - logging "## local build (keopscore version requirements is NOT hard-coded in pykeops)" - ;; - v) # enable verbosity - PYBUILD_VERBOSE=1 - logging "## verbose mode" - ;; - \?) # Invalid option - echo "Error: Invalid option" - exit 1 - ;; - esac -done - -################################################################################ -# script setup # -################################################################################ - -# project root directory -PROJDIR=$(git rev-parse --show-toplevel) - -# python exec -PYTHON="python3" - -# python environment for build -BUILD_VENV=${PROJDIR}/.build_venv +print_help() { + cat < /dev/null && pwd ) - -# python exec -PYTHON="python3" - -# python environment for test -DOC_VENV=${PROJDIR}/.doc_venv - -# python test requirements (names of packages to be installed with pip) -DOC_REQ="pip" - - -################################################################################ -# prepare python environment # -################################################################################ - -logging "-- Preparing python environment for doc build..." - -${PYTHON} -m venv --clear ${DOC_VENV} -source ${DOC_VENV}/bin/activate - -logging "---- Python version = $(python -V)" - -pip install -U ${DOC_REQ} - - -################################################################################ -# Installing keopscore # -################################################################################ - -logging "-- Installing keopscore..." - -pip install -e ${PROJDIR}/keopscore - -################################################################################ -# Installing pykeops # -################################################################################ - -logging "-- Installing pykeops..." - -pip install -e "${PROJDIR}/pykeops[full]" - - - -################################################################################ -# Building the doc # -################################################################################ - -printf "\n----------------------\n Building the doc \n----------------------\n\n" - -# go to the doc directory -CURRENT_DIR=$(pwd) -cd $PROJDIR/doc - -make clean -if [ $noplot = true ]; then - make html-noplot -else - make html -fi - -################################################################################ -# fixing doc link # -################################################################################ - -if [ $fix_link = true ]; then - printf "\n----------------------\n Fixing doc links \n----------------------\n\n" - # Fix some bad links due interaction between rtd-theme and sphinx-gallery - find . -path "*_auto_*" -name "plot_*.html" -exec sed -i "s/doc\/_auto_\(.*\)rst/pykeops\/pykeops\/\1py/" {} \; - find . -path "*_auto_*" -name "index.html" -exec sed -i "s/doc\/_auto_\(.*\)\/index\.rst/pykeops\/pykeops\/\1\//" {} \; -fi - -set -e - -# comes back to directory of -cd $CURRENT_DIR - +#!/usr/bin/env bash + +set -euo pipefail + +readonly PYTHON_BIN="python3" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +readonly DOC_VENV="${SCRIPT_DIR}/.doc_venv" +readonly DOC_REQUIREMENTS=(pip) +readonly DOC_DIR="${SCRIPT_DIR}/doc" + +FIX_LINKS=0 +NO_PLOT=0 +DOC_JOBS="" + +print_help() { + cat </dev/null + make clean + if [[ "${NO_PLOT}" -eq 1 ]]; then + make_target="html-noplot" + else + make_target="html" + fi + make "${make_args[@]}" "${make_target}" + popd >/dev/null +} + +fix_doc_links() { + if [[ "${FIX_LINKS}" -ne 1 ]]; then + return + fi + + log_step "" + log_step "----------------------" + log_step " Fixing doc links" + log_step "----------------------" + log_step "" + + pushd "${DOC_DIR}" >/dev/null + find . -path "*_auto_*" -name "plot_*.html" -exec \ + sed -i "s/doc\/_auto_\(.*\)rst/pykeops\/pykeops\/\1py/" {} \; + find . -path "*_auto_*" -name "index.html" -exec \ + sed -i "s/doc\/_auto_\(.*\)\/index\.rst/pykeops\/pykeops\/\1\//" {} \; + popd >/dev/null +} + +main() { + parse_options "$@" + prepare_python_environment + install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" + install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[full]" + build_doc + fix_doc_links +} + +main "$@" diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 9d08d6e12..441eec74c 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -1,48 +1,31 @@ +import sys import os -import keopscore - -############################################################## -# Verbosity level (we must do this before importing keopscore) -verbose = True -if os.getenv("PYKEOPS_VERBOSE") == "0": - verbose = False - os.environ["KEOPS_VERBOSE"] = "0" +import keopscore from . import config as pykeopsconfig -from keopscore import show_cuda_status -keops_get_build_folder = pykeopsconfig.pykeops_base.get_build_folder -from .config import pykeops_nvrtc_name -from .config import numpy_found, torch_found +########################################################### +# PykeOps version +__version__ = pykeopsconfig.get_version() -def set_verbose(val): - global verbose - verbose = val - keopscore.verbose = val +############################################################## +# Verbosity level (we must do this before importing keopscore) +verbose = pykeopsconfig.init_verbose() -########################################################### -# Set version -with open( - os.path.join(os.path.abspath(os.path.dirname(__file__)), "keops_version"), - encoding="utf-8", -) as v: - __version__ = v.read().rstrip() +def set_verbose(val): + pykeopsconfig.set_verbose(val) + sys.modules[__name__].verbose = pykeopsconfig.get_verbose() + ########################################################### # Utils default_device_id = 0 # default Gpu device number -if pykeopsconfig.pykeops_cuda.get_use_cuda(): - if not os.path.exists(pykeops_nvrtc_name(type="target")): - from .common.keops_io.LoadKeOps_nvrtc import compile_jit_binary - - compile_jit_binary() - def clean_pykeops(recompile_jit_binaries=True): r""" @@ -53,55 +36,81 @@ def clean_pykeops(recompile_jit_binaries=True): """ import pykeops - keopscore.clean_keops(recompile_jit_binary=recompile_jit_binaries) + keopscore.config.clean_keops(recompile_jit_binary=recompile_jit_binaries) keops_binder = pykeops.common.keops_io.keops_binder for key in keops_binder: keops_binder[key].reset() - if recompile_jit_binaries and pykeopsconfig.pykeops_cuda.get_use_cuda(): - pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() + if recompile_jit_binaries: + pykeops.common.keops_io.LoadKeOps_cpp.compile_jit_binary() + if pykeopsconfig.cuda.get_use_cuda(): + pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() -def check_health(): +def check_health(infos="all"): r""" Runs a complete sanity check of the KeOps installation within your system. This function verifies the setup and configuration of KeOps, including compilation flags, paths, .... Parameters: - config_type (str): The configuration to check. Options are: - 'base', 'cuda', 'openmp', 'platform', 'all'. + infos (str): The configuration to check. Options are: + 'cuda', 'cxx', 'openmp', 'platform', 'path', 'all'. Default is 'all'. Returns: None """ - import pykeops + keopscore.config.check_health(infos=infos) - keopscore.check_health() +def set_build_folder(path=None, reset_all=True): + from .common.keops_io import keops_binder + from .common.keops_io.cpp.LoadKeOps_cpp import ( + should_compile_binder as should_compile_cpp_binder, + ) -def set_build_folder(path=None): - import pykeops + # Set the build folder, reset keopscore cache and recompile JIT binaries if needed + keopscore.set_build_folder(path, reset_all=reset_all) - keopscore.set_build_folder(path) - keops_binder = pykeops.common.keops_io.keops_binder + # Reset the cache of all pykeops binders to ensure they will be reloaded from the new build folder for key in keops_binder: keops_binder[key].reset(new_save_folder=get_build_folder()) - if pykeopsconfig.pykeops_cuda.get_use_cuda() and not os.path.exists( - pykeops.config.pykeops_nvrtc_name(type="target") + + # Recompile pykeops binder binaries if needed + + if reset_all or should_compile_cpp_binder(): + from .common.keops_io.cpp import LoadKeOps_cpp + + LoadKeOps_cpp.compile_jit_binary() + + if reset_all or ( + pykeopsconfig.cuda.get_use_cuda() + and not os.path.exists(pykeopsconfig.pykeops_nvrtc_name(type="target")) ): - pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() + from .common.keops_io.nvrtc import LoadKeOps_nvrtc + + LoadKeOps_nvrtc.compile_jit_binary() def get_build_folder(): - return keops_get_build_folder() + return pykeopsconfig.get_build_folder() -if numpy_found: +if pykeopsconfig.numpy_found: from .numpy.test_install import test_numpy_bindings -if torch_found: + +if pykeopsconfig.torch_found: from .torch.test_install import test_torch_bindings -# next line is to ensure that cache file for formulas is loaded at import -from .common import keops_io + +# set the build folder and ensure it is in the python path +try: + set_build_folder(reset_all=False) +except Exception as e: + from .common.utils import pyKeOps_Warning + + pyKeOps_Warning( + f"An error occurred while setting up KeOps: {e}. Use pykeops.check_health() to get details on the current configuration.", + level=1, + ) diff --git a/pykeops/pykeops/benchmarks/benchmark_KNN.py b/pykeops/pykeops/benchmarks/benchmark_KNN.py index e4147ac4e..881e2b153 100644 --- a/pykeops/pykeops/benchmarks/benchmark_KNN.py +++ b/pykeops/pykeops/benchmarks/benchmark_KNN.py @@ -59,9 +59,6 @@ ) from dataset_utils import generate_samples -use_cuda = torch.cuda.is_available() - - ############################################## # We then specify the values of K that we will inspect: diff --git a/pykeops/pykeops/benchmarks/benchmark_utils.py b/pykeops/pykeops/benchmarks/benchmark_utils.py index 87e746a83..71b5c9b51 100644 --- a/pykeops/pykeops/benchmarks/benchmark_utils.py +++ b/pykeops/pykeops/benchmarks/benchmark_utils.py @@ -13,10 +13,11 @@ import numpy as np import torch +import pykeops.config # import jax -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() ################################################## # Utility functions: diff --git a/pykeops/pykeops/benchmarks/plot_accuracy.py b/pykeops/pykeops/benchmarks/plot_accuracy.py index c5c014a4e..c765536ec 100644 --- a/pykeops/pykeops/benchmarks/plot_accuracy.py +++ b/pykeops/pykeops/benchmarks/plot_accuracy.py @@ -13,15 +13,15 @@ output_filename = "accuracy" -import importlib import os import time import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() D = 3 @@ -145,7 +145,7 @@ def benchmark( ) ) else: - elapsed = np.NaN + elapsed = np.nan # accuracy ind = torch.randperm(y.shape[0]) @@ -159,11 +159,11 @@ def benchmark( mean_err = ( (out.double() - ref_out.double()).abs().mean() / ref_out.double().abs().mean() ).item() - mean_err = float("NaN") if mean_err == 0 else mean_err + mean_err = float("nan") if mean_err == 0 else mean_err max_err = ( (out.double() - ref_out.double()).abs().max() / ref_out.double().abs().mean() ).item() - max_err = float("NaN") if max_err == 0 else max_err + max_err = float("nan") if max_err == 0 else max_err print( "accuracy of an MxN convolution, with M = {}, N ={:7}: mean err={:.1e}, max err={:.1e}".format( M, N, mean_err, max_err diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py b/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py index 2ecdf516a..ff26ab9d3 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py @@ -22,12 +22,12 @@ import numpy as np import timeit import matplotlib +import pykeops.config from matplotlib import pyplot as plt from pykeops.numpy.utils import np_kernel from pykeops.torch.utils import torch_kernel from pykeops.torch import Vi, Vj, Pm - ###################################################################### # Benchmark specifications: # @@ -56,7 +56,7 @@ try: import torch - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" torchtype = torch.float32 if dtype == "float32" else torch.float64 diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py b/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py index cbe438cbc..312d71f33 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py @@ -23,12 +23,12 @@ import numpy as np import timeit import matplotlib +import pykeops.config from matplotlib import pyplot as plt from pykeops.numpy.utils import grad_np_kernel, chain_rules from pykeops.torch.utils import torch_kernel from pykeops.torch import Vi, Vj, Pm - ###################################################################### # Benchmark specifications: # @@ -62,7 +62,7 @@ try: import torch - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" torchtype = torch.float32 if dtype == "float32" else torch.float64 diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py b/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py index 67f07a515..cdba8c540 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py @@ -18,8 +18,6 @@ from benchmark_utils import random_normal, full_benchmark -use_cuda = torch.cuda.is_available() - ############################################## # Benchmark specifications: # diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py b/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py index e5bdc89a4..6565d3165 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py @@ -33,6 +33,7 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from scipy.sparse import diags @@ -45,7 +46,8 @@ from benchmark_utils import flatten, random_normal, unit_tensor, full_benchmark -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +torch_device = "cuda" if use_cuda else "cpu" if torch.__version__ >= "1.8": torchsolve = lambda A, B: torch.linalg.solve(A, B) @@ -72,7 +74,7 @@ # -def generate_samples(N, device="cuda", lang="torch", **kwargs): +def generate_samples(N, device=None, lang="torch", **kwargs): """Generates a point cloud x, a scalar signal b of size N and two regularization parameters. Args: @@ -83,6 +85,11 @@ def generate_samples(N, device="cuda", lang="torch", **kwargs): Returns: 3-uple of arrays: x, y, b """ + if device is None: + device = torch_device + if lang == "torch" and device == "cuda" and not use_cuda: + device = "cpu" + randn = random_normal(device=device, lang=lang) ones = unit_tensor(device=device, lang=lang) @@ -149,7 +156,7 @@ def Kinv_scipy(x, b, gamma, alpha, **kwargs): def Kinv_pytorch(x, b, gamma, alpha, **kwargs): - K_xx = alpha * torch.eye(x.shape[0], device=x.get_device()) + torch.exp( + K_xx = alpha * torch.eye(x.shape[0], device=x.device) + torch.exp( -gamma * sqdist_torch(x, x) ) res = torchsolve(K_xx, b) diff --git a/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py b/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py index 241eb5ea8..4c0dc816e 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py +++ b/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py @@ -21,11 +21,12 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from benchmark_utils import flatten, random_normal, full_benchmark -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() print( f"Running torch version {torch.__version__} with {'GPU' if use_cuda else 'CPU'}..." ) diff --git a/pykeops/pykeops/common/get_options.py b/pykeops/pykeops/common/get_options.py index 084371304..0f1b9de24 100644 --- a/pykeops/pykeops/common/get_options.py +++ b/pykeops/pykeops/common/get_options.py @@ -1,8 +1,8 @@ import re import numpy as np from collections import OrderedDict -import pykeops -import pykeops.config + +import pykeops.config as pykeopsconfig ############################################################ # define backend @@ -53,7 +53,7 @@ def define_tag_backend(self, backend, variables): # auto : infer everything if backend == "auto": return ( - int(pykeops.config.gpu_available), + int(pykeopsconfig.cuda.is_available()), self._find_grid(), self._find_mem(variables), ) @@ -84,7 +84,7 @@ def define_backend(self, backend, variables): @staticmethod def _find_dev(): - return int(pykeops.config.gpu_available) + return int(pykeopsconfig.cuda.is_available()) @staticmethod def _find_mem(variables): @@ -92,7 +92,7 @@ def _find_mem(variables): [type(var) is np.ndarray for var in variables] ): # Infer if we're working with numpy arrays or torch tensors: MemType = 0 - elif pykeops.config.torch_found: + elif pykeopsconfig.torch_found: import torch if all( diff --git a/pykeops/pykeops/common/gpu_utils.py b/pykeops/pykeops/common/gpu_utils.py index 100995c0a..4b3b7df78 100644 --- a/pykeops/pykeops/common/gpu_utils.py +++ b/pykeops/pykeops/common/gpu_utils.py @@ -1,5 +1,5 @@ -from keopscore.utils.gpu_utils import get_gpu_props +import pykeops.config as pykeopsconfig def get_gpu_number(): - return get_gpu_props()[0] + return pykeopsconfig.cuda.get_n_visible_devices() diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps.py b/pykeops/pykeops/common/keops_io/LoadKeOps.py index 57e00f6b5..26786b94a 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps.py @@ -57,9 +57,9 @@ def init( cat = 2 alias_args = var[3:-1].split(",") if len(alias_args) == 1: - ind, dim = k, eval(alias_args[0]) + ind, dim = k, int(alias_args[0]) elif len(alias_args) == 2: - ind, dim = eval(alias_args[0]), eval(alias_args[1]) + ind, dim = int(alias_args[0]), int(alias_args[1]) alias = f"{varname}=Var({ind},{dim},{cat})" aliases_new.append(alias) @@ -106,7 +106,6 @@ def init( ( self.params.tag, - self.params.source_name, self.params.low_level_code_file, self.params.tagI, self.params.tagZero, diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py deleted file mode 100644 index 7b632f3c9..000000000 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py +++ /dev/null @@ -1,194 +0,0 @@ -import os -import sysconfig - -import pykeops.config as pykeopsconfig - -get_build_folder = pykeopsconfig.pykeops_base.get_build_folder - -from keopscore.utils.Cache import Cache_partial -from pykeops.common.keops_io.LoadKeOps import LoadKeOps -from pykeops.common.utils import pyKeOps_Message -from keopscore.utils.misc_utils import KeOps_OS_Run -from pykeops.config import pykeops_cpp_name, python_includes - - -class LoadKeOps_cpp_class(LoadKeOps): - def __init__(self, *args, fast_init=False): - super().__init__(*args, fast_init=fast_init) - - def init_phase1(self): - srcname = pykeops_cpp_name(tag=self.params.tag, extension=".cpp") - - dllname = pykeops_cpp_name( - tag=self.params.tag, extension=sysconfig.get_config_var("EXT_SUFFIX") - ) - - if not os.path.exists(dllname): - f = open(srcname, "w") - f.write(self.get_pybind11_code()) - f.close() - compile_command = f"{pykeopsconfig.pykeops_base.get_cxx_compiler()} {pykeopsconfig.pykeops_base.get_cpp_flags()} {python_includes} {srcname} -o {dllname}" - pyKeOps_Message( - "Compiling pykeops cpp " + self.params.tag + " module ... ", - flush=True, - end="", - ) - KeOps_OS_Run(compile_command) - pyKeOps_Message("OK", use_tag=False, flush=True) - - def init_phase2(self): - import importlib - - mylib = importlib.import_module( - os.path.basename(pykeops_cpp_name(tag=self.params.tag)) - ) - - self.launch_keops_cpu = mylib.launch_pykeops_cpu - - def call_keops(self, nx, ny): - self.launch_keops_cpu( - self.params.dimy, - nx, - ny, - self.params.tagI, - self.params.tagZero, - self.params.use_half, - self.params.dimred, - self.params.use_chunk_mode, - self.params.indsi, - self.params.indsj, - self.params.indsp, - self.params.dim, - self.params.dimsx, - self.params.dimsy, - self.params.dimsp, - self.ranges_ptr_new, - self.outshape, - self.out_ptr, - self.args_ptr_new, - self.argshapes_new, - ) - - def get_pybind11_code(self): - return f""" -#include "{self.params.source_name}" - -#include -namespace py = pybind11; - -template < typename TYPE > -int launch_pykeops_{self.params.tag}_cpu(signed long int dimY, signed long int nx, signed long int ny, - int tagI, int tagZero, int use_half, - signed long int dimred, - int use_chunk_mode, - py::tuple py_indsi, py::tuple py_indsj, py::tuple py_indsp, - signed long int dimout, - py::tuple py_dimsx, py::tuple py_dimsy, py::tuple py_dimsp, - py::tuple py_ranges, - py::tuple py_shapeout, - long out_void, - py::tuple py_arg, - py::tuple py_argshape){{ - - /*------------------------------------*/ - /* Cast input args */ - /*------------------------------------*/ - - std::vector< int > indsi_v(py_indsi.size()); - for (auto i = 0; i < py_indsi.size(); i++) - indsi_v[i] = py::cast< int >(py_indsi[i]); - - - std::vector< int > indsj_v(py_indsj.size()); - for (auto i = 0; i < py_indsj.size(); i++) - indsj_v[i] = py::cast< int >(py_indsj[i]); - - - std::vector< int > indsp_v(py_indsp.size()); - for (auto i = 0; i < py_indsp.size(); i++) - indsp_v[i] = py::cast< int >(py_indsp[i]); - - - std::vector< signed long int > dimsx_v(py_dimsx.size()); - for (auto i = 0; i < py_dimsx.size(); i++) - dimsx_v[i] = py::cast< signed long int >(py_dimsx[i]); - - - std::vector< signed long int > dimsy_v(py_dimsy.size()); - for (auto i = 0; i < py_dimsy.size(); i++) - dimsy_v[i] = py::cast< signed long int >(py_dimsy[i]); - - - std::vector< signed long int > dimsp_v(py_dimsp.size()); - for (auto i = 0; i < py_dimsp.size(); i++) - dimsp_v[i] = py::cast< signed long int >(py_dimsp[i]); - - - // Cast the ranges arrays - std::vector< signed long int* > ranges_v(py_ranges.size()); - for (signed long int i = 0; i < py_ranges.size(); i++) - ranges_v[i] = (signed long int*) py::cast< signed long int >(py_ranges[i]); - signed long int **ranges = (signed long int**) ranges_v.data(); - - std::vector< signed long int > shapeout_v(py_shapeout.size()); - for (auto i = 0; i < py_shapeout.size(); i++) - shapeout_v[i] = py::cast< signed long int >(py_shapeout[i]); - - TYPE *out = (TYPE*) out_void; - // std::cout << "out_ptr : " << (long) out << std::endl; - - std::vector< TYPE* > arg_v(py_arg.size()); - for (int i = 0; i < py_arg.size(); i++) - arg_v[i] = (TYPE*) py::cast< long >(py_arg[i]); - TYPE **arg = (TYPE**) arg_v.data(); - - std::vector< std::vector< signed long int > > argshape_v(py_argshape.size()); - for (auto i = 0; i < py_argshape.size(); i++){{ - py::tuple tmp = py_argshape[i]; - std::vector< signed long int > tmp_v(tmp.size()); - for (auto j =0; j < tmp.size(); j++) - tmp_v[j] = py::cast< signed long int >(tmp[j]); - argshape_v[i] = tmp_v; - }} - - - return launch_keops_cpu_{self.params.tag}< TYPE >(dimY, - nx, - ny, - tagI, - tagZero, - use_half, - dimred, - use_chunk_mode, - indsi_v, - indsj_v, - indsp_v, - dimout, - dimsx_v, - dimsy_v, - dimsp_v, - ranges, - shapeout_v, - out, - arg, - argshape_v); - -}} - -PYBIND11_MODULE(pykeops_cpp_{self.params.tag}, m) {{ - m.doc() = "pyKeOps: KeOps for pytorch through pybind11 (pytorch flavour)."; - m.def("launch_pykeops_cpu", &launch_pykeops_{self.params.tag}_cpu < {cpp_dtype[self.params.dtype]} >, "Entry point to keops."); -}} - """ - - -LoadKeOps_cpp = Cache_partial( - LoadKeOps_cpp_class, use_cache_file=True, save_folder=get_build_folder() -) - -cpp_dtype = { - "float": "float", - "float32": "float", - "double": "double", - "float64": "double", -} diff --git a/pykeops/pykeops/common/keops_io/__init__.py b/pykeops/pykeops/common/keops_io/__init__.py index b8b2eb851..658775c37 100644 --- a/pykeops/pykeops/common/keops_io/__init__.py +++ b/pykeops/pykeops/common/keops_io/__init__.py @@ -1,13 +1,14 @@ import pykeops.config as pykeopsconfig +from .nvrtc import LoadKeOps_nvrtc -if pykeopsconfig.pykeops_cuda.get_use_cuda(): - from . import LoadKeOps_nvrtc, LoadKeOps_cpp +if pykeopsconfig.cuda.get_use_cuda(): + from .cpp import LoadKeOps_cpp keops_binder = { "nvrtc": LoadKeOps_nvrtc.LoadKeOps_nvrtc, "cpp": LoadKeOps_cpp.LoadKeOps_cpp, } else: - from . import LoadKeOps_cpp + from .cpp import LoadKeOps_cpp keops_binder = {"cpp": LoadKeOps_cpp.LoadKeOps_cpp} diff --git a/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py new file mode 100644 index 000000000..1eef91c66 --- /dev/null +++ b/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py @@ -0,0 +1,106 @@ +import pykeops.config as pykeopsconfig + +from keopscore.utils.Cache import Cache_partial +from keopscore.utils.system_utils import KeOps_OS_Run +from keopscore.utils.messages import KeOps_Error + +from pykeops.common.keops_io.LoadKeOps import LoadKeOps +from pykeops.common.utils import pyKeOps_Message + + +class LoadKeOps_cpp_class(LoadKeOps): + def __init__(self, *args, fast_init=False): + super().__init__(*args, fast_init=fast_init) + + def init_phase2(self): + import importlib + + if should_compile_binder(): + compile_jit_binary() + + pykeops_cpp = importlib.import_module("pykeops_cpp") + formula_library = self.params.low_level_code_file + if not formula_library: + raise RuntimeError( + "[pyKeOps] Error: empty CPU formula library path. " + "Please clean the KeOps cache and recompile formulas." + ) + + if self.params.c_dtype == "float": + self.launch_keops_cpu = pykeops_cpp.KeOps_module_float(formula_library) + elif self.params.c_dtype == "double": + self.launch_keops_cpu = pykeops_cpp.KeOps_module_double(formula_library) + else: + raise ValueError( + "The cpp backend only supports float32 and float64 inputs." + ) + + def call_keops(self, nx, ny): + self.launch_keops_cpu( + self.params.dimy, + nx, + ny, + self.params.tagI, + self.params.tagZero, + self.params.use_half, + self.params.dimred, + self.params.use_chunk_mode, + self.params.indsi, + self.params.indsj, + self.params.indsp, + self.params.dim, + self.params.dimsx, + self.params.dimsy, + self.params.dimsp, + self.ranges_ptr_new, + self.outshape, + self.out_ptr, + self.args_ptr_new, + self.argshapes_new, + ) + + +def should_compile_binder(): + return pykeopsconfig.path.target_needs_update( + pykeopsconfig.pykeops_cpp_binder_name(type="target"), + pykeopsconfig.pykeops_cpp_binder_name(type="src"), + ) + + +def compile_jit_binary(): + """ + Compile the reusable pybind11 entry point for the cpp backend. + """ + compile_command = ( + f"{pykeopsconfig.cxx.get_cxx_compiler()} " + f"{pykeopsconfig.cxx.get_compile_options()} " + f"{pykeopsconfig.path.get_include_options()} " + f"{pykeopsconfig.python_includes} " + f"{pykeopsconfig.pykeops_cpp_binder_name(type='src')} " + f"{pykeopsconfig.cxx.get_linking_options()} " + f"{pykeopsconfig.cxx.get_dynamic_loader_linking_options()} " + f"-o {pykeopsconfig.pykeops_cpp_binder_name(type='target')}" + ) + pyKeOps_Message("Compiling cpp binder for python", flush=True, end="", level=1) + pyKeOps_Message( + " in cache folder " + pykeopsconfig.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + pyKeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) + + out = KeOps_OS_Run(compile_command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling cpp binder {pykeopsconfig.pykeops_cpp_binder_name(type='target')} for python. See compiler output above." + ) + else: + pyKeOps_Message("OK", use_tag=False, flush=True, level=1) + + +LoadKeOps_cpp = Cache_partial( + LoadKeOps_cpp_class, + use_cache_file=True, + save_folder=pykeopsconfig.path.get_build_folder(), +) diff --git a/pykeops/pykeops/common/keops_io/cpp/__init__.py b/pykeops/pykeops/common/keops_io/cpp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp b/pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp new file mode 100644 index 000000000..e23d704bd --- /dev/null +++ b/pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp @@ -0,0 +1,273 @@ +#include "binders/cpp/keops_cpu_runtime.h" + +#include + +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +namespace py = pybind11; + +namespace pykeops_cpp_backend { + +class DynamicLibrary { +public: + explicit DynamicLibrary(const char *path) : path_(path), handle_(nullptr) { +#if defined(_WIN32) + handle_ = LoadLibraryA(path_.c_str()); + if (handle_ == nullptr) { + throw std::runtime_error("[pyKeOps] Error: could not load CPU formula library " + + path_ + ": " + last_error()); + } +#else + dlerror(); + handle_ = dlopen(path_.c_str(), RTLD_NOW | RTLD_LOCAL); + if (handle_ == nullptr) { + const char *error = dlerror(); + throw std::runtime_error("[pyKeOps] Error: could not load CPU formula library " + + path_ + ": " + (error ? error : "unknown error")); + } +#endif + } + + ~DynamicLibrary() { + if (handle_ != nullptr) { +#if defined(_WIN32) + FreeLibrary(static_cast(handle_)); +#else + dlclose(handle_); +#endif + } + } + + DynamicLibrary(const DynamicLibrary &) = delete; + DynamicLibrary &operator=(const DynamicLibrary &) = delete; + + void *symbol(const char *name) { +#if defined(_WIN32) + void *sym = reinterpret_cast( + GetProcAddress(static_cast(handle_), name)); + if (sym == nullptr) { + throw std::runtime_error("[pyKeOps] Error: could not find symbol " + + std::string(name) + " in " + path_ + ": " + + last_error()); + } + return sym; +#else + dlerror(); + void *sym = dlsym(handle_, name); + const char *error = dlerror(); + if (error != nullptr) { + throw std::runtime_error("[pyKeOps] Error: could not find symbol " + + std::string(name) + " in " + path_ + ": " + + error); + } + return sym; +#endif + } + +private: +#if defined(_WIN32) + static std::string last_error() { + DWORD error = GetLastError(); + if (error == 0) { + return "unknown error"; + } + + LPSTR buffer = nullptr; + size_t size = FormatMessageA( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | + FORMAT_MESSAGE_IGNORE_INSERTS, + nullptr, error, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + reinterpret_cast(&buffer), 0, nullptr); + + std::string message(buffer, size); + LocalFree(buffer); + return message; + } +#endif + + std::string path_; + void *handle_; +}; + +template class KeOps_module_python { +public: + explicit KeOps_module_python(const char *formula_library_name) + : library_(formula_library_name), launch_(nullptr) { + launch_ = reinterpret_cast( + library_.symbol("keops_cpu_launch")); + } + + int operator()(signed long int dimY, signed long int nx, signed long int ny, + int tagI, int tagZero, int use_half, + signed long int dimred, int use_chunk_mode, + py::tuple py_indsi, py::tuple py_indsj, py::tuple py_indsp, + signed long int dimout, py::tuple py_dimsx, + py::tuple py_dimsy, py::tuple py_dimsp, py::tuple py_ranges, + py::tuple py_shapeout, long out_void, py::tuple py_arg, + py::tuple py_argshape) { + + std::vector indsi_v(py_indsi.size()); + for (py::ssize_t i = 0; i < py_indsi.size(); i++) + indsi_v[i] = py::cast(py_indsi[i]); + + std::vector indsj_v(py_indsj.size()); + for (py::ssize_t i = 0; i < py_indsj.size(); i++) + indsj_v[i] = py::cast(py_indsj[i]); + + std::vector indsp_v(py_indsp.size()); + for (py::ssize_t i = 0; i < py_indsp.size(); i++) + indsp_v[i] = py::cast(py_indsp[i]); + + std::vector dimsx_v(py_dimsx.size()); + for (py::ssize_t i = 0; i < py_dimsx.size(); i++) + dimsx_v[i] = py::cast(py_dimsx[i]); + + std::vector dimsy_v(py_dimsy.size()); + for (py::ssize_t i = 0; i < py_dimsy.size(); i++) + dimsy_v[i] = py::cast(py_dimsy[i]); + + std::vector dimsp_v(py_dimsp.size()); + for (py::ssize_t i = 0; i < py_dimsp.size(); i++) + dimsp_v[i] = py::cast(py_dimsp[i]); + + std::vector ranges_v(py_ranges.size()); + for (py::ssize_t i = 0; i < py_ranges.size(); i++) + ranges_v[i] = + reinterpret_cast(py::cast(py_ranges[i])); + + std::vector shapeout_v(py_shapeout.size()); + for (py::ssize_t i = 0; i < py_shapeout.size(); i++) + shapeout_v[i] = py::cast(py_shapeout[i]); + + TYPE *out = reinterpret_cast(out_void); + + std::vector typed_arg_v(py_arg.size()); + std::vector arg_v(py_arg.size()); + for (py::ssize_t i = 0; i < py_arg.size(); i++) { + typed_arg_v[i] = reinterpret_cast(py::cast(py_arg[i])); + arg_v[i] = typed_arg_v[i]; + } + + std::vector> argshape_v(py_argshape.size()); + for (py::ssize_t i = 0; i < py_argshape.size(); i++) { + py::tuple tmp = py_argshape[i]; + std::vector tmp_v(tmp.size()); + for (py::ssize_t j = 0; j < tmp.size(); j++) + tmp_v[j] = py::cast(tmp[j]); + argshape_v[i] = tmp_v; + } + + for (int k = 0; k < (int)indsi_v.size(); k++) { + int idx = indsi_v[k]; + const auto &shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument( + "[pyKeOps] Error: Vi argument #" + std::to_string(idx) + + " requires at least 2 dimensions, got " + + std::to_string(shape.size()) + "."); + if (shape.back() != dimsx_v[k]) + throw std::invalid_argument( + "[pyKeOps] Error: Vi argument #" + std::to_string(idx) + + " has trailing dim " + std::to_string(shape.back()) + + ", expected " + std::to_string(dimsx_v[k]) + "."); + } + for (int k = 0; k < (int)indsj_v.size(); k++) { + int idx = indsj_v[k]; + const auto &shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument( + "[pyKeOps] Error: Vj argument #" + std::to_string(idx) + + " requires at least 2 dimensions, got " + + std::to_string(shape.size()) + "."); + if (shape.back() != dimsy_v[k]) + throw std::invalid_argument( + "[pyKeOps] Error: Vj argument #" + std::to_string(idx) + + " has trailing dim " + std::to_string(shape.back()) + + ", expected " + std::to_string(dimsy_v[k]) + "."); + } + for (int k = 0; k < (int)indsp_v.size(); k++) { + int idx = indsp_v[k]; + const auto &shape = argshape_v[idx]; + if ((int)shape.size() < 1) + throw std::invalid_argument( + "[pyKeOps] Error: Pm argument #" + std::to_string(idx) + + " requires at least 1 dimension (got a scalar)."); + if (shape.back() != dimsp_v[k]) + throw std::invalid_argument( + "[pyKeOps] Error: Pm argument #" + std::to_string(idx) + + " has trailing dim " + std::to_string(shape.back()) + + ", expected " + std::to_string(dimsp_v[k]) + "."); + } + + std::vector argshape_ptr_v(argshape_v.size()); + std::vector argshape_size_v(argshape_v.size()); + for (std::size_t i = 0; i < argshape_v.size(); ++i) { + argshape_ptr_v[i] = argshape_v[i].data(); + argshape_size_v[i] = static_cast(argshape_v[i].size()); + } + + KeOpsCpuArgs runtime_args; + runtime_args.dimY = dimY; + runtime_args.nx = nx; + runtime_args.ny = ny; + runtime_args.tagI = tagI; + runtime_args.tagZero = tagZero; + runtime_args.use_half = use_half; + runtime_args.dimred = dimred; + runtime_args.use_chunk_mode = use_chunk_mode; + runtime_args.n_indsi = static_cast(indsi_v.size()); + runtime_args.indsi = indsi_v.data(); + runtime_args.n_indsj = static_cast(indsj_v.size()); + runtime_args.indsj = indsj_v.data(); + runtime_args.n_indsp = static_cast(indsp_v.size()); + runtime_args.indsp = indsp_v.data(); + runtime_args.dimout = dimout; + runtime_args.n_dimsx = static_cast(dimsx_v.size()); + runtime_args.dimsx = dimsx_v.data(); + runtime_args.n_dimsy = static_cast(dimsy_v.size()); + runtime_args.dimsy = dimsy_v.data(); + runtime_args.n_dimsp = static_cast(dimsp_v.size()); + runtime_args.dimsp = dimsp_v.data(); + runtime_args.ranges = ranges_v.data(); + runtime_args.n_shapeout = static_cast(shapeout_v.size()); + runtime_args.shapeout = shapeout_v.data(); + runtime_args.out = out; + runtime_args.nargs = static_cast(arg_v.size()); + runtime_args.arg = arg_v.data(); + runtime_args.argshape = argshape_ptr_v.data(); + runtime_args.argshape_sizes = argshape_size_v.data(); + + py::gil_scoped_release release; + return launch_(&runtime_args); + } + +private: + DynamicLibrary library_; + KeOpsCpuLaunchFn launch_; +}; + +} // namespace pykeops_cpp_backend + +PYBIND11_MODULE(pykeops_cpp, m) { + m.doc() = "pyKeOps: reusable CPU frontend through pybind11."; + + py::class_>( + m, "KeOps_module_float", py::module_local()) + .def(py::init()) + .def("__call__", + &pykeops_cpp_backend::KeOps_module_python::operator()); + + py::class_>( + m, "KeOps_module_double", py::module_local()) + .def(py::init()) + .def("__call__", + &pykeops_cpp_backend::KeOps_module_python::operator()); +} diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py similarity index 69% rename from pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py rename to pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py index 95ef25489..5dd210046 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py @@ -1,14 +1,14 @@ -import os import sys -import pykeops +import pykeops.config as pykeopsconfig + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.utils.Cache import Cache_partial +from keopscore.utils.system_utils import KeOps_OS_Run +from keopscore.utils.messages import KeOps_Error + from pykeops.common.keops_io.LoadKeOps import LoadKeOps from pykeops.common.utils import pyKeOps_Message -from keopscore.utils.misc_utils import KeOps_OS_Run - -get_build_folder = pykeops.config.pykeops_base.get_build_folder class LoadKeOps_nvrtc_class(LoadKeOps): @@ -18,10 +18,10 @@ def __init__(self, *args, fast_init=False): def init_phase2(self): import importlib - if pykeops.get_build_folder() not in sys.path: + if pykeopsconfig.get_build_folder() not in sys.path: # The build folder is supposed to be in the python path, if not, # we add it - sys.path.append(pykeops.get_build_folder()) + sys.path.append(pykeopsconfig.get_build_folder()) pykeops_nvrtc = importlib.import_module("pykeops_nvrtc") @@ -80,17 +80,30 @@ def compile_jit_binary(): This function compile the main .so entry point to keops_nvrt binder... """ compile_command = Gpu_link_compile.get_compile_command( - extra_flags=pykeops.config.python_includes, - sourcename=pykeops.config.pykeops_nvrtc_name(type="src"), - dllname=pykeops.config.pykeops_nvrtc_name(type="target"), + pykeopsconfig.pykeops_nvrtc_name(type="src"), + pykeopsconfig.pykeops_nvrtc_name(type="target"), + extra_flags=pykeopsconfig.python_includes, ) - pyKeOps_Message("Compiling nvrtc binder for python ... ", flush=True, end="") - KeOps_OS_Run(compile_command) - pyKeOps_Message("OK", use_tag=False, flush=True) + pyKeOps_Message("Compiling nvrtc binder for python", flush=True, end="", level=1) + pyKeOps_Message( + " in cache folder " + pykeopsconfig.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + pyKeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) + + out = KeOps_OS_Run(compile_command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling nvrtc binder {pykeopsconfig.pykeops_nvrtc_name(type='target')} for python. See compiler output above." + ) + else: + pyKeOps_Message("OK", use_tag=False, flush=True, level=1) LoadKeOps_nvrtc = Cache_partial( LoadKeOps_nvrtc_class, use_cache_file=True, - save_folder=get_build_folder(), + save_folder=pykeopsconfig.get_build_folder(), ) diff --git a/pykeops/pykeops/common/keops_io/nvrtc/__init__.py b/pykeops/pykeops/common/keops_io/nvrtc/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp b/pykeops/pykeops/common/keops_io/nvrtc/pykeops_nvrtc.cpp similarity index 70% rename from pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp rename to pykeops/pykeops/common/keops_io/nvrtc/pykeops_nvrtc.cpp index 6805a9b3f..14feb33ab 100644 --- a/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp +++ b/pykeops/pykeops/common/keops_io/nvrtc/pykeops_nvrtc.cpp @@ -7,6 +7,7 @@ #include #include +#include namespace py = pybind11; @@ -76,6 +77,34 @@ template class KeOps_module_python : public KeOps_module { argshape_v[i] = tmp_v; } + /*------------------------------------*/ + /* Shape validation */ + /*------------------------------------*/ + for (int k = 0; k < (int)indsi_v.size(); k++) { + int idx = indsi_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " requires at least 3 dimensions, got " + std::to_string(shape.size()) + "."); + if (shape.back() != dimsx_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsx_v[k]) + "."); + } + for (int k = 0; k < (int)indsj_v.size(); k++) { + int idx = indsj_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " requires at least 2 dimensions, got " + std::to_string(shape.size()) + "."); + if (shape.back() != dimsy_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsy_v[k]) + "."); + } + for (int k = 0; k < (int)indsp_v.size(); k++) { + int idx = indsp_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 1) + throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " requires at least 1 dimension (got a scalar)."); + if (shape.back() != dimsp_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsp_v[k]) + "."); + } + return KeOps_module::launch_kernel( tagHostDevice, dimY, nx, ny, tagI, tagZero, use_half, tag1D2D, dimred, cuda_block_size, use_chunk_mode, indsi_v, indsj_v, indsp_v, dimout, diff --git a/pykeops/pykeops/common/lazy_tensor.py b/pykeops/pykeops/common/lazy_tensor.py index 791c5a2da..51356f800 100644 --- a/pykeops/pykeops/common/lazy_tensor.py +++ b/pykeops/pykeops/common/lazy_tensor.py @@ -5,7 +5,7 @@ import numpy as np -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from pykeops.common.utils import check_broadcasting @@ -328,28 +328,25 @@ def fixvariables(self): self.variables = newvars def separate_kwargs(self, kwargs): - # separating keyword arguments for Genred init vs Genred call... - # Currently the only additional optional keyword arguments that are passed to Genred init - # are accuracy options: dtype_acc, use_double_acc and sum_scheme, - # chunk mode option enable_chunks, - # use_fast_math option, - # and compiler option optional_flags. - kwargs_init = [] - kwargs_call = [] - for key in kwargs: - if key in ( - "dtype_acc", - "use_double_acc", - "sum_scheme", - "enable_chunks", - "use_fast_math", - "optional_flags", - ): - kwargs_init += [(key, kwargs[key])] - else: - kwargs_call += [(key, kwargs[key])] - kwargs_init = dict(kwargs_init) - kwargs_call = dict(kwargs_call) + """Split keyword arguments between ``Genred`` initialization and call. + + The initialization keywords are ``dtype_acc``, ``use_double_acc``, + ``sum_scheme``, ``enable_chunks``, ``use_fast_math`` and + ``optional_flags``. + """ + init_keys = { + "dtype_acc", + "use_double_acc", + "sum_scheme", + "enable_chunks", + "use_fast_math", + "optional_flags", + } + + kwargs_init = {key: kwargs.get(key) for key in init_keys if key in kwargs} + kwargs_call = { + key: value for key, value in kwargs.items() if key not in init_keys + } return kwargs_init, kwargs_call def promote(self, other, props, is_complex=False): diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index 950c1c7c9..2670f4649 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -1,5 +1,6 @@ import numpy as np +from keopscore.utils.messages import KeOps_Print, KeOps_Warning from pykeops.common.utils import get_tools @@ -87,46 +88,152 @@ def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): return out -def ConjugateGradientSolver(binding, linop, b, eps=1e-6): - # Conjugate gradient algorithm to solve linear system of the form - # Ma=b where linop is a linear operation corresponding - # to a symmetric and positive definite matrix +def ConjugateGradientSolver( + binding, linop, b, x0=None, eps=1e-6, maxiter=None, verbose=False +): + """ + Conjugate gradient algorithm to solve a linear system of the form + ``linop(x) = b``, where ``linop`` is a symmetric positive definite + linear operator. + + This implementation mostly follows SciPy's conjugate gradient solver. + + Parameters + ---------- + binding : str + Backend used by :func:`get_tools`, typically ``"numpy"`` or + ``"torch"``. + + linop : function + Function implementing the matrix-vector product associated with + the linear operator. Typically a PyKeOps routine. + + b : tensor + Right-hand side of the linear system. + + x0 : tensor, optional + Initial guess for the solution of the linear system. + + eps : float, optional + Relative tolerance used to define the absolute stopping threshold. + Internally, the solver stops when ``||r|| <= max(eps * ||b||, eps)``, + where ``r = b - linop(x)`` is the current residual. + + maxiter : int, optional + Maximum number of conjugate gradient iterations. Defaults to + ``10 * b.shape[0]``. + + verbose : bool, optional + If ``True``, print the convergence information dictionary before + returning. Defaults to ``False``. It contains the convergence + information with the following keys: + + - ``"status"``: ``"Converged"`` or ``"Maximum iterations reached"``. + - ``"niter"``: number of iterations performed. + - ``"residual_norm"``: final residual norm ``||r||``. + - ``"relative_residual_norm"``: final residual norm divided by ``||b||``. + - ``"atol"``: absolute tolerance used internally for stopping. + - ``"maxiter"``: effective maximum number of iterations. + - ``"x0_provided"``: whether a non-``None`` initial guess was given. + - ``"dtype"``: data type of the solution ``x``. + + Returns + ------- + x : tensor + Approximate solution returned by the conjugate gradient iterations. + + + """ + tools = get_tools(binding) - delta = tools.size(b) * eps**2 - a = 0 - r = tools.copy(b) + + # stopping criterion + atol, _ = _get_atol_rtol(tools.norm(b), eps) + maxiter = 10 * b.shape[0] if (maxiter is None) else maxiter + + x = tools.zeros_like(b) if (x0 is None) else tools.copy(x0) + r = tools.copy(b) if (x0 is None) else b - linop(x0) nr2 = (r**2).sum() - if nr2 < delta: - return 0 * r - p = tools.copy(r) - k = 0 - while True: - Mp = linop(p) - alp = nr2 / (p * Mp).sum() - a += alp * p - r -= alp * Mp - nr2new = (r**2).sum() - if nr2new < delta: - break - p = r + (nr2new / nr2) * p - nr2 = nr2new - k += 1 - return a + nr2new = nr2 + + if nr2 <= atol * atol: + it = -1 + else: + p = tools.copy(r) + + for it in range(maxiter): + Mp = linop(p) + alp = nr2 / (p * Mp).sum() + x += alp * p + r -= alp * Mp + nr2new = (r**2).sum() + if nr2new < atol * atol: + break + p = r + (nr2new / nr2) * p + nr2 = nr2new + + else: # for loop exhausted + # Return incomplete progress + KeOps_Warning( + "[KeOps CG]: Maximum iterations reached. Check convergence..." + ) + it = -maxiter - 1 + + if verbose: + nr = tools.sqrt(nr2new) + info = { + "status": ("Converged" if nr <= atol else "Maximum iterations reached"), + "niter": it + 1, + "dtype": tools.dtypename(x.dtype), + "residual_norm": float(nr), + "relative_residual_norm": float(nr / tools.norm(b)), + "atol": atol, + "maxiter": maxiter, + "x0_provided": x0 is not None, + } + + KeOps_Print(f"[KeOps CG]: {info}") + + return x + + +def _get_atol_rtol(b_norm, atol=0.0, rtol=1e-5): + """ + A helper function to handle tolerance normalization. See scipy.linalg.cg. + """ + atol = max(float(atol), float(rtol) * float(b_norm)) + + return atol, rtol def KernelLinearSolver( - binding, K, x, b, alpha=0, eps=1e-6, precond=False, precondKernel=None + binding, + K, + x, + b, + alpha=0, + eps=1e-6, + x0=None, + maxiter=None, + precond=False, + precondKernel=None, + verbose=False, ): tools = get_tools(binding) dtype = tools.dtype(x) - def PreconditionedConjugateGradientSolver(linop, b, invprecondop, eps=1e-6): + def PreconditionedConjugateGradientSolver( + linop, b, invprecondop, x0=None, eps=1e-6 + ): # Preconditioned conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix # invprecondop is linear operation corresponding to the inverse of the preconditioner matrix - a = 0 - r = tools.copy(b) + + atol, _ = _get_atol_rtol(tools.norm(b), eps) + + a = 0 if x0 is None else tools.copy(x0) + r = tools.copy(b) if x0 is None else b - linop(x0) z = invprecondop(r) p = tools.copy(z) rz = (r * z).sum() @@ -135,7 +242,7 @@ def PreconditionedConjugateGradientSolver(linop, b, invprecondop, eps=1e-6): alp = rz / (p * linop(p)).sum() a += alp * p r -= alp * linop(p) - if (r**2).sum() < eps**2: + if (r**2).sum() < atol * atol: break z = invprecondop(r) rznew = (r * z).sum() @@ -226,8 +333,18 @@ def f(x, y): if precond: invprecondop = NystromInversePreconditioner(K, precondKernel, x, alpha) - a = PreconditionedConjugateGradientSolver(KernelLinOp, b, invprecondop, eps) + a = PreconditionedConjugateGradientSolver( + KernelLinOp, b, invprecondop, x0=x0, eps=eps + ) else: - a = ConjugateGradientSolver(binding, KernelLinOp, b, eps=eps) + a = ConjugateGradientSolver( + binding, + KernelLinOp, + b, + eps=eps, + x0=x0, + maxiter=maxiter, + verbose=verbose, + ) return a diff --git a/pykeops/pykeops/common/utils.py b/pykeops/pykeops/common/utils.py index 104485ee9..8d4c6230a 100644 --- a/pykeops/pykeops/common/utils.py +++ b/pykeops/pykeops/common/utils.py @@ -95,19 +95,19 @@ def check_broadcasting(dims_1, dims_2): return max_tuple(padded_dims_1, padded_dims_2) -def pyKeOps_Print(message, **kwargs): - if pykeops.verbose: +def pyKeOps_Print(message, level=1, **kwargs): + if pykeops.verbose >= level: print(message, **kwargs) -def pyKeOps_Message(message, use_tag=True, **kwargs): - if pykeops.verbose: +def pyKeOps_Message(message, use_tag=True, level=1, **kwargs): + if pykeops.verbose >= level: tag = "[pyKeOps] " if use_tag else "" message = tag + message print(message, **kwargs) -def pyKeOps_Warning(message): - if pykeops.verbose: +def pyKeOps_Warning(message, level=1, **kwargs): + if pykeops.verbose >= level: message = "[pyKeOps] Warning : " + message - print(message) + print(message, **kwargs) diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index ab14cebea..2e9bdee93 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -1,45 +1,96 @@ import importlib.util +import os import sys -import sysconfig -from os.path import join, dirname, realpath -############################################################### +# Instantiating the keopscore.config main classes for pykeops +import keopscore.config + +cuda = keopscore.config.cuda +path = keopscore.config.path +openmp = keopscore.config.openmp +cxx = keopscore.config.cxx +debug = keopscore.config.debug + +get_build_folder = path.get_build_folder +# keep the old gpu_available variable for backward compatibility, but it is now recommended to use cuda.is_available() instead +gpu_available = cuda.is_available() + + # Initialize some variables: the values may be redefined later numpy_found = importlib.util.find_spec("numpy") is not None torch_found = importlib.util.find_spec("torch") is not None -import keopscore -from keopscore.config import * -# Instantiating the keopscore.config main classes for pykeops -pykeops_cuda = cuda_config -pykeops_openmp = openmp_config -pykeops_base = config +# Verbosity level +class _VerboseConfig: + def __init__(self, level=1): + self.level = level + + +_verbose_state = _VerboseConfig(1) + + +def set_verbose(val): + val = int(val) + if val not in (0, 1, 2): + raise ValueError("Verbose level must be 0, 1 or 2.") + + _verbose_state.level = val + os.environ["KEOPS_VERBOSE"] = str(val) + debug.set_verbose(val) + return _verbose_state.level + + +def init_verbose(): + env_val = os.getenv("PYKEOPS_VERBOSE") + if env_val in ("0", "1", "2"): + return set_verbose(int(env_val)) + + return set_verbose(1) -get_build_folder = pykeops_base.get_build_folder -gpu_available = pykeops_cuda.get_use_cuda() + +def get_verbose(): + return _verbose_state.level + + +# version + +base_dir_path = os.path.abspath(os.path.dirname(__file__)) + + +def read_version(version_file): + with open(version_file, encoding="utf-8") as v: + return v.read().rstrip() + + +_version = read_version(os.path.join(base_dir_path, "keops_version")) + + +def get_version(): + assert ( + _version == keopscore.__version__ + ), f"Version mismatch between pykeops and keopscore: {_version} vs {keopscore.__version__}" + return _version + + +# path +def get_pykeops_io_folder(): + return os.path.join(base_dir_path, "common", "keops_io") def pykeops_nvrtc_name(type="src"): basename = "pykeops_nvrtc" - extension = ".cpp" if type == "src" else sysconfig.get_config_var("EXT_SUFFIX") - return join( - ( - join(dirname(realpath(__file__)), "common", "keops_io") - if type == "src" - else config.get_build_folder() - ), - basename + extension, - ) - - -def pykeops_cpp_name(tag="", extension=""): - basename = "pykeops_cpp_" - return join( - config.get_build_folder(), - basename + tag + extension, - ) + if type == "src": + return os.path.join(get_pykeops_io_folder(), "nvrtc", basename + ".cpp") + return path.get_python_extension_path(basename) + + +def pykeops_cpp_binder_name(type="src"): + basename = "pykeops_cpp" + if type == "src": + return os.path.join(get_pykeops_io_folder(), "cpp", basename + ".cpp") + return path.get_python_extension_path(basename) python_includes = "$({python3} -m pybind11 --includes)".format(python3=sys.executable) diff --git a/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py b/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py index ed6b8c279..d3100edf2 100644 --- a/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py +++ b/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py @@ -63,7 +63,7 @@ # And on our GPUs, with copies between # the Host and Device memories: # -if pykeops.config.gpu_available: +if pykeops.config.cuda.is_available(): for gpuid in gpuids: d = my_routine(x, y, a, p, backend="GPU", device_id=gpuid) print( @@ -96,7 +96,7 @@ #################################################################### # And on the GPUs, with copies between the Host and Device memories: # -if pykeops.config.gpu_available: +if pykeops.config.cuda.is_available(): for gpuid in gpuids: d = ((p - aj) ** 2 * (xi + yj).exp()).sum( axis=1, backend="GPU", device_id=gpuid diff --git a/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py b/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py index 048808156..db4c6f21f 100755 --- a/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py +++ b/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py @@ -27,7 +27,7 @@ ######################################################################## # Define our dataset: two point clouds on the unit square. # -M, N = (5000, 5000) if pykeops.config.gpu_available else (2000, 2000) +M, N = (5000, 5000) if pykeops.config.cuda.is_available() else (2000, 2000) t = np.linspace(0, 2 * np.pi, M + 1)[:-1] x = np.stack((0.4 + 0.4 * (t / 7) * np.cos(t), 0.5 + 0.3 * np.sin(t)), 1) @@ -189,7 +189,7 @@ backends = ( (["CPU", "GPU"] if M * N < 4e8 else ["GPU"]) - if pykeops.config.gpu_available + if pykeops.config.cuda.is_available() else ["CPU"] ) for backend in backends: diff --git a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py index 4b460a452..264673511 100644 --- a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py +++ b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py @@ -27,7 +27,7 @@ # Define our dataset: # -N = 5000 if pykeops.config.gpu_available else 500 # Number of points +N = 5000 if pykeops.config.cuda.is_available() else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel diff --git a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py index 4f18b855e..2a7556f85 100644 --- a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py +++ b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py @@ -27,7 +27,7 @@ # Define our dataset: # -N = 5000 if pykeops.config.gpu_available else 500 # Number of points +N = 5000 if pykeops.config.cuda.is_available() else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel diff --git a/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py b/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py index 21327f745..0b4a4aeae 100644 --- a/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py +++ b/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py @@ -13,11 +13,14 @@ # First, the standard imports: import torch +import pykeops.config from pykeops.torch import Genred import matplotlib.pyplot as plt +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + # Choose the storage place for our data : CPU (host) or GPU (device) memory. -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +device = torch.device("cuda" if use_cuda else "cpu") #################################################################### # Then, the definition of our dataset: diff --git a/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py b/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py index 837cf964e..6dfcde53b 100644 --- a/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py +++ b/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py @@ -18,14 +18,16 @@ from matplotlib import pyplot as plt import matplotlib.cm as cm import torch +import pykeops.config from pykeops.torch import Vi, Vj, Pm, LazyTensor - ############################################## # Dataset: +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + # Choose the storage place for our data : CPU (host) or GPU (device) memory. -dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor +dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor # Three points in the plane R^2 y = torch.tensor([[0.2, 0.7], [0.5, 0.3], [0.7, 0.5]]).type(dtype) diff --git a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py index a7fd50a46..a5156358e 100644 --- a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py +++ b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py @@ -28,6 +28,7 @@ import matplotlib.pyplot as plt import torch +import pykeops.config from torch.autograd import grad from pykeops.torch import Genred @@ -39,7 +40,8 @@ N = 5000 # Choose the storage place for our data : CPU (host) or GPU (device) memory. -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = torch.device("cuda" if use_cuda else "cpu") dtype = "float32" # Could be 'float32' or 'float64' torchtype = torch.float32 if dtype == "float32" else torch.float64 diff --git a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py index 8b3a0dc0e..3ae45b873 100644 --- a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py +++ b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py @@ -28,6 +28,7 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from torch.autograd import grad @@ -148,7 +149,9 @@ # # Of course, this will only work if you own a Gpu... -if torch.cuda.is_available(): +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + +if use_cuda: # first transfer data on gpu pc, ac, xc, yc, ec = p.cuda(), a.cuda(), x.cuda(), y.cuda(), e.cuda() # then call the operations diff --git a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py index 08c7670bc..3db79107c 100644 --- a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py +++ b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py @@ -30,6 +30,7 @@ import matplotlib.pyplot as plt import torch +import pykeops.config from torch.autograd import grad from pykeops.torch import Genred @@ -162,7 +163,9 @@ # # Of course, this will only work if you own a Gpu... -if torch.cuda.is_available(): +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + +if use_cuda: # first transfer data on gpu pc, ac, xc, yc, bc, ec = p.cuda(), a.cuda(), x.cuda(), y.cuda(), b.cuda(), e.cuda() # then call the operations diff --git a/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py b/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py index 6865c501b..bcc6452e3 100644 --- a/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py +++ b/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py @@ -18,12 +18,13 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor nump = lambda t: t.cpu().numpy() -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() # dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor device = torch.device("cuda" if use_cuda else "cpu") dtype = torch.float32 # Standard float type for both CPU/GPU diff --git a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py index ac1be1424..83bffd0ca 100644 --- a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py +++ b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py @@ -19,6 +19,7 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import KernelSolve @@ -32,16 +33,21 @@ # Define our dataset: # -N = 5000 if torch.cuda.is_available() else 500 # Number of points +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = torch.device("cuda" if use_cuda else "cpu") + +N = 5000 if use_cuda else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel -x = torch.rand(N, D, requires_grad=True) -b = torch.rand(N, Dv) -g = torch.Tensor([0.5 / sigma**2]) # Parameter of the Gaussian RBF kernel +x = torch.rand(N, D, requires_grad=True, device=device) +b = torch.rand(N, Dv, device=device) +g = torch.tensor( + [0.5 / sigma**2], device=device +) # Parameter of the Gaussian RBF kernel -if torch.cuda.is_available(): +if use_cuda: sync = torch.cuda.synchronize else: diff --git a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py index 6fc302f7c..dbd74010c 100644 --- a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py +++ b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py @@ -17,6 +17,7 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor as keops @@ -26,14 +27,19 @@ # Define our dataset: # -N = 5000 if torch.cuda.is_available() else 500 # Number of points +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = torch.device("cuda" if use_cuda else "cpu") + +N = 5000 if use_cuda else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel -x = torch.rand(N, D, requires_grad=True) -b = torch.rand(N, Dv) -g = torch.Tensor([0.5 / sigma**2]) # Parameter of the Gaussian RBF kernel +x = torch.rand(N, D, requires_grad=True, device=device) +b = torch.rand(N, Dv, device=device) +g = torch.tensor( + [0.5 / sigma**2], device=device +) # Parameter of the Gaussian RBF kernel alpha = 0.01 # ridge regularization ############################################################################### @@ -59,7 +65,7 @@ # start = time.time() -K_xx = alpha * torch.eye(N) + torch.exp( +K_xx = alpha * torch.eye(N, device=device) + torch.exp( -torch.sum((x[:, None, :] - x[None, :, :]) ** 2, dim=2) / (2 * sigma**2) ) diff --git a/pykeops/pykeops/numpy/operations.py b/pykeops/pykeops/numpy/operations.py index b12876c0d..383d25eea 100644 --- a/pykeops/pykeops/numpy/operations.py +++ b/pykeops/pykeops/numpy/operations.py @@ -1,11 +1,10 @@ import numpy as np - +from pykeops import default_device_id from pykeops.common.get_options import get_tag_backend from pykeops.common.keops_io import keops_binder from pykeops.common.operations import ConjugateGradientSolver from pykeops.common.parse_type import get_sizes, complete_aliases, get_optional_flags from pykeops.common.utils import axis2cat -from pykeops import default_device_id from pykeops.common.utils import pyKeOps_Warning @@ -177,7 +176,16 @@ def __init__( self.optional_flags = optional_flags def __call__( - self, *args, backend="auto", device_id=-1, alpha=1e-10, eps=1e-6, ranges=None + self, + *args, + backend="auto", + device_id=-1, + ranges=None, + alpha=1e-10, + eps=1e-6, + x0=None, + maxiter=None, + verbose=False, ): r""" To apply the routine on arbitrary NumPy arrays. @@ -204,6 +212,21 @@ def __call__( **ridge regularization** parameter, added to the diagonal of the Kernel matrix :math:`K_{xx}`. + eps (float, default = 1e-6): Stopping criterion for the + conjugate gradient algorithm. + + x0 (2d array, default = None): Initial guess for the solution of + the linear system. Should be of the same shape as ``b``. + + maxiter (int, default = None): Maximum number of conjugate + gradient iterations. If ``None``, uses the default from + :func:`pykeops.common.operations.ConjugateGradientSolver`. + + verbose (bool, default = False): If ``True``, prints the + conjugate gradient convergence information dictionary + produced by + :func:`pykeops.common.operations.ConjugateGradientSolver`. + backend (string): Specifies the map-reduce scheme, as detailed in the documentation of the :class:`numpy.Genred ` module. @@ -272,4 +295,12 @@ def linop(var): res += alpha * var return res - return ConjugateGradientSolver("numpy", linop, varinv, eps=eps) + return ConjugateGradientSolver( + "numpy", + linop, + varinv, + eps=eps, + x0=x0, + maxiter=maxiter, + verbose=verbose, + ) diff --git a/pykeops/pykeops/numpy/test_install.py b/pykeops/pykeops/numpy/test_install.py index bd741fb56..e99d39adc 100644 --- a/pykeops/pykeops/numpy/test_install.py +++ b/pykeops/pykeops/numpy/test_install.py @@ -9,8 +9,7 @@ def test_numpy_bindings(): """ - This function try to compile a simple keops formula using the numpy binder. - + Try to compile a simple KeOps formula using the NumPy binder. """ x = np.arange(1, 10).reshape(-1, 3).astype("float32") y = np.arange(3, 9).reshape(-1, 3).astype("float32") @@ -18,7 +17,21 @@ def test_numpy_bindings(): import pykeops.numpy as pknp my_conv = pknp.Genred(formula, var) - if np.allclose(my_conv(x, y).flatten(), expected_res): - pyKeOps_Message("pyKeOps with numpy bindings is working!", use_tag=False) + + try: + keops_res = my_conv(x, y).flatten() + except Exception as e: + raise ValueError(f"Error during computation: {e}") + + if np.allclose(keops_res, expected_res): + pyKeOps_Message( + "pyKeOps with numpy bindings is working!", use_tag=False, level=1 + ) + return True else: - pyKeOps_Message("outputs wrong values...", use_tag=False) + pyKeOps_Message( + f"outputs wrong values: expected {expected_res} but get {keops_res}", + use_tag=False, + level=1, + ) + return False diff --git a/pykeops/pykeops/numpy/utils.py b/pykeops/pykeops/numpy/utils.py index cc0d015d4..035ade1a8 100644 --- a/pykeops/pykeops/numpy/utils.py +++ b/pykeops/pykeops/numpy/utils.py @@ -1,18 +1,17 @@ import numpy as np - +import pykeops.config as pykeopsconfig from pykeops.numpy import Genred, KernelSolve from pykeops.numpy.cluster import swap_axes as np_swap_axes -import pykeops.config class numpytools: - norm = np.linalg.norm - arraysum = np.sum - exp = np.exp - log = np.log + norm = staticmethod(np.linalg.norm) + arraysum = staticmethod(np.sum) + exp = staticmethod(np.exp) + log = staticmethod(np.log) Genred = Genred KernelSolve = KernelSolve - swap_axes = np_swap_axes + swap_axes = staticmethod(np_swap_axes) arraytype = np.ndarray float_types = [float, np.float16, np.float32, np.float64] @@ -107,6 +106,10 @@ def randn(m, n, dtype): def zeros(shape, dtype, device=None, requires_grad=None): return np.zeros(shape).astype(dtype) + @staticmethod + def zeros_like(x): + return np.zeros_like(x) + @staticmethod def empty(shape, dtype, device=None, requires_grad=None): return np.empty(shape, dtype=dtype) @@ -135,6 +138,10 @@ def device_type_index(x): def device_dict(x): return dict(cat="cpu") + @staticmethod + def sqrt(x): + return np.sqrt(x) + def squared_distances(x, y): x_norm = (x**2).sum(1).reshape(-1, 1) @@ -229,7 +236,7 @@ def WarmUpGpu(): from pykeops.common.utils import pyKeOps_Message pyKeOps_Message("Warming up the Gpu (numpy bindings) !!!") - if pykeops.config.gpu_available: + if pykeopsconfig.cuda.is_available(): formula = "Exp(-oos2*SqDist(x,y))*b" aliases = [ "x = Vi(1)", # First arg : i-variable, of size 1 diff --git a/pykeops/pykeops/sandbox/issue_327.py b/pykeops/pykeops/sandbox/issue_327.py index c37742a7d..9c6c01544 100644 --- a/pykeops/pykeops/sandbox/issue_327.py +++ b/pykeops/pykeops/sandbox/issue_327.py @@ -1,7 +1,6 @@ from pykeops.torch import Vi, Vj, LazyTensor import torch - xc = torch.randn(256, 5) xc.requires_grad_(True) diff --git a/pykeops/pykeops/sandbox/issue_335.py b/pykeops/pykeops/sandbox/issue_335.py index fef25a356..49b5b7d52 100644 --- a/pykeops/pykeops/sandbox/issue_335.py +++ b/pykeops/pykeops/sandbox/issue_335.py @@ -23,7 +23,7 @@ def covar_func(x1, x2=None): if __name__ == "__main__": - device = "cuda:0" + device = "cuda" train_x = torch.randn(30, 10000, 3, device=device) # covar_module = gpytorch.kernels.keops.MaternKernel(nu=2.5).to(device) diff --git a/pykeops/pykeops/sandbox/issue_335_alt.py b/pykeops/pykeops/sandbox/issue_335_alt.py index 85cd98853..fc4c2f578 100644 --- a/pykeops/pykeops/sandbox/issue_335_alt.py +++ b/pykeops/pykeops/sandbox/issue_335_alt.py @@ -25,7 +25,7 @@ def covar_func(x1, x2=None): B, M, N, D = 25, 10000, 10000, 3 if __name__ == "__main__": - device = "cuda:0" + device = "cuda" train_x = torch.randn(B, M, D, device=device) # covar_module = gpytorch.kernels.keops.MaternKernel(nu=2.5).to(device) diff --git a/pykeops/pykeops/sandbox/issue_353.py b/pykeops/pykeops/sandbox/issue_353.py index 7245ea04a..1017b1c1e 100644 --- a/pykeops/pykeops/sandbox/issue_353.py +++ b/pykeops/pykeops/sandbox/issue_353.py @@ -2,7 +2,6 @@ from pykeops.torch import Vi, Vj, LazyTensor, Genred import pykeops - d = 1 n, m = 1, 1 q_points = torch.rand((n, d)).requires_grad_(True).to("cuda") diff --git a/pykeops/pykeops/sandbox/test_chunks.py b/pykeops/pykeops/sandbox/test_chunks.py index 3a1530246..7bb38c084 100644 --- a/pykeops/pykeops/sandbox/test_chunks.py +++ b/pykeops/pykeops/sandbox/test_chunks.py @@ -12,7 +12,7 @@ test_grad = True test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand( diff --git a/pykeops/pykeops/sandbox/test_finalchunks.py b/pykeops/pykeops/sandbox/test_finalchunks.py index c697b7cb5..3861232f8 100644 --- a/pykeops/pykeops/sandbox/test_finalchunks.py +++ b/pykeops/pykeops/sandbox/test_finalchunks.py @@ -6,13 +6,13 @@ import torch from pykeops.torch import LazyTensor -M, N, D, DV = 10000, 10000, 3, 10000 +M, N, D, DV = 1000, 1000, 3, 1000 dtype = torch.float32 test_grad = True test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand( diff --git a/pykeops/pykeops/sandbox/test_finalchunks_ranges.py b/pykeops/pykeops/sandbox/test_finalchunks_ranges.py index aac85b2ac..aca386d4a 100644 --- a/pykeops/pykeops/sandbox/test_finalchunks_ranges.py +++ b/pykeops/pykeops/sandbox/test_finalchunks_ranges.py @@ -8,7 +8,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/sandbox/test_gpu_cpu2.py b/pykeops/pykeops/sandbox/test_gpu_cpu2.py index 048808156..db4c6f21f 100644 --- a/pykeops/pykeops/sandbox/test_gpu_cpu2.py +++ b/pykeops/pykeops/sandbox/test_gpu_cpu2.py @@ -27,7 +27,7 @@ ######################################################################## # Define our dataset: two point clouds on the unit square. # -M, N = (5000, 5000) if pykeops.config.gpu_available else (2000, 2000) +M, N = (5000, 5000) if pykeops.config.cuda.is_available() else (2000, 2000) t = np.linspace(0, 2 * np.pi, M + 1)[:-1] x = np.stack((0.4 + 0.4 * (t / 7) * np.cos(t), 0.5 + 0.3 * np.sin(t)), 1) @@ -189,7 +189,7 @@ backends = ( (["CPU", "GPU"] if M * N < 4e8 else ["GPU"]) - if pykeops.config.gpu_available + if pykeops.config.cuda.is_available() else ["CPU"] ) for backend in backends: diff --git a/pykeops/pykeops/sandbox/test_ifelse.py b/pykeops/pykeops/sandbox/test_ifelse.py index 3593a0ef6..58f8f0557 100644 --- a/pykeops/pykeops/sandbox/test_ifelse.py +++ b/pykeops/pykeops/sandbox/test_ifelse.py @@ -2,7 +2,6 @@ import pykeops from pykeops.torch import LazyTensor - ttypes = ( (torch.cuda.FloatTensor,) if torch.cuda.is_available() else (torch.FloatTensor,) ) diff --git a/pykeops/pykeops/sandbox/test_lazytensor_clamp.py b/pykeops/pykeops/sandbox/test_lazytensor_clamp.py index 2d021f3ac..005737d1d 100644 --- a/pykeops/pykeops/sandbox/test_lazytensor_clamp.py +++ b/pykeops/pykeops/sandbox/test_lazytensor_clamp.py @@ -10,7 +10,7 @@ test_grad = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py b/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py index 6486377a1..521aa40a3 100644 --- a/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py +++ b/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py @@ -10,7 +10,7 @@ dtype = torch.float16 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py b/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py index b7f88efcd..0a7bc16a6 100644 --- a/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py +++ b/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py @@ -12,7 +12,7 @@ test_grad = True test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand( diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py index 294607614..8cda610bd 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py @@ -10,7 +10,7 @@ test_grad = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py index 46cb12097..b0320bd1f 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py @@ -10,7 +10,7 @@ test_grad = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True @@ -122,7 +122,7 @@ def diss(i, j): #################################################################### from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, @@ -166,7 +166,7 @@ def DiffT(self, v, gradin): from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Zero import Zero -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py index 7e3df470c..fe9846327 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py @@ -10,7 +10,7 @@ # D is size of each sample M, N, D = 1000, 1000, 50 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True @@ -53,7 +53,7 @@ def SoftDTW_torch(x, y, gamma): ################################## from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, diff --git a/pykeops/pykeops/sandbox/test_torch_func_hessian.py b/pykeops/pykeops/sandbox/test_torch_func_hessian.py index f91422fc5..fcb6e503c 100644 --- a/pykeops/pykeops/sandbox/test_torch_func_hessian.py +++ b/pykeops/pykeops/sandbox/test_torch_func_hessian.py @@ -1,9 +1,8 @@ +import keopscore import torch from pykeops.torch import LazyTensor -import keopscore - -keopscore.auto_factorize = False +keopscore.config.reduction.set_auto_factorize(False) def fn_torch(x_i): diff --git a/pykeops/pykeops/test/__init__.py b/pykeops/pykeops/test/__init__.py index e69de29bb..7f70d7277 100644 --- a/pykeops/pykeops/test/__init__.py +++ b/pykeops/pykeops/test/__init__.py @@ -0,0 +1,30 @@ +def assert_torch_allclose(actual, expected, *, label=None, **kwargs): + + import torch + + # test dtype, convert long to float if needed + + if actual.dtype in [torch.int64, torch.int32]: + actual = actual.float() + + if expected.dtype in [torch.int64, torch.int32]: + expected = expected.float() + + ok = torch.allclose(actual, expected, **kwargs) + diff = torch.linalg.norm(actual - expected).item() + prefix = label if label is not None else "torch.allclose failed" + assert ( + ok + ), f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={torch.linalg.norm(expected).item():.6e}" + + +def assert_np_allclose(actual, expected, *, label=None, **kwargs): + + import numpy as np + + ok = np.allclose(actual, expected, **kwargs) + diff = float(np.linalg.norm(np.asarray(actual) - np.asarray(expected))) + prefix = label if label is not None else "np.allclose failed" + assert ( + ok + ), f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={float(np.linalg.norm(expected)):.6e}" diff --git a/pykeops/pykeops/test/conv2d.py b/pykeops/pykeops/test/conv2d.py deleted file mode 100644 index 1713f225f..000000000 --- a/pykeops/pykeops/test/conv2d.py +++ /dev/null @@ -1,56 +0,0 @@ -import math -import torch -from pykeops.torch import LazyTensor - -M, N, D, DV = 20000, 30000, 3, 1 - -dtype = torch.float32 - -torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" - -x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) -y = torch.rand(1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) -b = torch.randn(N, DV, requires_grad=True, device=device_id, dtype=dtype) - - -def fun(x, y, b, backend): - if "keops" in backend: - x = LazyTensor(x) - y = LazyTensor(y) - Dxy = ((x - y) ** 2).sum(dim=2) - Kxy = (-Dxy).exp() - if backend == "keops2D": - out = LazyTensor.__matmul__(Kxy, b, backend="GPU_2D") - else: - out = Kxy @ b - if device_id != "cpu": - torch.cuda.synchronize() - # print("out:",out) - return out - - -backends = ["keops2D", "torch"] - -out = [] -for backend in backends: - out.append(fun(x, y, b, backend).squeeze()) - -out_g = [] -for k, backend in enumerate(backends): - out_g.append(torch.autograd.grad((out[k] ** 2).sum(), [b], create_graph=True)[0]) - -out_g2 = [] -for k, backend in enumerate(backends): - out_g2.append(torch.autograd.grad((out_g[k] ** 2).sum(), [b])[0]) - - -class TestCase: - def test_conv2d_fw(self): - assert torch.allclose(out[0], out[1]) - - def test_conv2d_bw1(self): - assert torch.allclose(out_g[0], out_g[1]) - - def test_conv2d_bw2(self): - assert torch.allclose(out_g2[0], out_g2[1]) diff --git a/pykeops/pykeops/test/test_block_sparse_reduction.py b/pykeops/pykeops/test/test_block_sparse_reduction.py index 3ca6fe1e9..defa32824 100644 --- a/pykeops/pykeops/test/test_block_sparse_reduction.py +++ b/pykeops/pykeops/test/test_block_sparse_reduction.py @@ -1,9 +1,9 @@ import time import numpy as np import torch +import pykeops.config from pykeops.torch import LazyTensor - # Import clustering functions from KeOps from pykeops.torch.cluster import ( grid_cluster, @@ -14,7 +14,7 @@ def test_block_sparse_reduction(): - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") dtype = torch.float32 diff --git a/pykeops/pykeops/test/test_chunks.py b/pykeops/pykeops/test/test_chunks.py index b5b03d8d0..744f58f02 100644 --- a/pykeops/pykeops/test/test_chunks.py +++ b/pykeops/pykeops/test/test_chunks.py @@ -1,6 +1,8 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 200, 300, 300, 1 @@ -8,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -38,4 +41,4 @@ def fun(x, y, b, backend): def test_chunks(): - assert torch.allclose(out[0], out[1], atol=0.0001) + assert_torch_allclose(out[0], out[1], atol=0.0001, label="chunks") diff --git a/pykeops/pykeops/test/test_chunks_ranges.py b/pykeops/pykeops/test/test_chunks_ranges.py index 19d5c0e27..f6170381f 100644 --- a/pykeops/pykeops/test/test_chunks_ranges.py +++ b/pykeops/pykeops/test/test_chunks_ranges.py @@ -1,6 +1,8 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 2, 3, 200, 300, 300, 1 @@ -8,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -38,4 +41,6 @@ def fun(x, y, b, backend): def test_chunk_ranges(): - assert torch.allclose(out[0], out[1], rtol=0.0001, atol=0.0001) + assert_torch_allclose( + out[0], out[1], rtol=0.0001, atol=0.0001, label="chunks_ranges" + ) diff --git a/pykeops/pykeops/test/test_complex.py b/pykeops/pykeops/test/test_complex.py index 2378163f6..ae2c917f2 100644 --- a/pykeops/pykeops/test/test_complex.py +++ b/pykeops/pykeops/test/test_complex.py @@ -2,7 +2,9 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose dtype = torch.float32 dtype_c = torch.complex64 @@ -10,7 +12,8 @@ M, N, D = 1000, 1000, 1 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(1, N, D, dtype=dtype_c, requires_grad=True, device=device_id) @@ -40,7 +43,7 @@ def fun(x, p, f, backend): def test_complex_fw(): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="complex_fw") # out_g = [] diff --git a/pykeops/pykeops/test/test_complex_numpy.py b/pykeops/pykeops/test/test_complex_numpy.py index fb33a507c..fe4ce25d2 100644 --- a/pykeops/pykeops/test/test_complex_numpy.py +++ b/pykeops/pykeops/test/test_complex_numpy.py @@ -1,5 +1,6 @@ import numpy as np from pykeops.numpy import LazyTensor, ComplexLazyTensor +from pykeops.test import assert_np_allclose M, N, D = 1000, 1000, 3 @@ -24,4 +25,4 @@ def fun(x, y, backend): def test_complex_numpy(): - assert np.allclose(out[0], out[1]) + assert_np_allclose(out[0], out[1], label="complex_numpy") diff --git a/pykeops/pykeops/test/test_contiguous_numpy.py b/pykeops/pykeops/test/test_contiguous_numpy.py index aa96a5125..5f75a8a38 100644 --- a/pykeops/pykeops/test/test_contiguous_numpy.py +++ b/pykeops/pykeops/test/test_contiguous_numpy.py @@ -1,5 +1,6 @@ from pykeops.numpy import LazyTensor import numpy as np +from pykeops.test import assert_np_allclose np.random.seed(0) a1 = np.random.rand(2, 1000, 5) @@ -22,4 +23,4 @@ def test_contiguous_numpy(): - assert np.allclose(d2, d1) + assert_np_allclose(d2, d1, label="contiguous_numpy") diff --git a/pykeops/pykeops/test/test_contiguous_torch.py b/pykeops/pykeops/test/test_contiguous_torch.py index c8af66361..bf789dd83 100644 --- a/pykeops/pykeops/test/test_contiguous_torch.py +++ b/pykeops/pykeops/test/test_contiguous_torch.py @@ -1,5 +1,6 @@ from pykeops.torch import LazyTensor import torch +from pykeops.test import assert_torch_allclose torch.backends.cuda.matmul.allow_tf32 = False torch.manual_seed(0) @@ -24,4 +25,4 @@ def test_contiguous_torch(): - assert torch.allclose(d2, d1) + assert_torch_allclose(d2, d1, label="contiguous_torch") diff --git a/pykeops/pykeops/test/test_conv2d.py b/pykeops/pykeops/test/test_conv2d.py new file mode 100644 index 000000000..74afd0d3d --- /dev/null +++ b/pykeops/pykeops/test/test_conv2d.py @@ -0,0 +1,69 @@ +import math +import unittest +import torch +import pykeops.config +from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose + +M, N, D, DV = 2000, 3000, 3, 1 + +dtype = torch.float32 + +torch.manual_seed(42) + +torch.backends.cuda.matmul.allow_tf32 = False +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" + + +def fun(x, y, b, backend): + if "keops" in backend: + x = LazyTensor(x) + y = LazyTensor(y) + Dxy = ((x - y) ** 2).sum(dim=2) + Kxy = (-Dxy).exp() + if backend == "keops2D": + out = LazyTensor.__matmul__(Kxy, b, backend="GPU_2D") + else: + out = Kxy @ b + if device_id != "cpu": + torch.cuda.synchronize() + # print("out:",out) + return out + + +backends = ["keops2D", "torch"] + + +@unittest.skipUnless( + use_cuda, + "CUDA not available in torch or KeOps", +) +class TestCase(unittest.TestCase): + @classmethod + def setUpClass(cls): + x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) + y = torch.rand(1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) + b = torch.randn(N, DV, requires_grad=True, device=device_id, dtype=dtype) + cls.out = [] + for backend in backends: + cls.out.append(fun(x, y, b, backend).squeeze()) + + cls.out_g = [] + for k in range(len(backends)): + cls.out_g.append( + torch.autograd.grad((cls.out[k] ** 2).sum(), [b], create_graph=True)[0] + ) + + cls.out_g2 = [] + for k in range(len(backends)): + cls.out_g2.append(torch.autograd.grad((cls.out_g[k] ** 2).sum(), [b])[0]) + + def test_conv2d_fw(self): + assert_torch_allclose(self.out[0], self.out[1], label="conv2d_fw") + + def test_conv2d_bw1(self): + assert_torch_allclose(self.out_g[0], self.out_g[1], label="conv2d_bw1") + + def test_conv2d_bw2(self): + assert_torch_allclose(self.out_g2[0], self.out_g2[1], label="conv2d_bw2") diff --git a/pykeops/pykeops/test/test_finalchunks.py b/pykeops/pykeops/test/test_finalchunks.py index 5a7f98000..0bf164370 100644 --- a/pykeops/pykeops/test/test_finalchunks.py +++ b/pykeops/pykeops/test/test_finalchunks.py @@ -1,6 +1,8 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 200, 300, 3, 300 @@ -8,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -40,4 +43,4 @@ def fun(x, y, b, backend): def test_finalchunk(): # print(out[0] - out[1]) - assert torch.allclose(out[0], out[1], atol=0.0001) + assert_torch_allclose(out[0], out[1], atol=0.0001, label="finalchunks") diff --git a/pykeops/pykeops/test/test_finalchunks_ranges.py b/pykeops/pykeops/test/test_finalchunks_ranges.py index aac85b2ac..d24939629 100644 --- a/pykeops/pykeops/test/test_finalchunks_ranges.py +++ b/pykeops/pykeops/test/test_finalchunks_ranges.py @@ -1,6 +1,8 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 2, 3, 200, 300, 3, 300 @@ -8,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -38,4 +41,4 @@ def fun(x, y, b, backend): def test_finalchunks_ranges(): - assert torch.allclose(out[0], out[1], atol=0.0001) + assert_torch_allclose(out[0], out[1], atol=0.0001, label="finalchunks_ranges") diff --git a/pykeops/pykeops/test/test_float16.py b/pykeops/pykeops/test/test_float16.py index 161f1f9b7..4324d9424 100644 --- a/pykeops/pykeops/test/test_float16.py +++ b/pykeops/pykeops/test/test_float16.py @@ -1,14 +1,17 @@ # Test for Clamp operation using LazyTensors import pytest import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose dtype = torch.float16 M, N, D = 5, 5, 1 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.randn(M, 1, D, dtype=dtype, requires_grad=True, device=device_id) @@ -29,17 +32,25 @@ def fun(x, y, backend): class TestCase: out = [] - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") + @pytest.mark.skipif( + not use_cuda, + reason="Requires torch and KeOps CUDA support", + ) def test_float16_fw(self): for backend in ["torch", "keops"]: self.out.append(fun(x, y, backend).squeeze()) - assert torch.allclose(self.out[0], self.out[1], atol=0.001, rtol=0.001) + assert_torch_allclose( + self.out[0], self.out[1], atol=0.001, rtol=0.001, label="float16_fw" + ) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") + @pytest.mark.skipif( + not use_cuda, + reason="Requires torch and KeOps CUDA support", + ) def test_float16_bw(self): out_g = [] for k, backend in enumerate(["torch", "keops"]): out_g.append(torch.autograd.grad(self.out[k][0], [x])[0]) - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="float16_bw") diff --git a/pykeops/pykeops/test/test_gpu_cpu.py b/pykeops/pykeops/test/test_gpu_cpu.py index 505437f1a..ada773e83 100644 --- a/pykeops/pykeops/test/test_gpu_cpu.py +++ b/pykeops/pykeops/test/test_gpu_cpu.py @@ -1,8 +1,11 @@ import math import pytest import torch +import pykeops.config from pykeops.torch import LazyTensor +use_cuda = pykeops.config.cuda.is_available() + M, N, D, DV = 1000, 1000, 3, 1 dtype = torch.float32 @@ -39,8 +42,16 @@ def fun(x, y, b, backend): class TestCase: def test_torch_keops_cpu(self): - assert torch.allclose(out[0], out[1]) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") + assert torch.allclose( + out[0], out[1] + ), f"torch vs keops_cpu mismatch: ||ref-test||_2={torch.norm(out[0] - out[1]).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" + + @pytest.mark.skipif( + not use_cuda, + reason="Requires KeOps CUDA support", + ) def test_torch_keops_gpu(self): - assert torch.allclose(out[0], fun(x, y, b, ["keops_gpu"]).squeeze()) + out_gpu = fun(x, y, b, ["keops_gpu"]).squeeze() + assert torch.allclose( + out[0], out_gpu + ), f"torch vs keops_gpu mismatch: ||ref-test||_2={torch.norm(out[0] - out_gpu).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" diff --git a/pykeops/pykeops/test/test_kron.py b/pykeops/pykeops/test/test_kron.py index 29fcadb17..22104db09 100644 --- a/pykeops/pykeops/test/test_kron.py +++ b/pykeops/pykeops/test/test_kron.py @@ -2,6 +2,7 @@ import torch from pykeops.numpy import Genred +from pykeops.test import assert_np_allclose, assert_torch_allclose M = 11 N = 150 @@ -32,7 +33,9 @@ def test_kron_lazytensor_np(): Y = LazyTensor(y[None, :, :]) gamma_keops_lazytensor_np = (X.keops_kron(Y, dimfa, dimfb)).sum(axis=axis) - assert np.allclose(gamma_keops_lazytensor_np, gamma_py, atol=1e-6) + assert_np_allclose( + gamma_keops_lazytensor_np, gamma_py, atol=1e-6, label="kron_lazytensor_np" + ) ############################################################################ @@ -43,8 +46,11 @@ def test_kron_lazytensor_torch(): Y_t = LazyTensor_torch(torch.from_numpy(y)[None, :, :]) gamma_keops_lazytensor_torch = (X_t.keops_kron(Y_t, dimfa, dimfb)).sum(axis=axis) - assert torch.allclose( - gamma_keops_lazytensor_torch, torch.from_numpy(gamma_py), atol=1e-6 + assert_torch_allclose( + gamma_keops_lazytensor_torch, + torch.from_numpy(gamma_py), + atol=1e-6, + label="kron_lazytensor_torch", ) @@ -79,8 +85,12 @@ def test_kron_genred(): myconv2 = Genred(formula2, aliases, reduction_op="Sum", axis=axis) gamma_keops_genred_TensorDot = myconv2(x0, y0) - assert np.allclose(gamma_keops_genred_Kron, gamma_py2, atol=1e-6) - assert np.allclose(gamma_keops_genred_TensorDot, gamma_py2, atol=1e-6) + assert_np_allclose( + gamma_keops_genred_Kron, gamma_py2, atol=1e-6, label="kron_genred" + ) + assert_np_allclose( + gamma_keops_genred_TensorDot, gamma_py2, atol=1e-6, label="tensordot_genred" + ) ############################################################################ @@ -123,6 +133,10 @@ def test_kron_genred2(): myconv2 = Genred(formula2, aliases, reduction_op="Sum", axis=axis) gamma_keops_genred_TensorDot2 = myconv2(x2, y2) - assert np.allclose(gamma_keops_genred_Kron2, gamma_py3, atol=1e-6) - assert np.allclose(gamma_keops_genred_TensorDot2, gamma_py3, atol=1e-6) - assert np.allclose(gamma_py4, gamma_py3, atol=1e-6) + assert_np_allclose( + gamma_keops_genred_Kron2, gamma_py3, atol=1e-6, label="kron_genred2" + ) + assert_np_allclose( + gamma_keops_genred_TensorDot2, gamma_py3, atol=1e-6, label="tensordot_genred2" + ) + assert_np_allclose(gamma_py4, gamma_py3, atol=1e-6, label="numpy_kron_einsum") diff --git a/pykeops/pykeops/test/test_lazytensor_clamp.py b/pykeops/pykeops/test/test_lazytensor_clamp.py index fa3dbf456..7aae6c2db 100644 --- a/pykeops/pykeops/test/test_lazytensor_clamp.py +++ b/pykeops/pykeops/test/test_lazytensor_clamp.py @@ -1,10 +1,13 @@ import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D = 1000, 1000, 3 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.randn(M, 1, D, requires_grad=True, device=device_id) @@ -35,7 +38,9 @@ def fun(x, y, a, b, backend): class TestCase: def test_lazytensor_clamp_fw(self): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="clamp_fw") def test_lazytensor_clamp_bw(self): - assert torch.allclose(out_g[0], out_g[1], atol=0.01, rtol=0.001) + assert_torch_allclose( + out_g[0], out_g[1], atol=0.01, rtol=0.001, label="clamp_bw" + ) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py index b2c5c84cc..c6af5a374 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py @@ -1,11 +1,14 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 3, 4, 20, 25, 3, 2 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(1) x = torch.rand(1, B2, M, 1, D, device=device_id) / math.sqrt(D) @@ -41,7 +44,7 @@ def fun(x, y, b, backend): class TestCase: def test_lazytensor_gaussian_batch_fw(self): # print(out[0]- out[1]) - assert torch.allclose(out[0], out[1], atol=1e-6) + assert_torch_allclose(out[0], out[1], atol=1e-6, label="gaussian_batch_fw") def test_lazytensor_gaussian_batch_bw(self): - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="gaussian_batch_bw") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py b/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py index efa7712b1..f07f7be98 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 1000, 1000, 3, 1 @@ -48,10 +49,10 @@ def fun(x, y, b, backend): class TestClass: def test_lazytensor_gaussian_cpu(self): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="gaussian_cpu_fw") def test_lazytensor_gaussian_cpu_bw1(self): - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="gaussian_cpu_bw1") def test_lazytensor_gaussian_cpu_bw2(self): - assert torch.allclose(out_g2[0], out_g2[1]) + assert_torch_allclose(out_g2[0], out_g2[1], label="gaussian_cpu_bw2") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py index d78eca154..0c0c078d7 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py @@ -1,6 +1,8 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 2500, 2000, 3, 1 @@ -8,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = 0 if torch.cuda.is_available() else -1 +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = 0 if use_cuda else -1 torch.manual_seed(0) x = torch.rand(M, 1, D, dtype=dtype) / math.sqrt(D) @@ -38,4 +41,4 @@ def fun(x, y, b, backend): def test_lazytensor_gaussian_fromhost(): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="gaussian_fromhost") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py index 1a14d951b..eb75bc2b3 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py @@ -1,6 +1,8 @@ import math import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 20, 30, 3, 1 @@ -8,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -39,4 +42,4 @@ def fun(x, y, b, backend, out=None): def test_lazytensor_gaussian_inplace(): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="gaussian_inplace") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py b/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py index 8ea7bcbb2..180ddc751 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py @@ -1,6 +1,7 @@ import math import numpy as np from pykeops.numpy import LazyTensor +from pykeops.test import assert_np_allclose M, N, D, DV = 3000, 2000, 3, 1 @@ -35,4 +36,4 @@ def fun(x, y, b, backend, out=None): def test_lazytensor_gaussian_numpy_inplace(): - assert np.allclose(out[0], out[1]) + assert_np_allclose(out[0], out[1], label="gaussian_numpy_inplace") diff --git a/pykeops/pykeops/test/test_lazytensor_grad.py b/pykeops/pykeops/test/test_lazytensor_grad.py index 7010bfd64..2cf81bf31 100644 --- a/pykeops/pykeops/test/test_lazytensor_grad.py +++ b/pykeops/pykeops/test/test_lazytensor_grad.py @@ -1,13 +1,16 @@ import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 100, 100, 3, 1 dtype = torch.float32 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, requires_grad=True, device=device_id, dtype=dtype) @@ -45,10 +48,14 @@ def fun(x, y, b, backend): class TestClass: def test_lazytensor_grad(self): - assert torch.allclose(out[0], out[1], rtol=0.0001) + assert_torch_allclose(out[0], out[1], rtol=0.0001, label="lazytensor_grad_fw") def test_lazytensor_grad_bw1(self): - assert torch.allclose(out_g[0], out_g[1], rtol=0.0001) + assert_torch_allclose( + out_g[0], out_g[1], rtol=0.0001, label="lazytensor_grad_bw1" + ) def test_lazytensor_grad_bw2(self): - assert torch.allclose(out_g2[0], out_g2[1], rtol=0.001) + assert_torch_allclose( + out_g2[0], out_g2[1], rtol=0.001, label="lazytensor_grad_bw2" + ) diff --git a/pykeops/pykeops/test/test_lazytensor_tensordot.py b/pykeops/pykeops/test/test_lazytensor_tensordot.py index d1cf43348..c428962f8 100644 --- a/pykeops/pykeops/test/test_lazytensor_tensordot.py +++ b/pykeops/pykeops/test/test_lazytensor_tensordot.py @@ -1,11 +1,14 @@ import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N = 2, 10 # Matrix multiplication as a special case of Tensordot torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) a = torch.randn(4 * 7, requires_grad=True, device=device_id, dtype=torch.float64) @@ -18,4 +21,4 @@ def test_tensordot(): - assert torch.allclose(c.flatten(), C.flatten()) + assert_torch_allclose(c.flatten(), C.flatten(), label="lazytensor_tensordot") diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index 3b66ed731..28298c996 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -1,5 +1,7 @@ import os.path import sys +from contextlib import redirect_stdout +import io sys.path.append( os.path.join( @@ -20,39 +22,104 @@ import pykeops import pykeops.config +from pykeops.test import assert_np_allclose from pykeops.numpy.utils import ( - np_kernel, - grad_np_kernel, - differences, squared_distances, log_sum_exp, - np_kernel_sphere, ) +np.random.seed(42) + class NumpyUnitTestCase(unittest.TestCase): - A = int(4) # Batchdim 1 - B = int(6) # Batchdim 2 - M = int(10) - N = int(6) - D = int(3) - E = int(3) - nbatchdims = int(2) - - x = np.random.rand(M, D) - a = np.random.rand(M, E) - f = np.random.rand(M, 1) - y = np.random.rand(N, D) - b = np.random.rand(N, E) - g = np.random.rand(N, 1) - sigma = np.array([0.4]) - - X = np.random.rand(A, 1, M, D) - L = np.random.rand(1, B, M, 1) - Y = np.random.rand(1, B, N, D) - S = np.random.rand(A, B, 1) + 1 - - type_to_test = ["float32", "float64"] + + def setUp(self): + self.A = int(4) # Batchdim 1 + self.B = int(6) # Batchdim 2 + self.M = int(10) + self.N = int(6) + self.D = int(3) + self.E = int(3) + self.nbatchdims = int(2) + + self.x = np.random.rand(self.M, self.D) + self.a = np.random.rand(self.M, self.E) + self.f = np.random.rand(self.M, 1) + self.y = np.random.rand(self.N, self.D) + self.b = np.random.rand(self.N, self.E) + self.g = np.random.rand(self.N, 1) + self.sigma = np.array([0.4]) + + self.X = np.random.rand(self.A, 1, self.M, self.D) + self.L = np.random.rand(1, self.B, self.M, 1) + self.Y = np.random.rand(1, self.B, self.N, self.D) + self.S = np.random.rand(self.A, self.B, 1) + 1 + + self.type_to_test = ["float32", "float64"] + + ############################################################ + def test_numpytools_function_binding(self): + ############################################################ + from pykeops.numpy.utils import numpytools + + tools = numpytools() + x = self.x.astype(self.type_to_test[0]) + + self.assertTrue(np.array_equal(tools.copy(x), x)) + assert_np_allclose(tools.exp(x), np.exp(x)) + assert_np_allclose(tools.log(x + 1), np.log(x + 1)) + assert_np_allclose(tools.norm(x), np.linalg.norm(x)) + assert_np_allclose(tools.arraysum(x, axis=0), np.sum(x, axis=0)) + + ############################################################ + def test_cg_solver_stops_immediately_when_x0_is_good(self): + ############################################################ + + from pykeops.numpy import LazyTensor + + alpha = 2.0 + + x_i = LazyTensor(self.x[:, None, :]) + x_j = LazyTensor(self.x[None, :, :]) + K_xx = (((x_i - x_j).abs()).sum(-1)).exp() + + b = K_xx @ self.f + alpha * self.f + + x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) + assert_np_allclose(self.f, x) + + stream = io.StringIO() + with redirect_stdout(stream): + x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) + + assert_np_allclose(self.f, x) + output = stream.getvalue() + self.assertIn("'status': 'Converged'", output) + self.assertIn("'niter': 0", output) + self.assertIn("'x0_provided': True", output) + + ############################################################ + def test_cg_solver_verbose_prints_info(self): + ############################################################ + + from pykeops.numpy import LazyTensor + + alpha = 2.0 + + x_i = LazyTensor(self.x[:, None, :]) + x_j = LazyTensor(self.x[None, :, :]) + K_xx = (((x_i - x_j).abs()).sum(-1)).exp() + + b = K_xx @ self.f + alpha * self.f + + stream = io.StringIO() + with redirect_stdout(stream): + x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) + + assert_np_allclose(self.f, x) + output = stream.getvalue() + self.assertIn("'status': 'Converged'", output) + self.assertIn("'x0_provided': True", output) ############################################################ def test_generic_syntax_sum(self): @@ -63,7 +130,7 @@ def test_generic_syntax_sum(self): formula = "Square(p-a)*Exp(x+y)" axis = 1 # 0 means summation over i, 1 means over j - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -88,7 +155,7 @@ def test_generic_syntax_sum(self): ).T # compare output - self.assertTrue(np.allclose(gamma_keops, gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops, gamma_py, atol=1e-6) ############################################################ def test_generic_syntax_lse(self): @@ -98,7 +165,7 @@ def test_generic_syntax_lse(self): aliases = ["p=Pm(0,1)", "a=Vj(1,1)", "x=Vi(2,3)", "y=Vj(3,3)"] formula = "Square(p-a)*Exp(-SqNorm2(x-y))" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -123,7 +190,7 @@ def test_generic_syntax_lse(self): ) # compare output - self.assertTrue(np.allclose(gamma_keops.ravel(), gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops.ravel(), gamma_py, atol=1e-6) ############################################################ def test_generic_syntax_softmax(self): @@ -134,7 +201,7 @@ def test_generic_syntax_softmax(self): formula = "Square(p-a)*Exp(-SqNorm2(x-y))" formula_weights = "y" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -169,9 +236,7 @@ def np_softmax(x, w): ) # compare output - self.assertTrue( - np.allclose(gamma_keops.ravel(), gamma_py.ravel(), atol=1e-6) - ) + assert_np_allclose(gamma_keops.ravel(), gamma_py.ravel(), atol=1e-6) ############################################################ def test_non_contiguity(self): @@ -199,7 +264,7 @@ def test_non_contiguity(self): # check output self.assertFalse(yc_tmp.flags.c_contiguous) - self.assertTrue(np.allclose(gamma_keops1, gamma_keops2)) + assert_np_allclose(gamma_keops1, gamma_keops2) ############################################################ def test_heterogeneous_var_aliases(self): @@ -228,7 +293,7 @@ def test_heterogeneous_var_aliases(self): ) # compare output - self.assertTrue(np.allclose(gamma_keops.ravel(), gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops.ravel(), gamma_py, atol=1e-6) ############################################################ def test_formula_simplification(self): @@ -254,7 +319,7 @@ def test_formula_simplification(self): gamma_py = np.zeros_like(self.x) # compare output - self.assertTrue(np.allclose(gamma_keops, gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops, gamma_py, atol=1e-6) ############################################################ def test_argkmin(self): @@ -281,7 +346,7 @@ def test_argkmin(self): np.sum((self.x[:, np.newaxis, :] - self.y[np.newaxis, :, :]) ** 2, axis=2), axis=1, )[:, :3] - self.assertTrue(np.allclose(c.ravel(), cnp.ravel())) + assert_np_allclose(c.ravel(), cnp.ravel()) ############################################################ def test_LazyTensor_sum(self): @@ -326,7 +391,7 @@ def test_LazyTensor_sum(self): for res_keops, res_numpy in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_numpy.shape) - self.assertTrue(np.allclose(res_keops, res_numpy, atol=1e-3)) + assert_np_allclose(res_keops, res_numpy, atol=1e-3) if __name__ == "__main__": diff --git a/pykeops/pykeops/test/test_shape_regression.py b/pykeops/pykeops/test/test_shape_regression.py new file mode 100644 index 000000000..87f46d5f5 --- /dev/null +++ b/pykeops/pykeops/test/test_shape_regression.py @@ -0,0 +1,131 @@ +import re +import unittest +import numpy as np +import pykeops.config + +from pykeops.test import assert_np_allclose, assert_torch_allclose + + +class ShapeNumpyTestCase(unittest.TestCase): + + def setUp(self): + self.param = np.array([0.4], dtype=np.float32) + self.param_scalar = np.array(0.4, dtype=np.float32) + + self.x = np.array([[0.1], [0.2], [0.3]], dtype=np.float32) # shape (3,1) + self.x_vect = np.array([0.1, 0.2, 0.3], dtype=np.float32) # shape (3,) + + def test_numpy_pm1_shape_regression(self): + from pykeops.numpy import Genred + + out = Genred("param", ["param=Pm(1)"], axis=1)(self.param, backend="auto") + ref = np.full_like(out, self.param[0]) + assert_np_allclose(out, ref, atol=1e-7) + + def test_numpy_pm1_scalar_rejected(self): + from pykeops.numpy import Genred + + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Pm argument #0 requires at least 1 dimension (got a scalar)." + ), + ): + Genred("param", ["param=Pm(1)"], axis=1)(self.param_scalar, backend="auto") + + def test_numpy_vj_single_dim_formula_x(self): + from pykeops.numpy import Genred + + out = Genred("x", ["x=Vj(1)"], axis=1)(self.x, backend="auto") + ref = np.sum(self.x, axis=0, keepdims=True) + assert_np_allclose(out, ref, atol=1e-7) + + def test_numpy_vj_single_dim_formula_x_1d_rejected_python(self): + from pykeops.numpy import Genred + + # The error is raised in the Python code (parse_type.py), before even calling the C++ code + with self.assertRaisesRegex(IndexError, "tuple index out of range"): + Genred("y", ["y=Vj(1)"], axis=1)(self.x_vect, backend="auto") + + def test_numpy_vj_single_dim_formula_x_1d_rejected_cpp(self): + from pykeops.numpy import Genred + + # The error is raised in the C++ code (LoadKeOps_cpp.py or pykeops_nvrtc.py), after the Python code has accepted the input shape but before launching the kernel + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Vj argument #2 requires at least 2 dimensions, got 1." + ), + ): + Genred("y + z + l", ["y=Vj(1)", "z=Vi(1)", "l=Vj(1)"], axis=1)( + self.x, self.x, self.x_vect, backend="auto" + ) + + +try: + import torch + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + + +@unittest.skipUnless(HAS_TORCH, "torch not available") +class ShapeTorchTestCase(unittest.TestCase): + + def setUp(self): + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + device = "cuda" if use_cuda else "cpu" + self.param = torch.tensor([0.4], dtype=torch.float32, device=device) + self.param_scalar = torch.tensor(0.4, dtype=torch.float32, device=device) + self.x = torch.tensor([[0.1], [0.2], [0.3]], dtype=torch.float32, device=device) + self.x_vect = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, device=device) + + def test_torch_pm1_shape_regression(self): + from pykeops.torch import Genred + + out = Genred("param", ["param=Pm(1)"], axis=1)(self.param, backend="auto") + ref = torch.full_like(out, self.param[0]) + assert_torch_allclose(out, ref, atol=1e-7) + + def test_torch_pm1_scalar_rejected(self): + from pykeops.torch import Genred + + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Pm argument #0 requires at least 1 dimension (got a scalar)." + ), + ): + Genred("param", ["param=Pm(1)"], axis=1)(self.param_scalar, backend="auto") + + def test_torch_vj_single_dim_formula_x(self): + from pykeops.torch import Genred + + out = Genred("x", ["x=Vj(1)"], axis=1)(self.x, backend="auto") + ref = torch.sum(self.x, dim=0, keepdim=True) + assert_torch_allclose(out, ref, atol=1e-7) + + def test_torch_vj_single_dim_formula_x_1d_rejected_python(self): + from pykeops.torch import Genred + + with self.assertRaisesRegex(IndexError, "tuple index out of range"): + Genred("y", ["y=Vj(1)"], axis=1)(self.x_vect, backend="auto") + + def test_torch_vj_single_dim_formula_x_1d_rejected_cpp(self): + from pykeops.torch import Genred + + # The error is raised in the C++ code (LoadKeOps_cpp.py or pykeops_nvrtc.py), after the Python code has accepted the input shape but before launching the kernel + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Vj argument #2 requires at least 2 dimensions, got 1." + ), + ): + Genred("y + z + l", ["y=Vj(1)", "z=Vi(1)", "l=Vj(1)"], axis=1)( + self.x, self.x, self.x_vect, backend="auto" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/pykeops/pykeops/test/test_sinc.py b/pykeops/pykeops/test/test_sinc.py index 15589c265..9837504b4 100644 --- a/pykeops/pykeops/test/test_sinc.py +++ b/pykeops/pykeops/test/test_sinc.py @@ -2,10 +2,13 @@ import pytest import torch +import pykeops.config from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose -device = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = "cuda" if use_cuda else "cpu" x = torch.rand(5, 1, dtype=torch.float64) * 2 * math.pi y = x.data.clone() @@ -25,9 +28,9 @@ @pytest.mark.skipif(torch.__version__ < "1.8", reason="Requires torch>=1.8") def test_sinc(): - assert torch.allclose(s1, s2), torch.abs(s1 - s2) + assert_torch_allclose(s1, s2, label="sinc") @pytest.mark.skipif(torch.__version__ < "1.8", reason="Requires torch>=1.8") def test_sinc_grad(): - assert torch.allclose(x.grad, y.grad) + assert_torch_allclose(x.grad, y.grad, label="sinc_grad") diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index d6402cb00..7283d1ff5 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -1,5 +1,8 @@ import os.path import sys +from contextlib import redirect_stdout +import io +from math import prod sys.path.append( os.path.join( @@ -15,90 +18,123 @@ ) import unittest -import itertools -import numpy as np import pykeops import pykeops.config -from pykeops.numpy.utils import ( - squared_distances, - np_kernel, - log_np_kernel, - grad_np_kernel, - differences, - log_sum_exp, -) +from pykeops.test import assert_torch_allclose + +HAS_TORCH = False +use_cuda = False + +try: + import torch + + HAS_TORCH = True + + torch.manual_seed(42) + + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + device = "cuda" if use_cuda else "cpu" + + if use_cuda: + torch.backends.cuda.matmul.allow_tf32 = False + +except ImportError: + pass +@unittest.skipUnless(HAS_TORCH, "torch not available") class PytorchUnitTestCase(unittest.TestCase): - A = int(5) # Batchdim 1 - B = int(3) # Batchdim 2 - M = int(10) - N = int(6) - D = int(3) - E = int(3) - nbatchdims = int(2) - - x = np.random.rand(M, D) - a = np.random.rand(M, E) - e = np.random.rand(M, E) - f = np.random.rand(M, 1) - y = np.random.rand(N, D) - b = np.random.rand(N, E) - g = np.random.rand(N, 1) - p = np.random.rand(2) - sigma = np.array([0.4]) - alpha = np.array([0.1]) - - X = np.random.rand(A, B, M, D) - L = np.random.rand(A, 1, M, 1) - Y = np.random.rand(1, B, N, D) - S = np.random.rand(A, B, 1) + 1 - - try: + + def setUp(self): + self.A = int(5) # Batchdim 1 + self.B = int(3) # Batchdim 2 + self.M = int(10) + self.N = int(6) + self.D = int(3) + self.E = int(3) + self.nbatchdims = int(2) + + self.x64 = torch.rand( + (self.M, self.D), dtype=torch.float64, device=device, requires_grad=True + ) + self.a64 = torch.rand( + (self.M, self.E), dtype=torch.float64, device=device, requires_grad=False + ) + self.e64 = torch.rand( + (self.M, self.E), dtype=torch.float64, device=device, requires_grad=False + ) + self.f64 = torch.rand( + (self.M, 1), dtype=torch.float64, device=device, requires_grad=True + ) + self.y64 = torch.rand( + (self.N, self.D), dtype=torch.float64, device=device, requires_grad=False + ) + self.b64 = torch.rand( + (self.N, self.E), dtype=torch.float64, device=device, requires_grad=False + ) + self.g64 = torch.rand( + (self.N, 1), dtype=torch.float64, device=device, requires_grad=True + ) + self.p64 = torch.rand( + 2, dtype=torch.float64, device=device, requires_grad=False + ) + + self.sigma64 = torch.tensor([0.4], dtype=torch.float64, device=device) + self.alpha64 = torch.tensor([0.1], dtype=torch.float64, device=device) + + self.X64 = torch.rand( + (self.A, self.B, self.M, self.D), + dtype=torch.float64, + device=device, + requires_grad=True, + ) + self.L64 = torch.rand( + (self.A, 1, self.M, 1), + dtype=torch.float64, + device=device, + requires_grad=False, + ) + self.Y64 = torch.rand( + (1, self.B, self.N, self.D), + dtype=torch.float64, + device=device, + requires_grad=True, + ) + self.S64 = 1 + torch.rand( + (self.A, self.B, 1), dtype=torch.float64, device=device, requires_grad=True + ) + + self.x32 = self.x64.to(torch.float32).clone().requires_grad_(True) + self.a32 = self.a64.to(torch.float32).clone().requires_grad_(False) + self.e32 = self.e64.to(torch.float32).clone().requires_grad_(False) + self.f32 = self.f64.to(torch.float32).clone().requires_grad_(True) + self.y32 = self.y64.to(torch.float32).clone().requires_grad_(False) + self.b32 = self.b64.to(torch.float32).clone().requires_grad_(False) + self.g32 = self.g64.to(torch.float32).clone().requires_grad_(True) + self.p32 = self.p64.to(torch.float32).clone().requires_grad_(False) + + self.sigma32 = self.sigma64.to(torch.float32).clone().requires_grad_(False) + self.alpha32 = self.alpha64.to(torch.float32).clone().requires_grad_(False) + + self.X32 = self.X64.to(torch.float32).clone().requires_grad_(True) + self.L32 = self.L64.to(torch.float32).clone().requires_grad_(False) + self.Y32 = self.Y64.to(torch.float32).clone().requires_grad_(True) + self.S32 = self.S64.to(torch.float32).clone().requires_grad_(True) + + ############################################################ + def test_torchtools_function_binding(self): + ############################################################ + from pykeops.torch.utils import torchtools import torch - use_cuda = torch.cuda.is_available() - device = "cuda" if use_cuda else "cpu" - torch.backends.cuda.matmul.allow_tf32 = False + tools = torchtools() + x = self.x32.detach() - dtype = torch.float32 - xc = torch.tensor(x, dtype=dtype, device=device, requires_grad=True) - ac = torch.tensor(a, dtype=dtype, device=device, requires_grad=False) - ec = torch.tensor(e, dtype=dtype, device=device, requires_grad=False) - fc = torch.tensor(f, dtype=dtype, device=device, requires_grad=True) - yc = torch.tensor(y, dtype=dtype, device=device, requires_grad=False) - bc = torch.tensor(b, dtype=dtype, device=device, requires_grad=False) - gc = torch.tensor(g, dtype=dtype, device=device, requires_grad=True) - pc = torch.tensor(p, dtype=dtype, device=device, requires_grad=False) - sigmac = torch.tensor(sigma, dtype=dtype, device=device, requires_grad=False) - alphac = torch.tensor(alpha, dtype=dtype, device=device, requires_grad=False) - - Xc = torch.tensor(X, dtype=dtype, device=device, requires_grad=True) - Lc = torch.tensor(L, dtype=dtype, device=device, requires_grad=False) - Yc = torch.tensor(Y, dtype=dtype, device=device, requires_grad=True) - Sc = torch.tensor(S, dtype=dtype, device=device, requires_grad=True) - - dtype = torch.float64 - xcd = torch.tensor(x, dtype=dtype, device=device, requires_grad=False) - acd = torch.tensor(a, dtype=dtype, device=device, requires_grad=False) - ecd = torch.tensor(e, dtype=dtype, device=device, requires_grad=False) - fcd = torch.tensor(f, dtype=dtype, device=device, requires_grad=False) - ycd = torch.tensor(y, dtype=dtype, device=device, requires_grad=False) - bcd = torch.tensor(b, dtype=dtype, device=device, requires_grad=False) - gcd = torch.tensor(g, dtype=dtype, device=device, requires_grad=False) - pcd = torch.tensor(p, dtype=dtype, device=device, requires_grad=False) - sigmacd = torch.tensor(sigma, dtype=dtype, device=device, requires_grad=False) - alphacd = torch.tensor(alpha, dtype=dtype, device=device, requires_grad=False) - Xcd = torch.tensor(X, dtype=dtype, device=device, requires_grad=True) - Lcd = torch.tensor(L, dtype=dtype, device=device, requires_grad=False) - Ycd = torch.tensor(Y, dtype=dtype, device=device, requires_grad=True) - Scd = torch.tensor(S, dtype=dtype, device=device, requires_grad=True) - - print("Running Pytorch tests.") - except: - print("Pytorch could not be loaded. Skip tests.") - pass + self.assertTrue(torch.equal(tools.copy(x), x)) + assert_torch_allclose(tools.exp(x), torch.exp(x)) + assert_torch_allclose(tools.log(x + 1), torch.log(x + 1)) + assert_torch_allclose(tools.norm(x), torch.norm(x)) ############################################################ def test_generic_syntax_float(self): @@ -107,7 +143,7 @@ def test_generic_syntax_float(self): aliases = ["p=Pm(1)", "a=Vj(1)", "x=Vi(3)", "y=Vj(3)"] formula = "Square(p-a)*Exp(x+y)" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -116,17 +152,17 @@ def test_generic_syntax_float(self): with self.subTest(b=b): # Call cuda kernel gamma_keops = Genred(formula, aliases, axis=1)( - self.sigmac, self.gc, self.xc, self.yc, backend=b + self.sigma32, self.g32, self.x32, self.y32, backend=b ) - # Numpy version - gamma_py = np.sum( - (self.sigma - self.g) ** 2 - * np.exp((self.y.T[:, :, np.newaxis] + self.x.T[:, np.newaxis, :])), - axis=1, + # Torch reference in float64 for stable comparisons. + gamma_ref = torch.sum( + (self.sigma64 - self.g64) ** 2 + * torch.exp(self.y64.T[:, :, None] + self.x64.T[:, None, :]), + dim=1, ).T # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) + assert_torch_allclose( + gamma_keops.to(torch.float64), gamma_ref, atol=1e-6 ) ############################################################ @@ -136,7 +172,7 @@ def test_generic_syntax_double(self): aliases = ["p=Pm(1)", "a=Vj(1)", "x=Vi(3)", "y=Vj(3)"] formula = "Square(p-a)*Exp(x+y)" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -145,18 +181,16 @@ def test_generic_syntax_double(self): with self.subTest(b=b): # Call cuda kernel gamma_keops = Genred(formula, aliases, axis=1)( - self.sigmacd, self.gcd, self.xcd, self.ycd, backend=b + self.sigma64, self.g64, self.x64, self.y64, backend=b ) - # Numpy version - gamma_py = np.sum( - (self.sigma - self.g) ** 2 - * np.exp((self.y.T[:, :, np.newaxis] + self.x.T[:, np.newaxis, :])), - axis=1, + # Torch reference in float64. + gamma_ref = torch.sum( + (self.sigma64 - self.g64) ** 2 + * torch.exp(self.y64.T[:, :, None] + self.x64.T[:, None, :]), + dim=1, ).T # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) ############################################################ def test_generic_syntax_softmax(self): @@ -166,7 +200,7 @@ def test_generic_syntax_softmax(self): aliases = ["p=Pm(1)", "a=Vj(1)", "x=Vi(3)", "y=Vj(3)"] formula = "Square(p-a)*Exp(-SqNorm2(x-y))" formula_weights = "y" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -182,24 +216,23 @@ def test_generic_syntax_softmax(self): formula2=formula_weights, ) gamma_keops = myop( - self.sigmacd, self.gcd, self.xcd, self.ycd, backend=b + self.sigma64, self.g64, self.x64, self.y64, backend=b ) - # Numpy version - def np_softmax(x, w): - x -= np.max(x, axis=1)[:, None] # subtract the max for robustness - return np.exp(x) @ w / np.sum(np.exp(x), axis=1)[:, None] - - gamma_py = np_softmax( - (self.sigma - self.g.T) ** 2 - * np.exp(-squared_distances(self.x, self.y)), - self.y, + # Torch reference + sqdist = torch.sum( + (self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2 + ) + scores = (self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist) + scores = scores - torch.max(scores, dim=1, keepdim=True).values + gamma_ref = ( + torch.exp(scores) + @ self.y64 + / torch.sum(torch.exp(scores), dim=1, keepdim=True) ) # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) ############################################################ def test_generic_syntax_simple(self): @@ -209,14 +242,14 @@ def test_generic_syntax_simple(self): aliases = [ "P = Pm(2)", # 1st argument, a parameter, dim 2. "X = Vi(" - + str(self.xc.shape[1]) + + str(self.x64.shape[1]) + ") ", # 2nd argument, indexed by i, dim D. - "Y = Vj(" + str(self.yc.shape[1]) + ") ", + "Y = Vj(" + str(self.y64.shape[1]) + ") ", ] # 3rd argument, indexed by j, dim D. formula = "Pow((X|Y),2) * ((Elem(P,0) * X) + (Elem(P,1) * Y))" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -224,18 +257,16 @@ def test_generic_syntax_simple(self): for b in backend_to_test: with self.subTest(b=b): my_routine = Genred(formula, aliases, reduction_op="Sum", axis=1) - gamma_keops = my_routine(self.pc, self.xc, self.yc, backend=b) + gamma_keops = my_routine(self.p64, self.x64, self.y64, backend=b) - # Numpy version - scals = (self.x @ self.y.T) ** 2 # Memory-intensive computation! - gamma_py = self.p[0] * scals.sum(1).reshape(-1, 1) * self.x + self.p[ - 1 - ] * (scals @ self.y) + # Torch reference + scals = (self.x64 @ self.y64.T) ** 2 # Memory-intensive computation! + gamma_ref = self.p64[0] * scals.sum(1).reshape( + -1, 1 + ) * self.x64 + self.p64[1] * (scals @ self.y64) # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) ############################################################ def test_logSumExp_kernels_feature(self): @@ -265,18 +296,27 @@ def test_logSumExp_kernels_feature(self): for k in ["gaussian", "laplacian", "cauchy", "inverse_multiquadric"]: with self.subTest(k=k): # Call cuda kernel - gamma_lazy = kernels[k](self.xc, self.yc, self.sigmac) - gamma_lazy = gamma_lazy.logsumexp(dim=1, weight=Vj(self.gc.exp())).cpu() - # gamma = kernel_product(params, self.xc, self.yc, self.gc).cpu() + gamma_lazy = kernels[k](self.x64, self.y64, self.sigma64) + gamma_lazy = gamma_lazy.logsumexp(dim=1, weight=Vj(self.g64.exp())) - # Numpy version - log_K = log_np_kernel(self.x, self.y, self.sigma, kernel=k) - log_KP = log_K + self.g.T - gamma_py = log_sum_exp(log_KP, axis=1) + # Torch reference + sqdist = torch.sum( + (self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2 + ) + inv_s2 = 1 / (self.sigma64**2) + if k == "gaussian": + log_k = -(inv_s2 * sqdist) + elif k == "laplacian": + log_k = -torch.sqrt(inv_s2 * sqdist) + elif k == "cauchy": + log_k = -(1 + inv_s2 * sqdist).log() + else: + log_k = -0.5 * (1 + inv_s2 * sqdist).log() + gamma_ref = torch.logsumexp(log_k + self.g64.T, dim=1) # compare output - self.assertTrue( - np.allclose(gamma_lazy.data.numpy().ravel(), gamma_py, atol=1e-6) + assert_torch_allclose( + gamma_lazy.reshape(-1), gamma_ref.reshape(-1), atol=1e-6 ) ############################################################ @@ -288,37 +328,36 @@ def test_logSumExp_gradient_kernels_feature(self): aliases = [ "P = Pm(2)", # 1st argument, a parameter, dim 2. "X = Vi(" - + str(self.gc.shape[1]) + + str(self.g64.shape[1]) + ") ", # 2nd argument, indexed by i, dim D. - "Y = Vj(" + str(self.fc.shape[1]) + ") ", + "Y = Vj(" + str(self.f64.shape[1]) + ") ", ] # 3rd argument, indexed by j, dim D. formula = "(Elem(P,0) * X) + (Elem(P,1) * Y)" # Pytorch version my_routine = Genred(formula, aliases, reduction_op="LogSumExp", axis=1) - tmp = my_routine(self.pc, self.fc, self.gc, backend="auto") + tmp = my_routine(self.p64, self.f64, self.g64, backend="auto") res = torch.dot( torch.ones_like(tmp).view(-1), tmp.view(-1) ) # equivalent to tmp.sum() but avoiding contiguity pb - gamma_keops = torch.autograd.grad(res, [self.fc, self.gc], create_graph=False) - - # Numpy version - tmp = self.p[0] * self.f + self.p[1] * self.g.T - res_py = (np.exp(tmp)).sum(axis=1) - tmp2 = np.exp(tmp.T) / res_py.reshape(1, -1) - gamma_py = [np.ones(self.M) * self.p[0], self.p[1] * tmp2.T.sum(axis=0)] + gamma_keops = torch.autograd.grad(res, [self.f64, self.g64], create_graph=False) + + # Torch reference + tmp = self.p64[0] * self.f64 + self.p64[1] * self.g64.T + res_ref = torch.exp(tmp).sum(dim=1) + tmp2 = torch.exp(tmp.T) / res_ref.reshape(1, -1) + gamma_ref = [ + torch.ones(self.M, dtype=torch.float64, device=device) * self.p64[0], + self.p64[1] * tmp2.T.sum(dim=0), + ] # compare output - self.assertTrue( - np.allclose( - gamma_keops[0].cpu().data.numpy().ravel(), gamma_py[0], atol=1e-6 - ) + assert_torch_allclose( + gamma_keops[0].reshape(-1), gamma_ref[0].reshape(-1), atol=1e-6 ) - self.assertTrue( - np.allclose( - gamma_keops[1].cpu().data.numpy().ravel(), gamma_py[1], atol=1e-6 - ) + assert_torch_allclose( + gamma_keops[1].reshape(-1), gamma_ref[1].reshape(-1), atol=1e-6 ) ############################################################ @@ -329,45 +368,40 @@ def test_non_contiguity(self): aliases = [ "P = Pm(2)", # 1st argument, a parameter, dim 2. "X = Vi(" - + str(self.xc.shape[1]) + + str(self.x64.shape[1]) + ") ", # 2nd argument, indexed by i, dim D. - "Y = Vj(" + str(self.yc.shape[1]) + ") ", + "Y = Vj(" + str(self.y64.shape[1]) + ") ", ] # 3rd argument, indexed by j, dim D. formula = "Pow((X|Y),2) * ((Elem(P,0) * X) + (Elem(P,1) * Y))" my_routine = Genred(formula, aliases, reduction_op="Sum", axis=1) - yc_tmp = self.yc.t().contiguous().t() # create a non contiguous copy + yc_tmp = self.y64.t().contiguous().t() # create a non contiguous copy # check output self.assertFalse(yc_tmp.is_contiguous()) - my_routine(self.pc, self.xc, yc_tmp, backend="auto") + my_routine(self.p64, self.x64, yc_tmp, backend="auto") ############################################################ def test_heterogeneous_var_aliases(self): ############################################################ from pykeops.torch import Genred - from pykeops.numpy.utils import squared_distances aliases = ["p=Pm(0,1)", "x=Vi(1,3)", "y=Vj(2,3)"] formula = "Square(p-Var(3,1,1))*Exp(-SqNorm2(y-x))" # Call cuda kernel myconv = Genred(formula, aliases, reduction_op="Sum", axis=1) - gamma_keops = myconv(self.sigmac, self.xc, self.yc, self.gc, backend="auto") + gamma_keops = myconv(self.sigma64, self.x64, self.y64, self.g64, backend="auto") - # Numpy version - gamma_py = np.sum( - (self.sigma - self.g.T) ** 2 * np.exp(-squared_distances(self.x, self.y)), - axis=1, + # Torch reference + sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) + gamma_ref = torch.sum( + (self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist), dim=1 ) # compare output - self.assertTrue( - np.allclose( - gamma_keops.cpu().data.numpy().ravel(), gamma_py.ravel(), atol=1e-6 - ) - ) + assert_torch_allclose(gamma_keops.reshape(-1), gamma_ref.reshape(-1), atol=1e-6) ############################################################ def test_invkernel(self): @@ -385,33 +419,53 @@ def test_invkernel(self): Kinv = KernelSolve(formula, aliases, "b", axis=1) - c = Kinv(self.xc, self.xc, self.ac, self.sigmac, alpha=self.alphac) + c = Kinv(self.x64, self.x64, self.a64, self.sigma64, alpha=self.alpha64) if torch.__version__ >= "1.8": torchsolve = lambda A, B: torch.linalg.solve(A, B) else: torchsolve = lambda A, B: torch.solve(B, A)[0] c_ = torchsolve( - self.alphac * torch.eye(self.M, device=self.device) + self.alpha64 * torch.eye(self.M, device=device, dtype=torch.float64) + torch.exp( - -torch.sum((self.xc[:, None, :] - self.xc[None, :, :]) ** 2, dim=2) - * self.sigmac + -torch.sum((self.x64[:, None, :] - self.x64[None, :, :]) ** 2, dim=2) + * self.sigma64 ), - self.ac, + self.a64, ) - self.assertTrue( - np.allclose( - c.cpu().data.numpy().ravel(), c_.cpu().data.numpy().ravel(), atol=1e-4 - ) - ) + assert_torch_allclose(c.reshape(-1), c_.reshape(-1), atol=1e-4) - (u,) = torch.autograd.grad(c, self.xc, self.ec) - (u_,) = torch.autograd.grad(c_, self.xc, self.ec) - self.assertTrue( - np.allclose( - u.cpu().data.numpy().ravel(), u_.cpu().data.numpy().ravel(), atol=1e-4 - ) - ) + (u,) = torch.autograd.grad(c, self.x64, self.e64) + (u_,) = torch.autograd.grad(c_, self.x64, self.e64) + assert_torch_allclose(u.reshape(-1), u_.reshape(-1), atol=1e-4) + + ############################################################ + def test_cg_solver_stops_immediately_when_x0_is_good(self): + ############################################################ + + import torch + from pykeops.torch import LazyTensor + + alpha = 2.0 + + x_i = LazyTensor(self.x64[:, None, :]) + x_j = LazyTensor(self.x64[None, :, :]) + K_xx = (((x_i - x_j).abs()).sum(-1)).exp() + + b = K_xx @ self.f64 + alpha * self.f64 + + x = K_xx.solve(b, alpha=alpha, x0=self.f64, eps=1e-12) + assert_torch_allclose(self.f64, x) + + stream = io.StringIO() + with redirect_stdout(stream): + x = K_xx.solve(b, alpha=alpha, x0=self.f64, eps=1e-12, verbose=True) + + assert_torch_allclose(self.f64, x) + output = stream.getvalue() + self.assertIn("'status': 'Converged'", output) + self.assertIn("'niter': 0", output) + self.assertIn("'x0_provided': True", output) ############################################################ def test_softmax(self): @@ -436,22 +490,18 @@ def test_softmax(self): formula2=formula_weights, ) - c = softmax_op(self.xc, self.yc, self.bc) + c = softmax_op(self.x64, self.y64, self.b64) # compare with direct implementation cc = 0 for k in range(self.D): - xk = self.xc[:, k][:, None] - yk = self.yc[:, k][:, None] + xk = self.x64[:, k][:, None] + yk = self.y64[:, k][:, None] cc += (xk - yk.t()) ** 2 cc -= torch.max(cc, dim=1)[0][:, None] # subtract the max for robustness - cc = torch.exp(cc) @ self.bc / torch.sum(torch.exp(cc), dim=1)[:, None] + cc = torch.exp(cc) @ self.b64 / torch.sum(torch.exp(cc), dim=1)[:, None] - self.assertTrue( - np.allclose( - c.cpu().data.numpy().ravel(), cc.cpu().data.numpy().ravel(), atol=1e-6 - ) - ) + assert_torch_allclose(c.reshape(-1), cc.reshape(-1), atol=1e-6) ############################################################ def test_pickle(self): @@ -485,19 +535,17 @@ def test_LazyTensor_sum(self): results = [] # N.B.: We could loop over float32 and float64, but this would take longer... - for x, l, y, s in [(self.Xc, self.Lc, self.Yc, self.Sc)]: # Float32 + for x, l, y, s in [(self.X32, self.L32, self.Y32, self.S32)]: # Float32 x_i = x.unsqueeze(-2) l_i = l.unsqueeze(-2) y_j = y.unsqueeze(-3) s_p = s.unsqueeze(-2).unsqueeze(-2) if use_keops: - x_i, l_i, y_j, s_p = ( - LazyTensor(x_i), - LazyTensor(l_i), - LazyTensor(y_j), - LazyTensor(s_p), - ) + x_i = LazyTensor(x_i) + l_i = LazyTensor(l_i) + y_j = LazyTensor(y_j) + s_p = LazyTensor(s_p) D_ij = (0.5 * (l_i * x_i - y_j) ** 2 / s_p).sum(-1) K_ij = (-D_ij).exp() @@ -515,17 +563,8 @@ def test_LazyTensor_sum(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-3, - ), - "KeOps:\n" - + str(res_keops) - + "\nPyTorch:\n" - + str(res_torch) - + "\nMax error: {:.2e}".format((res_keops - res_torch).abs().max()), + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-3 ) ############################################################ @@ -539,19 +578,17 @@ def test_LazyTensor_logsumexp(self): results = [] # N.B.: We could loop over float32 and float64, but this would take longer... - for x, l, y, s in [(self.Xcd, self.Lcd, self.Ycd, self.Scd)]: # Float64 + for x, l, y, s in [(self.X64, self.L64, self.Y64, self.S64)]: # Float64 x_i = x.unsqueeze(-2) l_i = l.unsqueeze(-2) y_j = y.unsqueeze(-3) s_p = s.unsqueeze(-2).unsqueeze(-2) if use_keops: - x_i, l_i, y_j, s_p = ( - LazyTensor(x_i), - LazyTensor(l_i), - LazyTensor(y_j), - LazyTensor(s_p), - ) + x_i = LazyTensor(x_i) + l_i = LazyTensor(l_i) + y_j = LazyTensor(y_j) + s_p = LazyTensor(s_p) D_ij = ((l_i * x_i + y_j).relu() * s_p / 9).sum(-1) K_ij = -1 / (1 + D_ij) @@ -572,12 +609,8 @@ def test_LazyTensor_logsumexp(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-5, - ) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5 ) ############################################################ @@ -587,25 +620,23 @@ def test_LazyTensor_min_chunked(self): from pykeops.torch import LazyTensor import torch - X = np.random.rand(self.M, 990) - Xc = torch.tensor(X, dtype=self.dtype, device=self.device) - Y = np.random.rand(self.N, 990) - Yc = torch.tensor(Y, dtype=self.dtype, device=self.device) + X64 = torch.rand((self.M, 990), dtype=torch.float64, device=device) + Y64 = torch.rand((self.N, 990), dtype=torch.float64, device=device) + + X32 = X64.to(torch.float32).clone().requires_grad_(True) + Y32 = Y64.to(torch.float32).clone().requires_grad_(True) full_results = [] for use_keops in [True, False]: results = [] - # N.B.: We could loop over float32 and float64, but this would take longer... - for x, y in [(Xc, Yc)]: # Float32 + for x, y in [(X32, Y32), (X64, Y64)]: x_i = x.unsqueeze(-2) y_j = y.unsqueeze(-3) if use_keops: - x_i, y_j = ( - LazyTensor(x_i), - LazyTensor(y_j), - ) + x_i = LazyTensor(x_i) + y_j = LazyTensor(y_j) K_ij = ((-(((x_i + y_j)) ** 2)).exp()).sum(-1, keepdim=True) @@ -620,12 +651,8 @@ def test_LazyTensor_min_chunked(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-5, - ) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5 ) ############################################################ @@ -638,19 +665,17 @@ def test_LazyTensor_min(self): results = [] # N.B.: We could loop over float32 and float64, but this would take longer... - for x, l, y, s in [(self.Xc, self.Lc, self.Yc, self.Sc)]: # Float32 + for x, l, y, s in [(self.X32, self.L32, self.Y32, self.S32)]: # Float32 x_i = x.unsqueeze(-2) l_i = l.unsqueeze(-2) y_j = y.unsqueeze(-3) s_p = s.unsqueeze(-2).unsqueeze(-2) if use_keops: - x_i, l_i, y_j, s_p = ( - LazyTensor(x_i), - LazyTensor(l_i), - LazyTensor(y_j), - LazyTensor(s_p), - ) + x_i = LazyTensor(x_i) + l_i = LazyTensor(l_i) + y_j = LazyTensor(y_j) + s_p = LazyTensor(s_p) D_ij = ((1 + ((l_i * x_i + y_j).relu() * s_p) ** 2).log()).sum( -1, keepdim=True @@ -668,12 +693,8 @@ def test_LazyTensor_min(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-5, - ) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5 ) ############################################################ @@ -686,7 +707,9 @@ def my_tensordort_perm(a, b, dims=None, perm=None): return torch.tensordot(a, b, dims=dims).sum(3).permute(perm) def invert_permutation_numpy(permutation): - return np.arange(len(permutation))[np.argsort(permutation)] + return [ + idx for idx, _ in sorted(enumerate(permutation), key=lambda x: x[1]) + ] x = torch.randn(self.M, 2, 3, 2, 2, 4, requires_grad=True, dtype=torch.float64) y = torch.randn( @@ -699,18 +722,16 @@ def invert_permutation_numpy(permutation): perm_torch = (0,) + tuple([(i + 1) for i in invert_permutation_numpy(perm)]) sum_f_torch2 = my_tensordort_perm(x, y, dims=(contfa, contfb), perm=perm_torch) - f_keops = LazyTensor( - x.reshape(self.M, 1, int(np.array((dimfa)).prod())) - ).keops_tensordot( - LazyTensor(y.reshape(1, self.N, int(np.array(dimfb).prod()))), + f_keops = LazyTensor(x.reshape(self.M, 1, prod(dimfa))).keops_tensordot( + LazyTensor(y.reshape(1, self.N, prod(dimfb))), dimfa, dimfb, - tuple(np.array(contfa) - 1), - tuple(np.array(contfb) - 1), + tuple(i - 1 for i in contfa), + tuple(i - 1 for i in contfb), tuple(perm), ) sum_f_keops = f_keops.sum_reduction(dim=1) - self.assertTrue(torch.allclose(sum_f_keops.flatten(), sum_f_torch2.flatten())) + assert_torch_allclose(sum_f_keops.flatten(), sum_f_torch2.flatten()) e = torch.randn_like(sum_f_torch2) # checking gradients @@ -718,15 +739,11 @@ def invert_permutation_numpy(permutation): sum_f_keops, x, e.reshape(self.M, -1), retain_graph=True )[0] grad_torch = torch.autograd.grad(sum_f_torch2, x, e, retain_graph=True)[0] - self.assertTrue( - torch.allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) - ) + assert_torch_allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) grad_keops = torch.autograd.grad(sum_f_keops, y, e.reshape(self.M, -1))[0] grad_torch = torch.autograd.grad(sum_f_torch2, y, e)[0] - self.assertTrue( - torch.allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) - ) + assert_torch_allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) if __name__ == "__main__": diff --git a/pykeops/pykeops/test/test_torch_func.py b/pykeops/pykeops/test/test_torch_func.py index a6657c56a..476c7a4dc 100644 --- a/pykeops/pykeops/test/test_torch_func.py +++ b/pykeops/pykeops/test/test_torch_func.py @@ -1,11 +1,11 @@ +import keopscore import torch from pykeops.torch import LazyTensor - -import keopscore +from pykeops.test import assert_torch_allclose torch.manual_seed(0) -keopscore.auto_factorize = False +keopscore.config.reduction.set_auto_factorize(False) B1, B2, M, N, D = 5, 4, 10, 20, 2 @@ -38,36 +38,40 @@ class TestCase: def test_torch_func_vmap(self): res1 = torch.func.vmap(fn_torch)(x_i, y_j, b_j, p) res2 = torch.func.vmap(fn_keops)(x_i, y_j, b_j, p) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_vmap") def test_torch_func_grad(self): res1 = torch.func.grad(fn_torch, (0, 1, 2, 3))(x_i, y_j, b_j, p) res2 = torch.func.grad(fn_keops, (0, 1, 2, 3))(x_i, y_j, b_j, p) for k in range(4): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_grad[{k}]" + ) def test_torch_func_vjp(self): res1 = torch.func.vjp(fn_torch, x_i, y_j, b_j, p)[1](torch.tensor(1.0))[0] res2 = torch.func.vjp(fn_keops, x_i, y_j, b_j, p)[1](torch.tensor(1.0))[0] - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_vjp") def test_torch_func_jvp(self): _, res1 = torch.func.jvp(fn_torch, (x_i, y_j, b_j, p), (dx_i, dy_j, db_j, dp)) _, res2 = torch.func.jvp(fn_keops, (x_i, y_j, b_j, p), (dx_i, dy_j, db_j, dp)) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_jvp") def test_torch_func_jacrev(self): res1 = torch.func.jacrev(fn_torch, (0, 1, 2, 3))(x_i, y_j, b_j, p) res2 = torch.func.jacrev(fn_keops, (0, 1, 2, 3))(x_i, y_j, b_j, p) for k in range(4): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_jacrev[{k}]" + ) def test_torch_func_jacfwd(self): res1 = torch.func.jacfwd(fn_torch)(x_i, y_j, b_j, p) res2 = torch.func.jacfwd(fn_keops)(x_i, y_j, b_j, p) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_jacfwd") def test_torch_func_hessian(self): res1 = torch.func.hessian(fn_torch)(x_i, y_j, b_j, p) res2 = torch.func.hessian(fn_keops)(x_i, y_j, b_j, p) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_hessian") diff --git a/pykeops/pykeops/test/test_torch_func_logsumexp.py b/pykeops/pykeops/test/test_torch_func_logsumexp.py index 83dde275c..6795938c1 100644 --- a/pykeops/pykeops/test/test_torch_func_logsumexp.py +++ b/pykeops/pykeops/test/test_torch_func_logsumexp.py @@ -1,5 +1,6 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose torch.manual_seed(0) @@ -29,36 +30,40 @@ class TestCase: def test_torch_func_vmap(self): res1 = torch.func.vmap(fn_torch)(x_i, y_j, b_j) res2 = torch.func.vmap(fn_keops)(x_i, y_j, b_j) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_vmap") def test_torch_func_grad(self): res1 = torch.func.grad(fn_torch, (0, 1, 2))(x_i, y_j, b_j) res2 = torch.func.grad(fn_keops, (0, 1, 2))(x_i, y_j, b_j) for k in range(3): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_grad[{k}]" + ) def test_torch_func_vjp(self): res1 = torch.func.vjp(fn_torch, x_i, y_j, b_j)[1](torch.tensor(1.0))[0] res2 = torch.func.vjp(fn_keops, x_i, y_j, b_j)[1](torch.tensor(1.0))[0] - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_vjp") def test_torch_func_jvp(self): _, res1 = torch.func.jvp(fn_torch, (x_i, y_j, b_j), (dx_i, dy_j, db_j)) _, res2 = torch.func.jvp(fn_keops, (x_i, y_j, b_j), (dx_i, dy_j, db_j)) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_jvp") def test_torch_func_jacrev(self): res1 = torch.func.jacrev(fn_torch, (0, 1, 2))(x_i, y_j, b_j) res2 = torch.func.jacrev(fn_keops, (0, 1, 2))(x_i, y_j, b_j) for k in range(3): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_jacrev[{k}]" + ) def test_torch_func_jacfwd(self): res1 = torch.func.jacfwd(fn_torch)(x_i, y_j, b_j) res2 = torch.func.jacfwd(fn_keops)(x_i, y_j, b_j) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_jacfwd") def test_torch_func_hessian(self): res1 = torch.func.hessian(fn_torch)(x_i, y_j, b_j) res2 = torch.func.hessian(fn_keops)(x_i, y_j, b_j) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_hessian") diff --git a/pykeops/pykeops/test/test_transpose.py b/pykeops/pykeops/test/test_transpose.py index cd253b291..7c3f58f41 100644 --- a/pykeops/pykeops/test/test_transpose.py +++ b/pykeops/pykeops/test/test_transpose.py @@ -1,5 +1,6 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose x, y = torch.randn(1000, 3), torch.randn(2000, 3) x_i, y_j = LazyTensor(x[:, None, :]), LazyTensor(y[None, :, :]) @@ -16,4 +17,4 @@ def test_transpose(): - assert torch.allclose(K_keops.t() @ w, K_torch.t() @ w) + assert_torch_allclose(K_keops.t() @ w, K_torch.t() @ w, label="transpose") diff --git a/pykeops/pykeops/test/test_vmap.py b/pykeops/pykeops/test/test_vmap.py index 4f3f21e96..68444636f 100644 --- a/pykeops/pykeops/test/test_vmap.py +++ b/pykeops/pykeops/test/test_vmap.py @@ -1,5 +1,6 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose test_grad = True torch.manual_seed(0) @@ -31,13 +32,13 @@ def test_vmap_0(self): x_i, x_j, y_j, use_keops=False ) - assert torch.allclose(res_keops, res_torch) + assert_torch_allclose(res_keops, res_torch, label="vmap_0_fw") u = torch.rand(res_torch.shape) (res_torch_grad,) = torch.autograd.grad(res_torch, x_i, u) (res_keops_grad,) = torch.autograd.grad(res_keops, x_i, u) - assert torch.allclose(res_keops_grad, res_torch_grad) + assert_torch_allclose(res_keops_grad, res_torch_grad, label="vmap_0_bw") def test_vmap_1(self): x_i = torch.randn(10, 5, 1, 2, requires_grad=test_grad, dtype=torch.float64) @@ -52,13 +53,13 @@ def test_vmap_1(self): res_torch = torch.vmap(fn, in_dims=in_dims, out_dims=out_dims)( x_i, x_j, y_j, use_keops=False ) - assert torch.allclose(res_keops, res_torch) + assert_torch_allclose(res_keops, res_torch, label="vmap_1_fw") u = torch.rand(res_torch.shape) (res_torch_grad,) = torch.autograd.grad(res_torch, x_i, u) (res_keops_grad,) = torch.autograd.grad(res_keops, x_i, u) - assert torch.allclose(res_keops_grad, res_torch_grad) + assert_torch_allclose(res_keops_grad, res_torch_grad, label="vmap_1_bw") def test_vmap_2(self): x_i = torch.randn(10, 5, 1, 2, requires_grad=test_grad, dtype=torch.float64) @@ -73,10 +74,10 @@ def test_vmap_2(self): res_torch = torch.vmap(fn, in_dims=in_dims, out_dims=out_dims)( x_i, x_j, y_j, use_keops=False ) - assert torch.allclose(res_keops, res_torch) + assert_torch_allclose(res_keops, res_torch, label="vmap_2_fw") u = torch.rand(res_torch.shape) (res_torch_grad,) = torch.autograd.grad(res_torch, x_i, u) (res_keops_grad,) = torch.autograd.grad(res_keops, x_i, u) - assert torch.allclose(res_keops_grad, res_torch_grad) + assert_torch_allclose(res_keops_grad, res_torch_grad, label="vmap_2_bw") diff --git a/pykeops/pykeops/torch/generic/generic_red.py b/pykeops/pykeops/torch/generic/generic_red.py index 97e3eb898..edcabdfe2 100644 --- a/pykeops/pykeops/torch/generic/generic_red.py +++ b/pykeops/pykeops/torch/generic/generic_red.py @@ -45,7 +45,15 @@ def check_AD_supported(formula): def set_device(tagCPUGPU, tagHostDevice, device_id_request, *args): device_args = args[0].device - if tagCPUGPU == 1 & tagHostDevice == 1: + # Guard against unsupported configuration that can crash in native code: + # CPU backend cannot consume CUDA tensors. + if tagCPUGPU == 0 and tagHostDevice == 1: + raise ValueError( + "[KeOps] Incompatible configuration: CPU backend selected while input tensors are on CUDA. " + "Use CPU tensors or enable KeOps CUDA backend." + ) + + if tagCPUGPU == 1 and tagHostDevice == 1: for i in range(1, len(args)): if args[i].device.index != device_args.index: raise ValueError( diff --git a/pykeops/pykeops/torch/lazytensor/LazyTensor.py b/pykeops/pykeops/torch/lazytensor/LazyTensor.py index 656b7cbc2..69785283a 100644 --- a/pykeops/pykeops/torch/lazytensor/LazyTensor.py +++ b/pykeops/pykeops/torch/lazytensor/LazyTensor.py @@ -3,7 +3,6 @@ from pykeops.common.lazy_tensor import GenericLazyTensor, ComplexGenericLazyTensor from pykeops.torch.utils import torchtools - # Convenient aliases: diff --git a/pykeops/pykeops/torch/operations.py b/pykeops/pykeops/torch/operations.py index e6232e88b..9434abda6 100644 --- a/pykeops/pykeops/torch/operations.py +++ b/pykeops/pykeops/torch/operations.py @@ -1,6 +1,6 @@ -import torch import copy +import torch from pykeops.common.get_options import get_tag_backend from pykeops.common.keops_io import keops_binder from pykeops.common.operations import ConjugateGradientSolver @@ -11,13 +11,12 @@ get_optional_flags, ) from pykeops.common.utils import axis2cat +from pykeops.common.utils import pyKeOps_Warning from pykeops.torch.generic.generic_red import ( GenredAutograd_fun, Genred_parameters, set_device, ) -from pykeops import default_device_id -from pykeops.common.utils import pyKeOps_Warning class KernelSolveAutograd(torch.autograd.Function): @@ -78,10 +77,19 @@ def linop(var): res += params.alpha * var return res - result = ConjugateGradientSolver("torch", linop, varinv.data, params.eps) + result = ConjugateGradientSolver( + "torch", + linop, + varinv.data, + eps=params.eps, + x0=params.x0, + maxiter=params.maxiter, + verbose=params.verbose, + ) # relying on the 'ctx.saved_variables' attribute is necessary if you want to be able to differentiate the output # of the backward once again. It helps pytorch to keep track of 'who is who'. + # Here we do not save the extra info dict if it is in the output ctx.save_for_backward(*args, result) return result @@ -89,7 +97,6 @@ def linop(var): @staticmethod def backward(ctx, G): params = ctx.params - device_id = ctx.params myconv = ctx.myconv args = ctx.saved_tensors[:-1] # Unwrap the saved variables @@ -236,6 +243,7 @@ def __init__( that should be computed and reduced. The correct syntax is described in the :doc:`documentation <../../Genred>`, using appropriate :doc:`mathematical operations <../../../api/math-operations>`. + aliases (list of strings): A list of identifiers of the form ``"AL = TYPE(DIM)"`` that specify the categories and dimensions of the input variables. Here: @@ -250,16 +258,13 @@ def __init__( As described below, :meth:`__call__` will expect input Tensors whose shape are compatible with **aliases**. + varinvalias (string): The alphanumerical **alias** of the variable with respect to which we shall perform our conjugate gradient descent. **formula** is supposed to be linear with respect to **varinvalias**, but may be more sophisticated than a mere ``"K(x,y) * {varinvalias}"``. Keyword Args: - alpha (float, default = 1e-10): Non-negative - **ridge regularization** parameter, added to the diagonal - of the Kernel matrix :math:`K_{xx}`. - axis (int, default = 0): Specifies the dimension of the kernel matrix :math:`K_{x_ix_j}` that is reduced by our routine. The supported values are: @@ -332,7 +337,16 @@ def __init__( self.axis = axis def __call__( - self, *args, backend="auto", device_id=-1, alpha=1e-10, eps=1e-6, ranges=None + self, + *args, + backend="auto", + device_id=-1, + ranges=None, + alpha=1e-10, + eps=1e-6, + x0=None, + maxiter=None, + verbose=False, ): r""" Apply the routine on arbitrary torch Tensors. @@ -373,6 +387,24 @@ def __call__( If **None** (default), we simply use a **dense Kernel matrix** as we loop over all indices :math:`i\in[0,M)` and :math:`j\in[0,N)`. + alpha (float, default = 1e-10): Non-negative floating-point + **ridge regularization** parameter, added to the diagonal + of the Kernel matrix :math:`K_{xx}`. + + eps (float, default = 1e-6): Stopping criterion for the conjugate gradient algorithm. + + x0 (2d Tensor, default = None): Initial guess for the solution of the linear system. + should be of the same shape as b. + + maxiter (int, default = None): Maximum number of conjugate + gradient iterations. If ``None``, uses the default from + :func:`pykeops.common.operations.ConjugateGradientSolver`. + + verbose (bool, default = False): If ``True``, prints the + conjugate gradient convergence information dictionary + produced by + :func:`pykeops.common.operations.ConjugateGradientSolver`. + Returns: (M,D) or (N,D) Tensor: @@ -389,18 +421,25 @@ def __call__( nx, ny = get_sizes(self.aliases, *args) params = Genred_parameters() + params.formula = self.formula params.aliases = self.aliases params.varinvpos = self.varinvpos - params.alpha = alpha params.backend = backend params.dtype = dtype + params.device_id_request = device_id - params.eps = eps params.ranges = ranges params.optional_flags = self.optional_flags params.rec_multVar_highdim = self.rec_multVar_highdim params.nx = nx params.ny = ny + # parameter of the Conjugate Gradient + params.alpha = alpha + params.eps = eps + params.x0 = x0 + params.maxiter = maxiter + params.verbose = verbose + return KernelSolveAutograd.apply(params, *args) diff --git a/pykeops/pykeops/torch/test_install.py b/pykeops/pykeops/torch/test_install.py index 541ec3231..909054cc5 100644 --- a/pykeops/pykeops/torch/test_install.py +++ b/pykeops/pykeops/torch/test_install.py @@ -9,17 +9,30 @@ def test_torch_bindings(): """ - This function try to compile a simple keops formula using the pytorch binder. + Try to compile a simple KeOps formula using the PyTorch binder. """ + import pykeops + import pykeops.torch as pktorch + x = torch.arange(1, 10, dtype=torch.float32).view(-1, 3) y = torch.arange(3, 9, dtype=torch.float32).view(-1, 3) - import pykeops.torch as pktorch + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() my_conv = pktorch.Genred(formula, var) - if torch.allclose( - my_conv(x, y).view(-1), torch.tensor(expected_res).type(torch.float32) - ): + + try: + keops_res = my_conv(x, y).view(-1) + except Exception as e: + pyKeOps_Message(f"Error during computation: {e}", use_tag=False) + return False + + if torch.allclose(keops_res, torch.tensor(expected_res, dtype=torch.float32)): pyKeOps_Message("pyKeOps with torch bindings is working!", use_tag=False) + return True else: - pyKeOps_Message("outputs wrong values...", use_tag=False) + pyKeOps_Message( + f"outputs wrong values: expected {expected_res} but get {keops_res}", + use_tag=False, + ) + return False diff --git a/pykeops/pykeops/torch/utils.py b/pykeops/pykeops/torch/utils.py index 45fb1d509..c5df4a81d 100644 --- a/pykeops/pykeops/torch/utils.py +++ b/pykeops/pykeops/torch/utils.py @@ -3,7 +3,6 @@ from pykeops.torch import Genred, KernelSolve from pykeops.torch.cluster import swap_axes as torch_swap_axes - # from pykeops.torch.generic.generic_red import GenredLowlevel @@ -18,12 +17,12 @@ def is_on_device(x): class torchtools: - copy = torch.clone - exp = torch.exp - log = torch.log - norm = torch.norm + copy = staticmethod(torch.clone) + exp = staticmethod(torch.exp) + log = staticmethod(torch.log) + norm = staticmethod(torch.norm) - swap_axes = torch_swap_axes + swap_axes = staticmethod(torch_swap_axes) Genred = Genred KernelSolve = KernelSolve @@ -154,6 +153,10 @@ def zeros( *shape, dtype=dtype, device=device, requires_grad=requires_grad ) + @staticmethod + def zeros_like(x): + return torch.zeros_like(x) + @staticmethod def empty( shape, @@ -210,6 +213,10 @@ def device_type_index(x): def pointer(x): return x.data.data_ptr() + @staticmethod + def sqrt(x): + return torch.sqrt(x) + def squared_distances(x, y): x_norm = (x**2).sum(1).reshape(-1, 1) diff --git a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py index c0d5a1d7f..c4911111e 100644 --- a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py +++ b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py @@ -24,6 +24,8 @@ # clouds :math:`(x_i)_{i\in[1,M]}` and :math:`(y_j)_{j\in[1,N]}` in the unit square: import numpy as np +import torch +import pykeops.config M, N = 1000, 2000 x = np.random.rand(M, 2) @@ -55,9 +57,7 @@ # That's good! Going further, we can speed-up these computations # using the **CUDA routines** of the PyTorch library: -import torch - -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor x_i = tensor(x[:, None, :]) # (M, 1, 2) torch tensor diff --git a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py index fe1ef1e57..b6de58335 100644 --- a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py +++ b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py @@ -27,10 +27,11 @@ # formulas: import torch +import pykeops.config from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor M, N = (100000, 200000) if use_cuda else (1000, 2000) D = 3 diff --git a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py index 4d52a2a0b..043c75aa1 100644 --- a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py +++ b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py @@ -10,10 +10,11 @@ import time import torch +import pykeops.config from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor ########################################################################### diff --git a/pykeops/pykeops/tutorials/backends/plot_scipy.py b/pykeops/pykeops/tutorials/backends/plot_scipy.py index 99249d74b..8ec0ab717 100644 --- a/pykeops/pykeops/tutorials/backends/plot_scipy.py +++ b/pykeops/pykeops/tutorials/backends/plot_scipy.py @@ -45,7 +45,7 @@ ################################################################### # Create a toy dataset, a spiral in 2D sampled with 10,000 points: -N = 10000 if pykeops.config.gpu_available else 1000 +N = 10000 if pykeops.config.cuda.is_available() else 1000 t = np.linspace(0, 2 * np.pi, N + 1)[:-1] x = np.stack((0.4 + 0.4 * (t / 7) * np.cos(t), 0.5 + 0.3 * np.sin(t)), 1) x = x + 0.01 * np.random.randn(*x.shape) @@ -221,7 +221,7 @@ # let's generate a large "noisy Swiss roll" with **1,000,000 points** in the unit cube: # -N = 1000000 if pykeops.config.gpu_available else 1000 +N = 1000000 if pykeops.config.cuda.is_available() else 1000 t = np.linspace(0, 2 * np.pi, N + 1)[:-1] x = np.stack( ( @@ -238,7 +238,7 @@ # To **display** our toy dataset with the (not-so-efficient) PyPlot library, # we pick **10,000 points** at random: -N_display = 10000 if pykeops.config.gpu_available else N +N_display = 10000 if pykeops.config.cuda.is_available() else N indices_display = np.random.randint(0, N, N_display) _, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 8), subplot_kw=dict(projection="3d")) @@ -317,7 +317,7 @@ # pattern can be encoded in a small boolean matrix **keep** computed through: sigma = ( - 0.01 if pykeops.config.gpu_available else 0.1 + 0.01 if pykeops.config.cuda.is_available() else 0.1 ) # Standard deviation of our Gaussian kernel # Compute a coarse Boolean mask: D = np.sum((x_centroids[:, None, :] - x_centroids[None, :, :]) ** 2, 2) diff --git a/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py b/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py index faa502585..f082b115b 100644 --- a/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py +++ b/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py @@ -17,19 +17,20 @@ import matplotlib.cm as cm import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from torch.nn import Module from torch.nn.functional import softmax, log_softmax from pykeops.torch import Vi, Vj, LazyTensor - #################################################################### # Define our dataset: a collection of points :math:`(x_i)_{i\in[1,N]}` which describe a # spiral in the unit square. # Choose the storage place for our data : CPU (host) or GPU (device) memory. -dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor torch.manual_seed(0) N = 10000 # Number of samples t = torch.linspace(0, 2 * np.pi, N + 1)[:-1] @@ -114,7 +115,7 @@ def __init__(self, M, sparsity=0, D=2): def update_covariances(self): """Computes the full covariance matrices from the model's parameters.""" - (M, D, _) = self.A.shape + M, D, _ = self.A.shape self.params["gamma"] = (torch.matmul(self.A, self.A.transpose(1, 2))).view( M, D * D ) / 2 diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py index 0109c3988..f6f58b042 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py @@ -29,18 +29,18 @@ import time -import numpy as np from matplotlib import pyplot as plt -from pykeops.numpy import LazyTensor +import numpy as np import pykeops.config +from pykeops.numpy import LazyTensor ####################################################################### # Generate some data: dtype = "float64" -N = 10000 if pykeops.config.gpu_available else 1000 # Number of samples +N = 10000 if pykeops.config.cuda.is_available() else 1000 # Number of samples # Sampling locations: x = np.random.rand(N, 1).astype(dtype) @@ -79,14 +79,12 @@ def gaussian_kernel(x, y, sigma=0.1): start = time.time() K_xx = gaussian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 1D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 1D: {end - start:.5f}s" ) ####################################################################### @@ -146,19 +144,17 @@ def laplacian_kernel(x, y, sigma=0.1): # between a perfect fit (**alpha** = 0) and a # smooth interpolation (**alpha** = :math:`+\infty`): -alpha = 10 # Ridge regularization +alpha = 5 # Ridge regularization start = time.time() K_xx = laplacian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 2D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 2D: {end - start:.5f}s" ) ######################################################################## diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py index ea0a91b59..4c9cde8f0 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py @@ -30,14 +30,14 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt - from pykeops.torch import LazyTensor ############################################################################################### # Generate some data: -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor N = 10000 if use_cuda else 1000 # Number of samples @@ -80,14 +80,12 @@ def gaussian_kernel(x, y, sigma=0.1): start = time.time() K_xx = gaussian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 1D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 1D: {end - start:.5f}s" ) ############################################################################################### @@ -148,29 +146,26 @@ def laplacian_kernel(x, y, sigma=0.1): # between a perfect fit (**alpha** = 0) and a # smooth interpolation (**alpha** = :math:`+\infty`): -alpha = 10 # Ridge regularization +alpha = 5 # Ridge regularization start = time.time() K_xx = laplacian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 2D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 2D: {end - start:.5f}s" ) - ############################################################################################### # Display the (fitted) model on the unit square: # # Extrapolate on a uniform sample: X = Y = torch.linspace(0, 1, 101).type(dtype) -X, Y = torch.meshgrid(X, Y) +X, Y = torch.meshgrid(X, Y, indexing="ij") t = torch.stack((X.contiguous().view(-1), Y.contiguous().view(-1)), dim=1) K_tx = laplacian_kernel(t, x) diff --git a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py index 16a22b2b4..cd81c68d7 100644 --- a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py +++ b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py @@ -107,7 +107,7 @@ def KMeans(x, K=10, Niter=10, verbose=True): # ------------------------- # Second experiment with N=1,000,000 points in dimension D=100, with K=1,000 classes: -if pykeops.config.gpu_available: +if pykeops.config.cuda.is_available(): N, D, K = 1000000, 100, 1000 x = np.random.randn(N, D).astype(dtype) cl, c = KMeans(x, K) diff --git a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py index c09f87c7c..f3b469cd0 100644 --- a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py +++ b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py @@ -24,12 +24,13 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() dtype = torch.float32 if use_cuda else torch.float64 -device_id = "cuda:0" if use_cuda else "cpu" +device_id = "cuda" if use_cuda else "cpu" ######################################################################## # Simple implementation of the K-means algorithm: diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py b/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py index 92caaa137..1d47c9d82 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py @@ -22,11 +22,12 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor ###################################################################### diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py b/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py index a68fdcca2..74eec6690 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py @@ -32,7 +32,7 @@ # Dataset, in 2D: N, D = ( - 10000 if pykeops.config.gpu_available else 1000, + 10000 if pykeops.config.cuda.is_available() else 1000, 2, ) # Number of samples, dimension x = np.random.rand(N, D).astype(dtype) # Random samples on the unit square @@ -48,7 +48,7 @@ def fth(x): ############################# # Reference sampling grid, on the unit square: -M = 1000 if pykeops.config.gpu_available else 100 +M = 1000 if pykeops.config.cuda.is_available() else 100 tmp = np.linspace(0, 1, M).astype(dtype) g1, g2 = np.meshgrid(tmp, tmp) g = np.hstack((g1.reshape(-1, 1), g2.reshape(-1, 1))) diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_torch.py b/pykeops/pykeops/tutorials/knn/plot_knn_torch.py index 610cb48a5..e878ef14a 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_torch.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_torch.py @@ -22,11 +22,12 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor ###################################################################### @@ -48,7 +49,7 @@ def fth(x): M = 1000 if use_cuda else 100 tmp = torch.linspace(0, 1, M).type(dtype) -g2, g1 = torch.meshgrid(tmp, tmp) +g2, g1 = torch.meshgrid(tmp, tmp, indexing="ij") g = torch.cat((g1.contiguous().view(-1, 1), g2.contiguous().view(-1, 1)), dim=1) diff --git a/pykeops/setup.py b/pykeops/setup.py index d8154d91f..e82cd2c75 100644 --- a/pykeops/setup.py +++ b/pykeops/setup.py @@ -1,7 +1,7 @@ # Always prefer setuptools over distutils # To use a consistent encoding -from codecs import open import os +from codecs import open from os import path from setuptools import setup @@ -28,24 +28,29 @@ "Bug Reports": "https://github.com/getkeops/keops/issues", "Source": "https://github.com/getkeops/keops", }, - author="B. Charlier, J. Feydy, J. Glaunes", - author_email="benjamin.charlier@umontpellier.fr, jean.feydy@gmail.com, alexis.glaunes@parisdescartes.fr", + author="B. Charlier, J. Feydy, J. Glaunès", + author_email="benjamin.charlier@inrae.fr, jean.feydy@inria.com, alexis.glaunes@parisdescartes.fr", python_requires=">=3.8", classifiers=[ + "Topic :: Scientific/Engineering", "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", - "Topic :: Scientific/Engineering", - "License :: OSI Approved :: MIT License", + "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Operating System :: MacOS :: MacOS X", + "Programming Language :: C", "Programming Language :: C++", "Programming Language :: Python :: 3 :: Only", + "Environment :: GPU :: NVIDIA CUDA", + "License :: OSI Approved :: MIT License", ], keywords="kernels gpu autodiff", packages=[ "pykeops", "pykeops.common", "pykeops.common.keops_io", + "pykeops.common.keops_io.nvrtc", + "pykeops.common.keops_io.cpp", "pykeops.numpy", "pykeops.numpy.cluster", "pykeops.numpy.generic", @@ -61,7 +66,8 @@ "readme.md", "licence.txt", "keops_version", - "common/keops_io/pykeops_nvrtc.cpp", + "common/keops_io/cpp/pykeops_cpp.cpp", + "common/keops_io/nvrtc/pykeops_nvrtc.cpp", ], }, install_requires=["numpy", "pybind11", "keopscore"], @@ -70,21 +76,26 @@ "sphinx", "sphinx-gallery", "recommonmark", + "myst-parser", "sphinxcontrib-httpdomain", "sphinx_rtd_theme", "sphinx-prompt", - "breathe", "matplotlib", "imageio", "torch", "gpytorch", "scikit-learn", "multiprocess", - "faiss", "h5py", "jaxlib", "jax", + "plotly", + "si_prefix", + "pandas", ], "test:": ["pytest", "numpy", "torch"], + "cu12": ["cuda-toolkit[nvrtc,nvcc,cudart,cccl]==12.9.2"], + "cu13": ["cuda-toolkit[nvrtc,nvcc,cccl,cudart,crt]==13.*"], + "cuda": ["cuda-toolkit[nvrtc,nvcc,cccl,cudart,crt]"], }, ) diff --git a/pytest.sh b/pytest.sh index 3b56ed1e3..56d39c315 100755 --- a/pytest.sh +++ b/pytest.sh @@ -1,128 +1,195 @@ -#!/bin/bash - -# exit in case of any errors -set -e - -################################################################################ -# help # -################################################################################ -function print_help() { - # Display Help - echo "Test script for keopscore/pykeops packages." - echo - echo "Usage: $0 [option...]" - echo - echo " -h Print the help" - echo " -v Verbose mode" - echo - exit 1 -} +#!/usr/bin/env bash + +set -euo pipefail + +readonly PYTHON_BIN="python3" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +readonly TEST_VENV="${SCRIPT_DIR}/.test_venv_pytest" +readonly NO_TORCH_TEST_VENV="${SCRIPT_DIR}/.test_venv_pytest_no_torch" +readonly TEST_REQUIREMENTS=(pip) + +KEOPS_VERBOSE_LEVEL=-1 +PIP_CONSTRAINT_FILE="" + +print_help() { + cat < + Verbosity level forwarded to KEOPS_VERBOSE and PYKEOPS_VERBOSE + --pip-constraint + Constrain pip installs using the given constraints file. +EOF +} -# log with verbosity management -function logging() { - if [[ ${PYTEST_VERBOSE} == 1 ]]; then - echo -e $1 +log_verbose() { + if [[ "${KEOPS_VERBOSE_LEVEL}" -ge 1 ]]; then + printf '%b\n' "$1" fi } -################################################################################ -# process script options # -################################################################################ +parse_options() { + while [[ $# -gt 0 ]]; do + case "$1" in + -h) + print_help + exit 0 + ;; + -v) + shift + if [[ $# -eq 0 ]]; then + echo "Error: -v requires a level (0, 1 or 2)" + exit 1 + fi + if ! [[ "$1" =~ ^[0-2]$ ]]; then + echo "Error: Invalid -v value: $1 (expected 0, 1 or 2)" + exit 1 + fi + KEOPS_VERBOSE_LEVEL="$1" + ;; + --pip-constraint) + shift + if [[ $# -eq 0 ]]; then + echo "Error: --pip-constraint requires a file path" + exit 1 + fi + PIP_CONSTRAINT_FILE="$1" + ;; + --pip-constraint=*) + PIP_CONSTRAINT_FILE="${1#*=}" + ;; + *) + echo "Error: Invalid option: $1" + exit 1 + ;; + esac + shift + done + + if [[ -n "${PIP_CONSTRAINT_FILE}" && ! -f "${PIP_CONSTRAINT_FILE}" ]]; then + echo "Error: Constraint file not found: ${PIP_CONSTRAINT_FILE}" + exit 1 + fi -# default options -PYTEST_VERBOSE=0 + log_verbose "## verbose mode (level=${KEOPS_VERBOSE_LEVEL})" +} -# Get the options -while getopts 'hv' option; do - case $option in - h) # display Help - print_help - ;; - v) # enable verbosity - PYTEST_VERBOSE=1 - logging "## verbose mode" - ;; - \?) # Invalid option - echo "Error: Invalid option" - exit 1 - ;; - esac -done +run_with_keops_verbose() { + if [[ "${KEOPS_VERBOSE_LEVEL}" -ne -1 ]]; then + KEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" PYKEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" "$@" + else + "$@" + fi +} -################################################################################ -# script setup # -################################################################################ +pip_install() { + if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then + run_with_keops_verbose "${PYTHON_BIN}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" "$@" + return + fi -# project root directory -PROJDIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) + run_with_keops_verbose "${PYTHON_BIN}" -m pip install "$@" +} -# python exec -PYTHON="python3" +prepare_python_environment() { -# python environment for test -TEST_VENV=${PROJDIR}/.test_venv_pytest + local env_name="$1" -# python test requirements (names of packages to be installed with pip) -TEST_REQ="pip" + log_verbose "-- Preparing python environment for test..." + "${PYTHON_BIN}" -m venv --clear "${env_name}" -################################################################################ -# prepare python environment # -################################################################################ + # shellcheck disable=SC1091 + source "${env_name}/bin/activate" -logging "-- Preparing python environment for test..." + log_verbose "---- Python version = $(${PYTHON_BIN} -V)" + pip_install -U "${TEST_REQUIREMENTS[@]}" +} -${PYTHON} -m venv --clear ${TEST_VENV} -source ${TEST_VENV}/bin/activate +deactivate_python_environment() { -logging "---- Python version = $(python -V)" + log_verbose "-- Deactivate python environment..." -pip install -U ${TEST_REQ} + # `deactivate` is provided by the activation script; there is no bin/deactivate file. + if declare -F deactivate >/dev/null; then + deactivate + fi +} +install_editable_package() { + local package_name="$1" + local package_path="$2" -################################################################################ -# Installing keopscore # -################################################################################ + log_verbose "-- Installing ${package_name}..." + pip_install -e "${package_path}" +} -logging "-- Installing keopscore..." +run_python_outside_repo() { + local python_code="$1" + local python_bin="${2:-${PYTHON_BIN}}" + ( + cd /tmp + run_with_keops_verbose "${python_bin}" -c "${python_code}" + ) +} -pip install -e ${PROJDIR}/keopscore +clean_pykeops_cache() { + log_verbose "-- Cleaning pykeops..." + run_python_outside_repo 'import pykeops; pykeops.clean_pykeops()' +} -################################################################################ -# Installing pykeops # -################################################################################ +run_pykeops_health_check() { + echo "-- Running pykeops.check_health()..." + run_python_outside_repo 'import pykeops; pykeops.check_health()' +} -logging "-- Installing pykeops..." +run_pykeops_no_torch_smoke_test() { + local no_torch_python="${NO_TORCH_TEST_VENV}/bin/python" + local no_torch_smoke_code="import importlib.util;" + no_torch_smoke_code+=" assert importlib.util.find_spec('torch') is None," + no_torch_smoke_code+=" 'torch must not be installed in no-torch smoke env';" + no_torch_smoke_code+=" import pykeops;" + no_torch_smoke_code+=" assert pykeops.test_numpy_bindings()" -pip install -e "${PROJDIR}/pykeops[test]" + run_python_outside_repo "${no_torch_smoke_code}" "${no_torch_python}" +} -################################################################################ -# Cleaning cache # -################################################################################ +run_test_suite() { + local suite_name="$1" + local suite_path="$2" -logging "-- Cleaning pykeops..." + log_verbose "-- Running ${suite_name} tests..." + run_with_keops_verbose pytest -v "${suite_path}" +} -mkdir tmp -cd tmp -${PYTHON} -c "import pykeops; pykeops.clean_pykeops()" -cd .. -rmdir tmp +main() { + parse_options "$@" -################################################################################ -# Running keopscore tests # -################################################################################ + printf '%b' "****************************************************************************\n Start of pykeops no-torch smoke test\n****************************************************************************\n" + prepare_python_environment "${NO_TORCH_TEST_VENV}" + install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" + install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops" + clean_pykeops_cache -logging "-- Running keopscore tests..." + run_pykeops_health_check + run_pykeops_no_torch_smoke_test -pytest -v keopscore/keopscore/test/ + deactivate_python_environment + printf '%b' "****************************************************************************\n End of pykeops no-torch smoke test\n****************************************************************************\n\n\n\n\n\n" -################################################################################ -# Running pykeops tests # -################################################################################ + printf '%b' "****************************************************************************\n Start of pykeops tests\n****************************************************************************\n" + prepare_python_environment "${TEST_VENV}" + install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" + install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[test]" + run_pykeops_health_check + + clean_pykeops_cache + run_test_suite "keopscore" "keopscore/keopscore/test/" + run_test_suite "pykeops" "pykeops/pykeops/test/" + printf '%b' "****************************************************************************\n End of pykeops tests\n****************************************************************************\n\n\n\n\n\n" -logging "-- Running pykeops tests..." +} -pytest -v pykeops/pykeops/test/ +main "$@"