Merge pull request #64 from NVIDIA/release/v1.0.0

release: nvmath-python-1.0.0

Author: markjolah

.gitignore

@@ -9,6 +9,8 @@ __pycache__
 docs/_build
 docs/sphinx/**/generated
 docs/sphinx/generated
+internal/internal_docs/_build
+internal/internal_docs/sphinx/**/generated
 dist
 build
 wheelhouse

.pre-commit-config.yaml

@@ -41,9 +41,35 @@ repos:
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: "v1.18.1"
     hooks:
+      # We split the mypy run into two separate hooks to avoid a "Duplicate module named..."
+      # error. Some files in examples/device/numba_cuda_mlir/ have exactly the same name
+      # as files in examples/device/, and mypy treats them as identical modules if
+      # passed together.
       - id: mypy
-        # Envorce only one source of configuration.
+        name: mypy (main and examples)
+        # Enforce only one source of configuration.
         args: ["--config-file", "pyproject.toml"]
+        exclude: ^examples/device/numba_cuda_mlir/
+        additional_dependencies:
+          - cuda-core
+          - cuda-bindings>=12.9.2,<13
+          - cupy-cuda12x
+          - mpi4py>=4.1.0
+          - numba
+          - numba-cuda
+          - numpy
+          - pytest
+          - scipy
+          - torch
+          - types-cffi
+          - invoke
+          - cython>=3.0.4,!=3.1.0,!=3.1.1
+          - tomli
+      - id: mypy
+        name: mypy (numba-cuda-mlir examples)
+        # Enforce only one source of configuration.
+        args: ["--config-file", "pyproject.toml"]
+        files: ^examples/device/numba_cuda_mlir/
         additional_dependencies:
           - cuda-core
           - cuda-bindings>=12.9.2,<13
@@ -106,6 +132,15 @@ repos:
       # Ignore old internal README that will not be rendered as docs page
         args: ["--fix", "--ignore", "internal/gtc2024/README.md"]
 
+  - repo: local
+    hooks:
+      - id: check-lock-urls
+        name: Check that lock files use internal registries
+        entry: python .ci/check_locks.py
+        language: python
+        files: ^\.ci/locks/.*\.toml$
+        pass_filenames: true
+
   - repo: https://github.com/sphinx-contrib/sphinx-lint
     rev: v1.0.0
     hooks:

README.md

@@ -15,7 +15,7 @@ functionality that is missing from those frameworks.
 ## Some Examples
 
 Below are a few representative examples showcasing the three main categories of
-features nvmath-python offers: host, device, and distributed APIs.
+features nvmath-python offers: host, distributed host, and device APIs.
 
 ### Host APIs
 
