## Motivation and context
`tcgen05.ld/st` with a `.16x*b` atom accesses one 16-row half-slab from
each 32-row TMEM partition owned by a warp. Across a four-warp
warpgroup, the two physical half-slabs are:
| View | Physical TMEM lanes | PTX row immediate |
|---|---|---|
| lower half | `0..15, 32..47, 64..79, 96..111` | `0` |
| upper half | `16..31, 48..63, 80..95, 112..127` | `16` |
Layout D for an M=128 accumulator occupies both halves. Layout F exposes
an M=64 logical tile over one half, which is useful both for a native
M=64 accumulator and for reading either 64-row half of an existing
Layout D accumulator.
Before this PR, `tmem_datapath_layout("F", 64, cols)` could describe
only the lower half. The copy dispatcher classified a TMEM buffer only
as `"D"` or `"F"` and every M=64 `.16x*b` operation started at `row=0`.
As a result, there was no layout-preserving way to create a recognized
64-row view of the upper half of a Layout D accumulator: a normal Layout
F view still addressed the lower half, while a hand-written `+16@TLane`
layout was not recognized by the dispatcher.
The half-slab selection belongs in the buffer layout because layout is
the source of truth for physical placement in TIRx. It should not be an
out-of-band `copy_async` option. This PR therefore records the selection
in Layout F, carries it through datapath classification, and derives the
PTX row immediate from it.
A related invariant is that a default M=128 TMEM allocation must be
structurally identical to named Layout D. The dispatcher recognizes
datapaths structurally, so keeping a separate hand-written default
layout creates an unnecessary drift risk. This PR makes the default call
the public Layout D factory directly.
Finally, the register-side `tcgen05_atom_layout` must agree with the PTX
mapping from a logical `(row, col)` to `(laneid, wid_in_wg, register)`.
A self-consistent load/store round trip is not enough to prove that
mapping: raw PTX can move the same bits back even when the logical
layout label is wrong. Elementwise dispatch does consume that label, so
this PR adds direct mapping and elementwise compilation fences for all
`.16x*b` atom families. These are coverage additions; the production
atom-layout construction itself is unchanged here.
## Changes
- Add `sub_slab={0,1}` to `tmem_datapath_layout("F", ...)`.
- Encode the upper view as a `+16@TLane` offset and reject invalid
selectors, including nonzero selectors for Layout D.
- Classify TMEM layouts as `(datapath, sub_slab)` and emit `.16x*b` with
`row=(sub_slab + slab) * 16`.
- Preserve the existing M=128 behavior: Layout D with a 128-row `.16x*b`
fragment still emits two operations at rows `0` and `16`.
- Build the default M=128 TMEM layout through `tmem_datapath_layout("D",
...)`.
- Document the physical lane mapping and supported datapath/atom
combinations.
- Add direct atom-layout mapping coverage and a warpgroup elementwise
regression.
## Testing
- Static Layout F checks cover every logical row for both `sub_slab=0`
and `sub_slab=1`.
- B200 readback tests populate one Layout D accumulator, then verify
that lower and upper Layout F views reproduce the two corresponding
register halves for `.16x64b`, `.16x128b`, and `.16x256b`.
- Negative tests cover invalid sub-slab values and incompatible
datapath/atom pairings.
- Direct `(row, col) -> (laneid, wid_in_wg, register)` sweeps cover
supported `.16x*b` shapes and repetitions.
- Warpgroup elementwise codegen verifies that an atom-layout fragment
canonicalizes and slices correctly.
- Changed-files pre-commit checks.
TVM Documentation
This folder contains the source of TVM's documentation, hosted at https://tvm.apache.org/docs
Build Locally
With Docker (recommended)
-
Build TVM and the docs inside the tlcpack/ci-gpu image using the
ci.pyscript.# If this runs into errors, try cleaning your 'build' directory python tests/scripts/ci.py docs # See other doc building options python tests/scripts/ci.py docs --help -
Serve the docs and visit http://localhost:8000 in your browser
# Run an HTTP server you can visit to view the docs in your browser python tests/scripts/ci.py serve-docs
Native
-
Build TVM first in the repo root folder, then make it importable:
export TVM_HOME=/path-to-tvm export TVM_LIBRARY_PATH=$TVM_HOME/build pip install --target=$TVM_HOME/python $TVM_HOME/3rdparty/tvm-ffi export PYTHONPATH=$TVM_HOME/python:$PYTHONPATHdocs/conf.pyunconditionally importstvmat startup, so the build will fail immediately if TVM is not importable. -
Install dependencies
# Pillow on Ubuntu may require libjpeg-dev from apt ./docker/bash.sh ci_gpu -c \ 'python3 -m pip install --quiet sphinx-book-theme==1.1.4 && python3 -m pip freeze' > frozen-requirements.txt pip install -r frozen-requirements.txt -
Generate the docs
# TVM_TUTORIAL_EXEC_PATTERN=none skips the tutorial execution to the build # work on most environments (e.g. MacOS). export TVM_TUTORIAL_EXEC_PATTERN=none cd docs make html -
Run an HTTP server and visit http://localhost:8000 in your browser
cd _build/html && python3 -m http.server
Only Execute Specified Tutorials
The document build process will execute all the tutorials in the sphinx gallery.
This will cause failure in some cases when certain machines do not have necessary
environment. You can set TVM_TUTORIAL_EXEC_PATTERN to only execute
the path that matches the regular expression pattern.
For example, to only build tutorials under /get_started/tutorials, run
python tests/scripts/ci.py docs --tutorial-pattern=/get_started/tutorials
To only build one specific file, do
# The slash \ is used to get . in regular expression
python tests/scripts/ci.py docs --tutorial-pattern=file_name\.py
Helper Scripts
The following script mirrors the CI docs pipeline: it runs a sphinx pre-check
(when not running locally) and then performs a full make htmldepoly build,
including tutorial execution. You will need a GPU CI environment.
tests/scripts/task_python_docs.sh
To build docs locally without executing tutorials (fastest local iteration):
cd docs && TVM_TUTORIAL_EXEC_PATTERN=none make html
Note: the sphinx pre-check (warning validation) only runs in CI (IS_LOCAL=0).
python tests/scripts/ci.py docs always sets IS_LOCAL=1 and skips the
pre-check regardless of other flags.
To run the full build including tutorial executions:
python tests/scripts/ci.py docs --full
Define the Order of Tutorials
You can define the order of tutorials with subsection_order and
within_subsection_order in conf.py.
By default, the tutorials within one subsection are sorted by filename.
Google Colab Integration
All the TVM tutorials can be opened and used interactively in Google Colab by
clicking the button at the top of the page. To do this, sphinx-gallery builds
.ipynb files from each tutorial, which are automatically deployed to the
apache/tvm-site repo's
asf-site branch by @tvm-bot.
To make sure your tutorial runs correctly on Colab, any non-Python parts of
the tutorial (e.g. dependency installations) should be prefixed by an
IPython magic command.
These will not be included in the built HTML file. For example, to install
Pytorch in your tutorial, add a ReStructured Text block like the following:
######################################################################
# To run this tutorial, we must install PyTorch:
#
# .. code-block:: bash
#
# %%shell
# pip install torch
#
Interactive Bash Scripts
In stock IPython, the %%bash magic command should be used to run shell
commands. However, this command does not give real-time output - the
tutorial's user will not see any output until the entire cell finishes
running. When running commands that take several minutes (e.g. installing
dependencies), this is annoying.
Luckily, Google Colab has the %%shell magic command that does the same
thing as %%bash, but gives output in real time. This command is specific
to Colab, and its source code
is public. Thus, %%shell should be used instead of %%bash when writing
TVM tutorials.