wehub-resource-sync
252 lines
9.6 KiB
Python
252 lines
9.6 KiB
Python
# 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.
|
|
"""Reusable pipeline state and mbarrier helpers for SM100 kernels.
|
|
|
|
These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
|
|
instances are automatically treated as meta values inside @T.prim_func.
|
|
"""
|
|
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
@T.meta_class
|
|
class PipelineState:
|
|
"""Tracks stage and phase for a software-pipelined ring buffer.
|
|
|
|
This class does not know anything about full/empty barriers. Use it when
|
|
the kernel manually waits/signals barriers, or when the stage/phase drives
|
|
a ring not wrapped in a ``Pipeline``.
|
|
|
|
Parameters
|
|
----------
|
|
depth : int
|
|
Number of stages in the ring.
|
|
phase : int, optional
|
|
Initial phase. Omit when initialization should happen later.
|
|
"""
|
|
|
|
def __init__(self, depth: int, phase=None):
|
|
self.stage = T.local_scalar("int32")
|
|
self.phase = T.local_scalar("int32")
|
|
self.depth = depth
|
|
if phase is not None:
|
|
self.init(phase)
|
|
|
|
@T.inline
|
|
def init(self, phase):
|
|
self.stage = 0
|
|
self.phase = phase
|
|
|
|
@T.inline
|
|
def advance(self):
|
|
if self.depth > 1:
|
|
self.stage = self.stage + 1
|
|
if self.stage == self.depth:
|
|
self.stage = 0
|
|
self.phase = self.phase ^ 1
|
|
else:
|
|
self.phase = self.phase ^ 1
|
|
|
|
|
|
@T.meta_class
|
|
class MBarrier:
|
|
"""Mbarrier wrapper with regular ``mbarrier.arrive``.
|
|
|
|
Parameters
|
|
----------
|
|
pool : SMEMPool
|
|
Shared memory pool allocator.
|
|
depth : int
|
|
Number of barrier slots (one per pipeline stage).
|
|
phase_offset : int
|
|
XORed into the phase bit on every ``wait`` / ``arrive``.
|
|
leader : Expr, optional
|
|
Boolean predicate selecting the single thread that runs
|
|
``mbarrier.init``. Defaults to ``T.cuda.thread_rank() == 0`` --
|
|
thread 0 of the enclosing CTA, which always picks exactly one
|
|
thread regardless of which scope_id vars the caller declared.
|
|
Override only when you want a different CTA-local thread to do
|
|
the init.
|
|
|
|
Note: the default deliberately avoids ``T.warp_id()`` /
|
|
``T.lane_id()``. Those introduce deferred ``cta->warp`` /
|
|
``warp->thread`` ScopeIdDefs that the verifier cannot pin down
|
|
unless the kernel header declares the full warp/lane chain (e.g. a
|
|
single-CTA DSMEM kernel that only declares ``thread_id``). It also
|
|
avoids the synccheck false-deadlock on kernels that declare a
|
|
second warp-scope id. The generated CUDA is equivalent.
|
|
"""
|
|
|
|
def __init__(self, pool, depth, phase_offset=0, leader=None):
|
|
self.buf = pool.alloc((depth,), "uint64", align=8)
|
|
self.depth = depth
|
|
self.phase_offset = phase_offset
|
|
self.leader = leader if leader is not None else (T.cuda.thread_rank() == 0)
|
|
|
|
@T.inline
|
|
def init(self, count):
|
|
if self.leader:
|
|
for i in T.unroll(self.depth):
|
|
T.ptx.mbarrier.init(self.buf.ptr_to([i]), count)
|
|
|
|
@T.inline
|
|
def wait(self, stage, phase):
|
|
# Blocks: ``mbarrier.try_wait`` loops internally until the phase flips,
|
|
# so this returns only once the barrier has completed.
|
|
T.ptx.mbarrier.try_wait(self.buf.ptr_to([stage]), phase ^ self.phase_offset)
|
|
|
|
@T.inline
|
|
def arrive(self, stage, cta_id=None, pred=None, count=None):
|
|
# Default: local-CTA arrive — emits the simple
|
|
# ``mbarrier.arrive.shared.b64`` form. To arrive on a remote
|
|
# CTA's mbarrier in a cluster kernel, callers must pass
|
|
# ``cta_id=`` explicitly (e.g. ``bar.arrive(stage, cta_id=0)``)
|
|
# or use ``MBarrier.remote_view(rank).arrive(stage)``. Defaulting
|
|
# the cross-CTA path was both surprising (``bar.arrive(stage)``
|
|
# silently ``mapa`` ed across the cluster) and a per-call cost
|
|
# of ~3 PTX ops on every single-CTA kernel.
|
|
#
|
|
# ``count`` (cross-CTA path only) emits the explicit arrival-count
|
|
# operand, i.e. ``mbarrier.arrive.shared::cluster.b64 _, [addr], count``.
|
|
# When ``None`` the implicit count-of-1 form is emitted. Passing
|
|
# ``count=1`` is semantically identical but spells the count explicitly.
|
|
if cta_id is None:
|
|
T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
|
|
else:
|
|
actual_pred = True if pred is None else pred
|
|
T.ptx.mbarrier.arrive(
|
|
self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred, count=count
|
|
)
|
|
|
|
def ptr_to(self, idx):
|
|
return self.buf.ptr_to(idx)
|
|
|
|
def remote_view(self, rank):
|
|
"""Create a view of this barrier mapped to another CTA's shared memory.
|
|
|
|
Arrive-only: the returned view is built with ``object.__new__`` and
|
|
never copies ``self.leader``, so calling ``.init()`` on it would fail.
|
|
Use it solely to ``arrive`` on a remote CTA's mbarrier.
|
|
"""
|
|
from tvm.ir import PointerType, PrimType
|
|
from tvm.tirx import Var as TIRVar
|
|
|
|
expr = T.reinterpret("handle", T.ptx.map_shared_rank(self.buf.ptr_to([0]), rank))
|
|
ptr = TIRVar("remote_mbar_ptr", PointerType(PrimType("uint64")))
|
|
T.Bind(expr, var=ptr)
|
|
buf = T.decl_buffer([self.depth], "uint64", data=ptr, scope="shared")
|
|
remote = object.__new__(type(self))
|
|
remote.buf = buf
|
|
remote.depth = self.depth
|
|
remote.phase_offset = self.phase_offset
|
|
return remote
|
|
|
|
|
|
class TMABar(MBarrier):
|
|
"""Barrier signaled by TMA (mbarrier.arrive.expect_tx).
|
|
|
|
When ``tx_count`` is None, falls back to a remote mbarrier.arrive
|
|
(matching MBarrier.arrive defaults).
|
|
"""
|
|
|
|
@T.inline
|
|
def arrive(self, stage, tx_count=None, cta_id=None, pred=None):
|
|
# NOTE: this arrive() kwarg set intentionally differs from
|
|
# MBarrier.arrive (hardware necessity, LSP-incompatible by design).
|
|
# ``tx_count``: TMA byte count for ``mbarrier.arrive.expect_tx``.
|
|
# ``cta_id`` / ``pred``: forwarded to the underlying
|
|
# ``mbarrier.arrive`` (cluster path) when set; otherwise the
|
|
# arrive is local-CTA only. See ``MBarrier.arrive`` for the
|
|
# full default-local rationale.
|
|
if tx_count is not None:
|
|
T.ptx.mbarrier.arrive.expect_tx(self.buf.ptr_to([stage]), tx_count)
|
|
elif cta_id is None:
|
|
T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
|
|
else:
|
|
actual_pred = True if pred is None else pred
|
|
T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred)
|
|
|
|
|
|
class TCGen05Bar(MBarrier):
|
|
"""Barrier signaled by ``tcgen05`` commit.
|
|
|
|
The caller is responsible for ensuring only one thread issues the
|
|
commit, e.g. by wrapping the call in ``if T.ptx.elect_sync():``.
|
|
"""
|
|
|
|
@T.inline
|
|
def arrive(self, stage, cta_group=1, cta_mask=None):
|
|
# NOTE: this arrive() kwarg set intentionally differs from
|
|
# MBarrier.arrive (hardware necessity, LSP-incompatible by design).
|
|
if cta_mask is None and cta_group == 1:
|
|
T.ptx.tcgen05.commit(self.buf.ptr_to([stage]))
|
|
else:
|
|
T.ptx.tcgen05.commit(self.buf.ptr_to([stage]), cta_group=cta_group, cta_mask=cta_mask)
|
|
|
|
|
|
# Barrier-type tags accepted by Pipeline's ``full=`` / ``empty=`` arguments.
|
|
_BAR_KINDS = {"tma": TMABar, "tcgen05": TCGen05Bar, "mbar": MBarrier}
|
|
|
|
|
|
@T.meta_class
|
|
class Pipeline:
|
|
"""A full/empty mbarrier pair for a software-pipelined data flow.
|
|
|
|
Pass barrier-type tags and ``Pipeline`` constructs and ``init``\\ s the
|
|
barriers itself. Tags: ``"tma"`` (TMABar), ``"tcgen05"`` (TCGen05Bar),
|
|
``"mbar"`` (MBarrier). The barrier type and arrival count of each event
|
|
stay explicit at the call site -- e.g. ``Pipeline(pool, n, full="tma",
|
|
empty="tcgen05", init_empty=NUM_CONSUMER)``.
|
|
|
|
Both signals are required: a ``Pipeline`` is a *pair*. For a one-way event
|
|
(a pure "X happened" signal with no slot to recycle) use a bare barrier
|
|
(``TMABar``/``TCGen05Bar``/``MBarrier``) directly -- it has no empty side.
|
|
|
|
Parameters
|
|
----------
|
|
pool : SMEMPool
|
|
Shared memory pool allocator.
|
|
stages : int
|
|
Number of pipeline stages (barrier slots).
|
|
full, empty : str
|
|
Barrier-type tag for the full / empty signal (see above).
|
|
init_full, init_empty : int
|
|
Expected arrival count for the full / empty barrier.
|
|
empty_phase_offset : int
|
|
XORed into the empty barrier's phase bit on every wait / arrive.
|
|
leader : Expr, optional
|
|
Propagated to both barriers; defaults to thread 0 of the CTA.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pool,
|
|
stages,
|
|
*,
|
|
full,
|
|
empty,
|
|
init_full=1,
|
|
init_empty=1,
|
|
empty_phase_offset=0,
|
|
leader=None,
|
|
):
|
|
self.stages = stages
|
|
self.full = _BAR_KINDS[full](pool, stages, leader=leader)
|
|
self.full.init(init_full)
|
|
self.empty = _BAR_KINDS[empty](pool, stages, phase_offset=empty_phase_offset, leader=leader)
|
|
self.empty.init(init_empty)
|