@@ -102,6 +102,83 @@ s = cp.fft.fftn(a, axes=[-1], norm="ortho")
 assert cp.allclose(r, s)
 ```
 
+### Distributed Host APIs
+
+Distributed Host APIs are called from host code but execute on a distributed
+(multi-node multi-GPU) system. The following example shows the use of the
+function-form distributed FFT with CuPy ndarrays:
+
+```python
+import cupy as cp
+from mpi4py import MPI
+
+import nvmath.distributed
+from nvmath.distributed.distribution import Slab
+
+# Initialize nvmath.distributed.
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+nranks = comm.Get_size()
+device_id = rank % cp.cuda.runtime.getDeviceCount()
+nvmath.distributed.initialize(device_id, comm, backends=["nvshmem"])
+
+# The global 3-D FFT size is (512, 256, 512).
+# In this example, the input data is distributed across processes according to
+# the cuFFTMp Slab distribution on the X axis.
+shape = 512 // nranks, 256, 512
+
+# cuFFTMp uses the NVSHMEM PGAS model for distributed computation, which requires GPU
+# operands to be on the symmetric heap.
+a = nvmath.distributed.allocate_symmetric_memory(shape, cp, dtype=cp.complex128)
+# a is a cupy ndarray and can be operated on using in-place cupy operations.
+with cp.cuda.Device(device_id):
+    a[:] = (
+        cp.random.rand(*shape, dtype=cp.float64)
+        + 1j * cp.random.rand(*shape, dtype=cp.float64)
+    )
+
+# Forward FFT.
+# In this example, the forward FFT operand is distributed according
+# to Slab.X distribution. With redistribute=False, the FFT result will be
+# distributed according to Slab.Y distribution.
+b = nvmath.distributed.fft.fft(a, distribution=Slab.X, options={"redistribute": False})
+
+# Distributed FFT performs computations in-place. The result is stored in the same
+# buffer as operand a. Note, however, that operand b has a different shape (due
+# to Slab.Y distribution).
+if rank == 0:
+    print(f"Shape of a on rank {rank} is {a.shape}")
+    print(f"Shape of b on rank {rank} is {b.shape}")
+
+# Inverse FFT.
+# Recall from previous transform that the inverse FFT operand is distributed according
+# to Slab.Y. With redistribute=False, the inverse FFT result will be distributed
+# according to Slab.X distribution.
+c = nvmath.distributed.fft.ifft(b, distribution=Slab.Y, options={"redistribute": False})
+
+# The shape of c is the same as a (due to Slab.X distribution). Once again, note that
+# a, b and c are sharing the same symmetric memory buffer (distributed FFT operations
+# are in-place).
+if rank == 0:
+    print(f"Shape of c on rank {rank} is {c.shape}")
+
+# Synchronize the default stream
+with cp.cuda.Device(device_id):
+    cp.cuda.get_current_stream().synchronize()
+
+if rank == 0:
+    print(f"Input type = {type(a)}, device = {a.device}")
+    print(f"FFT output type = {type(b)}, device = {b.device}")
+    print(f"IFFT output type = {type(c)}, device = {c.device}")
+
+# GPU operands on the symmetric heap are not garbage-collected and the user is
+# responsible for freeing any that they own (this deallocation is a collective
+# operation that must be called by all processes at the same point in the execution).
+# All cuFFTMp operations are inplace (a, b, and c share the same memory buffer), so
+# we take care to only free the buffer once.
+nvmath.distributed.free_symmetric_memory(a)
+```
+
 ### Device-side APIs
 
 nvmath-python exposes NVIDIA's device-side (Dx) APIs. This allows developers to call NVIDIA
@@ -195,89 +272,12 @@ if __name__ == "__main__":
     main()
 ```
 
-### Distributed APIs
-
-Distributed APIs are called from host code but execute on a distributed
-(multi-node multi-GPU) system. The following example shows the use of the
-function-form distributed FFT with CuPy ndarrays:
-
-```python
-import cupy as cp
-from mpi4py import MPI
-
-import nvmath.distributed
-from nvmath.distributed.distribution import Slab
-
-# Initialize nvmath.distributed.
-comm = MPI.COMM_WORLD
-rank = comm.Get_rank()
-nranks = comm.Get_size()
-device_id = rank % cp.cuda.runtime.getDeviceCount()
-nvmath.distributed.initialize(device_id, comm, backends=["nvshmem"])
-
-# The global 3-D FFT size is (512, 256, 512).
-# In this example, the input data is distributed across processes according to
-# the cuFFTMp Slab distribution on the X axis.
-shape = 512 // nranks, 256, 512
-
-# cuFFTMp uses the NVSHMEM PGAS model for distributed computation, which requires GPU
-# operands to be on the symmetric heap.
-a = nvmath.distributed.allocate_symmetric_memory(shape, cp, dtype=cp.complex128)
-# a is a cupy ndarray and can be operated on using in-place cupy operations.
-with cp.cuda.Device(device_id):
-    a[:] = (
-        cp.random.rand(*shape, dtype=cp.float64)
-        + 1j * cp.random.rand(*shape, dtype=cp.float64)
-    )
-
-# Forward FFT.
-# In this example, the forward FFT operand is distributed according
-# to Slab.X distribution. With reshape=False, the FFT result will be
-# distributed according to Slab.Y distribution.
-b = nvmath.distributed.fft.fft(a, distribution=Slab.X, options={"reshape": False})
-
-# Distributed FFT performs computations in-place. The result is stored in the same
-# buffer as operand a. Note, however, that operand b has a different shape (due
-# to Slab.Y distribution).
-if rank == 0:
-    print(f"Shape of a on rank {rank} is {a.shape}")
-    print(f"Shape of b on rank {rank} is {b.shape}")
-
-# Inverse FFT.
-# Recall from previous transform that the inverse FFT operand is distributed according
-# to Slab.Y. With reshape=False, the inverse FFT result will be distributed according
-# to Slab.X distribution.
-c = nvmath.distributed.fft.ifft(b, distribution=Slab.Y, options={"reshape": False})
-
-# The shape of c is the same as a (due to Slab.X distribution). Once again, note that
-# a, b and c are sharing the same symmetric memory buffer (distributed FFT operations
-# are in-place).
-if rank == 0:
-    print(f"Shape of c on rank {rank} is {c.shape}")
-
-# Synchronize the default stream
-with cp.cuda.Device(device_id):
-    cp.cuda.get_current_stream().synchronize()
-
-if rank == 0:
-    print(f"Input type = {type(a)}, device = {a.device}")
-    print(f"FFT output type = {type(b)}, device = {b.device}")
-    print(f"IFFT output type = {type(c)}, device = {c.device}")
-
-# GPU operands on the symmetric heap are not garbage-collected and the user is
-# responsible for freeing any that they own (this deallocation is a collective
-# operation that must be called by all processes at the same point in the execution).
-# All cuFFTMp operations are inplace (a, b, and c share the same memory buffer), so
-# we take care to only free the buffer once.
-nvmath.distributed.free_symmetric_memory(a)
-```
-
 ## License
 
 All files hosted in this repository are subject to the [Apache 2.0](./LICENSE) license.
 
 ## Disclaimer
 
