CUTLASS 3.1 (#915)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>

python/docs_src/source/_templates/layout.html

@@ -0,0 +1,94 @@
+{% extends "!layout.html" %}
+  {% block sidebartitle %} {{ super() }}
+
+  <style>
+    /* Sidebar header (and topbar for mobile) */
+    .wy-side-nav-search, .wy-nav-top {
+      background: #76b900;
+    }
+
+    .wy-menu > p > span.caption-text {
+      color: #76b900;
+    }
+
+    .wy-menu-vertical p {
+      height: 32px;
+      line-height: 32px;
+      padding: 0 1.618em;
+      margin: 12px 0 0;
+      display: block;
+      font-weight: 700;
+      text-transform: uppercase;
+      font-size: 85%;
+      white-space: nowrap;
+    }
+
+    .wy-side-nav-search a:link, .wy-nav-top a:link {
+      color: #fff;
+    }
+    .wy-side-nav-search a:visited, .wy-nav-top a:visited {
+      color: #fff;
+    }
+    .wy-side-nav-search a:hover, .wy-nav-top a:hover {
+      color: #fff;
+    }
+
+    .wy-menu-vertical a:link, .wy-menu-vertical a:visited {
+      color: #d9d9d9
+    }
+
+    .wy-menu-vertical a:active {
+      background-color: #76b900
+    }
+
+    .wy-side-nav-search>div.version {
+      color: rgba(0, 0, 0, 0.3)
+    }
+
+    .wy-nav-content {
+      max-width: 1000px;
+    }
+
+    /* override table width restrictions */
+    .wy-table-responsive table td, .wy-table-responsive table th {
+        /* !important prevents the common CSS stylesheets from
+          overriding this as on RTD they are loaded after this stylesheet */
+        white-space: normal !important;
+    }
+
+    .wy-table-responsive {
+        overflow: visible !important;
+    }
+
+  </style>
+  {% endblock %}
+
+  {% block footer %} {{ super() }}
+
+  <style>
+  a:link, a:visited {
+    color: #76b900;
+  }
+
+  a:hover {
+    color: #8c0;
+  }
+
+  html.writer-html4 .rst-content dl:not(.docutils)>dt, html.writer-html5 .rst-content dl[class]:not(.option-list):not(.field-list):not(.footnote):not(.citation):not(.glossary):not(.simple)>dt {
+    background: rgba(118, 185, 0, 0.1);
+    color: rgba(59,93,0,1);
+    border-top: solid 3px rgba(59,93,0,1);
+  }
+
+  html.writer-html4 .rst-content dl:not(.docutils) .property, html.writer-html5 .rst-content dl[class]:not(.option-list):not(.field-list):not(.footnote):not(.glossary):not(.simple) .property {
+    text-transform: capitalize;
+    display: inline-block;
+    padding-right: 8px;
+  }
+  </style>
+
+  {%- if nvidia_analytics_id %}
+  <script type="text/javascript">_satellite.pageBottom();</script>
+  {%- endif %}
+
+  {% endblock %}

python/docs_src/source/conf.py

@@ -0,0 +1,100 @@
+# 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
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath('..'))
+sys.path.insert(0, os.path.abspath('../..'))
+sys.path.insert(0, os.path.abspath('../../media/docs'))
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = 'CUTLASS Python interface'
+copyright = '2023, NVIDIA'
+author = 'NVIDIA'
+release = '3.1.0'
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+        'myst_parser',
+        'nbsphinx',
+        'nbsphinx_link', 
+        'sphinx_copybutton',
+        'sphinx.ext.autodoc',
+        'sphinx.ext.autosectionlabel',
+        'sphinx.ext.autosummary',
+        'sphinx.ext.coverage',
+        'sphinx.ext.extlinks',
+        'sphinx.ext.ifconfig',
+        'sphinx.ext.intersphinx',
+        'sphinx.ext.mathjax',
+        'sphinx.ext.napoleon',
+        'sphinx.ext.viewcode',
+        'sphinx_inline_tabs',
+        ]
+
+source_suffix = {
+    '.rst': 'restructuredtext',
+    '.md': 'markdown',
+}
+
+autodoc_typehints = 'description'
+
+pygments_style = "sphinx"
+pygments_dark_style = "monokai"
+
+templates_path = ['_templates']
+exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+
+# Ignore errors when converting notebooks
+nbsphinx_allow_errors = True
+
+language = 'en'
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_static_path = ['_static']
+
+html_title = "CUTLASS Python"
+html_baseurl = 'docs'
+html_theme = 'furo'
+html_theme_options = {
+	"light_logo": "cutlass-logo-small.png",
+	"dark_logo": "cutlass-logo-small.png",
+    "light_css_variables": {
+        "color-brand-primary": "#76B900",
+        "color-brand-content": "#76B900",
+    },
+    "dark_css_variables": {
+        "color-brand-primary": "#76B900",
+        "color-brand-content": "#76B900",
+    },
+    "footer_icons": [
+        {
+            "name": "GitHub",
+            "url": "https://github.com/NVIDIA/cutlass",
+            "html": """
+                <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 16 16">
+                    <path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
+                </svg>
+            """,
+            "class": "",
+        },
+    ],
+}

