jaxdecomp 0.2.0__tar.gz

jaxdecomp-0.2.0/.clang-format

@@ -0,0 +1,15 @@
+---
+BasedOnStyle: LLVM
+ColumnLimit: 120
+CommentPragmas:  '^\\.+'
+DerivePointerAlignment: false
+Language: Cpp
+PointerAlignment: Left
+UseTab: Never
+AlignAfterOpenBracket: Align
+AlignTrailingComments: true
+AllowShortBlocksOnASingleLine: true
+AllowShortCaseLabelsOnASingleLine : true
+AllowShortIfStatementsOnASingleLine: true
+AllowShortLoopsOnASingleLine: true
+...

jaxdecomp-0.2.0/.github/workflows/formatting.yml

@@ -0,0 +1,21 @@
+name: Code Formatting
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v3
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip isort
+        python -m pip install pre-commit
+    - name: Run pre-commit
+      run: python -m pre_commit run --all-files

jaxdecomp-0.2.0/.github/workflows/github-deploy.yml

@@ -0,0 +1,68 @@
+name: Build and upload to PyPI
+
+on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches:
+    - main
+  release:
+     types:
+       - published
+
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        # macos-13 is an intel runner, macos-14 is apple silicon
+        os: [ubuntu-latest]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build wheels
+        uses: pypa/cibuildwheel@v2.21.3
+        env:
+         CIBW_BUILD: "cp310-* cp311-* cp312-*"
+         CIBW_BUILD_VERBOSITY: 2
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
+          path: ./wheelhouse/*.whl
+
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build sdist
+        run: pipx run build --sdist
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-sdist
+          path: dist/*.tar.gz
+
+  upload_pypi:
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    # if: github.event_name == 'release' && github.event.action == 'published'
+    # or, alternatively, upload to PyPI on every tag starting with 'v' (remove on: release above to use this)
+    # if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          # unpacks all CIBW artifacts into dist/
+          pattern: cibw-*
+          path: dist
+          merge-multiple: true
+
+      - uses: pypa/gh-action-pypi-publish@release/v1
+      #with:
+      #    repository-url: https://test.pypi.org/legacy/

jaxdecomp-0.2.0/.github/workflows/joss-paper-pdf.yml

@@ -0,0 +1,28 @@
+name: Draft PDF
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  paper:
+    runs-on: ubuntu-latest
+    name: Paper Draft
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Build draft PDF
+        uses: openjournals/openjournals-draft-action@master
+        with:
+          journal: joss
+          # This should be the path to the paper within your repo.
+          paper-path: joss-paper/paper.md
+      - name: Upload
+        uses: actions/upload-artifact@v3
+        with:
+          name: paper
+          # This is the output path where Pandoc will write the compiled
+          # PDF. Note, this should be the same directory as the input
+          # paper.md
+          path: joss-paper/paper.pdf

jaxdecomp-0.2.0/.github/workflows/tests.yml

@@ -0,0 +1,39 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.10.4]
+
+    steps:
+    - name: Checkout Source
+      uses: actions/checkout@v2.3.1
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install jax[cpu]
+        pip install .[test]
+
+    - name: Run tests
+      run: |
+        cd tests
+        export JAX_PLATFORM_NAME=cpu
+        export XLA_FLAGS='--xla_force_host_platform_device_count=8'
+        pytest -v

jaxdecomp-0.2.0/.gitignore

@@ -0,0 +1,137 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+.vscode/
+
+scripts/experimental
+compile_commands.json
+CMakeFiles
+traces*
+notes.txt

jaxdecomp-0.2.0/.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/cuDecomp"]
+	path = third_party/cuDecomp
+	url = https://github.com/NVIDIA/cuDecomp.git

jaxdecomp-0.2.0/.pre-commit-config.yaml

@@ -0,0 +1,28 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+-   repo: https://github.com/google/yapf
+    rev: v0.40.2
+    hooks:
+    - id: yapf
+      args: ['--parallel', '--in-place']
+-   repo: https://github.com/pycqa/isort
+    rev: 5.13.2
+    hooks:
+      - id: isort
+        name: isort (python)
+-   repo: https://github.com/pre-commit/mirrors-clang-format
+    rev: v18.1.4
+    hooks:
+      - id: clang-format
+        files: '\.(c|cc|cpp|h|hpp|cxx|hh|cu|cuh)$'
+        exclude: '^third_party/|/pybind11/'
+        name: clang-format
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.13.0
+    hooks:
+    -   id: mypy

jaxdecomp-0.2.0/.style.yapf

@@ -0,0 +1,2 @@
+[style]
+based_on_style = yapf

jaxdecomp-0.2.0/CHANGELOG.md

@@ -0,0 +1,37 @@
+# Change log
+
+## jaxdecomp 0.2.0
+
+* Changes
+  * jaxDecomp works without MPI and using only JAX as backend
+  * with mesh is no longer required (will be deprecated by JAX)
+  * Added support for fftfreq
+  * Added testing for all functions
+  * Added static typing and checked with mypy
+
+## jaxdecomp 0.1.0
+
+* Changes
+  * Fixed bug with Halo
+  * Added joss paper
+
+
+## jaxdecomp 0.0.1
+
+* Changes
+  * New version compatible with JAX 0.4.30
+  * jaxDecomp now works in a multi-host environment
+  * Added custom partitioning for FFTs
+  * Added custom partitioning for halo exchange
+  * Added custom partitioning for slice_pad and slice_unpad
+  * Add example for multi-host FFTs in `examples/jaxdecomp_lpt.py`
+
+
+## jaxdecomp 0.0.1rc2
+* Changes
+  * Added utility to run autotuning
+
+
+## jaxdecomp 0.0.1rc1 (Nov. 25th 2022)
+
+Initial pre-release, include support for parallel ffts and halo exchange

jaxdecomp-0.2.0/CMakeLists.txt

@@ -0,0 +1,97 @@
+cmake_minimum_required(VERSION 3.19...3.25)
+
+
+
+project(jaxdecomp LANGUAGES CXX)
+
+# NVCC 12 does not support C++20
+set(CMAKE_CXX_STANDARD  17)
+set(CMAKE_CUDA_STANDARD 17)
+
+option(JD_CUDECOMP_BACKEND "Use cuDecomp backend" OFF)
+# Set default build type to Release
+if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+  set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
+  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
+endif()
+
+set(PYBIND11_FINDPYTHON ON)
+find_package(pybind11 CONFIG REQUIRED)
+
+# Check for CUDA
+include(CheckLanguage)
+check_language(CUDA)
+
+if(CMAKE_CUDA_COMPILER AND JD_CUDECOMP_BACKEND)
+  enable_language(CUDA)
+
+  # Latest JAX v0.4.26 no longer supports cuda 11.8
+  find_package(CUDAToolkit REQUIRED VERSION 12)
+  set(NVHPC_CUDA_VERSION ${CUDAToolkit_VERSION_MAJOR}.${CUDAToolkit_VERSION_MINOR})
+
+  message(STATUS "Using CUDA ${NVHPC_CUDA_VERSION}")
+
+  add_subdirectory(third_party/cuDecomp)
+
+  option(CUDECOMP_BUILD_FORTRAN "Build Fortran bindings" OFF)
+  option(CUDECOMP_ENABLE_NVSHMEM "Enable NVSHMEM" OFF)
+  option(CUDECOMP_BUILD_EXTRAS "Build benchmark, examples, and tests" OFF)
+
+  # 70: Volta, 80: Ampere, 89: RTX 4060
+  set(CUDECOMP_CUDA_CC_LIST "70;80;89" CACHE STRING "List of CUDA compute capabilities to build cuDecomp for.")
+
+  find_package(NVHPC REQUIRED COMPONENTS MATH MPI NCCL)
+
+  string(REPLACE "/lib64" "/include" NVHPC_MATH_INCLUDE_DIR ${NVHPC_MATH_LIBRARY_DIR})
+  string(REPLACE "/lib64" "/include" NVHPC_CUDA_INCLUDE_DIR ${NVHPC_CUDA_LIBRARY_DIR})
+
+
+  find_library(NCCL_LIBRARY
+      NAMES nccl
+      HINTS ${NVHPC_NCCL_LIBRARY_DIR}
+    )
+  string(REPLACE "/lib" "/include" NCCL_INCLUDE_DIR ${NVHPC_NCCL_LIBRARY_DIR})
+
+
+  message(STATUS "Using NCCL library: ${NCCL_LIBRARY}")
+  message(STATUS "NVHPC NCCL lib dir: ${NVHPC_NCCL_LIBRARY_DIR}")
+  message(STATUS "NCCL include dir: ${NCCL_INCLUDE_DIR}")
+
+  # Add _jaxdecomp modulei
+  pybind11_add_module(_jaxdecomp
+                        src/csrc/halo.cu
+                        src/csrc/jaxdecomp.cc
+                        src/csrc/grid_descriptor_mgr.cc
+                        src/csrc/fft.cu
+                        src/csrc/transpose.cu
+  )
+
+  set_target_properties(_jaxdecomp PROPERTIES CUDA_ARCHITECTURES "${CUDECOMP_CUDA_CC_LIST}")
+
+  target_include_directories(_jaxdecomp
+    PRIVATE
+    ${CMAKE_CURRENT_LIST_DIR}/src/csrc/include
+    ${CMAKE_CURRENT_SOURCE_DIR}/third_party/cuDecomp/include
+    ${NVHPC_CUDA_INCLUDE_DIR}
+    ${MPI_CXX_INCLUDE_DIRS}
+    ${NVHPC_MATH_INCLUDE_DIR}
+    ${NCCL_INCLUDE_DIR}
+  )
+
+  target_link_libraries(_jaxdecomp PRIVATE MPI::MPI_CXX)
+  target_link_libraries(_jaxdecomp PRIVATE NVHPC::CUFFT)
+  target_link_libraries(_jaxdecomp PRIVATE NVHPC::CUTENSOR)
+  target_link_libraries(_jaxdecomp PRIVATE NVHPC::CUDA)
+  target_link_libraries(_jaxdecomp PRIVATE ${NCCL_LIBRARY})
+  target_link_libraries(_jaxdecomp PRIVATE cudecomp)
+  target_link_libraries(_jaxdecomp PRIVATE stdc++fs)
+  set_target_properties(_jaxdecomp PROPERTIES LINKER_LANGUAGE CXX)
+  target_compile_definitions(_jaxdecomp PRIVATE JD_CUDECOMP_BACKEND)
+else()
+  pybind11_add_module(_jaxdecomp src/csrc/jaxdecomp.cc)
+  target_include_directories(_jaxdecomp PRIVATE ${CMAKE_CURRENT_LIST_DIR}/src/csrc/include)
+  target_compile_definitions(_jaxdecomp PRIVATE JD_JAX_BACKEND)
+endif()
+
+set_target_properties(_jaxdecomp PROPERTIES INSTALL_RPATH "$ORIGIN/lib")
+install(TARGETS _jaxdecomp LIBRARY DESTINATION . PUBLIC_HEADER DESTINATION .)

jaxdecomp-0.2.0/CONTRIBUTING.md

@@ -0,0 +1,46 @@
+# Code and Contribution guidelines
+
+## Code formatting
+
+Formatting is enforced using [yapf](https://github.com/google/yapf) and automatically applied using pre-commit hooks. To manually format the code, run the following command:
+
+```shell
+yapf -i -r .
+```
+but we highly recommend using the pre-commit hooks to ensure consistent formatting across the codebase, see below.
+
+### Pre-commit
+
+We use pre-commit to enforce code formatting and quality standards. Follow these steps to install and use pre-commit:
+
+1. Make sure you have Python installed on your system.
+
+2. Install pre-commit by running the following command in your terminal:
+
+    ```shell
+    pip install pre-commit
+    ```
+
+3. Navigate to the root directory of your project.
+
+4. Run the following command to initialize pre-commit:
+
+    ```shell
+    pre-commit install
+    ```
+
+5. Now, whenever you make a commit, pre-commit will automatically run the configured hooks on the files you modified. If any issues are found, pre-commit will prevent the commit from being made.
+
+    You can also manually run pre-commit on all files by running the following command:
+
+    ```shell
+    pre-commit run --all-files
+    ```
+
+    This is useful if you want to check all files before making a commit.
+
+6. Customize the pre-commit configuration by creating a `.pre-commit-config.yaml` file in the root directory of your project. This file allows you to specify which hooks should be run and how they should be configured.
+
+    For more information on configuring pre-commit, refer to the [pre-commit documentation](https://pre-commit.com/#configuration).
+
+That's it! You now have pre-commit set up to automatically enforce code formatting and quality standards in your project.

jaxdecomp-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Differentiable Universe Initiative
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.