v4.2 tag release. (#2638)
This commit is contained in:
@@ -72,6 +72,55 @@ All loop indices must be |Constexpr|.
|
||||
for i in cutlass.range(bound, unroll=2):
|
||||
cute.printf("%d\\n", i)
|
||||
|
||||
Software Pipelining
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Software pipelining is a technique used to optimize loops. Typically, this involves writing a prefetch loop and a main loop.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def example():
|
||||
...
|
||||
# build a circular buffer
|
||||
buffer = ...
|
||||
|
||||
# prefetch loop
|
||||
for i in range(prefetch_stages):
|
||||
cute.copy(atom, gmem[i], buffer[i], ...)
|
||||
|
||||
# main loop
|
||||
for i in range(bound):
|
||||
if i + prefetch_stages < bound:
|
||||
cute.copy(atom, gmem[i + prefetch_stages], buffer[(i + prefetch_stages) % total_stages], ...)
|
||||
|
||||
use(buffer[i % total_stages])
|
||||
|
||||
...
|
||||
|
||||
This can be tedious to write and tune. |DSL| provides a loop attribute to ask the compiler to do this.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def example():
|
||||
...
|
||||
# build a circular buffer
|
||||
buffer = ...
|
||||
|
||||
for i in cutlass.range(bound, prefetch_stages=prefetch_stages):
|
||||
# Compiler automatically handles the pipelining:
|
||||
# - Generates prefetch loop for initial stages
|
||||
# - In main loop, prefetches future data while using current data
|
||||
cute.copy(atom, gmem[i], buffer[i % total_stages], ...)
|
||||
use(buffer[i % total_stages]) # Uses data from previous iterations
|
||||
|
||||
...
|
||||
|
||||
Compiler will automatically generate the prefetch loop with `prefetch_stages` iterations and a corresponding main loop.
|
||||
|
||||
This feature is experimental and only supported on sm90 and above.
|
||||
|
||||
|
||||
If-Else Statements
|
||||
------------------
|
||||
|
||||
Reference in New Issue
Block a user