wehub-resource-sync
137 lines
4.7 KiB
ReStructuredText
137 lines
4.7 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 → local
|
|
=================
|
|
|
|
The ``local`` variant lowers a reduction (``sum`` / ``max`` / ``min``) when **both
|
|
source and destination are register** (``local``) buffers. At thread scope it is a
|
|
plain sequential reduction over each thread's own elements; at warp scope, if the
|
|
destination layout carries a ``laneid`` replica, it also folds across lanes with a
|
|
``__shfl_xor`` tree. Source:
|
|
``python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py``.
|
|
|
|
What it accepts
|
|
---------------
|
|
|
|
.. code-block:: python
|
|
|
|
@register_dispatch(op_name, "cuda", variant="local", priority=10, when=[
|
|
predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["local"]),
|
|
predicate("local_valid", validate_reduction_local),
|
|
])
|
|
|
|
.. list-table::
|
|
:header-rows: 1
|
|
:widths: 22 78
|
|
|
|
* - Property
|
|
- Requirement
|
|
* - target / priority
|
|
- ``cuda``; priority ``10``
|
|
* - operand scope
|
|
- src **and** dst in ``local`` (registers), equal dtype
|
|
* - exec scope
|
|
- ``thread`` (always valid — pure thread-local); ``warp`` / ``warpgroup``
|
|
require a valid (non-swizzled) ``TileLayout``; ``warp`` may additionally
|
|
cross-lane reduce when ``thread_reduce`` and a ``laneid`` shard→replica
|
|
pattern are present
|
|
* - shape
|
|
- dst spatial dims match src; reduced dims have ``local_extent == 1`` on dst
|
|
|
|
Demonstration program
|
|
----------------------
|
|
|
|
A single thread reduces a 4-element ``float32`` register vector to a scalar
|
|
(thread-wise path, from ``test_reduction.py``):
|
|
|
|
.. code-block:: python
|
|
|
|
@T.prim_func
|
|
def test_func(A_ptr: T.handle, B_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, [4], "float32", layout=TileLayout(S[(4,)]))
|
|
B = T.match_buffer(B_ptr, [1], "float32", layout=TileLayout(S[(1,)]))
|
|
T.device_entry(); T.cta_id([1]); T.thread_id([1])
|
|
A_local = T.alloc_buffer([4], "float32", scope="local")
|
|
B_local = T.alloc_buffer([1], "float32", scope="local")
|
|
for i in T.serial(4): A_local[i] = A[i]
|
|
Tx.sum(B_local, A_local, accum=False) # reduction local dispatch
|
|
B[0] = B_local[0]
|
|
|
|
(4 < 8 elements, so this stays on ``local`` rather than the
|
|
:doc:`sm100_packed` fast path.)
|
|
|
|
Algorithm
|
|
---------
|
|
|
|
**Thread-wise** (``_emit_reduction_local_thread_wise``): a spatial loop over the
|
|
output positions, each initialized to the op's identity (unless ``accum``), then a
|
|
reduction loop accumulating the source — no cross-thread communication:
|
|
|
|
.. code-block:: python
|
|
|
|
for spa in range(spatial_len):
|
|
if not accum: dst[spa] = identity
|
|
for red in range(reduction_len):
|
|
dst[spa] = op(dst[spa], src[spa, red])
|
|
|
|
**Warp-shuffle** (``_gen_warp_shuffle_reduce``): when the dst layout has a
|
|
``laneid`` replica, each lane first reduces its own elements, then
|
|
``T.cuda.warp_reduce`` folds across lanes — a ``__shfl_xor`` tree over the **full**
|
|
``0xFFFFFFFF`` mask. (This differs from :doc:`shared`, which uses explicit
|
|
``tvm_warp_shuffle_xor`` steps over ``__activemask()`` at the *group* width.)
|
|
|
|
Generated TIRx IR
|
|
-----------------
|
|
|
|
For the 4-element thread reduction:
|
|
|
|
.. code-block:: python
|
|
|
|
for spa in range(1):
|
|
for red in range(4):
|
|
dst[...] = dst[...] + src[...] # op = sum
|
|
|
|
Generated CUDA
|
|
--------------
|
|
|
|
.. code-block:: c++
|
|
|
|
for (int red = 0; red < 4; ++red)
|
|
B_local_ptr[0] = B_local_ptr[0] + A_local_ptr[red];
|
|
|
|
(Verified on ``sm_100a`` — ``B == sum(A)``.)
|
|
|
|
How inputs change the algorithm
|
|
-------------------------------
|
|
|
|
.. list-table::
|
|
:header-rows: 1
|
|
:widths: 28 72
|
|
|
|
* - input
|
|
- effect
|
|
* - op
|
|
- ``sum`` → ``+``, ``max`` → ``max``, ``min`` → ``min`` (and the identity)
|
|
* - exec scope
|
|
- ``thread`` → sequential; ``warp`` with a ``laneid`` replica → adds a
|
|
``__shfl_xor`` cross-lane tree
|
|
* - axes / shape
|
|
- set the spatial vs reduction loop extents
|
|
* - accum
|
|
- ``True`` reuses the old dst value instead of the identity
|