Skip to content

[Perf] gemm_bf16 gfx1250: cross-tile carry to hide K-tile prologue ds… #1125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
amd-hhashemi wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Perf] gemm_bf16 gfx1250: cross-tile carry to hide K-tile prologue ds… #1125

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 41 additions & 12 deletions kernels/gemm/gemm_bf16_gfx1250.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,16 @@ def load_b(buf, wn, ks):
for cf in c_frags:
cf.store(Vec.filled(8, 0.0, fx.Float32))

# Cross-tile carry: each tile defers its LAST K-step's WMMAs to the next tile so the
# next tile's step-0 read latency hides behind them. carry_* hold that deferred K-step's
# A/B fragments; zeroed so the first tile's carry-WMMA adds 0 (no prologue peel needed).
carry_act = [fx.make_rmem_tensor(8, fx.Int32) for _ in range_constexpr(wmma_m_rep)]
carry_wt = [fx.make_rmem_tensor(8, fx.Int32) for _ in range_constexpr(wmma_n_rep)]
for t in carry_act:
t.store(Vec.filled(8, 0, fx.Int32))
for t in carry_wt:
t.store(Vec.filled(8, 0, fx.Int32))

def _rmem(n, v):
t = fx.make_rmem_tensor(n, fx.Int32)
t.store(v)
Expand All @@ -213,20 +223,37 @@ def _mma_ks(state):
# B is the instruction's A operand: the accumulator's fast dim is N.
fx.gemm(wmma_atom, c_frags[idx], wt[wn], act[wm], c_frags[idx])

def compute_ktile(buf, prefetch_kt):
cur = _load_ks(buf, 0)
for ks in range_constexpr(K_WS):
nxt = _load_ks(buf, ks + 1) if const_expr(ks + 1 < K_WS) else None
rocdl.s_wait_dscnt(KS_DS if const_expr(nxt is not None) else 0)
if const_expr(ks == 0 and prefetch_kt is not None and wmma_m_rep > 1):
def _mma_carry():
for wm in range_constexpr(wmma_m_rep):
for wn_raw in range_constexpr(wmma_n_rep):
wn = (wmma_n_rep - 1 - wn_raw) if (wm % 2 == 1) else wn_raw
idx = wm * wmma_n_rep + wn
fx.gemm(wmma_atom, c_frags[idx], carry_wt[wn], carry_act[wm], c_frags[idx])

def _load_ks_carry(buf):
# Issue this tile's LAST K-step reads directly into the persistent carry tensors
# (no copy); consumed by the next tile's _mma_carry after its s_wait_dscnt.
for wm in range_constexpr(wmma_m_rep):
carry_act[wm].store(load_a(buf, wm, K_WS - 1))
for wn in range_constexpr(wmma_n_rep):
carry_wt[wn].store(load_b(buf, wn, K_WS - 1))

def compute_head(buf, prefetch_kt):
cur = _load_ks(buf, 0) # this tile's step-0 reads
rocdl.s_wait_dscnt(KS_DS) # prev tile's carry reads done; step-0 in flight
_mma_carry() # prev tile's deferred K-step -> hides step-0 latency
for ks in range_constexpr(K_WS - 1): # compute steps 0 .. K_WS-2
if const_expr(ks < K_WS - 2):
nxt = _load_ks(buf, ks + 1) # next head step (fresh)
else:
_load_ks_carry(buf) # step K_WS-1 -> carry (deferred to next tile)
nxt = None
rocdl.s_wait_dscnt(KS_DS)
if const_expr(ks == 0 and prefetch_kt is not None):
rocdl.sched_barrier(0)
issue(prefetch_kt % num_buffers, prefetch_kt)
rocdl.sched_barrier(0)
_mma_ks(cur)
if const_expr(ks == 0 and prefetch_kt is not None and wmma_m_rep == 1):
rocdl.sched_barrier(0)
issue(prefetch_kt % num_buffers, prefetch_kt)
rocdl.sched_barrier(0)
if const_expr(nxt is not None):
cur = nxt
rocdl.sched_dsrd(KS_DS) # prologue group
Expand All @@ -244,15 +271,17 @@ def compute_ktile(buf, prefetch_kt):
for kt in range(n_steady):
buf = _bidx(_buf_ptr(kt % num_buffers))
pipeline_fence(outstanding=TDM_PW * (num_buffers - 2), use_cluster=False)
compute_ktile(buf, kt + (num_buffers - 1))
compute_head(buf, kt + (num_buffers - 1))
if const_expr(use_cluster) and kt % num_buffers == num_buffers - 1:
cluster.cluster_barrier()
for j in range_constexpr(num_buffers - 1):
kt = n_steady + j
buf = _bidx(_buf_ptr(kt % num_buffers))
pipeline_fence(outstanding=TDM_PW * (num_buffers - 2 - j), use_cluster=False)
compute_ktile(buf, None)
compute_head(buf, None)

rocdl.s_wait_dscnt(0) # drain the final tile's deferred K-step
_mma_carry()
accs = [c_frags[idx].load() for idx in range_constexpr(n_acc)]
pipeline_fence(outstanding=0, use_cluster=use_cluster)
for wm in range_constexpr(wmma_m_rep):
Expand Down
Loading