-nvmath-python is in a Beta state. Beta products may not be fully functional, may contain
-errors or design flaws, and may be changed at any time without notice. We appreciate your
-feedback to improve and iterate on our Beta products.
+nvmath-python contains features marked as experimental. Experimental features may not be
+fully functional, may contain errors or design flaws, and may be changed at any time without
+notice. We appreciate your feedback to improve and iterate on our experimental features.

docs/sphinx/_ext/experimental.css

@@ -0,0 +1,75 @@
+/* Experimental API styling */
+
+/* Shared vars for every experimental left bar. */
+:root {
+    --experimental-bar-width: 6px;
+    --experimental-bar-pad: 10px;
+}
+
+/* Add left border to the entire method/function/class/attribute marked experimental. */
+/* This covers both the signature and docstring. */
+/* `attribute` covers dataclass option fields, which napoleon renders
+/* as `.. attribute::` blocks rather than Parameters rows. */
+/* Uses theme's attention color variables that automatically adapt to light/dark mode. */
+/* The negative margin equals border + padding so the content keeps its normal */
+/* alignment with non-experimental siblings; only the bar falls into the left */
+/* gutter. It is derived from the bar constants so it stays in sync. */
+dl.py.method.experimental,
+dl.py.function.experimental,
+dl.py.class.experimental,
+dl.py.attribute.experimental {
+    border-left: var(--experimental-bar-width) solid var(--pst-color-attention-bg);
+    padding-left: var(--experimental-bar-pad);
+    margin-left: calc(-1 * (var(--experimental-bar-width) + var(--experimental-bar-pad)));
+}
+
+/* Style the experimental marker box. */
+/* This is a simple container created by .. experimental:: directive. */
+/* Using CSS variables from the theme's attention admonition style. */
+/* These variables automatically change with the theme switcher. */
+.experimental-marker {
+    background-color: var(--pst-color-attention-bg);
+    padding: 8px 12px;
+    margin: 12px 0;
+    border-radius: 4px;
+}
+
+/* Style the paragraph inside the container. */
+.experimental-marker > p {
+    color: var(--pst-color-attention-text);
+    font-weight: 700;
+    margin: 0;
+    padding-left: 8px;
+    font-size: 0.95em;
+}
+
+/* The module banner names its module in a `code` span. Keep it monospace but */
+/* let it inherit the banner's attention colors instead of the theme's default */
+/* code background/border, which would clash with the banner fill. */
+.experimental-marker > p code {
+    background-color: transparent;
+    color: inherit;
+    border: none;
+    padding: 0;
+}
+
+/* Highlight an experimental parameter row ("Parameters" list) or attribute row */
+/* (":ivar:"/"Variables" list). Both are set by mark_experimental_apis on the */
+/* row's block container -- the <li> for several entries, or the field body <dd> */
+/* when a lone entry collapses the list -- so the bar spans the whole row.
+/* (An attribute written as a ".. attribute::" object instead gets the object
+/* border via dl.py.attribute.experimental above.) */
+.experimental-param,
+.experimental-ivar {
+    border-left: var(--experimental-bar-width) solid var(--pst-color-attention-bg);
+    padding-left: var(--experimental-bar-pad);
+    border-radius: 4px;
+}
+
+/* Whole-module experimental: one bar spanning the entire module page. */
+/* Set by mark_experimental_apis on the `compound` div wrapping the page body. */
+/* Mirrors the per-entry treatment above (border-left + padding-left). */
+.experimental-module {
+    border-left: var(--experimental-bar-width) solid var(--pst-color-attention-bg);
+    padding-left: var(--experimental-bar-pad);
+}