This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 40d297af8 [DOCS] Add HTML diagrams and streamline matrix instruction
guides (#209)
40d297af8 is described below
commit 40d297af8c9b13a5bc890e0e5622a578c21469b2
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Nov 5 23:58:42 2026 -0500
[DOCS] Add HTML diagrams and streamline matrix instruction guides (#209)
This pr replaces the ASCII diagrams in the Warp Matrix Instructions and
Blackwell Tensor Cores guides with HTML diagrams hosted on tvm-site,
with automatic sizing and theme synchronization.
and also streamlines the explanations of register fragments, MMA
execution, and TMEM lifetime, removing details repeated in the diagrams
and code.
---
.../tile/blackwell_tensor_cores.rst | 141 ++++++++++++---------
docs/programming_guides/tile/warp_matrix.rst | 114 +++++++++--------
2 files changed, 138 insertions(+), 117 deletions(-)
diff --git a/docs/programming_guides/tile/blackwell_tensor_cores.rst
b/docs/programming_guides/tile/blackwell_tensor_cores.rst
index 6ed98cbcb..fd5d899a6 100644
--- a/docs/programming_guides/tile/blackwell_tensor_cores.rst
+++ b/docs/programming_guides/tile/blackwell_tensor_cores.rst
@@ -27,38 +27,49 @@ shared memory.
The Accumulator's Lifetime
--------------------------
-The example computes ``Out = A @ B.T`` with A of shape ``(128, 64)`` and
-B of shape ``(64, 64)``. One CTA of 128 threads performs these stages:
-
-.. code-block:: text
-
- allocate TMEM -> stage A/B in shared memory -> issue MMA
- |
- commit + wait for completion
- |
- release TMEM <- write Out <- load TMEM into registers + wait
-
-Different stages require different participants. The conditions in the
-kernel select them explicitly:
-
-.. list-table::
- :header-rows: 1
-
- * - Participants
- - Operations
- * - ``warp == 0``: all 32 threads of warp 0
- - Allocate TMEM, relinquish the allocation permit, and deallocate TMEM.
- * - ``tid == 0``: thread 0
- - Initialize the barrier, issue the MMA instructions, and commit them.
- * - All four warps
- - Stage inputs, synchronize, load the completed result from TMEM, and
write Out.
-
-``mma_shared_layout`` constructs the swizzled shared layouts for A and B.
-Mode 3 selects a 128-byte swizzle; its 1024-byte period determines the
-shared allocation alignment used below.
-``tmem_mma_operand_layout`` maps the float32 accumulator to TMEM lanes
-and columns for the chosen instruction. The ``128 x 64`` accumulator
-uses 64 columns, so the allocation and deallocation both specify 64.
+The example computes ``Out = A @ B.T`` with A of shape ``(128, 64)``,
+B of shape ``(64, 64)``, and Out of shape ``(128, 64)``. One CTA of 128
+threads performs these stages:
+
+.. raw:: html
+
+ <iframe data-tirx-matrix-dataflow
+
src="https://tvm.apache.org/images/docs/tirx/blackwell_tensor_cores.html";
+ title="Blackwell matrix multiplication: TMEM allocation, MMA
completion, and output"
+ style="display: block; width: 100%; height: 740px; border: 0;"
+ loading="lazy"></iframe>
+ <script>
+ // Keep both matrix diagrams sized and themed with their documentation
pages.
+ (() => {
+ const dark = matchMedia('(prefers-color-scheme: dark)');
+
document.querySelectorAll('iframe[data-tirx-matrix-dataflow]').forEach((frame)
=> {
+ if (frame.dataset.ready) return;
+ frame.dataset.ready = 'true';
+ const origin = new URL(frame.src, document.baseURI).origin;
+ const syncTheme = () => {
+ const selected = document.documentElement.dataset.theme;
+ const theme = ['light', 'dark'].includes(selected) ? selected :
(dark.matches ? 'dark' : 'light');
+ frame.contentWindow.postMessage({ type:
'tirx-matrix-dataflow:theme', theme }, origin);
+ };
+ addEventListener('message', (event) => {
+ if (event.source !== frame.contentWindow || event.origin !==
origin) return;
+ const data = event.data;
+ if (data?.type !== 'tirx-matrix-dataflow:height') return;
+ if (!Number.isFinite(data.height) || data.height < 100 ||
data.height > 3000) return;
+ frame.style.height = Math.ceil(data.height) + 'px';
+ });
+ new MutationObserver(syncTheme).observe(document.documentElement, {
+ attributes: true, attributeFilter: ['data-theme']
+ });
+ dark.addEventListener('change', syncTheme);
+ frame.addEventListener('load', syncTheme);
+ syncTheme();
+ });
+ })();
+ </script>
+
+The kernel uses ``d`` for the TMEM accumulator and ``frag`` for the output
+values held in each thread's registers:
.. code-block:: python
@@ -133,35 +144,38 @@ uses 64 columns, so the allocation and deallocation both
specify 64.
T.ptx.tcgen05.relinquish_alloc_permit.cta_group__1.sync.aligned()
T.ptx.tcgen05.dealloc.cta_group__1.sync.aligned.b32(address[0],
T.uint32(64))
-The allocation writes a TMEM base address to the shared ``address`` slot.
-After the CTA barrier publishes the address, ``decl_tmem`` gives the allocated
-storage a tensor shape and layout. It does not perform another allocation.
-
-Each thread stages input values in shared memory and executes a proxy
-fence. The following CTA barrier lets the issuing thread proceed only
-after all writers have finished this handoff to the asynchronous reader.
-The A/B views split K=64 into four K=16 slices. For example,
-``av`` has shape ``(1, 4, 128, 16)``: one M tile, four K repeats, and a
-``128 x 16`` A slice per issue. B likewise has one N tile and four K
-repeats. Consequently, the tile call issues four MMA instructions.
-``accum=False`` starts a fresh result on the first slice; subsequent
-slices accumulate into D.
-
-The elected thread commits these operations to ``mbar``, initialized with
-one expected arrival. All threads wait for phase 0 and execute
-``tcgen05.fence::after_thread_sync`` before reading TMEM. This completion
-protocol uses a barrier arrival, rather than the expected-byte count
-used by :doc:`TMA copies <tma>`.
-
-The register fragment assigns one output row to each thread. Its
-``(4, 32, 64)`` view lets each of the four warps load its 32 TMEM lanes.
-``tcgen05.wait::ld`` then makes the loaded registers ready for the output
-stores. It waits for these loads; the preceding barrier and fence are
-what establish that MMA has finished. The final CTA barrier keeps TMEM
-alive until every warp has finished using it, after which warp 0 gives
-up further allocations and releases the columns.
-
-On a supported GPU, we can check the product as follows:
+All 32 threads of warp 0 collectively allocate 64 TMEM columns for the
+``128 x 64`` float32 accumulator. The allocation writes its base address
+to ``address[0]``.
+
+The inputs are staged in shared-memory buffers ``a`` and ``b``.
+``mma_shared_layout`` mode ``3`` selects a 128-byte swizzle with
+1024-byte base alignment. Each thread executes ``fence.proxy.async``
+after its copies; the CTA barrier then ensures all inputs are ready
+for MMA and the TMEM address is visible to all threads.
+``T.cuda.decl_tmem`` binds that allocation to tensor ``d``.
+
+The views ``av`` and ``bv`` split K=64 into four K=16 slices, so thread 0's
+single ``T.cuda.tile.apply`` call expands into four MMA instructions.
+With ``accum=False``, the first instruction writes its product into ``d``;
+the next three add their products to it.
+
+Thread 0 calls ``tcgen05.commit`` to signal ``mbar`` once all four MMA
+instructions complete, matching the barrier's arrival count of 1.
+All threads wait for phase 0, then execute
+``tcgen05.fence::after_thread_sync`` before loading the result from TMEM.
+
+``register_layout`` assigns one output row of 64 float32 values to each
+thread. All four warps use ``tcgen05.ld`` to load their rows into ``frag``.
+``tcgen05.wait::ld`` waits for these loads to complete before each thread
+writes its row to ``Out``.
+
+The final CTA barrier ensures every warp has finished using TMEM.
+Warp 0 then collectively relinquishes the CTA's allocation permit and
+deallocates the 64 columns.
+
+On a supported GPU, we can compare the result with a PyTorch float32
+matrix product:
.. code-block:: python
@@ -202,6 +216,11 @@ have different participants and completion rules:
- One elected thread issues the copy and commits it to an
``mbarrier``; consumers wait and use the thread-sync fence.
+For ``tcgen05.mma`` and ``tcgen05.cp``, a commit signals completion through
+a barrier arrival. :doc:`TMA loads <tma>` instead track completion using
+an expected-byte count. The barrier setup must match the operation being
+tracked.
+
``tcgen05.cp`` is useful for operands that must reside in TMEM, such as
the scale factors of a block-scaled MMA. Its shared source must be ready
for asynchronous access, just as the MMA inputs above must be. The
diff --git a/docs/programming_guides/tile/warp_matrix.rst
b/docs/programming_guides/tile/warp_matrix.rst
index bb5225099..ca035f7d5 100644
--- a/docs/programming_guides/tile/warp_matrix.rst
+++ b/docs/programming_guides/tile/warp_matrix.rst
@@ -30,46 +30,46 @@ From Shared Memory to a Matrix Product
All 32 lanes cooperate on one ``16 x 8`` output tile. The data follows
this path:
-.. code-block:: text
-
- float16 inputs -> shared memory -> ldmatrix -> register fragments
- |
- mma.sync
- |
- float32 registers -> Out
-
-Each lane holds part of a matrix, called its register fragment. In this
-example, the fragments have the following sizes:
-
-.. list-table::
- :header-rows: 1
-
- * - Registers in each lane
- - Contents
- - Complete matrix across the warp
- * - ``ar``: four uint32 words
- - Eight float16 elements, packed in pairs
- - A: ``(16, 16)``
- * - ``br``: two uint32 words
- - Four float16 elements, packed in pairs
- - B: ``(8, 16)``
- * - ``dr``: four float32 values
- - Four output elements
- - Out: ``(16, 8)``
-
-The instruction fixes which elements belong to each lane. The
-``ldmatrix`` views use coordinates ``(matrix, row, packed_word)``: each
-``8 x 8`` half matrix has eight rows of four words. A contains four such
-matrices and B contains two, giving view shapes ``(4, 8, 4)`` and
-``(2, 8, 4)``. The load layouts place their words in ``ar`` and ``br``.
-
-The ``*_mma_layout`` views describe those same registers with MMA operand
-coordinates. For example, A's ``(1, 1, 16, 8)`` view has one M repeat,
-one K repeat, 16 rows, and eight packed words per row. These views let
-both instructions use the same fragments without moving values between
-lanes. The
-:ref:`warp MMA reference <cuda-tile-contract-mma-sync>` gives the fixed
-lane mappings represented by these layouts.
+.. raw:: html
+
+ <iframe data-tirx-matrix-dataflow
+ src="https://tvm.apache.org/images/docs/tirx/warp_matrix.html";
+ title="Warp matrix multiplication: shared memory, register
fragments, and output"
+ style="display: block; width: 100%; height: 600px; border: 0;"
+ loading="lazy"></iframe>
+ <script>
+ // Keep both matrix diagrams sized and themed with their documentation
pages.
+ (() => {
+ const dark = matchMedia('(prefers-color-scheme: dark)');
+
document.querySelectorAll('iframe[data-tirx-matrix-dataflow]').forEach((frame)
=> {
+ if (frame.dataset.ready) return;
+ frame.dataset.ready = 'true';
+ const origin = new URL(frame.src, document.baseURI).origin;
+ const syncTheme = () => {
+ const selected = document.documentElement.dataset.theme;
+ const theme = ['light', 'dark'].includes(selected) ? selected :
(dark.matches ? 'dark' : 'light');
+ frame.contentWindow.postMessage({ type:
'tirx-matrix-dataflow:theme', theme }, origin);
+ };
+ addEventListener('message', (event) => {
+ if (event.source !== frame.contentWindow || event.origin !==
origin) return;
+ const data = event.data;
+ if (data?.type !== 'tirx-matrix-dataflow:height') return;
+ if (!Number.isFinite(data.height) || data.height < 100 ||
data.height > 3000) return;
+ frame.style.height = Math.ceil(data.height) + 'px';
+ });
+ new MutationObserver(syncTheme).observe(document.documentElement, {
+ attributes: true, attributeFilter: ['data-theme']
+ });
+ dark.addEventListener('change', syncTheme);
+ frame.addEventListener('load', syncTheme);
+ syncTheme();
+ });
+ })();
+ </script>
+
+A register fragment is the part of a matrix held in one lane's registers.
+The kernel below uses ``ar`` and ``br`` for the input fragments and ``dr``
+for the output fragment:
.. code-block:: python
@@ -126,22 +126,24 @@ lane mappings represented by these layouts.
for r in range(4):
Out[lane // 4 + 8 * (r // 2), 2 * (lane % 4) + r % 2] = dr[r]
-Each thread first copies part of A and B to shared memory. The CTA barrier
-ensures those stores are visible before ``ldmatrix`` reads them. The
-``uint32`` views pack two float16 elements into each instruction operand;
-they preserve the bits rather than convert the values to integers.
-
-``ldmatrix.x4`` loads the four A matrices, and ``ldmatrix.x2`` loads the
-two B matrices. Both instructions require all 32 lanes to participate.
-
-The MMA call takes ``D, A, B, C`` in that order. Its views include repeat
-axes of length one before the matrix dimensions. The selected
-``m16n8k16`` instruction computes the whole product in one issue, with
-``T.float32(0)`` as the initial C value. Each lane receives four output
-elements in ``dr``; the final loop writes them to their matrix positions.
-For lane 0, ``dr[0]`` through ``dr[3]`` belong to ``Out[0, 0]``,
-``Out[0, 1]``, ``Out[8, 0]``, and ``Out[8, 1]``. The other lanes use the
-same coordinate formula to cover the remaining output elements.
+The threads copy ``A`` and ``B`` into shared-memory buffers ``a`` and ``b``,
+then synchronize before ``ldmatrix`` reads them.
+
+The ``uint32`` views reinterpret pairs of float16 values as 32-bit words.
+``rearrange`` expresses the shared-memory views in
+``(matrix, row, packed_word)`` coordinates for ``ldmatrix``.
+All 32 lanes participate in ``ldmatrix.x4`` to load A into ``ar`` and
+``ldmatrix.x2`` to load B into ``br``.
+
+The load and MMA layouts describe the same register fragments in each
+instruction's coordinates. This lets ``mma.sync`` reuse the loaded values
+without moving them between lanes. The :ref:`warp MMA reference
+<cuda-tile-contract-mma-sync>` gives the fixed lane mappings.
+
+The MMA call takes ``D, A, B, C``, with ``T.float32(0)`` as the initial C
+value. The ``m16n8k16`` instruction computes the product in one issue,
+leaving four float32 results per lane in ``dr``. The final loop stores
+them at their corresponding positions in ``Out``.
We can compile the kernel and compare its result with a float32 matrix
product:
