wehub-resource-sync
134 lines
4.9 KiB
ReStructuredText
134 lines
4.9 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.
|
||
|
||
elementwise → reg
|
||
=================
|
||
|
||
The ``reg`` variant lowers an elementwise op (``sqrt``, ``exp``, ``add``,
|
||
``fma``, …) when **all operands are register** (``local``) buffers. Like the copy
|
||
:doc:`../copy/reg` variant the partition is *induced* by the operands' register
|
||
layout — the thread axes are dropped, leaving each thread its private bundle — and
|
||
the op is applied to every register in that bundle. Source:
|
||
``python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py``.
|
||
|
||
What it accepts
|
||
---------------
|
||
|
||
``is_reg_ewise(spec)`` builds the predicate:
|
||
|
||
.. code-block:: python
|
||
|
||
def check(op_call, sctx):
|
||
if not sctx.is_target("cuda"): return False, "non-cuda target"
|
||
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"): ...
|
||
ok, reason = _all_threads_active(sctx)
|
||
plan, msg = spec.parse(op_call)
|
||
for br in buffer_regions(plan):
|
||
if br.buffer.scope() != "local": # every operand register
|
||
return False, f"operand scope {br.buffer.scope()} != local"
|
||
if br.buffer.layout is None: ...
|
||
# + spec.check_extras (dtype rules), pick_anchor + _validate_anchor_layout,
|
||
# _validate_scope_level_anchor
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 22 78
|
||
|
||
* - Property
|
||
- Requirement
|
||
* - target / scope / priority
|
||
- ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta`` (all active);
|
||
priority ``10``
|
||
* - operands
|
||
- **every** operand in ``local`` (registers)
|
||
* - op
|
||
- any registry op (unary / binary / ``fma``); ``spec.check_extras`` validates
|
||
the dtype combo
|
||
* - register layout
|
||
- the anchor register layout must validate, and its thread axis must match the
|
||
scope (it induces the partition)
|
||
|
||
Demonstration program
|
||
----------------------
|
||
|
||
A warp takes the elementwise ``sqrt`` of a ``32×8`` ``float32`` register tile
|
||
(register layout ``S[(32,8):(1@laneid,1)]`` — lane ``i`` owns row ``i``):
|
||
|
||
.. code-block:: python
|
||
|
||
from tvm.tirx.layout import S, TileLayout, laneid
|
||
|
||
r_layout = TileLayout(S[(32, 8) : (1 @ laneid, 1)]); fs = (slice(0, 32), slice(0, 8))
|
||
|
||
@T.prim_func
|
||
def k(A_ptr: T.handle, B_ptr: T.handle):
|
||
A = T.match_buffer(A_ptr, (32, 8), "float32"); B = T.match_buffer(B_ptr, (32, 8), "float32")
|
||
T.device_entry(); T.cta_id([1]); T.lane_id([32]); tid = T.thread_id([32])
|
||
A_smem = T.alloc_buffer((32, 8), "float32", scope="shared", layout=TileLayout(S[(32, 8)]))
|
||
Tx.warp.copy(A_smem[fs], A[fs]); T.cuda.cta_sync()
|
||
R = T.alloc_buffer((32, 8), "float32", scope="local", layout=r_layout)
|
||
Tx.warp.copy(R[fs], A_smem[fs])
|
||
Tx.warp.sqrt(R[fs], R[fs]) # elementwise reg dispatch
|
||
Tx.warp.copy(A_smem[fs], R[fs]); T.cuda.cta_sync()
|
||
Tx.warp.copy(B[fs], A_smem[fs])
|
||
|
||
Algorithm
|
||
---------
|
||
|
||
**1. Parse and check.** ``spec.parse`` builds the op plan; the predicate confirms
|
||
every operand is a register buffer and the anchor register layout is valid.
|
||
|
||
**2. Induce the partition** from the anchor's thread axis (``laneid`` here): drop
|
||
the thread iters, leaving each thread its private bundle (8 elements per lane).
|
||
|
||
**3. Apply the op per register** — a per-thread loop over the bundle:
|
||
|
||
Generated TIRx IR
|
||
-----------------
|
||
|
||
.. code-block:: python
|
||
|
||
buffer[f] = T.sqrt(buffer_1[f]) # over each register f in the lane's bundle
|
||
|
||
Generated CUDA
|
||
--------------
|
||
|
||
.. code-block:: c++
|
||
|
||
r_local_ptr[f_2] = sqrtf(r_local_ptr[f_2]); // per-register, private to the lane
|
||
|
||
(Verified on ``sm_100a`` — the result equals ``sqrt(A)``.)
|
||
|
||
How inputs change the algorithm
|
||
-------------------------------
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 28 72
|
||
|
||
* - input
|
||
- effect
|
||
* - op
|
||
- unary → ``sqrtf`` / ``expf`` / … per register; binary → ``a + b`` per
|
||
register; ``fma`` → ``a * b + c``
|
||
* - dtype
|
||
- the register element type (``sqrtf`` vs ``hsqrt`` etc.); the bundle size is
|
||
the per-lane element count
|
||
* - register layout
|
||
- the anchor's thread axis sets the partition; a wider per-lane bundle means a
|
||
longer loop
|