derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
1219 lines
62 KiB
Python
1219 lines
62 KiB
Python
# Copyright Advanced Micro Devices, Inc., or its affiliates.
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
################################################################################
|
|
# GR (global read) emit and alloc dispatch.
|
|
#
|
|
# singledispatch over GR tag sentinels (GRTag_1x2, GRTag_TLU1, etc.).
|
|
# ABGRTile calls these via self.config.tag as the dispatch key.
|
|
#
|
|
# Structure:
|
|
# 1. Dispatch bases — @singledispatch declarations
|
|
# 2. Implementations — logic functions decorated with @register
|
|
#
|
|
# To add a new GR shape: define a new tag in SubtileGeometry.py, add geometry
|
|
# instances with that tag, and register a new implementation below.
|
|
# LR emit lives in a separate file (SubtileLREmit.py).
|
|
################################################################################
|
|
|
|
import math
|
|
from functools import singledispatch
|
|
|
|
from rocisa.code import Module
|
|
from rocisa.container import DPPModifiers, EXEC, MUBUFModifiers, VCC, vgpr, sgpr, mgpr
|
|
from rocisa.enum import RegisterType
|
|
from rocisa.instruction import (
|
|
BufferLoadB128,
|
|
SAddCU32, SAddU32, SAddU64, SAndB32, SMaxI32, SMinU32, SMovB32, SMovB64, SMulI32,
|
|
SNop, SOrB32, SSubI32, SXorB32,
|
|
SCBranchSCC1, SCmpEQU32, SEndpgm,
|
|
SLShiftLeftB64, SLShiftRightB32,
|
|
VAddU32, VAndB32, VCmpXEqU32,
|
|
VLShiftLeftB32, VLShiftRightB32, VMovB32,
|
|
TensorLoadToLds,
|
|
VMulLOU32, VReadfirstlaneB32, VSubU32, VXorB32,
|
|
)
|
|
|
|
from .SubtileGeometry import (
|
|
RegList,
|
|
GRTag_1x1, GRTag_1x2, GRTag_2x2, GRTag_TLU1,
|
|
)
|
|
from .SubtileScaleEmit import emitScaleGRLDSSwap
|
|
|
|
from math import ceil, log, log2, prod
|
|
from rocisa.code import Label
|
|
from ...Common import INDEX_CHARS, clusterEnabled
|
|
from ...Common.DataType import DataType
|
|
|
|
|
|
################################################################################
|
|
# 1. Dispatch bases
|
|
################################################################################
|
|
|
|
@singledispatch
|
|
def _emitGlobalReadOffset(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"emitGlobalReadOffset not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _emitGlobalRead(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"emitGlobalRead not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _emitLocalWrite(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"emitLocalWrite not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _allocGROffsetRegisters(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"allocGROffsetRegisters not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _deallocGROffsetRegisters(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"deallocGROffsetRegisters not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _emitDTLInit(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"emitDTLInit not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _emitGRLDSBufferSwap(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"emitGRLDSBufferSwap not implemented for {type(tag).__name__}")
|
|
|
|
@singledispatch
|
|
def _emitGRPtrUpdate(tag, tile, ti, writer, kernel):
|
|
raise NotImplementedError(f"emitGRPtrUpdate not implemented for {type(tag).__name__}")
|
|
|
|
# Stubs for tags not yet implemented.
|
|
_stub = lambda tag, tile, ti, writer, kernel: None
|
|
_emitGlobalReadOffset.register(GRTag_TLU1)(_stub)
|
|
_allocGROffsetRegisters.register(GRTag_TLU1)(_stub)
|
|
_deallocGROffsetRegisters.register(GRTag_TLU1)(_stub)
|
|
_emitGlobalRead.register(GRTag_TLU1)(_stub)
|
|
_emitDTLInit.register(GRTag_TLU1)(_stub)
|
|
_emitGRLDSBufferSwap.register(GRTag_TLU1)(_stub)
|
|
_emitGRPtrUpdate.register(GRTag_TLU1)(_stub)
|
|
for _tag in (GRTag_1x1, GRTag_1x2, GRTag_2x2, GRTag_TLU1):
|
|
_emitLocalWrite.register(_tag)(_stub)
|
|
|
|
|
|
################################################################################
|
|
# 2. Implementations — TLU=0 (shared by GRTag_1x2 and GRTag_2x2)
|
|
################################################################################
|
|
|
|
@_emitGlobalReadOffset.register(GRTag_1x1)
|
|
@_emitGlobalReadOffset.register(GRTag_1x2)
|
|
@_emitGlobalReadOffset.register(GRTag_2x2)
|
|
def _emitGROffset_TLU0(tag, tile, ti, writer, kernel):
|
|
return Module(f"GR Offset TLU0 ({ti.tc})") # STUB — legacy path in graTileAssignment
|
|
"""GR offset for row-major (TLU=0) geometry with swizzling and rotation.
|
|
|
|
Ported from legacy graTileAssignment. Operates on a single tensor component.
|
|
|
|
1. Compute waveId, laneId, colId, rowId from Serial (v0).
|
|
2. Swizzle colId via DPP quad_perm to avoid LDS bank conflicts.
|
|
3. Intra-wave rotation: shift colId based on LDS row parity.
|
|
4. Inter-wave rotation: additional shift from waveId (when waves_coop > 1).
|
|
5. Unified wave partition: localRow + partitionRow from waveId.
|
|
6. Compute byte offsets for each GR load into sharedVgprGROffset[].
|
|
7. Compute subtile perpendicular soffsets.
|
|
"""
|
|
module = Module(f"GR Offset TLU0 ({ti.tc})")
|
|
tc = ti.tc
|
|
loadWidth = ti.loadWidthGR
|
|
subIterKBytes = ti.subIterKBytes
|
|
blockSize = subIterKBytes // loadWidth
|
|
wavesize = kernel["WavefrontSize"]
|
|
bpe = ti.bpe
|
|
bpeBits = int(8 * bpe)
|
|
strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
|
|
ldsRowBankSize = writer.states.archCaps["LDSBankCount"] * writer.states.archCaps["LDSBankWidth"]
|
|
|
|
wg_m = ti.waveGroupSize
|
|
numWaves = ti.numWaves
|
|
waves_coop = numWaves // wg_m
|
|
numRowsPerWave = wavesize // blockSize
|
|
numRowsPerLDSBanks = ldsRowBankSize // subIterKBytes
|
|
|
|
tmpVgpr = writer.vgprPool.checkOut(4, tag="_emitGROffset_TLU0_tmpVgpr")
|
|
colId = tmpVgpr
|
|
rowId = tmpVgpr + 1
|
|
waveId = tmpVgpr + 2
|
|
localRow = tmpVgpr + 3
|
|
tmpSgpr = writer.sgprPool.checkOut(1, tag="_emitGROffset_TLU0_tmpSgpr", preventOverflow=False)
|
|
|
|
# --- 1. waveId, laneId, colId, rowId ---
|
|
module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1),
|
|
src=vgpr("Serial"), comment=f"{tc}: waveId"))
|
|
module.add(VAndB32(dst=vgpr(localRow), src0=vgpr("Serial"), src1=wavesize-1,
|
|
comment=f"{tc}: laneId"))
|
|
module.add(VAndB32(dst=vgpr(colId), src0=vgpr("Serial"), src1=blockSize-1,
|
|
comment=f"{tc}: colId for {loadWidth}B load"))
|
|
module.add(VLShiftRightB32(dst=vgpr(rowId), shiftHex=hex(blockSize.bit_length()-1),
|
|
src=vgpr(localRow), comment=f"{tc}: rowId within wave"))
|
|
|
|
# --- 2. Swizzle: DPP quad_perm swap colId pairs on even LDS rows ---
|
|
tmpSwz = writer.vgprPool.checkOut(2, tag="_emitGROffset_TLU0_tmpSwz")
|
|
ldsRowId = tmpSwz
|
|
swzTmp = tmpSwz + 1
|
|
|
|
module.addComment0(f"{tc}: Swizzling")
|
|
module.add(VLShiftRightB32(dst=vgpr(ldsRowId), shiftHex=hex(blockSize.bit_length()-1),
|
|
src=vgpr(localRow), comment=f"{tc}: row id within wave"))
|
|
module.add(VLShiftRightB32(dst=vgpr(ldsRowId), shiftHex=hex(numRowsPerLDSBanks.bit_length()-1),
|
|
src=vgpr(ldsRowId), comment=f"{tc}: lds row id"))
|
|
module.add(VAndB32(dst=vgpr(swzTmp), src0=vgpr(ldsRowId), src1=hex(1),
|
|
comment=f"{tc}: lds row id %% 2"))
|
|
module.add(VCmpXEqU32(dst=VCC(), src0=0, src1=vgpr(swzTmp),
|
|
comment=f"{tc}: lds row id %% 2 == 0?"))
|
|
module.add(VMovB32(dst=vgpr(colId), src=vgpr(colId), dpp=DPPModifiers(quad_perm=[1,0,3,2]),
|
|
comment=f"{tc}: swap colId pairs"))
|
|
module.add(SMovB64(dst=EXEC(), src=-1))
|
|
|
|
# --- 3. Intra-wave rotation: blockSize - (ldsRowId // 2) * 2 ---
|
|
module.addComment0(f"{tc}: Intra-wave rotation")
|
|
module.add(VLShiftRightB32(dst=vgpr(swzTmp), shiftHex=hex(1), src=vgpr(ldsRowId)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(swzTmp), shiftHex=hex(1), src=vgpr(swzTmp),
|
|
comment=f"{tc}: (ldsRowId // 2) * 2"))
|
|
module.add(VSubU32(dst=vgpr(swzTmp), src0=hex(blockSize), src1=vgpr(swzTmp),
|
|
comment=f"{tc}: rotation = blockSize - (ldsRowId//2)*2"))
|
|
|
|
# --- 4. Inter-wave rotation (when waves cooperate on a subtile) ---
|
|
if waves_coop > 1:
|
|
waveRotation = writer.vgprPool.checkOut(1, tag="_emitGROffset_TLU0_waveRotation")
|
|
module.addComment0(f"{tc}: Inter-wave rotation")
|
|
module.add(VAndB32(dst=vgpr(waveRotation), src0=vgpr(waveId), src1=hex(1)))
|
|
module.add(VLShiftLeftB32(dst=vgpr(waveRotation),
|
|
shiftHex=hex((2*numRowsPerLDSBanks).bit_length() - 1), src=vgpr(waveRotation)))
|
|
module.add(VSubU32(dst=vgpr(waveRotation), src0=vgpr(swzTmp), src1=vgpr(waveRotation)))
|
|
module.add(VAddU32(dst=vgpr(colId), src0=vgpr(waveRotation), src1=vgpr(colId)))
|
|
writer.vgprPool.checkIn(waveRotation)
|
|
else:
|
|
module.add(VAddU32(dst=vgpr(colId), src0=vgpr(swzTmp), src1=vgpr(colId)))
|
|
|
|
module.add(VAndB32(dst=vgpr(colId), src0=vgpr(colId), src1=hex(blockSize-1),
|
|
comment=f"{tc}: (col + rotation) %% blockSize"))
|
|
writer.vgprPool.checkIn(tmpSwz)
|
|
|
|
# --- 5. Unified wave partition ---
|
|
rowOffset = writer.vgprPool.checkOut(1, tag="_emitGROffset_TLU0_rowOffset")
|
|
partitionStride = ti.mmaTileShape[0] * int(ti.localSubtileGrid[0])
|
|
waves_coop_shift = max(0, waves_coop.bit_length() - 1) if waves_coop > 0 else 0
|
|
module.add(VAndB32(dst=vgpr(localRow), src0=hex(waves_coop - 1), src1=vgpr(waveId),
|
|
comment=f"{tc}: waveId %% {waves_coop}"))
|
|
module.add(VLShiftRightB32(dst=vgpr(rowOffset), shiftHex=hex(waves_coop_shift),
|
|
src=vgpr(waveId), comment=f"{tc}: waveId // {waves_coop}"))
|
|
module.add(VLShiftLeftB32(dst=vgpr(localRow), shiftHex=hex(numRowsPerWave.bit_length()-1),
|
|
src=vgpr(localRow), comment=f"{tc}: local row * {numRowsPerWave}"))
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=partitionStride,
|
|
comment=f"{tc}: partition stride"))
|
|
module.add(VMulLOU32(dst=vgpr(rowOffset), src0=sgpr(tmpSgpr), src1=vgpr(rowOffset),
|
|
comment=f"{tc}: partition row offset"))
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=vgpr(localRow), src1=vgpr(rowOffset),
|
|
comment=f"{tc}: + local row"))
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=vgpr(rowId), src1=vgpr(rowOffset),
|
|
comment=f"{tc}: + lane rowId"))
|
|
|
|
# --- 6. Compute byte offsets for each GR load ---
|
|
tmpVgpr2 = writer.vgprPool.checkOut(2, tag="_emitGROffset_TLU0_tmpVgpr2")
|
|
colBytes = tmpVgpr2 + 1
|
|
for i in range(ti.numGRPerSubtile):
|
|
useColId = colId
|
|
# For numGRPerSubtile > 1 with single-wave subtiles: rotate colId between loads
|
|
if i > 0 and waves_coop == 1 and ti.numGRPerSubtile > 1:
|
|
rotatedCol = writer.vgprPool.checkOut(1, tag="_emitGROffset_TLU0_rotatedCol")
|
|
colRotation = blockSize // 2
|
|
module.add(VAddU32(dst=vgpr(rotatedCol), src0=colRotation, src1=vgpr(colId),
|
|
comment=f"{tc}: rotate col for GR {i}"))
|
|
module.add(VAndB32(dst=vgpr(rotatedCol), src0=vgpr(rotatedCol), src1=hex(blockSize-1),
|
|
comment=f"{tc}: (col + {colRotation}) %% blockSize"))
|
|
useColId = rotatedCol
|
|
|
|
module.add(VLShiftLeftB32(dst=vgpr(colBytes), shiftHex=hex(loadWidth.bit_length()-1),
|
|
src=vgpr(useColId), comment=f"{tc}: colId * {loadWidth}"))
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr2), src0=sgpr(strideRef), src1=vgpr(rowOffset),
|
|
comment=f"{tc}: rowOffset * stride"))
|
|
module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr2), shiftHex=hex(bpeBits.bit_length()-1),
|
|
src=vgpr(tmpVgpr2), comment=f"{tc}: * bpe"))
|
|
module.add(VLShiftRightB32(dst=vgpr(tmpVgpr2), shiftHex=hex(3), src=vgpr(tmpVgpr2),
|
|
comment=f"{tc}: bits to bytes"))
|
|
module.add(VAddU32(dst=vgpr(tile.sharedVgprGROffset[i]), src0=vgpr(colBytes), src1=vgpr(tmpVgpr2),
|
|
comment=f"{tc}: GR offset {i}"))
|
|
|
|
if i > 0 and waves_coop == 1 and ti.numGRPerSubtile > 1:
|
|
writer.vgprPool.checkIn(rotatedCol)
|
|
|
|
if i + 1 < ti.numGRPerSubtile:
|
|
advance = ti.subtileShape[0] * ti.mmaTileShape[0] // ti.numGRPerSubtile
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=advance, src1=vgpr(rowOffset),
|
|
comment=f"{tc}: advance row for GR {i+1}"))
|
|
writer.vgprPool.checkIn(tmpVgpr2)
|
|
|
|
# --- 7. Subtile perpendicular soffsets ---
|
|
subtileRowElements = ti.subtileShape[0] * ti.mmaTileShape[0]
|
|
s_stride_bpe = int(subtileRowElements * bpe)
|
|
for reg_idx in range(len(ti.localSubtilesRegister)):
|
|
rl = ti.localSubtilesRegister[reg_idx]
|
|
if len(rl) == 0:
|
|
continue
|
|
if rl.is_sgpr:
|
|
module.add(SMulI32(dst=rl.ref(0), src0=hex(s_stride_bpe * reg_idx),
|
|
src1=sgpr(strideRef), comment=f"{tc}: subtile row {reg_idx} soffset"))
|
|
else:
|
|
stmp = writer.sgprPool.checkOut(1, tag="_emitGROffset_TLU0_stmp")
|
|
for i, reg in enumerate(rl):
|
|
module.add(SMulI32(dst=sgpr(stmp), src0=hex(s_stride_bpe * reg_idx),
|
|
src1=sgpr(strideRef), comment=f"{tc}: subtile row {reg_idx} soffset"))
|
|
module.add(VAddU32(dst=vgpr(reg), src0=vgpr(tile.sharedVgprGROffset[i]), src1=sgpr(stmp),
|
|
comment=f"{tc}: bake soffset into vgpr"))
|
|
writer.sgprPool.checkIn(stmp)
|
|
|
|
writer.vgprPool.checkIn(rowOffset)
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
return module
|
|
|
|
|
|
@_allocGROffsetRegisters.register(GRTag_1x1)
|
|
@_allocGROffsetRegisters.register(GRTag_1x2)
|
|
@_allocGROffsetRegisters.register(GRTag_2x2)
|
|
def _allocGROffsetRegs_TLU0(tag, tile, ti, writer, kernel):
|
|
"""Allocate GR offset registers for TLU=0 shapes.
|
|
|
|
Two register groups are allocated:
|
|
|
|
1. sharedVgprGROffset[]: one VGPR per GR load within a subtile.
|
|
These hold per-lane byte offsets for buffer_load (colId * loadWidth +
|
|
rowOffset * stride * bpe). Shared across all subtile rows — only the
|
|
soffset changes between rows.
|
|
|
|
2. localSubtilesRegister[]: one RegList per perpendicular subtile row.
|
|
Each entry holds the constant M-direction offset (soffset) that shifts
|
|
the shared VGPR offset to the correct subtile row.
|
|
|
|
Row 0 needs no offset (soffset=0), so its RegList is left empty.
|
|
Row 1+ gets either:
|
|
- 1 SGPR (preferred): used as the soffset field in buffer_load.
|
|
The shared VGPR offset is reused as-is across rows.
|
|
- numGRPerSubtile VGPRs (fallback when SGPRs exhausted): each VGPR
|
|
has the shared offset + row offset baked in, replacing soffset.
|
|
"""
|
|
# TDM handles global reads without per-lane VGPRs
|
|
hasTDM = kernel.get("enableTDMA", False) and kernel.get("enableTDMB", False)
|
|
if hasTDM:
|
|
tile.sharedVgprGROffset = []
|
|
ti.localSubtilesRegister = []
|
|
return
|
|
|
|
# Per-lane byte offsets: one VGPR per GR load within a subtile
|
|
tile.sharedVgprGROffset = []
|
|
for i in range(ti.numGRPerSubtile):
|
|
tile.sharedVgprGROffset.append(writer.vgprPool.checkOut(1, tag="_allocGROffsetRegs_TLU0_sharedVgprGROffset"))
|
|
|
|
# Per-subtile-row soffset registers.
|
|
# perpDimSize = how many GR subtile shapes tile the perpendicular (M) dimension
|
|
# per wave. Each position needs its own soffset register.
|
|
ti.localSubtilesRegister = []
|
|
# perpDimSize: distinct soffset positions in M = how many localSubtile rows
|
|
# need their own soffset register. localGRGranularity[0] tells how many
|
|
# consecutive localSubtile rows one GR load covers (>1 only for bc==1 with
|
|
# wave-cooperative expansion, i.e. loadRatioGR > 1).
|
|
localSubtileRowCount = int(ti.localSubtileGrid[0])
|
|
gran = tile.localGRGranularity(ti.numWaves)
|
|
perpDimSize = math.ceil(localSubtileRowCount / gran[0])
|
|
tmpSgprBuffer = 3
|
|
sgprLimit = writer.states.regCaps["MaxSgpr"] - tmpSgprBuffer
|
|
|
|
for reg_idx in range(perpDimSize):
|
|
useSgpr = writer.sgprPool.size() < sgprLimit
|
|
if useSgpr:
|
|
rl = RegList(writer.sgprPool, RegisterType.Sgpr)
|
|
else:
|
|
rl = RegList(writer.vgprPool, RegisterType.Vgpr)
|
|
ti.localSubtilesRegister.append(rl)
|
|
# Row 0 is the base position — no soffset needed, RegList stays empty.
|
|
if reg_idx == 0:
|
|
continue
|
|
if useSgpr:
|
|
# SGPR path: 1 register for soffset, shared VGPR offset reused.
|
|
rl.alloc(preventOverflow=False)
|
|
else:
|
|
# VGPR fallback: one VGPR per GR load, each with soffset baked in.
|
|
for i in range(ti.numGRPerSubtile):
|
|
rl.alloc(preventOverflow=False)
|
|
|
|
|
|
@_deallocGROffsetRegisters.register(GRTag_1x1)
|
|
@_deallocGROffsetRegisters.register(GRTag_1x2)
|
|
@_deallocGROffsetRegisters.register(GRTag_2x2)
|
|
def _deallocGROffsetRegs_TLU0(tag, tile, ti, writer, kernel):
|
|
"""Deallocate GR offset registers for TLU=0 shapes."""
|
|
if isinstance(tile.sharedVgprGROffset, list):
|
|
for voff in tile.sharedVgprGROffset:
|
|
writer.vgprPool.checkIn(voff)
|
|
tile.sharedVgprGROffset = []
|
|
if isinstance(ti.localSubtilesRegister, list):
|
|
for rl in ti.localSubtilesRegister:
|
|
rl.dealloc()
|
|
ti.localSubtilesRegister = []
|
|
|
|
|
|
# --- GR load emit (TLU=0) ---------------------------------------------------
|
|
|
|
@_emitGlobalRead.register(GRTag_1x1)
|
|
@_emitGlobalRead.register(GRTag_1x2)
|
|
@_emitGlobalRead.register(GRTag_2x2)
|
|
def _emitGR_TLU0(tag, tile, ti, writer, kernel):
|
|
"""Emit buffer_load_dwordx4 (DTL) for all subtiles in the local grid.
|
|
|
|
For each subtile (sId0, sId1):
|
|
- Computes LDS write address (m0) from LocalWriteBaseAddr + subtile offset.
|
|
- Emits buffer_load_b128 with lds=True (direct-to-LDS).
|
|
- Uses soffset (SGPR path) or baked VGPR offset for the subtile row.
|
|
|
|
When loadRatioGR > 1, multiple subtiles share one GR load; only the first
|
|
subtile in each group emits the load.
|
|
"""
|
|
module = Module(f"GR Load TLU0 ({ti.tc})")
|
|
tc = ti.tc
|
|
isGlc = bool(kernel.get(f"NonTemporal{tc}", 0) & 0x1)
|
|
isSlc = bool(kernel.get(f"NonTemporal{tc}", 0) & 0x2)
|
|
isNT = bool(kernel.get(f"NonTemporal{tc}", 0) & 0x4)
|
|
|
|
perpDimSize = len(ti.localSubtilesRegister)
|
|
|
|
# TODO: Remove legacy TileInfo dependency after full migration.
|
|
# Currently uses legacy's grid/sizes because subtileShape expansion in for_kernel
|
|
# changes subtileSize/localSubtileGrid/loadRatioGR, which must match the LDS
|
|
# layout computed from legacy values.
|
|
legacyTi = getattr(writer.states, tc.lower()).tileInfo
|
|
localGrid0 = int(legacyTi.localSubtileGrid[0])
|
|
localGrid1 = int(legacyTi.localSubtileGrid[1])
|
|
legacyLoadRatio = legacyTi.loadRatioGR
|
|
legacySubtileSize = int(legacyTi.subtileSize)
|
|
|
|
for j in range(localGrid1):
|
|
for i in range(localGrid0):
|
|
slowId = i
|
|
if legacyLoadRatio == 2.0:
|
|
slowId = int(i // legacyLoadRatio)
|
|
reg_idx = slowId
|
|
|
|
# Skip duplicate loads when loadRatio > 1
|
|
if legacyLoadRatio > 1:
|
|
linearId = j * localGrid0 + i
|
|
grBaseId = int(linearId // legacyLoadRatio)
|
|
firstInGroup = int(grBaseId * legacyLoadRatio)
|
|
if linearId != firstInGroup:
|
|
continue
|
|
|
|
rl = ti.localSubtilesRegister[min(reg_idx, perpDimSize - 1)]
|
|
offsetK = j * int(ti.mmaTileShape[1] * ti.subtileShape[1] * ti.bpe)
|
|
|
|
module.addComment0(f"GR load {tc} subtile [{i},{j}]")
|
|
|
|
subtileOffset = int(math.ceil(legacyLoadRatio * legacySubtileSize)) if legacyLoadRatio else legacySubtileSize
|
|
WriteBaseAddr = f"LocalWriteBaseAddr{tc}"
|
|
|
|
for gr_idx in range(legacyTi.numGRPerSubtile):
|
|
m0Offset = gr_idx * subtileOffset + (i + j * int(legacyTi.globalSubtileGrid[0])) * legacySubtileSize
|
|
module.add(SAddU32(dst=mgpr(0), src0=sgpr(WriteBaseAddr), src1=(m0Offset - offsetK)))
|
|
mubuf = MUBUFModifiers(offen=True, offset12=offsetK, glc=isGlc, slc=isSlc, nt=isNT, lds=True)
|
|
|
|
use_sgpr = rl.is_sgpr if len(rl) > 0 else True
|
|
soffset = rl.ref(0) if len(rl) > 0 and use_sgpr else 0
|
|
voff = tile.sharedVgprGROffset[gr_idx] if use_sgpr or len(rl) == 0 else rl.indices[gr_idx]
|
|
module.add(BufferLoadB128(dst=None, vaddr=vgpr(voff), saddr=sgpr(f"Srd{tc}", 4),
|
|
soffset=soffset, mubuf=mubuf, comment=f"GR{gr_idx} [{i},{j}]"))
|
|
|
|
return module
|
|
|
|
|
|
# --- DTL init (TLU=0) -------------------------------------------------------
|
|
|
|
@_emitDTLInit.register(GRTag_1x1)
|
|
@_emitDTLInit.register(GRTag_1x2)
|
|
@_emitDTLInit.register(GRTag_2x2)
|
|
def _emitDTLInit_TLU0(tag, tile, ti, writer, kernel):
|
|
return Module(f"DTL Init ({ti.tc})") # STUB — legacy path in globalReadDTLInitCommonSgpr
|
|
"""Compute LocalWriteBaseAddr and Swap SGPR for one tensor component.
|
|
|
|
The DTL (direct-to-LDS) buffer_load writes data at m0 = LocalWriteBaseAddr + subtile offset.
|
|
LocalWriteBaseAddr is the wave's base LDS position, derived from the wave partition.
|
|
Swap holds the XOR mask to toggle between double-buffer halves.
|
|
|
|
For double-buffering: LocalWriteBaseAddr XOR Swap flips to the other buffer.
|
|
|
|
Requires sgprs: LocalWriteBaseAddr{tc}, Swap{tc} (must be pre-allocated by caller).
|
|
"""
|
|
module = Module(f"DTL Init ({ti.tc})")
|
|
tc = ti.tc
|
|
wavesize = kernel["WavefrontSize"]
|
|
wg_m = ti.waveGroupSize
|
|
numWaves = ti.numWaves
|
|
waves_coop = numWaves // wg_m
|
|
|
|
vgprWaveId = writer.vgprPool.checkOut(1, tag="_emitDTLInit_TLU0_vgprWaveId")
|
|
rowOffset = writer.vgprPool.checkOut(1, tag="_emitDTLInit_TLU0_rowOffset")
|
|
|
|
module.add(VLShiftRightB32(dst=vgpr(vgprWaveId), shiftHex=hex(wavesize.bit_length()-1),
|
|
src=vgpr("Serial"), comment=f"{tc}: waveId"))
|
|
|
|
# Wave partition: same unified formula as GR offset step 5
|
|
numRowsPerWave = wavesize // (ti.subIterKBytes // ti.loadWidthGR)
|
|
partitionStride = ti.mmaTileShape[0] * int(ti.localSubtileGrid[0])
|
|
waves_coop_shift = max(0, waves_coop.bit_length() - 1) if waves_coop > 0 else 0
|
|
|
|
module.add(VLShiftRightB32(dst=vgpr(rowOffset), shiftHex=hex(waves_coop_shift),
|
|
src=vgpr(vgprWaveId), comment=f"{tc}: partitionRow = waveId // {waves_coop}"))
|
|
tmpSgpr = writer.sgprPool.checkOut(1, tag="_emitDTLInit_TLU0_tmpSgpr", preventOverflow=False)
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=partitionStride))
|
|
module.add(VMulLOU32(dst=vgpr(rowOffset), src0=sgpr(tmpSgpr), src1=vgpr(rowOffset),
|
|
comment=f"{tc}: partition row offset"))
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
# Scale by subIterKBytes to get LDS byte offset
|
|
module.add(VLShiftLeftB32(dst=vgpr(rowOffset),
|
|
shiftHex=hex(ti.subIterKBytes.bit_length()-1), src=vgpr(rowOffset),
|
|
comment=f"{tc}: * subIterKBytes"))
|
|
|
|
# Move to SGPR via readfirstlane (uniform across wave)
|
|
module.add(SNop(waitState=0, comment="wait for VGPR"))
|
|
WriteBaseAddr = f"LocalWriteBaseAddr{tc}"
|
|
Swap = f"Swap{tc}"
|
|
module.add(VReadfirstlaneB32(dst=sgpr(WriteBaseAddr), src=vgpr(rowOffset),
|
|
comment=f"{tc}: base LDS offset"))
|
|
|
|
# Add global LDS start offset for B (B data follows A in LDS)
|
|
ldsStartOffset = getattr(writer, f'ldsStartOffset{tc}', 0)
|
|
if ldsStartOffset:
|
|
module.add(SAddU32(dst=sgpr(WriteBaseAddr), src0=sgpr(WriteBaseAddr),
|
|
src1=hex(ldsStartOffset), comment=f"{tc}: + ldsStartOffset"))
|
|
|
|
# Swap mask: XOR(base, base + ldsTotalSize) toggles between buffer halves
|
|
module.add(SAddU32(dst=sgpr(Swap), src0=sgpr(WriteBaseAddr), src1=writer.ldsTotalSize))
|
|
module.add(SXorB32(dst=sgpr(Swap), src0=sgpr(WriteBaseAddr), src1=sgpr(Swap)))
|
|
|
|
writer.vgprPool.checkIn(vgprWaveId)
|
|
writer.vgprPool.checkIn(rowOffset)
|
|
return module
|
|
|
|
|
|
# --- GR LDS buffer swap (TLU=0) ---------------------------------------------
|
|
|
|
@_emitGRLDSBufferSwap.register(GRTag_1x1)
|
|
@_emitGRLDSBufferSwap.register(GRTag_1x2)
|
|
@_emitGRLDSBufferSwap.register(GRTag_2x2)
|
|
def _emitGRLDSSwap_TLU0(tag, tile, ti, writer, kernel):
|
|
"""Toggle GR DTL write target between double-buffer halves.
|
|
|
|
XOR LocalWriteBaseAddr with Swap to flip to the other LDS buffer.
|
|
"""
|
|
module = Module()
|
|
tc = ti.tc
|
|
module.addComment0("Emit code to swap %s GR m0 offsets"%tc)
|
|
module.add(SXorB32(dst=sgpr(f"LocalWriteBaseAddr{tc}"),
|
|
src0=sgpr(f"LocalWriteBaseAddr{tc}"), src1=sgpr(f"Swap{tc}"),
|
|
comment=""))
|
|
return module
|
|
|
|
|
|
# --- GR pointer update (TLU=0) ----------------------------------------------
|
|
|
|
@_emitGRPtrUpdate.register(GRTag_1x1)
|
|
@_emitGRPtrUpdate.register(GRTag_1x2)
|
|
@_emitGRPtrUpdate.register(GRTag_2x2)
|
|
def _emitGRPtrUpdate_TLU0(tag, tile, ti, writer, kernel):
|
|
"""Advance SRD base pointer by one depthU iteration (depthU * bpe bytes)."""
|
|
tc = ti.tc
|
|
# TDM path: advance Address{tc} and sync the TDM descriptor instead of SRD.
|
|
if kernel.get("enableTDM%s" % tc, False):
|
|
module = Module(f"TDM GR Ptr Update ({tc})")
|
|
inc = int(ti.depthUBytes)
|
|
module.addComment0("TDM addr update: %s += %u" % (tc, inc))
|
|
module.add(SAddU64(dst=sgpr("Address%s" % tc, 2), src0=sgpr("Address%s" % tc, 2), src1=inc))
|
|
group0 = "tdm%sGroup0" % tc
|
|
module.add(SMovB64(dst=sgpr("%s+2" % group0, 2), src=sgpr("Address%s" % tc, 2), comment="sync descriptor global addr"))
|
|
module.add(SOrB32(dst=sgpr("%s+3" % group0), src0=sgpr("%s+3" % group0), src1=hex(2 << 30), comment="restore type field"))
|
|
return module
|
|
|
|
module = Module(f"GR Ptr Update ({tc})")
|
|
inc = int(ti.depthUBytes)
|
|
module.add(SAddU32(dst=sgpr(f"Srd{tc}"), src0=sgpr(f"Srd{tc}"), src1=inc,
|
|
comment=f"{tc}: advance SRD by {inc} bytes"))
|
|
module.add(SAddCU32(dst=sgpr(f"Srd{tc}+1"), src0=sgpr(f"Srd{tc}+1"), src1=0,
|
|
comment=f"{tc}: carry"))
|
|
return module
|
|
|
|
|
|
################################################################################
|
|
# Legacy GR emit functions (moved from SubtileBasedKernel.py)
|
|
################################################################################
|
|
|
|
##################################################
|
|
# Subroutine to generate GR offset calculation code
|
|
#
|
|
def graInitPointer(writer, kernel):
|
|
module = Module()
|
|
module.addComment0("REMOVE WHEN IMPLEMNTED: Placeholder for GR base pointer init")
|
|
for i in range(8):
|
|
module.addComment("")
|
|
|
|
return module
|
|
|
|
|
|
##################################################
|
|
# Compute GR offset for a single matrix (A or B)
|
|
#
|
|
def _grComputeOffset(module, writer, tileInfo, colId, rowId, output):
|
|
tc = tileInfo.tc
|
|
bpeBits = int(8*tileInfo.bpe)
|
|
|
|
tmpVgpr = writer.vgprPool.checkOut(2, tag="_grComputeOffset_tmpVgpr")
|
|
colBytes = tmpVgpr + 1
|
|
loadWidth = tileInfo.loadWidthGR
|
|
|
|
module.add(VLShiftLeftB32(dst=vgpr(colBytes), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(colId), comment="scale col_id by load_width"))
|
|
MT0 = tileInfo.globalMMATileGrid[0] * tileInfo.mmaTileShape[0]
|
|
subtileSize = tileInfo.subtileShape[0]*tileInfo.mmaTileShape[0]
|
|
strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(strideRef), src1=vgpr(rowId), comment="%s: rowId * stride"%tc))
|
|
module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr), shiftHex=hex(bpeBits.bit_length()-1), src=vgpr(tmpVgpr), comment="%s: rowId*stride*bpe"%tc))
|
|
module.add(VLShiftRightB32(dst=vgpr(tmpVgpr), shiftHex=hex(3), src=vgpr(tmpVgpr), comment="to bytes"))
|
|
module.add(VAddU32(dst=vgpr(output), src0=vgpr(colBytes), src1=vgpr(tmpVgpr), comment="%s: GR row_offset"%tc))
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
|
|
##################################################
|
|
# Compute subtile perpendicular offsets for a single matrix
|
|
#
|
|
# TODO: need to generalize this to support TLU=1
|
|
def _grComputeSubtileOffsets(writer, module, tileInfo):
|
|
tc = tileInfo.tc
|
|
strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
|
|
subtile_size = tileInfo.subtileShape[0]*tileInfo.mmaTileShape[0]
|
|
# rowOffset between 2 subtiles offset, ie how many consecutive subtile covered by a single subtileOffset.
|
|
# rowOffset = numGRPerSubtile * (local load ratio * subtile size)
|
|
rowOffset = math.ceil(tileInfo.numGRPerSubtile*tileInfo.loadRatioGR*subtile_size)
|
|
s_stride = int(rowOffset * tileInfo.bpe)
|
|
|
|
for regId in range(len(tileInfo.localSubtilesRegister)):
|
|
rl = tileInfo.localSubtilesRegister[regId]
|
|
for i, reg in enumerate(rl):
|
|
if rl.is_sgpr:
|
|
module.add(SMulI32(dst=sgpr(reg), src0=hex(s_stride * regId), src1=sgpr(strideRef), comment="%s: %u rows offset, stride %u, %u"%(tc, rowOffset, s_stride, regId)))
|
|
else:
|
|
stmp = writer.sgprPool.checkOut(1, tag="_grComputeSubtileOffsets_stmp")
|
|
module.add(SMulI32(dst=sgpr(stmp), src0=hex(s_stride * regId), src1=sgpr(strideRef), comment="%s: %u rows offset, stride %u, %u"%(tc, rowOffset, s_stride, regId)))
|
|
module.add(VAddU32(dst=vgpr(reg), src0=vgpr(tileInfo.sharedVgprGROffset[i]), src1=sgpr(stmp)))
|
|
writer.sgprPool.checkIn(stmp)
|
|
|
|
# Compute wave partition offset for a single tile (A or B)
|
|
#
|
|
def _grComputeRowPartition(module, kernel, writer, tileInfo, waveId, rowOffset):
|
|
subIterKBytes = tileInfo.subIterKBytes
|
|
wavesize = kernel["WavefrontSize"]
|
|
loadWidth = tileInfo.loadWidthGR
|
|
numRowsPerWave = wavesize // (subIterKBytes // loadWidth)
|
|
tc = tileInfo.tc
|
|
tmpVgpr = writer.vgprPool.checkOut(2, tag="_grComputeRowPartition_tmpVgpr")
|
|
tmpSgpr = writer.sgprPool.checkOut(1, tag="_grComputeRowPartition_tmpSgpr", preventOverflow=False)
|
|
localRow = tmpVgpr
|
|
partitionRow = tmpVgpr+1
|
|
partitionOffset = tileInfo.mmaTileShape[0]*tileInfo.localSubtileGrid[0]
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=partitionOffset, comment="%s: row offset"%tc))
|
|
|
|
if tileInfo.loadRatioGR == 1.0:
|
|
module.add(VAndB32(dst=vgpr(localRow), src0=hex(1), src1=vgpr(waveId), comment="%s: waveId %% 2"%tc))
|
|
module.add(VLShiftRightB32(dst=vgpr(partitionRow), shiftHex=hex(1), src=vgpr(waveId), comment="%s: waveId / 2"%tc))
|
|
elif tileInfo.loadRatioGR == 0.5:
|
|
module.add(VMovB32(dst=vgpr(localRow), src=0, comment="%s"%tc))
|
|
module.add(VMovB32(dst=vgpr(partitionRow), src=vgpr(waveId), comment="%s"%tc))
|
|
elif tileInfo.loadRatioGR == 2.0:
|
|
module.add(VMovB32(dst=vgpr(localRow), src=vgpr(waveId), comment="%s"%tc))
|
|
module.add(VMovB32(dst=vgpr(partitionRow), src=0, comment="%s"%tc))
|
|
else:
|
|
raise NotImplementedError("Unsupported loadRatioGR for wave partition: %s"%str(tileInfo.loadRatioGR))
|
|
|
|
module.add(VLShiftLeftB32(dst=vgpr(localRow), shiftHex=hex(numRowsPerWave.bit_length()-1), src=vgpr(localRow), comment="%s: local row offset"%tc))
|
|
module.add(VMulLOU32(dst=vgpr(partitionRow), src0=sgpr(tmpSgpr), src1=vgpr(partitionRow), comment="%s: wave row offset"%tc))
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=vgpr(localRow), src1=vgpr(partitionRow), comment="%s: row offset"%tc))
|
|
|
|
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
##################################################
|
|
# Compute GR offsets for all subtiles of a single matrix (A or B)
|
|
#
|
|
def _grComputeAllOffsets(module, writer, tileInfo, colId, rowId, rowOffset):
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=vgpr(rowId), src1=vgpr(rowOffset), comment="%s: row offset"%tileInfo.tc))
|
|
_grComputeOffset(module, writer, tileInfo, colId, rowOffset, tileInfo.sharedVgprGROffset[0])
|
|
for i in range(1, len(tileInfo.sharedVgprGROffset)):
|
|
subtileSize = tileInfo.subtileShape[0] * tileInfo.mmaTileShape[0]
|
|
offset = math.ceil(subtileSize * tileInfo.loadRatioGR)
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=offset, src1=vgpr(rowOffset), comment="%s: advance row for GR offset %u"%(tileInfo.tc, i)))
|
|
|
|
# Apply Rotation on entire wave. Only applies to 4x case as a subtile is loaded by a single wave in 2 steps. (waveId rotation not applied)
|
|
rotatedcolId = writer.vgprPool.checkOut(1, tag="_grComputeAllOffsets_rotatedcolId")
|
|
loadWidth = tileInfo.loadWidthGR
|
|
if tileInfo.loadRatioGR == 0.5:
|
|
blockSize = tileInfo.subIterKBytes // loadWidth
|
|
colRotation = blockSize // 2
|
|
module.add(VAddU32(dst=vgpr(rotatedcolId), src0=colRotation, src1=vgpr(colId), comment="%s: rotate col for GR offset %u"%(tileInfo.tc, i)))
|
|
module.add(VAndB32(dst=vgpr(rotatedcolId), src0=vgpr(rotatedcolId), src1=hex(blockSize-1), comment="(col + %d) %% block_size"%colRotation))
|
|
else:
|
|
module.add(VMovB32(dst=vgpr(rotatedcolId), src=vgpr(colId), comment=""))
|
|
|
|
_grComputeOffset(module, writer, tileInfo, rotatedcolId, rowOffset, tileInfo.sharedVgprGROffset[i])
|
|
writer.vgprPool.checkIn(rotatedcolId)
|
|
|
|
##################################################
|
|
# Apply swizzling and rotation to col IDs for GR offset calculation.
|
|
#
|
|
# Swizzling reorders column indices to avoid LDS bank conflicts.
|
|
# Two levels of rotation are applied to the column IDs:
|
|
# 1. Intra-wave rotation: rotates colId based on the LDS row id within
|
|
# a single wave. The rotation offset is: blockSize - (ldsRowId // 2) * 2.
|
|
# This ensures consecutive rows access different LDS banks.
|
|
# 2. Inter-wave rotation: an additional per-wave offset derived from waveId
|
|
# shifts the column further so that different waves also avoid bank
|
|
# conflicts with each other. Only applied when loadRatioGR != 0.5
|
|
# (i.e. when multiple waves share the same subtile region).
|
|
#
|
|
##################################################
|
|
# Subroutine to generate GR offset calculation code
|
|
#
|
|
def graTileAssignment(writer, kernel, useSwizzling=True):
|
|
return _graTileAssignment_legacy(writer, kernel, useSwizzling)
|
|
|
|
# --- Legacy interleaved A+B GR offset (temporary, matches reference exactly) ---
|
|
|
|
def _grComputeOffset_legacy(module, writer, tileInfo, colId, rowId, output):
|
|
tc = tileInfo.tc
|
|
bpeBits = int(8*tileInfo.bpe)
|
|
tmpVgpr = writer.vgprPool.checkOut(2, tag="_grComputeOffset_legacy_tmpVgpr")
|
|
colBytes = tmpVgpr + 1
|
|
loadWidth = tileInfo.loadWidthGR
|
|
module.add(VLShiftLeftB32(dst=vgpr(colBytes), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(colId), comment="scale col_id by load_width"))
|
|
strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
|
|
module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(strideRef), src1=vgpr(rowId), comment="%s: rowId * stride"%tc))
|
|
module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr), shiftHex=hex(bpeBits.bit_length()-1), src=vgpr(tmpVgpr), comment="%s: rowId*stride*bpe"%tc))
|
|
module.add(VLShiftRightB32(dst=vgpr(tmpVgpr), shiftHex=hex(3), src=vgpr(tmpVgpr), comment="to bytes"))
|
|
module.add(VAddU32(dst=vgpr(output), src0=vgpr(colBytes), src1=vgpr(tmpVgpr), comment="%s: GR row_offset"%tc))
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
|
|
def _grComputeSubtileOffsets_legacy(writer, module, tileInfo):
|
|
tc = tileInfo.tc
|
|
strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
|
|
subtile_size = tileInfo.subtileShape[0]*tileInfo.mmaTileShape[0]
|
|
rowOffset = math.ceil(tileInfo.numGRPerSubtile*tileInfo.loadRatioGR*subtile_size)
|
|
s_stride = int(rowOffset * tileInfo.bpe)
|
|
for regId in range(len(tileInfo.localSubtilesRegister)):
|
|
rl = tileInfo.localSubtilesRegister[regId]
|
|
for i, reg in enumerate(rl):
|
|
if rl.is_sgpr:
|
|
module.add(SMulI32(dst=sgpr(reg), src0=hex(s_stride * regId), src1=sgpr(strideRef), comment="%s: %u rows offset, stride %u, %u"%(tc, rowOffset, s_stride, regId)))
|
|
else:
|
|
stmp = writer.sgprPool.checkOut(1, tag="_grComputeSubtileOffsets_legacy_stmp")
|
|
module.add(SMulI32(dst=sgpr(stmp), src0=hex(s_stride * regId), src1=sgpr(strideRef), comment="%s: %u rows offset, stride %u, %u"%(tc, rowOffset, s_stride, regId)))
|
|
module.add(VAddU32(dst=vgpr(reg), src0=vgpr(tileInfo.sharedVgprGROffset[i]), src1=sgpr(stmp)))
|
|
writer.sgprPool.checkIn(stmp)
|
|
|
|
def _grComputeRowPartition_legacy(module, kernel, writer, tileInfo, waveId, rowOffset):
|
|
subIterKBytes = tileInfo.subIterKBytes
|
|
wavesize = kernel["WavefrontSize"]
|
|
loadWidth = tileInfo.loadWidthGR
|
|
numRowsPerWave = wavesize // (subIterKBytes // loadWidth)
|
|
tc = tileInfo.tc
|
|
tmpVgpr = writer.vgprPool.checkOut(2, tag="_grComputeRowPartition_legacy_tmpVgpr")
|
|
tmpSgpr = writer.sgprPool.checkOut(1, tag="_grComputeRowPartition_legacy_tmpSgpr", preventOverflow=False)
|
|
localRow = tmpVgpr
|
|
partitionRow = tmpVgpr+1
|
|
partitionOffset = tileInfo.mmaTileShape[0]*tileInfo.localSubtileGrid[0]
|
|
module.add(SMovB32(dst=sgpr(tmpSgpr), src=partitionOffset, comment="%s: row offset"%tc))
|
|
if tileInfo.loadRatioGR == 1.0:
|
|
module.add(VAndB32(dst=vgpr(localRow), src0=hex(1), src1=vgpr(waveId), comment="%s: waveId %% 2"%tc))
|
|
module.add(VLShiftRightB32(dst=vgpr(partitionRow), shiftHex=hex(1), src=vgpr(waveId), comment="%s: waveId / 2"%tc))
|
|
elif tileInfo.loadRatioGR == 0.5:
|
|
module.add(VMovB32(dst=vgpr(localRow), src=0, comment="%s"%tc))
|
|
module.add(VMovB32(dst=vgpr(partitionRow), src=vgpr(waveId), comment="%s"%tc))
|
|
elif tileInfo.loadRatioGR == 2.0:
|
|
module.add(VMovB32(dst=vgpr(localRow), src=vgpr(waveId), comment="%s"%tc))
|
|
module.add(VMovB32(dst=vgpr(partitionRow), src=0, comment="%s"%tc))
|
|
else:
|
|
raise NotImplementedError("Unsupported loadRatioGR for wave partition: %s"%str(tileInfo.loadRatioGR))
|
|
module.add(VLShiftLeftB32(dst=vgpr(localRow), shiftHex=hex(numRowsPerWave.bit_length()-1), src=vgpr(localRow), comment="%s: local row offset"%tc))
|
|
module.add(VMulLOU32(dst=vgpr(partitionRow), src0=sgpr(tmpSgpr), src1=vgpr(partitionRow), comment="%s: wave row offset"%tc))
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=vgpr(localRow), src1=vgpr(partitionRow), comment="%s: row offset"%tc))
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
writer.sgprPool.checkIn(tmpSgpr)
|
|
|
|
def _grComputeAllOffsets_legacy(module, writer, tileInfo, colId, rowId, rowOffset):
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=vgpr(rowId), src1=vgpr(rowOffset), comment="%s: row offset"%tileInfo.tc))
|
|
_grComputeOffset_legacy(module, writer, tileInfo, colId, rowOffset, tileInfo.sharedVgprGROffset[0])
|
|
for i in range(1, len(tileInfo.sharedVgprGROffset)):
|
|
subtileSize = tileInfo.subtileShape[0] * tileInfo.mmaTileShape[0]
|
|
offset = math.ceil(subtileSize * tileInfo.loadRatioGR)
|
|
module.add(VAddU32(dst=vgpr(rowOffset), src0=offset, src1=vgpr(rowOffset), comment="%s: advance row for GR offset %u"%(tileInfo.tc, i)))
|
|
rotatedcolId = writer.vgprPool.checkOut(1, tag="_grComputeAllOffsets_legacy_rotatedcolId")
|
|
loadWidth = tileInfo.loadWidthGR
|
|
if tileInfo.loadRatioGR == 0.5:
|
|
if tileInfo.bpe == 1: # FP8: intra-block K_group +2 rotation, preserving block bit
|
|
tmpBlock = writer.vgprPool.checkOut(1, tag="_grComputeAllOffsets_legacy_tmpBlock")
|
|
module.add(VAndB32(dst=vgpr(tmpBlock), src0=vgpr(colId), src1=hex(4), comment="%s: block_bit = colId & 4"%tileInfo.tc))
|
|
module.add(VAndB32(dst=vgpr(rotatedcolId), src0=vgpr(colId), src1=hex(3), comment="%s: K_group = colId & 3"%tileInfo.tc))
|
|
module.add(VAddU32(dst=vgpr(rotatedcolId), src0=vgpr(rotatedcolId), src1=hex(2), comment="%s: K_group + 2"%tileInfo.tc))
|
|
module.add(VAndB32(dst=vgpr(rotatedcolId), src0=vgpr(rotatedcolId), src1=hex(3), comment="%s: (K_group+2) %% 4"%tileInfo.tc))
|
|
module.add(VAddU32(dst=vgpr(rotatedcolId), src0=vgpr(rotatedcolId), src1=vgpr(tmpBlock), comment="%s: K_group_rot + block_bit"%tileInfo.tc))
|
|
writer.vgprPool.checkIn(tmpBlock)
|
|
else: # FP4/FP16: half-block rotation
|
|
blockSize = tileInfo.subIterKBytes // loadWidth
|
|
colRotation = blockSize // 2
|
|
module.add(VAddU32(dst=vgpr(rotatedcolId), src0=colRotation, src1=vgpr(colId), comment="%s: rotate col for GR offset %u"%(tileInfo.tc, i)))
|
|
module.add(VAndB32(dst=vgpr(rotatedcolId), src0=vgpr(rotatedcolId), src1=hex(blockSize-1), comment="(col + %d) %% block_size"%colRotation))
|
|
else:
|
|
module.add(VMovB32(dst=vgpr(rotatedcolId), src=vgpr(colId), comment=""))
|
|
_grComputeOffset_legacy(module, writer, tileInfo, rotatedcolId, rowOffset, tileInfo.sharedVgprGROffset[i])
|
|
writer.vgprPool.checkIn(rotatedcolId)
|
|
|
|
def _grSwizzleColIds_legacy(module, writer, tileInfoA, tileInfoB, blockSize, numRowsPerLDSBanks,
|
|
laneId, colIdA, colIdB, waveId):
|
|
tmpVgpr = writer.vgprPool.checkOut(3, tag="_grSwizzleColIds_legacy_tmpVgpr")
|
|
ldsRowId = tmpVgpr
|
|
tmp = tmpVgpr + 1
|
|
waveRotation = tmpVgpr + 2
|
|
half = blockSize // 2
|
|
module.addComment0("Swizzling")
|
|
module.add(VLShiftRightB32(dst=vgpr(ldsRowId), shiftHex=hex(blockSize.bit_length()-1), src=vgpr(laneId), comment="row id within wave"))
|
|
module.add(VLShiftRightB32(dst=vgpr(ldsRowId), shiftHex=hex(numRowsPerLDSBanks.bit_length()-1), src=vgpr(ldsRowId), comment="lds row id"))
|
|
module.add(VAndB32(dst=vgpr(tmp), src0=vgpr(ldsRowId), src1=hex(1), comment="swap_bit = ldsRowId & 1"))
|
|
if tileInfoA.bpe == 1: # FP8: step1=block-swap, step2=wave K_group rotation
|
|
# Step 1: block-swap (XOR blockSize//2 for odd ldsRowId)
|
|
module.add(VLShiftLeftB32(dst=vgpr(tmp), shiftHex=hex(int(math.log2(half))), src=vgpr(tmp),
|
|
comment=f"swap_bit * {half}"))
|
|
module.add(VXorB32(dst=vgpr(colIdA), src0=vgpr(colIdA), src1=vgpr(tmp),
|
|
comment="FP8 step1: block-swap colIdA"))
|
|
module.add(VMovB32(dst=vgpr(colIdB), src=vgpr(colIdA), comment="colIdB = colIdA"))
|
|
# Step 2: K_group rotation = (waveId & 1) * 2 (only for loadRatioGR != 0.5)
|
|
module.add(VAndB32(dst=vgpr(tmp), src0=vgpr(waveId), src1=hex(1), comment="wave_half = waveId & 1"))
|
|
module.add(VLShiftLeftB32(dst=vgpr(tmp), shiftHex=hex(1), src=vgpr(tmp), comment="rotation = wave_half * 2"))
|
|
for tInfo, cId in [(tileInfoA, colIdA), (tileInfoB, colIdB)]:
|
|
if tInfo.loadRatioGR != 0.5:
|
|
module.add(VAndB32(dst=vgpr(waveRotation), src0=vgpr(cId), src1=hex(4), comment="FP8 step2: block_bit = colId & 4"))
|
|
module.add(VAndB32(dst=vgpr(cId), src0=vgpr(cId), src1=hex(3), comment="K_group = colId & 3"))
|
|
module.add(VAddU32(dst=vgpr(cId), src0=vgpr(cId), src1=vgpr(tmp), comment="K_group + rotation"))
|
|
module.add(VAndB32(dst=vgpr(cId), src0=vgpr(cId), src1=hex(3), comment="(K_group+rotation) % 4"))
|
|
module.add(VAddU32(dst=vgpr(cId), src0=vgpr(cId), src1=vgpr(waveRotation), comment="K_group_rot + block_bit"))
|
|
else: # FP4/FP16: pair-swap (even ldsRowId) + intra/inter-wave rotation
|
|
module.add(VCmpXEqU32(dst=VCC(), src0=0, src1=vgpr(tmp), comment="lds row id % 2 == 0 ?"))
|
|
module.add(VMovB32(dst=vgpr(colIdA), src=vgpr(colIdA), dpp=DPPModifiers(quad_perm=[1,0,3,2]), comment="swap colId pairs for swizzling"))
|
|
module.add(SMovB64(dst=EXEC(), src=-1))
|
|
module.add(VMovB32(dst=vgpr(colIdB), src=vgpr(colIdA), comment=""))
|
|
module.addComment0("Rotation within a single wave")
|
|
module.add(VLShiftRightB32(dst=vgpr(tmp), shiftHex=hex(1), src=vgpr(ldsRowId), comment=""))
|
|
module.add(VLShiftLeftB32(dst=vgpr(tmp), shiftHex=hex(1), src=vgpr(tmp), comment="(ldsRowId //2) * 2"))
|
|
module.add(VSubU32(dst=vgpr(tmp), src0=hex(blockSize), src1=vgpr(tmp), comment="rotation offset : blockSize - (ldsRowId//2)*2"))
|
|
for tInfo, cId in [(tileInfoA, colIdA), (tileInfoB, colIdB)]:
|
|
if tInfo.loadRatioGR != 0.5:
|
|
module.addComment0("Rotation per wave")
|
|
module.add(VAndB32(dst=vgpr(waveRotation), src0=vgpr(waveId), src1=hex(1), comment=""))
|
|
module.add(VLShiftLeftB32(dst=vgpr(waveRotation), shiftHex=hex((2*numRowsPerLDSBanks).bit_length() - 1), src=vgpr(waveRotation), comment=""))
|
|
module.add(VSubU32(dst=vgpr(waveRotation), src0=vgpr(tmp), src1=vgpr(waveRotation), comment=""))
|
|
module.add(VAddU32(dst=vgpr(cId), src0=vgpr(waveRotation), src1=vgpr(cId), comment=""))
|
|
else:
|
|
module.add(VAddU32(dst=vgpr(cId), src0=vgpr(tmp), src1=vgpr(cId), comment=""))
|
|
module.add(VAndB32(dst=vgpr(colIdA), src0=vgpr(colIdA), src1=hex(blockSize-1), comment="(col + offset) % block_size"))
|
|
module.add(VAndB32(dst=vgpr(colIdB), src0=vgpr(colIdB), src1=hex(blockSize-1), comment="(col + offset) % block_size"))
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
|
|
def _graTileAssignment_legacy(writer, kernel, useSwizzling=True):
|
|
module = Module()
|
|
module.addComment0("GR Offset Calculation for Subtile Based Tiling")
|
|
tileInfoA = writer.states.a.tileInfo
|
|
tileInfoB = writer.states.b.tileInfo
|
|
subIterKBytes = tileInfoA.subIterKBytes
|
|
wavesize = kernel["WavefrontSize"]
|
|
ldsRowBankSize = writer.states.archCaps["LDSBankCount"] * writer.states.archCaps["LDSBankWidth"]
|
|
loadWidth = tileInfoA.loadWidthGR
|
|
assert subIterKBytes % loadWidth == 0
|
|
assert subIterKBytes <= ldsRowBankSize
|
|
blockSize = subIterKBytes // loadWidth
|
|
numRowsPerLDSBanks = ldsRowBankSize // subIterKBytes
|
|
tmpVgpr = writer.vgprPool.checkOut(7, tag="_graTileAssignment_legacy_tmpVgpr")
|
|
colIdA = tmpVgpr
|
|
colIdB = tmpVgpr + 1
|
|
rowId = tmpVgpr + 2
|
|
rowOffsetA = tmpVgpr + 3
|
|
rowOffsetB = tmpVgpr + 4
|
|
waveId = tmpVgpr + 5
|
|
laneId = tmpVgpr + 6
|
|
module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment="Wave Id"))
|
|
module.add(VAndB32(dst=vgpr(laneId), src0=vgpr("Serial"), src1=wavesize-1, comment=""))
|
|
module.add(VAndB32(dst=vgpr(colIdA), src0=vgpr("Serial"), src1=(blockSize-1), comment="get col_id in wave for %uB load"%loadWidth))
|
|
module.add(VLShiftRightB32(dst=vgpr(rowId), shiftHex=hex(blockSize.bit_length()-1), src=vgpr(laneId), comment="row id within wave"))
|
|
_grSwizzleColIds_legacy(module, writer, tileInfoA, tileInfoB, blockSize, numRowsPerLDSBanks,
|
|
laneId, colIdA, colIdB, waveId)
|
|
_grComputeRowPartition_legacy(module, kernel, writer, tileInfoA, waveId, rowOffsetA)
|
|
_grComputeRowPartition_legacy(module, kernel, writer, tileInfoB, waveId, rowOffsetB)
|
|
_grComputeAllOffsets_legacy(module, writer, tileInfoA, colIdA, rowId, rowOffsetA)
|
|
_grComputeAllOffsets_legacy(module, writer, tileInfoB, colIdB, rowId, rowOffsetB)
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
_grComputeSubtileOffsets_legacy(writer, module, tileInfoA)
|
|
_grComputeSubtileOffsets_legacy(writer, module, tileInfoB)
|
|
return module
|
|
|
|
##################################################
|
|
# Subroutine to generate GR load code
|
|
#
|
|
def emitSingleBufferLoad(tileInfo, kernel, sId0, sId1):
|
|
"""Emit buffer_load instructions for a single subtile (sId0, sId1).
|
|
|
|
When loadRatioGR > 1, multiple local subtiles share the same global read.
|
|
Only the first subtile in each group emits the load; others return empty.
|
|
|
|
Args:
|
|
tileInfo: TileInfo or TileInfo for the tensor component
|
|
sId0: Subtile row index
|
|
sId1: Subtile column index (K-dimension)
|
|
"""
|
|
module = Module()
|
|
|
|
# TDM path: emit one tensor_load_to_lds per tensor, skip all per-subtile DTL loads
|
|
if kernel.get("enableTDM%s" % tileInfo.tc[0], False):
|
|
if sId0 == 0 and sId1 == 0:
|
|
tc = tileInfo.tc
|
|
group0 = "tdm%sGroup0" % tc
|
|
group1 = "tdm%sGroup1" % tc
|
|
module.add(TensorLoadToLds(sgpr(group0, 4), sgpr(group1, 8), None, None,
|
|
comment="TDM: global->LDS for %s" % tc))
|
|
return module
|
|
|
|
linearId = tileInfo.getLocalSubtileLinearId(sId0, sId1)
|
|
grBaseId = int(math.floor(linearId / tileInfo.loadRatioGR))
|
|
|
|
if tileInfo.loadRatioGR > 1:
|
|
firstInGroup = int(grBaseId * tileInfo.loadRatioGR)
|
|
if linearId != firstInGroup:
|
|
return module
|
|
|
|
tc = tileInfo.tc
|
|
isGlc = bool(kernel["NonTemporal%s"%tc] & 0x1)
|
|
isSlc = bool(kernel["NonTemporal%s"%tc] & 0x2)
|
|
isNT = bool(kernel["NonTemporal%s"%tc] & 0x4)
|
|
|
|
regListIdx = tileInfo.grRegGroupForSubtileRow(sId0)
|
|
regList = tileInfo.localSubtilesRegister[regListIdx]
|
|
useSgpr = regList.is_sgpr
|
|
|
|
offsetK = sId1 * int(tileInfo.mmaTileShape[1] * tileInfo.subtileShape[1] * tileInfo.bpe)
|
|
|
|
subtileOffset = int(math.ceil(tileInfo.loadRatioGR*tileInfo.subtileSize))
|
|
WriteBaseAddr = "LocalWriteBaseAddr%s"%tc
|
|
for i in range(tileInfo.numGRPerSubtile):
|
|
m0Offset = int(i * subtileOffset + (sId0 + sId1 * tileInfo.globalSubtileGrid[0]) * tileInfo.subtileSize)
|
|
module.add(SAddU32(dst=mgpr(0), src0=sgpr(WriteBaseAddr), src1=(m0Offset - offsetK)))
|
|
mubuf = MUBUFModifiers(offen=True, offset12=offsetK, glc=isGlc, slc=isSlc, nt=isNT, lds=True)
|
|
|
|
soffset = regList.ref(0) if len(regList) > 0 and useSgpr else 0
|
|
voff = tileInfo.sharedVgprGROffset[i] if useSgpr or len(regList) == 0 else regList.indices[i]
|
|
module.add(BufferLoadB128(dst=None, vaddr=vgpr(voff), saddr=sgpr("Srd%s"%tc, 4), soffset=soffset, mubuf=mubuf, comment="grBaseId = %u, i= %u"%(grBaseId , i)))
|
|
|
|
return module
|
|
|
|
|
|
def emitSubtileBufferLoad(tc, writer, kernel, subtileId):
|
|
tileInfo = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo
|
|
return emitSingleBufferLoad(tileInfo, kernel, subtileId[0], subtileId[1])
|
|
|
|
##################################################
|
|
# Subroutine to generate GR load code
|
|
# Initial idea: maybe store asm in modules in a separate obj?
|
|
#
|
|
def globalReadDoSubtile(tc, writer, kernel):
|
|
module = Module()
|
|
|
|
tileInfo = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo
|
|
|
|
for j in range(tileInfo.localSubtileGrid[1]):
|
|
for i in range(tileInfo.localSubtileGrid[0]):
|
|
module.addComment0("Emit load for %s subtile: [%u, %u]"%(tc, i, j))
|
|
module.add(emitSubtileBufferLoad(tc, writer, kernel, [i, j]))
|
|
|
|
return module
|
|
|
|
##################################################
|
|
# Subroutine to generate DTL M0 LDS buffer swap
|
|
#
|
|
def globalReadDTLInitCommonSgpr(writer, kernel):
|
|
return _globalReadDTLInitCommonSgpr_legacy(writer, kernel)
|
|
|
|
def _globalReadDTLInitCommonSgpr_legacy(writer, kernel):
|
|
module = Module()
|
|
tileInfoA = writer.states.a.tileInfo
|
|
tileInfoB = writer.states.b.tileInfo
|
|
wavesize = kernel["WavefrontSize"]
|
|
vgprWaveId = writer.vgprPool.checkOut(1, tag="_globalReadDTLInitCommonSgpr_legacy_vgprWaveId")
|
|
module.addComment0("Compute shared offsets used by m0 in DTL loads")
|
|
module.add(VLShiftRightB32(dst=vgpr(vgprWaveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment="Wave Id"))
|
|
tmpVgpr = writer.vgprPool.checkOut(2, tag="_globalReadDTLInitCommonSgpr_legacy_tmpVgpr")
|
|
rowOffsetA = tmpVgpr
|
|
rowOffsetB = tmpVgpr + 1
|
|
_grComputeRowPartition_legacy(module, kernel, writer, tileInfoA, vgprWaveId, rowOffsetA)
|
|
_grComputeRowPartition_legacy(module, kernel, writer, tileInfoB, vgprWaveId, rowOffsetB)
|
|
subIterKBytes = tileInfoA.subIterKBytes
|
|
module.add(VLShiftLeftB32(dst=vgpr(rowOffsetA), shiftHex=hex((subIterKBytes).bit_length()-1), src=vgpr(rowOffsetA), comment="Apply wave-specific offset for A"))
|
|
module.add(VLShiftLeftB32(dst=vgpr(rowOffsetB), shiftHex=hex((subIterKBytes).bit_length()-1), src=vgpr(rowOffsetB), comment="Apply wave-specific offset for B"))
|
|
module.add(SNop(waitState=0, comment="Wait for VGPR to be ready"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr("LocalWriteBaseAddrA"), src=vgpr(rowOffsetA), comment="Store base LDS offset, will be modified"))
|
|
module.add(VReadfirstlaneB32(dst=sgpr("LocalWriteBaseAddrB"), src=vgpr(rowOffsetB), comment="Store base LDS offset, will be modified"))
|
|
module.add(SAddU32(dst=sgpr("LocalWriteBaseAddrB"), src0=sgpr("LocalWriteBaseAddrB"), src1=hex(writer.ldsStartOffsetB), comment=""))
|
|
module.add(SAddU32(dst=sgpr("SwapA"), src0=sgpr("LocalWriteBaseAddrA"), src1=writer.ldsTotalSize, comment=""))
|
|
module.add(SXorB32(dst=sgpr("SwapA"), src0=sgpr("LocalWriteBaseAddrA"), src1=sgpr("SwapA"), comment=""))
|
|
module.add(SAddU32(dst=sgpr("SwapB"), src0=sgpr("LocalWriteBaseAddrB"), src1=writer.ldsTotalSize, comment=""))
|
|
module.add(SXorB32(dst=sgpr("SwapB"), src0=sgpr("LocalWriteBaseAddrB"), src1=sgpr("SwapB"), comment=""))
|
|
writer.vgprPool.checkIn(vgprWaveId)
|
|
writer.vgprPool.checkIn(tmpVgpr)
|
|
return module
|
|
|
|
##################################################
|
|
# Subroutine to generate DTL M0 LDS buffer swap
|
|
#
|
|
def globalReadLDSBufferSwap(tc, writer, kernel):
|
|
if tc in ['A', 'B']:
|
|
ti_ = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo
|
|
if kernel.get("enableTDM%s" % tc, False):
|
|
ldsAddrSgpr = "tdmLdsAddr%s" % tc
|
|
swapSgpr = "tdmLdsSwapMask%s" % tc
|
|
module = Module()
|
|
module.addComment0("TDM: swap %s LDS buffer (XOR with per-tensor swap mask)" % tc)
|
|
module.add(SXorB32(dst=sgpr(ldsAddrSgpr), src0=sgpr(ldsAddrSgpr), src1=sgpr(swapSgpr), comment=""))
|
|
group0 = "tdm%sGroup0" % tc
|
|
module.add(SMovB32(dst=sgpr("%s+1" % group0), src=sgpr(ldsAddrSgpr), comment="sync descriptor LDS addr"))
|
|
return module
|
|
return ti_.emitGRLDSBufferSwap(writer, kernel)
|
|
else:
|
|
ti_ = writer.states.mxsa.tileInfo if tc == 'MXSA' else writer.states.mxsb.tileInfo
|
|
return emitScaleGRLDSSwap(ti_, writer, kernel)
|
|
|
|
|
|
|
|
################################################################################
|
|
# TDM subtile functions (global offset, descriptor init, StreamK offset)
|
|
################################################################################
|
|
|
|
def tdmGlobalOffsetSubtile(writer, kernel, tP):
|
|
"""Per-wave global address for subtile TDM.
|
|
|
|
All waves cooperatively load the tile: wave w covers M-rows
|
|
[w*mt/numWaves, (w+1)*mt/numWaves) across the full wave count, rather
|
|
than only this tensor's wave axis. Splitting over every wave avoids the
|
|
duplicate loads the axis-only split issued for waves sharing an axis id.
|
|
The LDS tile end-state (identity map global-row r -> LDS-row r) is
|
|
unchanged; the barrier before local reads (WaitGROp has_sync) makes
|
|
every wave's rows visible to all consumers.
|
|
"""
|
|
tc = tP["tensorChar"]
|
|
ti = tP["idx"]
|
|
bpe = tP["bpeGR"]
|
|
tlu = tP["tlu"]
|
|
mt = kernel[f"MacroTile{ti}"]
|
|
wavelen = kernel["WavefrontSize"]
|
|
numWaves = prod(kernel["MIWaveGroup"])
|
|
mod = Module(f"TDM Global Offset Subtile {tc}")
|
|
|
|
with writer.allocTmpSgpr(3) as tmpSgprRes:
|
|
tmp = tmpSgprRes.idx
|
|
waveOff = tmpSgprRes.idx + 2
|
|
|
|
tileStride = writer.strideRef(tc, ti)
|
|
mod.add(SMulI32(dst=sgpr(tmp), src0=tileStride, src1=int(mt * bpe),
|
|
comment=f"stride * MT({mt}) * bpe({bpe})"))
|
|
mod.add(SMulI32(dst=sgpr(tmp), src0=sgpr(tmp), src1=sgpr(f"WorkGroup{ti}"),
|
|
comment="*= wgId"))
|
|
|
|
if numWaves > 1:
|
|
mod.add(VReadfirstlaneB32(dst=sgpr(waveOff), src=vgpr("Serial"), comment="first tId"))
|
|
mod.add(SLShiftRightB32(dst=sgpr(waveOff), src=sgpr(waveOff),
|
|
shiftHex=hex(int(ceil(log2(wavelen)))), comment=f"wId = tId / {wavelen}"))
|
|
tileStrideSep = writer.strideRef(tc, 3) if tlu else writer.strideRef(tc, ti)
|
|
mod.add(SMulI32(dst=sgpr(waveOff), src0=sgpr(waveOff), src1=int(mt // numWaves * bpe),
|
|
comment=f"waveOff = waveId * {mt // numWaves} * {bpe}"))
|
|
mod.add(SMulI32(dst=sgpr(waveOff), src0=sgpr(waveOff), src1=tileStrideSep,
|
|
comment="waveOff *= stride"))
|
|
mod.add(SAddU32(dst=sgpr(tmp), src0=sgpr(tmp), src1=sgpr(waveOff), comment="+= waveOff"))
|
|
|
|
mod.add(SAddU32(dst=sgpr(f"Address{tc}"), src0=sgpr(f"Address{tc}"), src1=sgpr(tmp),
|
|
comment=f"+= offset(lo)"))
|
|
mod.add(SAddCU32(dst=sgpr(f"Address{tc}+1"), src0=sgpr(f"Address{tc}+1"), src1=0,
|
|
comment=f"+= offset(hi)"))
|
|
|
|
if kernel["ProblemType"]["Batched"] and kernel["ProblemType"]["StridedBatched"]:
|
|
ia = tP["ia"]
|
|
batchStrideName = f"Stride{tc}{writer.states.indexChars[ia[2]]}"
|
|
mod.addModuleAsFlatItems(writer.s_mul_u64_u32(sgpr(tmp), sgpr(tmp+1),
|
|
sgpr(batchStrideName), sgpr("WorkGroup2"),
|
|
comment="Batch: Stride*WG"))
|
|
mod.add(SLShiftLeftB64(dst=sgpr(tmp, 2), src=sgpr(tmp, 2),
|
|
shiftHex=int(log2(bpe)), comment="scale by bpe"))
|
|
mod.add(SAddU64(dst=sgpr(f"Address{tc}", 2), src0=sgpr(tmp, 2), src1=sgpr(f"Address{tc}", 2),
|
|
comment="+= batch"))
|
|
|
|
return mod
|
|
|
|
|
|
def initTDMDescriptorSubtile(writer, kernel, tP):
|
|
"""Subtile variant of initTDMDescriptor()."""
|
|
from ...Components.TensorDataMover import TensorDataMoverLoad
|
|
comp = TensorDataMoverLoad.find(writer)
|
|
tc = tP['tensorChar']
|
|
ti = tP["idx"]
|
|
tileChar = tP["tileChar"]
|
|
mod = Module(f"Init TDM Descriptor Subtile {tc}")
|
|
|
|
def descSgprName(idx):
|
|
assert idx < 2
|
|
return f"tdm{tc}Group{idx}"
|
|
|
|
def strideRefName():
|
|
return f"Stride{tc}{tileChar}"
|
|
|
|
def sizeRefName(idx):
|
|
idxChar = INDEX_CHARS[idx]
|
|
return f"Size{idxChar}"
|
|
|
|
dtype = kernel["ProblemType"][f"DataType{tc}"]
|
|
mt = kernel[f"MacroTile{ti}"]
|
|
du = kernel["DepthU"]
|
|
bpe = tP["bpeGR"]
|
|
numWaves = prod(kernel["MIWaveGroup"])
|
|
wavelen = kernel["WavefrontSize"]
|
|
|
|
# Use subtile LDS offsets from writer state (not kernel["LdsOffset{tc}"])
|
|
ldsOffsetMap = {
|
|
'A': writer.ldsStartOffsetA,
|
|
'B': writer.ldsStartOffsetB,
|
|
}
|
|
ldsConstOffset = ldsOffsetMap.get(tc, 0)
|
|
|
|
sizeTile0, sizeTile1 = du, mt
|
|
# TDM D# Group1 pad fields
|
|
# padAmountBytes -> pad_amount [31:25], bytes inserted per pad event
|
|
# padIntervalBytes -> pad_interval [24:22], bytes written between pads
|
|
# Sourced from TileInfo.ldsRowPadBytes so GR and LR
|
|
# see the same value.
|
|
tileInfoForTc = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo
|
|
padAmountBytes = int(getattr(tileInfoForTc, "ldsRowPadBytes", 0))
|
|
padIntervalBytes = int(du * bpe) if padAmountBytes else 0
|
|
|
|
mod.add(comp.initOperands(descSgprName(0), descSgprName(1), None, None))
|
|
mod.add(comp.setDataType(dtype, descSgprName(1)))
|
|
mod.add(comp.setGlobalAddr(descSgprName(0), f"Address{tc}"))
|
|
# OR the per-tensor broadcast mask into the descriptor for TDM multicast.
|
|
# Subtile loads both A and B on every wave, so it uses split masks
|
|
# (MulticastMask{tc}), not the non-subtile single parity mask.
|
|
if kernel["Multicast"] and clusterEnabled(kernel["ClusterDim"]):
|
|
mod.add(comp.setMulticastMask(descSgprName(1), f"MulticastMask{tc}", writer))
|
|
|
|
with writer.allocTmpSgpr(1) as tmpSgprRes:
|
|
waveOffsetSgprIdx = tmpSgprRes.idx
|
|
mod.add(VReadfirstlaneB32(sgpr(waveOffsetSgprIdx), vgpr("Serial"), "first tId"))
|
|
mod.add(SLShiftRightB32(sgpr(waveOffsetSgprIdx), ceil(log2(wavelen)), sgpr(waveOffsetSgprIdx), "wId=fTid // wavelen"))
|
|
# Each wave writes its mt/numWaves rows to a distinct LDS region,
|
|
# matching the cooperative full-wave global split in
|
|
# tdmGlobalOffsetSubtile. The union over all waves covers the whole
|
|
# mt-row tile (identity map global-row r -> LDS-row r).
|
|
if padIntervalBytes != 0 and padAmountBytes != 0:
|
|
tileBytes = round(mt // numWaves * du * bpe)
|
|
padBytes = tileBytes // padIntervalBytes * padAmountBytes
|
|
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), tileBytes + padBytes,
|
|
f"woffset = wId * ({tileBytes}+{padBytes})"))
|
|
else:
|
|
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), round(mt // numWaves * du * bpe),
|
|
"woffset = wId * (mt // numWaves * du * bpe)"))
|
|
mod.add(SAddU32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), ldsConstOffset,
|
|
f"ldsOffset = woffset + {ldsConstOffset} (subtile LDS offset for {tc})"))
|
|
mod.add(comp.setLdsAddr(descSgprName(0), sgpr(waveOffsetSgprIdx)))
|
|
# Save LDS offset to tracking SGPR for runtime double-buffer swap
|
|
ldsTrackSgpr = f"tdmLdsAddr{tc}"
|
|
mod.add(SMovB32(dst=sgpr(ldsTrackSgpr), src=sgpr(waveOffsetSgprIdx), comment=f"init {ldsTrackSgpr} for buffer tracking"))
|
|
# Compute swap mask: swapMask = addr XOR (addr + ldsTotalSize)
|
|
# Used by globalReadLDSBufferSwap to toggle between buffer 0 and buffer 1.
|
|
swapMaskSgpr = f"tdmLdsSwapMask{tc}"
|
|
ldsTotalSize = writer.ldsTotalSize
|
|
mod.add(SAddU32(dst=sgpr(swapMaskSgpr), src0=sgpr(waveOffsetSgprIdx), src1=ldsTotalSize, comment=f"addr + ldsTotalSize({ldsTotalSize})"))
|
|
mod.add(SXorB32(dst=sgpr(swapMaskSgpr), src0=sgpr(waveOffsetSgprIdx), src1=sgpr(swapMaskSgpr), comment=f"swapMask = addr XOR (addr + ldsTotalSize)"))
|
|
sizeShifter = 1 if dtype.isFloat4() else 0
|
|
sizeShifterDim = sizeShifter
|
|
|
|
mod.add(comp.setIterationEnabled(descSgprName(1), False))
|
|
mod.add(comp.setPadding(descSgprName(1), padIntervalBytes, padAmountBytes))
|
|
mod.add(comp.setTensorDim0(descSgprName(1), sizeRefName(3), writer, sizeShifterDim))
|
|
mod.add(comp.setTensorDim1(descSgprName(1), sizeRefName(ti), writer))
|
|
|
|
sizeShifterTile = sizeShifter
|
|
mod.add(comp.setTensorTile0(descSgprName(1), sizeTile0, writer, sizeShifterTile))
|
|
|
|
# Clamp each wave's Tile1 (free-dim-1) load extent to the valid remainder.
|
|
# tdmGlobalOffsetSubtile bases wave w at row w*(mt//numWaves), but the
|
|
# descriptor Dim1 does not bound the walk, so an edge tile (free dim < mt)
|
|
# reads past the tensor. setTensorTile1 takes a compile-time int, so write its
|
|
# field (+4[15:0]) with a runtime clamp. No-op when the tile fits.
|
|
perWaveRows = sizeTile1 // numWaves
|
|
if numWaves > 1:
|
|
with writer.allocTmpSgpr(2) as tileClampRes:
|
|
validRows = tileClampRes.idx
|
|
waveRowStart = tileClampRes.idx + 1
|
|
mod.add(VReadfirstlaneB32(sgpr(waveRowStart), vgpr("Serial"), "first tId"))
|
|
mod.add(SLShiftRightB32(sgpr(waveRowStart), ceil(log2(wavelen)), sgpr(waveRowStart),
|
|
"wId = fTid // wavelen"))
|
|
mod.add(SMulI32(sgpr(waveRowStart), sgpr(waveRowStart), perWaveRows,
|
|
f"waveGlobalRowStart = wId * {perWaveRows}"))
|
|
mod.add(SSubI32(dst=sgpr(validRows), src0=sgpr(sizeRefName(ti)), src1=sgpr(waveRowStart),
|
|
comment="Size_free - waveGlobalRowStart"))
|
|
mod.add(SMaxI32(dst=sgpr(validRows), src0=sgpr(validRows), src1=0,
|
|
comment="saturate negative remainder to 0"))
|
|
mod.add(SMinU32(dst=sgpr(validRows), src0=sgpr(validRows), src1=perWaveRows,
|
|
comment=f"clamp to per-wave rows ({perWaveRows})"))
|
|
mod.add(SAndB32(sgpr(f"{descSgprName(1)}+4"), sgpr(f"{descSgprName(1)}+4"),
|
|
hex(0xFFFF0000), "clear tile1 field"))
|
|
mod.add(SOrB32(sgpr(f"{descSgprName(1)}+4"), sgpr(f"{descSgprName(1)}+4"),
|
|
sgpr(validRows), "set tile1 = clamped validRows"))
|
|
else:
|
|
mod.add(comp.setTensorTile1(descSgprName(1), perWaveRows, writer))
|
|
mod.add(comp.setTensorStride0(descSgprName(1), strideRefName(), sizeShifterTile))
|
|
return mod
|
|
|
|
|
|
def tdmApplyStreamKOffsetSubtile(writer, kernel, tP):
|
|
"""Apply the StreamK K-offset to the subtile TDM descriptor.
|
|
|
|
StreamK=3 DP-partial work items have a nonzero StreamKLocalStart and must read
|
|
their own K-slice. Advance Address{tc} by StreamKLocalStart unroll iterations
|
|
(inc = ti.depthUBytes, matching _emitGRPtrUpdate_TLU0's per-iteration advance)
|
|
and re-sync the descriptor. No-op when StreamKLocalStart == 0.
|
|
"""
|
|
tc = tP["tensorChar"]
|
|
ti = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo
|
|
inc = int(ti.depthUBytes) # per-unroll TDM advance; same source as _emitGRPtrUpdate_TLU0
|
|
group0 = f"tdm{tc}Group0"
|
|
mod = Module(f"TDM StreamK K-offset subtile {tc}")
|
|
with writer.allocTmpSgpr(2, alignment=2, tag="tdmSkOffset") as tmpSgprRes:
|
|
o = tmpSgprRes.idx
|
|
mod.add(SMulI32(dst=sgpr(o), src0=sgpr("StreamKLocalStart"), src1=inc,
|
|
comment=f"SK K-start * depthU*bpe ({inc})"))
|
|
mod.add(SMovB32(dst=sgpr(o + 1), src=0, comment="SK K-start offset hi = 0"))
|
|
mod.add(SAddU32(dst=sgpr(f"Address{tc}+0"), src0=sgpr(f"Address{tc}+0"), src1=sgpr(o),
|
|
comment="Address += SK K-start offset (lo)"))
|
|
mod.add(SAddCU32(dst=sgpr(f"Address{tc}+1"), src0=sgpr(f"Address{tc}+1"), src1=sgpr(o + 1),
|
|
comment="Address += SK K-start offset (hi, carry)"))
|
|
mod.add(SMovB64(dst=sgpr(f"{group0}+2", 2), src=sgpr(f"Address{tc}", 2),
|
|
comment="sync descriptor global addr"))
|
|
mod.add(SOrB32(dst=sgpr(f"{group0}+3"), src0=sgpr(f"{group0}+3"), src1=hex(2 << 30),
|
|
comment="restore descriptor type field"))
|
|
return mod
|
|
|
|
##################################################
|
|
# Subroutine to update ptrs
|
|
#
|
|
def globalReadPtrUpdates(tc, writer, kernel):
|
|
ti_ = writer.states.a.tileInfo if tc == 'A' else writer.states.b.tileInfo
|
|
return ti_.emitGRPtrUpdate(writer, kernel)
|