This is an automated email from the ASF dual-hosted git repository.
tqchen 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 391dec40a [DOCS] Clarify programming guides and CUDA tile examples
(#201)
391dec40a is described below
commit 391dec40a5cd667c7a40b7e963af3f7e7c411fba
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Nov 5 09:17:44 2026 -0500
[DOCS] Clarify programming guides and CUDA tile examples (#201)
Clarify the Programming Guides with consistent type notation and more
explicit examples.
- Standardize dtype arguments on `T.*` types and show `T.float32x4`
explicitly in the vector-access example.
- Add a `break`/`continue` example and improve introductory text.
- Explain the differences between `tile.apply` and `tile.ewise`, their
operand shapes, and equivalent add-one implementations.
---
docs/programming_guides/control_flow.rst | 25 ++--
docs/programming_guides/cuda_ptx_interfaces.rst | 12 +-
docs/programming_guides/data_layout.rst | 6 +-
docs/programming_guides/metaprogramming.rst | 9 +-
docs/programming_guides/tile/arithmetic.rst | 18 +--
docs/programming_guides/tile/async_copies.rst | 6 +-
.../tile/blackwell_tensor_cores.rst | 24 ++--
docs/programming_guides/tile/data_movement.rst | 10 +-
docs/programming_guides/tile/index.rst | 22 ++--
docs/programming_guides/tile/tile_patterns.rst | 127 +++++++++++++++------
docs/programming_guides/tile/tma.rst | 8 +-
docs/programming_guides/tile/warp_matrix.rst | 20 ++--
docs/programming_guides/types_tensors_memory.rst | 48 ++++----
examples/programming_guides/control_flow.py | 34 ++++++
14 files changed, 228 insertions(+), 141 deletions(-)
diff --git a/docs/programming_guides/control_flow.rst
b/docs/programming_guides/control_flow.rst
index e963e8f4b..b3a070a3b 100644
--- a/docs/programming_guides/control_flow.rst
+++ b/docs/programming_guides/control_flow.rst
@@ -18,18 +18,13 @@
Control Flow
============
-.. note::
-
- All the code in this guide is in
- `examples/programming_guides/control_flow.py
-
<https://github.com/apache/tvm-tirx/blob/main/examples/programming_guides/control_flow.py>`_.
-
This guide shows how to write loops and conditions in TIRx kernels. TIRx uses
Python's own control flow syntax: ``if``, ``elif``, ``else``, ``for``,
``while``, ``break``, and ``continue`` all mean what they mean in Python, and
each thread executes them independently with its own values.
-The examples use the following imports and a symbolic length ``n``:
+The kernels below illustrate these control-flow constructs. They use the
+following imports and a symbolic tensor length ``n``:
.. literalinclude:: ../../examples/programming_guides/control_flow.py
:language: python
@@ -39,8 +34,8 @@ The examples use the following imports and a symbolic length
``n``:
``for`` Loops
-------------
-A ``for`` loop over ``range(...)`` runs its body once per iteration, in order,
-within each thread. ``T.unroll(...)`` iterates in the same way but asks the
+A ``for`` loop over ``range(...)`` executes its iterations sequentially within
+each thread. ``T.unroll(...)`` iterates in the same way but asks the
compiler to expand the loop into one copy of its body per iteration. The
following kernel uses both to sum each row of an ``n x 64`` matrix, one row
per thread:
@@ -93,6 +88,18 @@ local scalars, so they keep their updates from one iteration
to the next.
iteration, just as in Python. Both work in ``for`` and ``while`` loops, and
each thread decides on its own when to break or continue.
+The following kernel scans each row, skips negative values, and stops at the
+first zero. Each thread sums the positive values encountered in its own row:
+
+.. literalinclude:: ../../examples/programming_guides/control_flow.py
+ :language: python
+ :start-after: # [break_continue.begin]
+ :end-before: # [break_continue.end]
+
+For a row ``[2, -1, 3, 0, 9, 9, 9, 9]``, ``continue`` skips ``-1`` and
+``break`` stops the loop at ``0``, so the thread writes ``5``. The store
+``B[row] = total`` still executes after ``break``; only the loop ends.
+
``if`` Statements
-----------------
diff --git a/docs/programming_guides/cuda_ptx_interfaces.rst
b/docs/programming_guides/cuda_ptx_interfaces.rst
index 29f59cefb..1706542c5 100644
--- a/docs/programming_guides/cuda_ptx_interfaces.rst
+++ b/docs/programming_guides/cuda_ptx_interfaces.rst
@@ -57,11 +57,11 @@ For example, a warp can sum its 32 input values with
from tvm_tirx.script import tirx as T
@T.function
- def warp_reduce(A: T.Tensor((32,), "float32", align=16)):
+ def warp_reduce(A: T.Tensor((32,), T.float32, align=16)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
lane_id = T.cuda.lane_id()
- v = T.alloc_local((1,), "float32")
- i = T.alloc_local((1,), "int32")
+ v = T.alloc_local((1,), T.float32)
+ i = T.alloc_local((1,), T.int32)
v[0] = A[lane_id]
i[0] = 16
while i[0] >= 1:
@@ -100,7 +100,7 @@ This kernel uses ``add.rn.f32`` to add two float32 values
in each thread;
):
with T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32)):
tid = T.cuda.thread_idx("x")
- result = T.local_scalar("float32")
+ result = T.local_scalar(T.float32)
T.ptx.add.rn.f32(result, A[tid], B[tid])
C[tid] = result
@@ -136,10 +136,10 @@ otherwise:
@T.function
- def relu_cuda(A: T.Tensor((256,), "float32"), B: T.Tensor((256,),
"float32")):
+ def relu_cuda(A: T.Tensor((256,), T.float32), B: T.Tensor((256,),
T.float32)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=256))
tx = T.cuda.thread_idx("x")
- B[tx] = T.cuda.func_call("my_relu", A[tx], SRC, ty="float32")
+ B[tx] = T.cuda.func_call("my_relu", A[tx], SRC, ty=T.float32)
The generated CUDA source includes ``my_relu`` and calls it for each
element:
diff --git a/docs/programming_guides/data_layout.rst
b/docs/programming_guides/data_layout.rst
index 33118e362..7f40edf9a 100644
--- a/docs/programming_guides/data_layout.rst
+++ b/docs/programming_guides/data_layout.rst
@@ -77,10 +77,10 @@ add-one kernel launches one warp and gives each lane four
values:
from tvm_tirx.script import tirx as T
@T.function
- def add_one_layout(A: T.Tensor((32, 4), "float32"), B: T.Tensor((32, 4),
"float32")):
+ def add_one_layout(A: T.Tensor((32, 4), T.float32), B: T.Tensor((32, 4),
T.float32)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
lane = T.cuda.lane_id()
- fragment = T.alloc_local((32, 4), "float32", layout=fragment_layout)
+ fragment = T.alloc_local((32, 4), T.float32, layout=fragment_layout)
lane_values = fragment.view(4, layout="default")
for j in T.unroll(4):
lane_values[j] = A[lane, j] + 1
@@ -145,7 +145,7 @@ Inside a device region, attach the layout when allocating
shared memory:
.. code-block:: python
- shared = T.alloc_shared((8, 64), "float16", layout=swizzled, align=1024)
+ shared = T.alloc_shared((8, 64), T.float16, layout=swizzled, align=1024)
The eight 128-byte rows form one 1024-byte swizzle period, matching the
allocation's alignment. Instructions such as TMA and MMA must interpret
diff --git a/docs/programming_guides/metaprogramming.rst
b/docs/programming_guides/metaprogramming.rst
index 086b0cc6a..fb987d7ca 100644
--- a/docs/programming_guides/metaprogramming.rst
+++ b/docs/programming_guides/metaprogramming.rst
@@ -18,12 +18,6 @@
Metaprogramming
===============
-.. note::
-
- All the code in this guide is in
- `examples/programming_guides/metaprogramming.py
-
<https://github.com/apache/tvm-tirx/blob/main/examples/programming_guides/metaprogramming.py>`_.
-
Metaprogramming means writing code that constructs or specializes other code.
In TIRx, ordinary Python functions can generate kernel definitions, inline
helpers can insert reusable statements, and compile-time parameters can
@@ -34,7 +28,8 @@ This guide shows these three techniques: capturing Python
values in kernel
generators, reusing code with ``@T.inline``, and specializing ``@T.jit``
functions.
-The examples use the following imports:
+To define and specialize kernels using these techniques, start with the
+following imports:
.. literalinclude:: ../../examples/programming_guides/metaprogramming.py
:language: python
diff --git a/docs/programming_guides/tile/arithmetic.rst
b/docs/programming_guides/tile/arithmetic.rst
index 0c522c187..7fc1e0b85 100644
--- a/docs/programming_guides/tile/arithmetic.rst
+++ b/docs/programming_guides/tile/arithmetic.rst
@@ -51,11 +51,11 @@ values to each of 64 threads:
@T.function
def affine_relu(
- A: T.Tensor((2, 32, 4), "float32", align=16),
- B: T.Tensor((2, 32, 4), "float32", align=16),
+ A: T.Tensor((2, 32, 4), T.float32, align=16),
+ B: T.Tensor((2, 32, 4), T.float32, align=16),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
- registers = T.alloc_local((2, 32, 4), "float32",
layout=register_layout)
+ registers = T.alloc_local((2, 32, 4), T.float32,
layout=register_layout)
T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, A, scope="cta")
scalars = registers.view(2, 32, 4, 1)
T.cuda.tile.apply(
@@ -130,19 +130,19 @@ register container, exposed below through a uint16 view:
@T.function
def to_float16(
- A: T.Tensor((2, 32, 4), "float32", align=16),
- B: T.Tensor((2, 32, 4), "float16", align=8),
+ A: T.Tensor((2, 32, 4), T.float32, align=16),
+ B: T.Tensor((2, 32, 4), T.float16, align=8),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
- registers = T.alloc_local((2, 32, 4), "float32",
layout=register_layout)
- converted = T.alloc_local((2, 32, 4), "float16",
layout=register_layout)
- words = converted.view("uint16")
+ registers = T.alloc_local((2, 32, 4), T.float32,
layout=register_layout)
+ converted = T.alloc_local((2, 32, 4), T.float16,
layout=register_layout)
+ words = converted.view(T.uint16.dtype)
T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, A, scope="cta")
T.cuda.tile.apply(
T.ptx["cvt.rn.f16.f32"],
words.view(2, 32, 4, 1), registers.view(2, 32, 4, 1), scope="cta",
)
- T.cuda.tile.apply(T.ptx["st.global.v4.b16"], B.view("uint16"), words,
scope="cta")
+ T.cuda.tile.apply(T.ptx["st.global.v4.b16"], B.view(T.uint16.dtype),
words, scope="cta")
``cvt`` performs the numeric conversion. The uint16 views expose the
resulting bits for the PTX operands and reuse the float16 storage. The
diff --git a/docs/programming_guides/tile/async_copies.rst
b/docs/programming_guides/tile/async_copies.rst
index c34068a2e..fd1885228 100644
--- a/docs/programming_guides/tile/async_copies.rst
+++ b/docs/programming_guides/tile/async_copies.rst
@@ -41,12 +41,12 @@ then reads the next thread's row:
@T.function
def async_read_next_row(
- A: T.Tensor((64, 4), "float32", align=16),
- B: T.Tensor((64, 4), "float32", align=16),
+ A: T.Tensor((64, 4), T.float32, align=16),
+ B: T.Tensor((64, 4), T.float32, align=16),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
tid = T.cuda.thread_idx("x")
- shared = T.alloc_shared((64, 4), "float32", align=16)
+ shared = T.alloc_shared((64, 4), T.float32, align=16)
T.cuda.tile.apply(
T.ptx["cp.async.cg.shared.global"],
shared.sub[tid, :], A.sub[tid, :], chunk_bytes=16,
diff --git a/docs/programming_guides/tile/blackwell_tensor_cores.rst
b/docs/programming_guides/tile/blackwell_tensor_cores.rst
index 1bf3f46c4..6ed98cbcb 100644
--- a/docs/programming_guides/tile/blackwell_tensor_cores.rst
+++ b/docs/programming_guides/tile/blackwell_tensor_cores.rst
@@ -67,25 +67,25 @@ uses 64 columns, so the allocation and deallocation both
specify 64.
from tvm_tirx.backend.cuda.layout import mma_shared_layout,
tmem_mma_operand_layout
from tvm_tirx.tirx.layout import S, TileLayout
- a_layout = mma_shared_layout("float16", 3, (128, 64))
- b_layout = mma_shared_layout("float16", 3, (64, 64))
- d_layout = tmem_mma_operand_layout("D", (128, 64), "float32", M=128,
cta_group=1)
+ a_layout = mma_shared_layout(T.float16, 3, (128, 64))
+ b_layout = mma_shared_layout(T.float16, 3, (64, 64))
+ d_layout = tmem_mma_operand_layout("D", (128, 64), T.float32, M=128,
cta_group=1)
register_layout = TileLayout(S[(128, 64) : (1 @ tid_in_wg, 1)])
@T.function
def blackwell_matmul(
- A: T.Tensor((128, 64), "float16"),
- B: T.Tensor((64, 64), "float16"),
- Out: T.Tensor((128, 64), "float32"),
+ A: T.Tensor((128, 64), T.float16),
+ B: T.Tensor((64, 64), T.float16),
+ Out: T.Tensor((128, 64), T.float32),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
tid = T.cuda.thread_idx("x")
warp = T.cuda.warp_id()
- a = T.alloc_shared((128, 64), "float16", layout=a_layout, align=1024)
- b = T.alloc_shared((64, 64), "float16", layout=b_layout, align=1024)
- address = T.alloc_shared((1,), "uint32", layout="default")
- mbar = T.alloc_shared((1,), "uint64", layout="default", align=8)
+ a = T.alloc_shared((128, 64), T.float16, layout=a_layout, align=1024)
+ b = T.alloc_shared((64, 64), T.float16, layout=b_layout, align=1024)
+ address = T.alloc_shared((1,), T.uint32, layout="default")
+ mbar = T.alloc_shared((1,), T.uint64, layout="default", align=8)
if warp == 0:
T.ptx.tcgen05.alloc.cta_group__1.sync.aligned.shared__cta.b32(
T.address_of(address), T.uint32(64)
@@ -99,7 +99,7 @@ uses 64 columns, so the allocation and deallocation both
specify 64.
T.ptx.fence.proxy.async_.shared__cta()
T.cuda.cta_sync()
d = T.cuda.decl_tmem(
- address[0], ty_args=[T.Tensor((128, 64), "float32", scope="tmem",
layout=d_layout)]
+ address[0], ty_args=[T.Tensor((128, 64), T.float32, scope="tmem",
layout=d_layout)]
)
av = a.rearrange("m (kr k) -> kr m k", k=16).view(1, 4, 128, 16)
bv = b.rearrange("n (kr k) -> kr n k", k=16).view(1, 4, 64, 16)
@@ -117,7 +117,7 @@ uses 64 columns, so the allocation and deallocation both
specify 64.
)
T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
T.ptx.tcgen05.fence__after_thread_sync()
- frag = T.alloc_local((128, 64), "float32", layout=register_layout)
+ frag = T.alloc_local((128, 64), T.float32, layout=register_layout)
T.cuda.tile.apply(
T.ptx["tcgen05.ld.sync.aligned.32x32b.x64.b32"],
frag.view(4, 32, 64),
diff --git a/docs/programming_guides/tile/data_movement.rst
b/docs/programming_guides/tile/data_movement.rst
index be04efe88..1bba749fa 100644
--- a/docs/programming_guides/tile/data_movement.rst
+++ b/docs/programming_guides/tile/data_movement.rst
@@ -57,13 +57,13 @@ other threads, then reads the next thread's row:
@T.function
def read_next_row(
- A: T.Tensor((64, 4), "float32", align=16),
- B: T.Tensor((64, 4), "float32", align=16),
+ A: T.Tensor((64, 4), T.float32, align=16),
+ B: T.Tensor((64, 4), T.float32, align=16),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
tid = T.cuda.thread_idx("x")
- registers = T.alloc_local((4,), "float32", layout="default")
- shared = T.alloc_shared((64, 4), "float32", align=16)
+ registers = T.alloc_local((4,), T.float32, layout="default")
+ shared = T.alloc_shared((64, 4), T.float32, align=16)
T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, A.sub[tid, :])
T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], shared.sub[tid, :],
registers)
T.cuda.cta_sync()
@@ -108,7 +108,7 @@ for another read. Insert the following statements after its
global load:
.. code-block:: python
- saved = T.alloc_local((4, 1), "float32", layout="default")
+ saved = T.alloc_local((4, 1), T.float32, layout="default")
T.cuda.tile.apply(T.ptx["mov.b32"], saved, registers.view(4, 1))
The view exposes four repeats, each ending in a single-register atom
diff --git a/docs/programming_guides/tile/index.rst
b/docs/programming_guides/tile/index.rst
index 43c7e32ec..01781e59b 100644
--- a/docs/programming_guides/tile/index.rst
+++ b/docs/programming_guides/tile/index.rst
@@ -20,20 +20,20 @@
Tile Operations
===============
-This guide shows how to use CUDA tile operations to apply a computation to
-tensor regions. ``T.cuda.tile.apply`` repeats a specified PTX instruction, such
-as a vector load or matrix multiplication. ``T.cuda.tile.ewise`` applies a
-scalar TIRx function to register elements. You choose the operand views,
+TIRx provides two CUDA tile operations: ``T.cuda.tile.apply`` and
+``T.cuda.tile.ewise``. ``apply`` repeats a specified PTX instruction, such as
+a vector load or matrix multiplication, while ``ewise`` applies a scalar
+TIRx function to register elements. You choose the operand views,
instruction variants, and synchronization required by the kernel.
-Start with :doc:`tile_patterns` for a load, compute, and store example and
-an explanation of how operand axes divide the work across threads. The
-:doc:`instruction reference </reference/python/backend.cuda.tile.instructions>`
-lists each family's parameters, operand layouts, and alignment requirements.
+The guides below cover data movement, arithmetic, asynchronous transfers,
+and matrix instructions. Start with :doc:`tile_patterns` for a load,
+compute, and store example that shows how operand axes divide the work
+across threads.
-The guides below show how to use these interfaces for data movement,
-arithmetic, asynchronous transfers, and matrix instructions. Asynchronous
-examples include the completion protocol required by their instruction family.
+For each instruction family's parameters, operand layouts, and alignment
+requirements, see the
+:doc:`instruction reference
</reference/python/backend.cuda.tile.instructions>`.
.. toctree::
:maxdepth: 1
diff --git a/docs/programming_guides/tile/tile_patterns.rst
b/docs/programming_guides/tile/tile_patterns.rst
index 008f6ad31..dbfceffae 100644
--- a/docs/programming_guides/tile/tile_patterns.rst
+++ b/docs/programming_guides/tile/tile_patterns.rst
@@ -26,44 +26,46 @@ A CUDA tile operation repeats a chosen computation over a
tensor region.
operand views describe which values each thread handles and how many times
it repeats that computation.
-Reading an Operand's Shape
---------------------------
-
-Suppose each of 64 threads processes twelve float32 values. A
-``ld.global.v4.b32`` instruction loads four 32-bit words at a time, so
-each thread needs three loads. We can describe this work with shape
-``(2, 32, 3, 4)``:
+Choosing a Tile Operation
+-------------------------
-.. code-block:: text
-
- (2, 32, 3, 4)
- └───┘ │ └─ atom: four words used by one instruction
- group └──── repeat: three instructions per thread
+Both operations repeat a computation over the values described by the
+operands. They differ in how you specify that computation:
.. list-table::
:header-rows: 1
- :widths: 20 80
-
- * - Part
- - Meaning in this example
- * - Group
- - ``(2, 32)`` selects one of two warps and a lane within that warp.
- * - Repeat
- - ``(3,)`` repeats the chosen instruction three times for that thread.
- * - Atom
- - ``(4,)`` selects the four words transferred by ``.v4.b32``.
-
-The instruction fixes the atom shape. The leading axes describe its
-repetitions and the threads that execute them. The register layout below
-maps the first two coordinates to ``warpid`` and ``laneid``, leaving
-``3 * 4 = 12`` local positions per thread.
+ :widths: 15 40 45
+
+ * - Operation
+ - You supply
+ - Each repetition
+ * - ``apply``
+ - A PTX instruction, such as ``T.ptx["ld.global.v4.b32"]`` or
+ ``T.ptx["add.rn.f32"]``.
+ - Executes that instruction on the values required by its operand form.
+ * - ``ewise``
+ - A scalar TIRx function, such as ``plus_one(value)``.
+ - Inlines the function for one element, potentially producing several
+ instructions.
+
+Use ``apply`` when you want a particular instruction, including its
+rounding or other variants. It supports arithmetic and matrix instructions
+as well as data movement. Use ``ewise`` to express a computation as a
+scalar function over register elements.
A Load, Compute, and Store Example
----------------------------------
-The kernel uses this layout to load the values, add one to each, and
-store the results. ``scope="cta"`` tells the tile calls to use the
-leading ``(2, 32)`` axes for the threads in the block:
+Consider an add-one kernel with 64 threads, each handling twelve float32
+values. It makes three tile calls: ``apply`` to load the values into
+registers, ``ewise`` to add one to each value, and ``apply`` to store the
+results. The load and store use the ``.v4.b32`` instruction forms, which
+transfer four 32-bit words at a time.
+
+The tensors have shape ``(2, 32, 3, 4)``: two warps, 32 lanes per
+warp, and three groups of four values per thread. The register layout
+maps the first two axes to ``warpid`` and ``laneid``. ``scope="cta"``
+tells the tile calls to use those leading axes for the threads in the block:
.. code-block:: python
@@ -79,11 +81,11 @@ leading ``(2, 32)`` axes for the threads in the block:
@T.function
def add_one_tile(
- A: T.Tensor((2, 32, 3, 4), "float32", align=16),
- B: T.Tensor((2, 32, 3, 4), "float32", align=16),
+ A: T.Tensor((2, 32, 3, 4), T.float32, align=16),
+ B: T.Tensor((2, 32, 3, 4), T.float32, align=16),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
- registers = T.alloc_local((2, 32, 3, 4), "float32",
layout=register_layout)
+ registers = T.alloc_local((2, 32, 3, 4), T.float32,
layout=register_layout)
T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, A, scope="cta")
T.cuda.tile.ewise(plus_one, registers, registers, scope="cta")
T.cuda.tile.apply(T.ptx["st.global.v4.b32"], B, registers, scope="cta")
@@ -106,6 +108,37 @@ We can compile and check the result:
compiled(a, b)
torch.testing.assert_close(b, a + 1)
+Reading an Operand's Shape
+--------------------------
+
+For ``apply``, operand axes describe the execution **group**, the
+instruction **repeats**, and the instruction's **atom**. The atom is the
+set of values used by one instruction; here it is four consecutive
+32-bit words. The same register view has a different interpretation
+for ``ewise``, which applies a scalar function to every element:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 20 25 30
+
+ * - Operation in this kernel
+ - Group
+ - Repeat
+ - Atom
+ * - ``apply`` with ``ld/st.global.v4.b32``
+ - ``(2, 32)``: warp and lane
+ - ``(3,)``: three instructions per thread
+ - ``(4,)``: four words per instruction
+ * - ``ewise(plus_one, ...)``
+ - ``(2, 32)``: warp and lane
+ - ``(3, 4)``: twelve scalar applications per thread
+ - One scalar element, with no trailing atom axis
+
+The chosen instruction fixes ``apply``'s atom shape. The operand view
+specifies how often to repeat it, and the layout maps the coordinates to
+threads and storage. ``ewise`` uses every coordinate after the group
+prefix for a scalar function application.
+
Execution Scope and Scalar Functions
------------------------------------
@@ -116,9 +149,31 @@ has no group prefix. Other scopes and warp-collective
instructions use
the group shapes listed in the :ref:`coordinate reference
<cuda-tile-operand-coordinates>`.
-``ewise`` uses one scalar value per invocation, so it has no trailing
-instruction atom. For the same register view, ``(2, 32)`` still selects a
-thread, while ``(3, 4)`` now specifies twelve scalar invocations.
+``ewise``'s function parameters and return type must match the input and
+destination elements. In ``ewise(plus_one, registers, registers, ...)``,
+the first ``registers`` is the destination and the second is the input;
+the result replaces each input value in place. A function can contain
+multiple scalar operations, as the SiLU example in
+:doc:`arithmetic` demonstrates.
+
+Using apply for Arithmetic
+--------------------------
+
+The add-one computation can also use an explicitly selected PTX
+instruction. In the kernel above, replace the ``ewise`` call with:
+
+.. code-block:: python
+
+ scalars = registers.view(2, 32, 3, 4, 1)
+ T.cuda.tile.apply(
+ T.ptx["add.rn.f32"], scalars, scalars, T.float32(1), scope="cta"
+ )
+
+``add.rn.f32`` uses a single-register atom, represented by the trailing
+``1``. The ``(3, 4)`` repeat axes issue twelve additions per thread,
+reusing the scalar constant ``1``. The view reuses the same register
+storage. Here ``apply`` selects the addition instruction and its rounding
+mode directly; ``ewise`` expresses the computation through ``plus_one``.
Preparing Operand Views
-----------------------
diff --git a/docs/programming_guides/tile/tma.rst
b/docs/programming_guides/tile/tma.rst
index 7e18e5afa..c39834569 100644
--- a/docs/programming_guides/tile/tma.rst
+++ b/docs/programming_guides/tile/tma.rst
@@ -44,13 +44,13 @@ all 32 lanes update one column each:
@T.function
def tma_add_one(
- A: T.Tensor((8, 32), "float32", align=16),
- B: T.Tensor((8, 32), "float32", align=16),
+ A: T.Tensor((8, 32), T.float32, align=16),
+ B: T.Tensor((8, 32), T.float32, align=16),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
lane = T.cuda.lane_id()
- shared = T.alloc_shared((8, 32), "float32", layout="default",
align=128)
- mbar = T.alloc_shared((1,), "uint64", align=8)
+ shared = T.alloc_shared((8, 32), T.float32, layout="default",
align=128)
+ mbar = T.alloc_shared((1,), T.uint64, align=8)
if lane == 0:
T.ptx.mbarrier.init.shared.b64(mbar.ptr_to([0]), T.uint32(1))
T.ptx.mbarrier.arrive.expect_tx.shared.b64(mbar.ptr_to([0]),
T.uint32(1024))
diff --git a/docs/programming_guides/tile/warp_matrix.rst
b/docs/programming_guides/tile/warp_matrix.rst
index df164ab6c..bb5225099 100644
--- a/docs/programming_guides/tile/warp_matrix.rst
+++ b/docs/programming_guides/tile/warp_matrix.rst
@@ -86,14 +86,14 @@ lane mappings represented by these layouts.
@T.function
def warp_matmul(
- A: T.Tensor((16, 16), "float16"),
- B: T.Tensor((8, 16), "float16"),
- Out: T.Tensor((16, 8), "float32"),
+ A: T.Tensor((16, 16), T.float16),
+ B: T.Tensor((8, 16), T.float16),
+ Out: T.Tensor((16, 8), T.float32),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
lane = T.cuda.lane_id()
- a = T.alloc_shared((16, 16), "float16", layout="default", align=16)
- b = T.alloc_shared((8, 16), "float16", layout="default", align=16)
+ a = T.alloc_shared((16, 16), T.float16, layout="default", align=16)
+ b = T.alloc_shared((8, 16), T.float16, layout="default", align=16)
for q in range(8):
index = lane + 32 * q
a[index // 16, index % 16] = A[index // 16, index % 16]
@@ -101,13 +101,13 @@ lane mappings represented by these layouts.
index = lane + 32 * q
b[index // 16, index % 16] = B[index // 16, index % 16]
T.cuda.cta_sync()
- a_words = a.view("uint32").rearrange(
+ a_words = a.view(T.uint32.dtype).rearrange(
"(group row) (chunk word) -> (chunk group) row word", group=2,
chunk=2, word=4
)
- b_words = b.view("uint32").rearrange("row (chunk word) -> chunk row
word", chunk=2)
- ar = T.alloc_local((4,), "uint32", layout="default")
- br = T.alloc_local((2,), "uint32", layout="default")
- dr = T.alloc_local((4,), "float32", layout="default")
+ b_words = b.view(T.uint32.dtype).rearrange("row (chunk word) -> chunk
row word", chunk=2)
+ ar = T.alloc_local((4,), T.uint32, layout="default")
+ br = T.alloc_local((2,), T.uint32, layout="default")
+ dr = T.alloc_local((4,), T.float32, layout="default")
T.cuda.tile.apply(
T.ptx["ldmatrix.sync.aligned.m8n8.x4.shared.b16"],
ar.view(4, 8, 4, layout=a_load_layout), a_words, scope="warp",
diff --git a/docs/programming_guides/types_tensors_memory.rst
b/docs/programming_guides/types_tensors_memory.rst
index 39753f6fb..186a6baaf 100644
--- a/docs/programming_guides/types_tensors_memory.rst
+++ b/docs/programming_guides/types_tensors_memory.rst
@@ -22,7 +22,7 @@
Types, Tensors, and Memory
==========================
-This guide looks at how CUDA kernels represent values and manage their storage.
+This guide covers data types, tensors, and memory spaces in CUDA kernels.
A tensor's dtype describes its elements, and its scope selects the memory
space. We will extend the add-one computation with vector accesses, shared
storage, and thread-private values, then create views of existing data.
@@ -55,12 +55,8 @@ in generated CUDA source:
* - ``float32x4``
- ``float4``
-A ``float32x4`` value contains four float32 components, called lanes.
-CUDA C++ names this type ``float4``: each ``float`` is 32 bits, and ``4``
-is the number of components.
-
-We can extend the add-one kernel with vector accesses so each thread
-reads and updates four elements as a group:
+In the following add-one kernel, each thread reads four ``float32``
+elements into a ``float32x4`` value and updates them together:
.. code-block:: python
@@ -68,19 +64,19 @@ reads and updates four elements as a group:
@T.function
def add_one_vector(
- A: T.Tensor((256,), "float32", align=16),
- B: T.Tensor((256,), "float32", align=16),
+ A: T.Tensor((256,), T.float32, align=16),
+ B: T.Tensor((256,), T.float32, align=16),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
tx = T.cuda.thread_idx("x")
- vector_value = A[T.Ramp(tx * 4, 1, 4)]
+ vector_value: T.float32x4 = A[T.Ramp(tx * 4, 1, 4)]
B[T.Ramp(tx * 4, 1, 4)] = vector_value + 1
``T.Ramp(tx * 4, 1, 4)`` selects the four consecutive indices ``4*tx``,
``4*tx + 1``, ``4*tx + 2``, and ``4*tx + 3``. Its arguments specify the
-starting index, stride, and lane count. Because ``A`` has float32 elements,
-the read at these four positions produces ``vector_value`` with inferred
-dtype ``float32x4``.
+starting index, stride, and number of elements (lanes). Because ``A`` has
+float32 elements, the read produces a four-lane value, matching the
+``T.float32x4`` annotation on ``vector_value``.
``vector_value + 1`` adds one to each lane. The store uses the same indices
to write the four results to the corresponding float32 elements in ``B``.
@@ -93,7 +89,7 @@ vector loads and stores.
Declaring Tensors
-----------------
-A parameter such as ``A: T.Tensor((128,), "float32")`` describes a tensor
+A parameter such as ``A: T.Tensor((128,), T.float32)`` describes a tensor
passed by the caller. Inside a kernel, ``T.alloc_shared`` and
``T.alloc_local`` allocate scratch tensors. Both are conveniences for
``T.alloc_tensor``, whose tensor type specifies the shape, dtype, and scope.
@@ -125,9 +121,9 @@ one shared tensor and declare another view of the same
elements:
.. code-block:: python
- shared = T.alloc_shared((128,), "float32")
+ shared = T.alloc_shared((128,), T.float32)
shared_view = T.decl_tensor(
- shared.data, ty_args=[T.Tensor((128,), "float32", scope="shared")]
+ shared.data, ty_args=[T.Tensor((128,), T.float32, scope="shared")]
)
``shared.data`` is the base pointer for ``shared_view``. Both names refer
@@ -152,10 +148,10 @@ neighbor, wrapping the last thread back to the first:
.. code-block:: python
@T.function
- def shared_shift(A: T.Tensor((128,), "float32"), B: T.Tensor((128,),
"float32")):
+ def shared_shift(A: T.Tensor((128,), T.float32), B: T.Tensor((128,),
T.float32)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
tx = T.cuda.thread_idx("x")
- shared = T.alloc_shared((128,), "float32")
+ shared = T.alloc_shared((128,), T.float32)
shared[tx] = A[tx]
T.cuda.cta_sync()
B[tx] = shared[(tx + 1) % 128]
@@ -184,12 +180,12 @@ uses it for the shared tensor:
@T.function
def shared_shift_dynamic(
- A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")
+ A: T.Tensor((128,), T.float32), B: T.Tensor((128,), T.float32)
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
tx = T.cuda.thread_idx("x")
pool = SMEMPool()
- shared = pool.alloc((128,), "float32")
+ shared = pool.alloc((128,), T.float32)
pool.commit()
shared[tx] = A[tx]
T.cuda.cta_sync()
@@ -216,10 +212,10 @@ adding one and writing them back:
.. code-block:: python
@T.function
- def add_one_local(A: T.Tensor((256,), "float32"), B: T.Tensor((256,),
"float32")):
+ def add_one_local(A: T.Tensor((256,), T.float32), B: T.Tensor((256,),
T.float32)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
tx = T.cuda.thread_idx("x")
- values = T.alloc_local((4,), "float32")
+ values = T.alloc_local((4,), T.float32)
for k in T.unroll(4):
values[k] = A[tx * 4 + k]
for k in T.unroll(4):
@@ -255,7 +251,7 @@ and writes one output for each sum:
.. code-block:: python
@T.function
- def sum_four(A: T.Tensor((256,), "float32"), B: T.Tensor((64,),
"float32")):
+ def sum_four(A: T.Tensor((256,), T.float32), B: T.Tensor((64,),
T.float32)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
tx = T.cuda.thread_idx("x")
base: T.let[T.int32] = tx * 4
@@ -303,7 +299,7 @@ same elements:
- A view with the axes reordered
* - ``A.rearrange(pattern, **sizes)``
- A view that splits, reorders, or merges axes using a named pattern
- * - ``A.view("uint32")``
+ * - ``A.view(T.uint32.dtype)``
- A view of the same bytes as 32-bit unsigned words
For a contiguous one-dimensional tensor ``A`` with 256 elements, we can
@@ -339,8 +335,8 @@ a device region, eight float16 elements can be viewed as
four uint32 words:
.. code-block:: python
- half_values = T.alloc_local((8,), "float16")
- packed_words = half_values.view("uint32")
+ half_values = T.alloc_local((8,), T.float16)
+ packed_words = half_values.view(T.uint32.dtype)
``packed_words`` has shape ``(4,)``. Each word contains the bits of two
adjacent float16 elements. This is the representation used by instructions
diff --git a/examples/programming_guides/control_flow.py
b/examples/programming_guides/control_flow.py
index c37e0582c..e06a63513 100644
--- a/examples/programming_guides/control_flow.py
+++ b/examples/programming_guides/control_flow.py
@@ -109,6 +109,40 @@ tvm_tirx.compile(search_sorted, target="cuda")(s, q, b)
torch.testing.assert_close(b, torch.searchsorted(s, q).to(torch.int32))
+# [break_continue.begin]
[email protected]
+def sum_until_zero(A: T.Tensor((n, 8), T.int32), B: T.Tensor((n,), T.int32)):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=T.ceildiv(n, 128),
block=128))
+ row = T.cuda.block_idx("x") * 128 + T.cuda.thread_idx("x")
+ if row < n:
+ total: T.int32 = 0
+ for j in range(8):
+ value = A[row, j]
+ if value < 0:
+ continue
+ if value == 0:
+ break
+ total += value
+ B[row] = total
+ # [break_continue.end]
+
+
+a = torch.tensor(
+ [
+ [2, -1, 3, 0, 9, 9, 9, 9],
+ [0, 4, 5, 6, 7, 8, 9, 10],
+ [-2, 1, -3, 2, 0, 4, 5, 6],
+ [1, 2, 3, 4, 5, 6, 7, 8],
+ [-1, -2, -3, -4, -5, -6, -7, -8],
+ ],
+ dtype=torch.int32,
+ device="cuda",
+)
+b = torch.empty(a.shape[0], dtype=torch.int32, device="cuda")
+tvm_tirx.compile(sum_until_zero, target="cuda")(a, b)
+torch.testing.assert_close(b, torch.tensor([5, 0, 3, 36, 0],
dtype=torch.int32, device="cuda"))
+
+
# [return.begin]
@T.function
def scale(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
