wehub-resource-sync
137 lines
5.3 KiB
ReStructuredText
137 lines
5.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.
|
||
|
||
reduction → shared
|
||
==================
|
||
|
||
The ``shared`` variant lowers a reduction (``sum`` / ``max`` / ``min``) when
|
||
**both source and destination are shared** memory. At CTA / warpgroup / warp scope
|
||
it partitions the threads into groups — one group per output position — has each
|
||
thread gather a chunk of the reduction axis, then folds the group with an adaptive
|
||
``__shfl_xor`` tree. Source:
|
||
``python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py``.
|
||
|
||
What it accepts
|
||
---------------
|
||
|
||
.. code-block:: python
|
||
|
||
@register_dispatch(op_name, "cuda", variant="shared", priority=10, when=[
|
||
predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["shared*"]),
|
||
predicate("shared_valid", validate_reduction_shared),
|
||
])
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 22 78
|
||
|
||
* - Property
|
||
- Requirement
|
||
* - target / priority
|
||
- ``cuda``; priority ``10``
|
||
* - operand scope
|
||
- src **and** dst in ``shared*``, equal dtype
|
||
* - exec scope
|
||
- ``cta`` / ``warpgroup`` / ``warp`` (shuffle tree) or ``thread`` (sequential)
|
||
* - thread binding
|
||
- ``threadIdx.x`` present and **1-D** (no ``threadIdx.y`` / ``z``)
|
||
* - shape
|
||
- ``dst`` size equals the source's spatial extent (product of the non-reduced
|
||
dims)
|
||
|
||
Demonstration program
|
||
----------------------
|
||
|
||
A 32-thread CTA reduces each row of a ``4×8`` ``float32`` shared tile (reduce
|
||
axis ``-1``) to a ``4``-vector (from ``test_reduction.py``):
|
||
|
||
.. code-block:: python
|
||
|
||
@T.prim_func
|
||
def test_reduction(A_ptr: T.handle, B_ptr: T.handle):
|
||
A = T.match_buffer(A_ptr, (4, 8), "float32", layout=TileLayout(S[(4, 8)]))
|
||
B = T.match_buffer(B_ptr, (4,), "float32", layout=TileLayout(S[(4,)]))
|
||
T.device_entry(); T.cta_id([1]); T.thread_id([32])
|
||
A_smem = T.alloc_buffer((4, 8), "float32", scope="shared", layout=TileLayout(S[(4, 8)]))
|
||
B_smem = T.alloc_buffer((4,), "float32", scope="shared", layout=TileLayout(S[(4,)]))
|
||
Tx.cta.copy(A_smem, A); T.cuda.cta_sync()
|
||
Tx.cta.sum(B_smem, A_smem, axes=(-1,), accum=False) # reduction shared dispatch
|
||
T.cuda.cta_sync()
|
||
Tx.cta.copy(B, B_smem)
|
||
|
||
Algorithm
|
||
---------
|
||
|
||
**1. Choose the group size.** ``group_size = min(next_power_of_2(reduction_len),
|
||
32, thread_cnt)`` — here ``reduction_len = 8`` ⇒ ``group_size = 8``. Each group of 8
|
||
lanes cooperatively reduces one row; the CTA processes the 4 rows in parallel.
|
||
|
||
**2. Gather + shuffle tree.** Each lane loads its slice of the reduction axis into a
|
||
register, then ``log2(group_size)`` ``shfl_xor`` steps (masks ``1, 2, 4``) fold the
|
||
group; lane 0 of each group writes the result, followed by a barrier:
|
||
|
||
.. code-block:: python
|
||
|
||
mask = T.tvm_warp_activemask()
|
||
for i in range(n_shuffles): # n_shuffles = log2(group_size)
|
||
thread_data[0] = op(thread_data[0],
|
||
T.tvm_warp_shuffle_xor(mask, thread_data[0], 1 << i, group_size, 32))
|
||
|
||
(``warp`` uses ``warp_sync``; ``warpgroup`` ``warpgroup_sync(8)``; ``cta``
|
||
``cta_sync``. Thread scope is instead the sequential loop of :doc:`local`.)
|
||
|
||
Generated TIRx IR
|
||
-----------------
|
||
|
||
.. code-block:: python
|
||
|
||
thread_data[0] = thread_data[0] + T.tvm_warp_shuffle_xor(T.tvm_warp_activemask(), thread_data[0], 1, 8, 32)
|
||
thread_data[0] = thread_data[0] + T.tvm_warp_shuffle_xor(T.tvm_warp_activemask(), thread_data[0], 2, 8, 32)
|
||
thread_data[0] = thread_data[0] + T.tvm_warp_shuffle_xor(T.tvm_warp_activemask(), thread_data[0], 4, 8, 32)
|
||
|
||
Generated CUDA
|
||
--------------
|
||
|
||
.. code-block:: c++
|
||
|
||
thread_data_ptr[0] = thread_data_ptr[0] + __shfl_xor_sync(__activemask(), thread_data_ptr[0], 1, 8);
|
||
thread_data_ptr[0] = thread_data_ptr[0] + __shfl_xor_sync(__activemask(), thread_data_ptr[0], 2, 8);
|
||
thread_data_ptr[0] = thread_data_ptr[0] + __shfl_xor_sync(__activemask(), thread_data_ptr[0], 4, 8);
|
||
|
||
(Verified on ``sm_100a`` — each ``B[r] == sum(A[r, :])``. The shuffle width ``8``
|
||
is the group size, not the full warp.)
|
||
|
||
How inputs change the algorithm
|
||
-------------------------------
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 28 72
|
||
|
||
* - input
|
||
- effect
|
||
* - op
|
||
- ``sum`` → ``+`` shuffle tree; ``max`` / ``min`` → the corresponding combine
|
||
* - reduction length / thread count
|
||
- set ``group_size = min(next_pow2(reduction_len), 32, thread_cnt)`` and hence
|
||
the number of shuffle steps
|
||
* - exec scope
|
||
- ``cta`` / ``warpgroup`` / ``warp`` → shuffle tree (different sync); ``thread``
|
||
→ sequential loop
|
||
* - accum
|
||
- ``True`` combines the reduced value with the old dst
|