Sphinx Documentation, First Pass (#113)

* Started documentation.

* Documentation is beginning to work.

* Documentation generator is working.

* Began writing documentation.

* Documentation is shaping up.

* Making progress on writing documentation.

* Adding environment configurations.

* Added description of tests and benchmarks.

* More progress on documentation.

* Finished first pass of documentation.

* Almost ready to deploy docs.

* Reverted README for now.

* Removed comments from workflow file.

* Removed Makefile.

* Reformatted.

* Minor documentation updates.

* Linted.

Author: vbharadwaj-bk

.github/workflows/docs.yaml

@@ -0,0 +1,49 @@
+name: Deploy documentation to Github Pages 
+on:
+  workflow_dispatch:
+
+permissions: write-all
+
+concurrency:
+  group: "pages"
+  cancel-in-progress: false
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Setup Pages
+        uses: actions/configure-pages@v3
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+         python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install sphinx furo
+      - name: Build website
+        run: |
+          sphinx-build -M dirhtml docs docs/_build
+          
+      - name: Fix permissions
+        run: |
+          chmod -c -R +rX "docs/_build/dirhtml/" | while read line; do
+            echo "::warning title=Invalid file permissions automatically fixed::$line"
+          done
+      - name: Upload artifact
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: './docs/_build/dirhtml'
+  deploy:
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4

.gitignore

@@ -16,6 +16,9 @@ experimental/*
 MACE_models/*
 triton_autotuning/*
 sandbox/*
+docs/_build
+docs/_static
+docs/_templates
 
 # bash scripts
 get_gpu_node.sh

README.md

@@ -348,4 +348,4 @@ Copyright (c) 2025, The Regents of the University of California, through Lawrenc
 
 If you have questions about your rights to use or distribute this software, please contact Berkeley Lab's Intellectual Property Office at IPO@lbl.gov.
 
-NOTICE. This Software was developed under funding from the U.S. Department of Energy and the U.S. Government consequently retains certain rights. As such, the U.S. Government has been granted for itself and others acting on its behalf a paid-up, nonexclusive, irrevocable, worldwide license in the Software to reproduce, distribute copies to the public, prepare derivative works, and perform publicly and display publicly, and to permit others to do so.
+NOTICE. This Software was developed under funding from the U.S. Department of Energy and the U.S. Government consequently retains certain rights. As such, the U.S. Government has been granted for itself and others acting on its behalf a paid-up, nonexclusive, irrevocable, worldwide license in the Software to reproduce, distribute copies to the public, prepare derivative works, and perform publicly and display publicly, and to permit others to do so.

docs/api.rst

@@ -0,0 +1,47 @@
+OpenEquivariance API
+==============================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Contents:
+
+OpenEquivariance exposes two key classes: :py:class:`openequivariance.TensorProduct`, which replaces
+``o3.TensorProduct`` from e3nn, and :py:class:`openequivariance.TensorProductConv`, which fuses
+the CG tensor product with a subsequent graph convolution. Initializing either class triggers
+JIT compilation of a custom kernel, which can take a few seconds. 
+
+Both classes require a configuration object specified 
+by :py:class:`openequivariance.TPProblem`, which has a constructor
+almost identical to ``o3.TensorProduct``. 
+We recommend reading the `e3nn documentation <https://docs.e3nn.org/en/latest/>`_ before
+trying our code. OpenEquivariance cannot accelerate all tensor products; see 
+:doc:`this page </supported_ops>` for a list of supported configurations.
+
+.. autoclass:: openequivariance.TensorProduct
+    :members:
+    :undoc-members:
+    :exclude-members: name
+
+.. autoclass:: openequivariance.TensorProductConv
+    :members:
+    :undoc-members:
+    :exclude-members: name
+
+.. autoclass:: openequivariance.TPProblem
+    :members:
+    :undoc-members:
+
+.. autofunction:: openequivariance.torch_to_oeq_dtype
+
+.. autofunction:: openequivariance.torch_ext_so_path
+
+API Identical to e3nn
+---------------------
+
+These remaining API members are identical to the corresponding
+objects in ``e3nn.o3``. You can freely mix these objects from
+both packages. 
+
+.. autoclass:: openequivariance.Irreps
+    :members:
+    :undoc-members:

docs/conf.py

@@ -0,0 +1,38 @@
+import sys
+from pathlib import Path
+
+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = "OpenEquivariance"
+copyright = "2025, The Regents of the University of California, through Lawrence Berkeley National Laboratory."
+author = "Vivek Bharadwaj, Austin Glover, Aydin Buluc, James Demmel"
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = []
+
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = "furo"
+# html_static_path = ["_static"]
+
+extensions = [
+    "sphinx.ext.autodoc",
+]
+
+sys.path.insert(0, str(Path("..").resolve()))
+
+autodoc_mock_imports = ["torch", "openequivariance.extlib", "jinja2"]
+autodoc_typehints = "description"

docs/index.rst

@@ -0,0 +1,28 @@
+.. OpenEquivariance documentation master file, created by
+   sphinx-quickstart on Tue Jun  3 00:20:54 2025.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+OpenEquivariance
+==============================
+
+OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon
+tensor product, a key kernel in equivariant graph neural networks. We offer
+an identical interface to e3nn and produce the same results 
+(up to numerical roundoff). Our package exhibits up to an order of magnitude
+speedup over e3nn and competitive performance with NVIDIA's cuEquivariance. 
+
+Here, you can find our API reference, installation instructions, 
+and troubleshooting guide. We support for both NVIDIA and AMD GPUs through
+our PyTorch interface, including support for JITScript compilation accessible
+from C++.
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Contents:
+
+   installation
+   api
+   supported_ops
+   tests_and_benchmarks
+   models 

docs/installation.rst

@@ -0,0 +1,95 @@
+Installation
+==============================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Contents:
+
+You need the following to install OpenEquivariance:
+
+- A Linux system equipped with an NVIDIA / AMD graphics card.
+- PyTorch >= 2.4 (>= 2.8 for AOTI and export).
+- GCC 9+ and the CUDA / HIP toolkit. The command
+  ``c++ --version`` should return >= 9.0; see below for details on 
+  setting an alternate compiler.
+
+Installation is one easy command, followed by import verification: 
+
+.. code-block:: bash 
+
+    pip install git+https://github.com/PASSIONLab/OpenEquivariance 
+    python -c "import openequivariance"
+
+The second line triggers a build of the C++ extension we use to compile
+kernels, which can take a couple of minutes. Subsequent imports are
+much faster since this extension is cached. 
+
+
+Compiling the Integrated PyTorch Extension
+------------------------------------------
+To support ``torch.compile``, ``torch.export``, and
+JITScript, OpenEquivariance needs to compile a C++ extension
+tightly integrated with PyTorch. If you see a warning that
+this extension could not be compiled, first check:
+
+.. code-block:: bash 
+
+    c++ --version 
+    
+To build the extension with an alternate compiler, set the 
+``CC`` and ``CXX``
+environment variable and retry the import:
+
+.. code-block:: bash
+
+    export CCC=/path/to/your/gcc
+    export CXX=/path/to/your/g++
+    python -c "import openequivariance"
+
+These configuration steps are required only ONCE after 
+installation (or upgrade) with pip. 
+
+Configurations on Major Platforms 
+---------------------------------
+OpenEquivariance has been tested on both supercomputers and lab clusters.
+Here are some tested environment configuration files. If use OpenEquivariance
+on a widely-used platform, send us a pull request to add your configuration! 
+
+NERSC Perlmutter (NVIDIA A100)
+""""""""""""""""""""""""""""""
+
+.. code-block:: bash
+    :caption: env.sh (last updated June 2024)
+
+    module load gcc 
+    module load conda
+
+    # Deactivate any base environments
+    for i in $(seq ${CONDA_SHLVL}); do 
+        conda deactivate
+    done
+
+    conda activate <your-conda-env>
+
+
+OLCF Frontier (AMD MI250x)
+""""""""""""""""""""""""""
+You need to install a HIP-enabled verison of PyTorch to use our package. 
+To do this, follow the steps `here <https://docs.olcf.ornl.gov/software/analytics/pytorch_frontier.html>`_.
+
+
+.. code-block:: bash
+    :caption: env.sh (last updated June 2024) 
+
+    module load PrgEnv-gnu/8.6.0
+    module load miniforge3/23.11.0-0
+    module load rocm/6.4.0
+    module load craype-accel-amd-gfx90a
+
+    for i in $(seq ${CONDA_SHLVL}); do
+        conda deactivate
+    done
+
+    conda activate <your-conda-env>
+    export CC=cc
+    export CXX=CC

docs/livebuild.sh

@@ -0,0 +1 @@
+sphinx-autobuild -M dirhtml . _build --watch ../openequivariance 

docs/models.rst

@@ -0,0 +1,46 @@
+Running MACE and Nequip
+==============================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Contents:
+
+MACE
+----
+
+We have modified MACE to use our accelerated kernels instead
+of the standard e3nn backend. Here are the steps to replicate
+our MACE benchmark:
+
+1. Install ``oeq`` and our modified version of MACE via
+
+   .. code-block:: bash
+
+      pip uninstall mace-torch
+      pip install git+https://github.com/vbharadwaj-bk/mace_oeq_integration.git@oeq_experimental
+
+2. Download the ``carbon.xyz`` data file, available at 
+   `<https://portal.nersc.gov/project/m1982/equivariant_nn_graphs/>`_.
+
+   This graph has 158K edges. With the original e3nn backend, you would need a GPU with 80GB
+   of memory to run the experiments. ``oeq`` provides a memory-efficient equivariant convolution,
+   so we expect the test to succeed.
+
+3. Benchmark OpenEquivariance:
+
+   .. code-block:: bash
+
+      python tests/mace_driver.py carbon.xyz -o outputs/mace_tests -i oeq
+
+4. If you have a GPU with 80GB of memory *or* supply a smaller molecular graph
+   as the input file, you can run the full benchmark that includes ``e3nn`` and ``cue``:
+
+   .. code-block:: bash
+
+      python tests/mace_driver.py carbon.xyz -o outputs/mace_tests -i e3nn cue oeq
+
+Nequip
+------
+See the 
+`official Nequip documentation <https://nequip.readthedocs.io/en/latest/guide/accelerations/openequivariance.html>`_
+to use OpenEquivariance with Nequip.