python/docs_src/source/contribute.md

@@ -0,0 +1,9 @@
+# Contributing
+
+Thank you for your interest in contributing to the CUTLASS Python interface. Based on the type of contribution, it will fall into two categories:
+
+1. You want to report a bug, feature request, or documentation issue
+    - File an [issue](https://github.com/NVIDIA/cutlass/issues/new/choose) describing what you encountered or what you want to see changed.
+    - The CUTLASS team will evaluate the issues and triage them, scheduling them for a release. If you believe the issue needs priority attention, comment on the issue to notify the team.
+2. You want to implement a feature or bug-fix
+    - We welcome contributions from the community. We recommend that you contribute via a [pull request](https://github.com/NVIDIA/cutlass/pulls). If you have questions about CUTLASS, consider asking a question via the [Discussions](https://github.com/NVIDIA/cutlass/discussions) tab. Please be sure to search through both existing issues and discussions to see whether your question has already been answered.

python/docs_src/source/cutlass.emit.rst

@@ -0,0 +1,18 @@
+Emitters
+========
+
+Common
+------
+
+.. automodule:: cutlass.emit.common
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+PyTorch
+-------
+
+.. automodule:: cutlass.emit.pytorch
+   :members:
+   :undoc-members:
+   :show-inheritance:

python/docs_src/source/cutlass.op.rst

@@ -0,0 +1,26 @@
+Operations
+==========
+
+GEMM
+----
+
+.. automodule:: cutlass.op.gemm
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Grouped GEMM
+------------
+
+.. automodule:: cutlass.op.gemm_grouped
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Operation
+---------
+
+.. automodule:: cutlass.op.op
+   :members:
+   :undoc-members:
+   :show-inheritance:

python/docs_src/source/cutlass.rst

@@ -0,0 +1,36 @@
+CUTLASS
+=======
+
+Subpackages
+-----------
+
+.. toctree::
+   :maxdepth: 1
+
+   cutlass.emit
+   cutlass.op
+   cutlass.utils
+
+Epilogue
+--------
+
+.. automodule:: cutlass.epilogue
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Library Defaults
+----------------
+
+.. automodule:: cutlass.library_defaults
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Swizzle
+----------
+
+.. automodule:: cutlass.swizzle
+   :members:
+   :undoc-members:
+   :show-inheritance:

python/docs_src/source/cutlass.utils.rst

@@ -0,0 +1,18 @@
+Utilities
+=========
+
+Checks
+------
+
+.. automodule:: cutlass.utils.check
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Data Types
+----------
+
+.. automodule:: cutlass.utils.datatypes
+   :members:
+   :undoc-members:
+   :show-inheritance:

python/docs_src/source/examples.rst

@@ -0,0 +1,9 @@
+Examples
+==================
+
+.. toctree::
+   :maxdepth: 5
+
+   Basic GEMM <externals/00_basic_gemm.nblink>
+   Epilogue <externals/01_epilogue.nblink>
+   PyTorch Extension <externals/02_pytorch_extension_grouped_gemm.nblink>

python/docs_src/source/externals/00_basic_gemm.nblink

@@ -0,0 +1,3 @@
+{
+  "path": "./../../../../examples/python/00_basic_gemm.ipynb"
+}

python/docs_src/source/externals/01_epilogue.nblink

@@ -0,0 +1,3 @@
+{
+  "path": "./../../../../examples/python/01_epilogue.ipynb"
+}

python/docs_src/source/externals/02_pytorch_extension_grouped_gemm.nblink

@@ -0,0 +1,3 @@
+{
+  "path": "./../../../../examples/python/02_pytorch_extension_grouped_gemm.ipynb"
+}

python/docs_src/source/index.rst

@@ -0,0 +1,55 @@
+.. CUTLASS Python interface documentation master file, created by
+   sphinx-quickstart on Mon Feb 13 17:57:39 2023.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+.. include:: ../../README.md
+   :start-line: 1
+   :parser: markdown
+
+.. toctree::
+   :hidden:
+
+   Home <self>
+
+.. toctree::
+   :hidden:
+   :caption: Getting Started:
+
+   install.md
+   Getting Started <externals/00_basic_gemm.nblink>
+   contribute.md
+
+.. toctree::
+   :hidden:
+   :caption: Python Documentation:
+
+   modules.rst
+
+.. toctree::
+   :hidden:
+   :caption: Examples and Tutorials:
+
+   examples.rst
+
+.. toctree::
+   :hidden:
+   :caption: Advanced:
+
+.. toctree::
+   :hidden:
+   :caption: FAQ:
+
+.. toctree::
+   :hidden:
+   :caption: Reference:
+
+   Github <https://github.com/NVIDIA/cutlass>
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

python/docs_src/source/install.md

@@ -0,0 +1,37 @@
+# Installation
+
+## Installing from source
+
+Installing from source requires the latest CUDA Toolkit that matches the major.minor of CUDA Python installed.
+
+Prior to installing the CUTLASS Python interface, one may optionally set the following environment variables:
+* `CUTLASS_PATH`: the path to the cloned CUTLASS repository
+* `CUDA_INSTALL_PATH`: the path to the installation of CUDA
+
+If these environment variables are not set, the installation process will infer them to be the following:
+* `CUTLASS_PATH`: one directory level above the current directory (i.e., `$(pwd)/..`)
+* `CUDA_INSTALL_PATH`: the directory holding `/bin/nvcc` for the first version of `nvcc` on `$PATH` (i.e., `which nvcc | awk -F'/bin/nvcc' '{print $1}'`)
+
+**NOTE:** The version of `cuda-python` installed must match the CUDA version in `CUDA_INSTALL_PATH`.
+
+### Installing a developer-mode package
+The CUTLASS Python interface can currently be installed via:
+```bash
+python setup.py develop --user
+```
+This will allow changes to the Python interface source to be reflected when using the Python interface.
+
+We plan to add support for installing via `python setup.py install` in a future release.
+
+## Docker
+To ensure that you have all of the necessary Python modules for running the examples using the
+CUTLASS Python interface, we recommend using one of the Docker images for CUDA [11.8](../../../python/docker/Dockerfile-cuda11.8-pytorch)
+and [12.0](../../../python/docker/Dockerfile-cuda12.0-pytorch) are located in the docker directory.
+
+For example, to build and launch a container that uses CUDA 12.0 via an NGC PyTorch container, run:
+```bash
+docker build -t cutlass-cuda12.0:latest -f docker/Dockerfile-cuda12.0-pytorch .
+docker run --gpus all -it --rm cutlass-cuda12.0:latest
+```
+
+The CUTLASS Python interface has been tested with CUDA 11.8 and CUDA 12.0 on Python 3.8.10 and 3.9.7.

python/docs_src/source/modules.rst

@@ -0,0 +1,7 @@
+CUTLASS Python API
+==================
+
+.. toctree::
+   :maxdepth: 5
+
+   cutlass