wehub-resource-sync
113 lines
4.3 KiB
ReStructuredText
113 lines
4.3 KiB
ReStructuredText
.. Licensed to the Apache Software Foundation (ASF) under one
|
|
or more contributor license agreements. See the NOTICE file
|
|
distributed with this work for additional information
|
|
regarding copyright ownership. The ASF licenses this file
|
|
to you under the Apache License, Version 2.0 (the
|
|
"License"); you may not use this file except in compliance
|
|
with the License. You may obtain a copy of the License at
|
|
|
|
.. http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
.. Unless required by applicable law or agreed to in writing,
|
|
software distributed under the License is distributed on an
|
|
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
KIND, either express or implied. See the License for the
|
|
specific language governing permissions and limitations
|
|
under the License.
|
|
|
|
Compiling and inspecting
|
|
========================
|
|
|
|
Wrap the ``PrimFunc`` in an ``IRModule`` and compile with
|
|
``tvm.compile(mod, target=..., tir_pipeline="tirx")``; it runs the TIRx lowering
|
|
pipeline and returns an ``Executable`` you call directly. The arch (e.g.
|
|
``sm_100a``) is auto-detected from the device, so the target ``"cuda"`` is enough.
|
|
|
|
.. code-block:: python
|
|
|
|
target = tvm.target.Target("cuda")
|
|
exe = tvm.compile(tvm.IRModule({"main": scale}), target=target, tir_pipeline="tirx")
|
|
|
|
``tir_pipeline="tirx"`` selects the TIRx lowering pipeline (``LowerTIRx`` →
|
|
tile-primitive dispatch → host/device split → finalize). Compiling inside a
|
|
``with target:`` block also works and lets the kernel pick up the target context.
|
|
|
|
Inspecting the result
|
|
---------------------
|
|
|
|
Read the IR with ``.show()`` / ``.script()``, and read the generated CUDA from the
|
|
compiled module.
|
|
|
|
.. code-block:: python
|
|
|
|
scale.show() # pretty-print the TIRx (TVMScript)
|
|
print(scale.script()) # ... the same, as a string
|
|
|
|
# the generated CUDA C source, from the compiled Executable:
|
|
print(exe.mod.imports[0].inspect_source())
|
|
|
|
Debug aids: ``T.print_buffer(C.data, "float32", False, False, 1, (M,))`` emits a
|
|
runtime ``printf`` of a buffer into the kernel; ``T.hint("message")`` (statement
|
|
or ``with`` block) attaches structured hints that survive a script round-trip.
|
|
|
|
From simple to complex
|
|
----------------------
|
|
|
|
A natural native progression, each rung adding one capability:
|
|
|
|
#. **Elementwise** — ``device_entry`` + ``thread_id`` + a guarded store (the first
|
|
kernel).
|
|
#. **Shared-memory reduction** — stage into ``T.alloc_shared``, then a
|
|
``cta_sync``-separated tree (shown in full below). Adds shared memory and a
|
|
block barrier.
|
|
#. **Warp / block reduction** — ``T.tvm_warp_shuffle_xor`` or ``T.cuda.cta_sum``
|
|
to combine partial results across lanes/warps (the warp all-reduce in
|
|
:doc:`threads_sync`).
|
|
#. **Async pipeline** — ``T.ptx.cp_async`` (or TMA ``cp_async.bulk.tensor``) with
|
|
``T.ptx.mbarrier.*`` to overlap loads with compute.
|
|
|
|
Rung 2 in full — a 256-element block sum via a shared-memory tree reduction
|
|
(shared buffer, ``cta_sync``, a ``while`` loop, and a thread predicate):
|
|
|
|
.. code-block:: python
|
|
|
|
@T.prim_func
|
|
def block_sum(A_ptr: T.handle, out_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (256,), "float32")
|
|
out = T.match_buffer(out_ptr, (1,), "float32")
|
|
|
|
T.device_entry()
|
|
bx = T.cta_id([1])
|
|
tx = T.thread_id([256])
|
|
|
|
sm = T.alloc_shared((256,), "float32")
|
|
sm[tx] = A[tx]
|
|
T.cuda.cta_sync()
|
|
|
|
s = T.alloc_local((1,), "int32")
|
|
s[0] = 128
|
|
while s[0] >= 1:
|
|
if tx < s[0]:
|
|
sm[tx] += sm[tx + s[0]]
|
|
T.cuda.cta_sync()
|
|
s[0] = s[0] // 2
|
|
|
|
if tx == 0:
|
|
out[0] = sm[0]
|
|
|
|
exe = tvm.compile(tvm.IRModule({"main": block_sum}),
|
|
target=tvm.target.Target("cuda"), tir_pipeline="tirx")
|
|
a = torch.arange(256, device="cuda", dtype=torch.float32)
|
|
out = torch.zeros(1, device="cuda")
|
|
exe(a, out) # out[0] == 32640.0
|
|
|
|
The full tile-level GEMM/attention ladder (sync → TMA → warp specialization →
|
|
2-CTA cluster) is built on top of these and the dispatchable tile primitives in
|
|
:doc:`../../tile_primitives`.
|
|
|
|
Next steps
|
|
----------
|
|
|
|
- :doc:`../../layout` — how buffers map to physical resources (``TileLayout``).
|
|
- :doc:`../../tile_primitives` — the dispatchable ops these native idioms lower to.
|