Skip to content

Commit cee42b5

Browse files
mengzcaibriannwumengzcaicycheng
authored
Enable Cluster Barrier on gfx1250 (ROCm#7477)
[AIHPBLAS-1333](https://amd-hub.atlassian.net/browse/AIHPBLAS-1333) ## Motivation After ROCm#7324 (cluster launch + TDM multicast). Add cluster barrier support for gfx1250. Cluster barrier enables inter-workgroup synchronization within a cluster by splitting s_barrier into separate signal and wait instructions, allowing workgroups in the same cluster to coordinate data sharing more efficiently during TDM multicast operations. <!-- Explain the purpose of this PR and the goals it aims to achieve. --> ## Technical Details 1. Rocisa SBarrier split 2. Cluster Barrier support 3. sgprWaveIdx allocation 4. Reschedule cluster_barrier signal <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan Reuse tdm_multicast_gfx1250.yaml test suite (introduced in the PR#7324) which covers cluster launch + multicast + barrier scenarios. <!-- Explain any relevant testing done to verify this PR. --> ## Test Result All TDM multicast tests passed on gfx1250. No regressions observed on unrelated kernels. <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. [AIHPBLAS-1333]: https://amd-hub.atlassian.net/browse/AIHPBLAS-1333?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ --------- Co-authored-by: briannwu <briannwu@amd.com> Co-authored-by: mengzcai <mengzcai@amd.com> Co-authored-by: Chuang-Yu Cheng <cycheng2@amd.com>
1 parent a47f86f commit cee42b5

18 files changed

Lines changed: 353 additions & 115 deletions

File tree

‎projects/hipblaslt/tensilelite/Tensile/Common/GlobalParameters.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -524,7 +524,15 @@
524524
# code evicted from the I-cache before it runs. Software prefetch helps keep instruction fetch
525525
# ahead of execution. False: no SGPR reserved; Stinky prefetch pass disabled for that kernel.
526526
{"SwInstructionPrefetch": [True]},
527+
# ClusterDim — workgroup cluster dimensions [x, y] for clustered kernel launch.
528+
# [1, 1] disables clustering. Non-[1, 1] enables Multicast so workgroups within
529+
# a cluster can share data loaded via TDM-multicast, reducing redundant global reads.
527530
{"ClusterDim": [[1, 1]]},
531+
# ClusterBarrier — True: emit split signal/wait cluster_barrier instructions
532+
# so workgroups in a cluster synchronize before/after consuming shared
533+
# TDM-multicast data. Requires ClusterDim != [1, 1] and TDMInst != 0;
534+
# False: standard per-WG barriers, no inter-WG synchronization.
535+
{"ClusterBarrier": [ False ]},
528536
{"HalfPLR": [0]}
529537
]
530538

‎projects/hipblaslt/tensilelite/Tensile/Common/RequiredParameters.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,5 +122,6 @@ def getRequiredParametersMin() -> set:
122122
'MinGRIncPerMfma',
123123
'UsePLRPack',
124124
'UseSubtileImpl',
125+
'ClusterBarrier',
125126
'ClusterDim',
126127
})

‎projects/hipblaslt/tensilelite/Tensile/Common/ValidParameters.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1026,6 +1026,8 @@ def makeValidMatrixInstructions():
10261026
# "Auto": triggers defaulting in Solution.assignDerivedParameters. The default is
10271027
# TDM iff TDMInst != 0, otherwise BufferLoad.
10281028
"MXLoadInst": ["Auto", "TDM", "BufferLoad", "GlobalLoad"],
1029+
# Enable cluster barrier.
1030+
"ClusterBarrier": [False, True],
10291031
# Cluster dimension. Clusters have up to 16 work-groups in a cluster, but each work-group in a
10301032
# cluster runs on a separate WGP.
10311033
"ClusterDim": validClusterDimensions,

‎projects/hipblaslt/tensilelite/Tensile/Components/GSU.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1149,9 +1149,10 @@ def partialWriteBatch(self, writer, kernel, ss, batchIdx, beta, edge, gwvw, batc
11491149
# for the top-left corner this thread will write. These are not changed
11501150
# across all the store loop iters.
11511151
if writer.db["ConservativeWaitCnt"] & 0x10:
1152-
module.add(SBarrier("debug"))
1152+
module.add(SBarrier(comment="debug"))
11531153
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
1154-
module.add(SBarrier("debug"))
1154+
module.add(SBarrier(comment="debug"))
1155+
11551156
if not edge and writer.db["ForceEdgeStores"]>=2:
11561157
module.add(writer.getBomb()) # should not get here
11571158
if edge and writer.db["AssertNoEdge"]:
@@ -1305,9 +1306,10 @@ def reductionBatch(self, writer, kernel, ss, batchIdx, alpha, beta, edge, atomic
13051306
# for the top-left corner this thread will write. These are not changed
13061307
# across all the store loop iters.
13071308
if writer.db["ConservativeWaitCnt"] & 0x10:
1308-
module.add(SBarrier("debug"))
1309+
module.add(SBarrier(comment="debug"))
13091310
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
1310-
module.add(SBarrier("debug"))
1311+
module.add(SBarrier(comment="debug"))
1312+
13111313
if not edge and writer.db["ForceEdgeStores"]>=2:
13121314
module.add(writer.getBomb()) # should not get here
13131315
if edge and writer.db["AssertNoEdge"]:
@@ -1374,9 +1376,9 @@ def reductionBatch(self, writer, kernel, ss, batchIdx, alpha, beta, edge, atomic
13741376
module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))
13751377

13761378
if writer.db["ConservativeWaitCnt"] & 0x40:
1377-
module.add(SBarrier("debug"))
1379+
module.add(SBarrier(comment="debug"))
13781380
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
1379-
module.add(SBarrier("debug"))
1381+
module.add(SBarrier(comment="debug"))
13801382

13811383
# return registers to pool:
13821384
lastDataD = -1

‎projects/hipblaslt/tensilelite/Tensile/Components/GlobalWriteBatch.py‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ def emit(self) -> Module:
190190
(self.parentWriter.states.useBias != DataDirection.NONE or \
191191
self.kernel["ProblemType"].get("UseScaleAlphaVec", 0)):
192192
module.add(SWaitCnt(dscnt=0, comment="drain bias/SAV LDS reads"))
193-
module.add(SBarrier("sync waves before subtile paired stores"))
193+
module.add(SBarrier(comment="sync waves before subtile paired stores"))
194194
self._epilog(module)
195195
return module
196196

@@ -371,9 +371,10 @@ def _prolog(self, module: Module):
371371
# for the top-left corner this thread will write. These are not changed
372372
# across all the store loop iters.
373373
if self.debugConfig["ConservativeWaitCnt"] & 0x10:
374-
module.add(SBarrier("debug"))
374+
module.add(SBarrier(comment="debug"))
375375
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
376-
module.add(SBarrier("debug"))
376+
module.add(SBarrier(comment="debug"))
377+
377378
if not self.edge and self.debugConfig["ForceEdgeStores"] >= 2:
378379
module.add(self.parentWriter.getBomb()) # should not get here
379380
if self.edge and self.debugConfig["AssertNoEdge"]:
@@ -521,13 +522,13 @@ def addEpilogueLoad(modGwvw, ldName: str, addrVecVgpr, addrVec, dataVec, loadedD
521522
# Group bias load with C input to
522523
if isSingleKernel and (not self.isLocalBarrierInit):
523524
loadInputCode.add(SWaitCnt(dscnt=0, comment="Wait for LDS write"))
524-
loadInputCode.add(SBarrier("LDS write barrier"))
525+
loadInputCode.add(SBarrier(comment="LDS write barrier"))
525526
self.isLocalBarrierInit = True
526527
loadInputCode.add(self.parentWriter.addLdsLoad(self.kernel["ProblemType"]["ComputeDataType"], dataVec, ldsAddrVgpr, vecOffset, gwvw, comment=comment))
527528
else:
528529
if isSingleKernel and (not self.isLocalBarrierInit):
529530
module.add(SWaitCnt(dscnt=0, comment="Wait for LDS write"))
530-
module.add(SBarrier("LDS write barrier"))
531+
module.add(SBarrier(comment="LDS write barrier"))
531532
self.isLocalBarrierInit = True
532533
module.add(self.parentWriter.addLdsLoad(self.kernel["ProblemType"]["ComputeDataType"], dataVec, ldsAddrVgpr, vecOffset, gwvw, comment=comment))
533534
loadedDataVec[dataVec] = ceil(self.kernel["ProblemType"]["ComputeDataType"].numBytes() * gwvw / 16)
@@ -1549,9 +1550,9 @@ def applyScaleVec(vecModule, addressStr, dataScaleVec, factorDim, isGlobal=True)
15491550
module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))
15501551

15511552
if self.parentWriter.db["ConservativeWaitCnt"] & 0x40:
1552-
module.add(SBarrier("debug"))
1553+
module.add(SBarrier(comment="debug"))
15531554
module.add(SWaitCnt(vscnt=0, comment="ConservativeWaitCnt"))
1554-
module.add(SBarrier("debug"))
1555+
module.add(SBarrier(comment="debug"))
15551556

15561557
def _emitSubtilePackedPermute(self, vPack: int, vPermAddr: int, addrWhilePermuting=None) -> Module:
15571558
"""Shuffle four packed dwords across wave halves for a subtile dwordx4 store.

‎projects/hipblaslt/tensilelite/Tensile/Components/PersistentLoop.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,11 @@
2020
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
2121
################################################################################
2222

23+
from math import ceil, log2
24+
2325
from rocisa.code import Module, Label
2426
from rocisa.container import vgpr, sgpr
25-
from rocisa.instruction import VMovB32, SCmpGeU32, SMulI32
27+
from rocisa.instruction import VMovB32, SCmpGeU32, SMulI32, SLShiftRightB32, VReadfirstlaneB32
2628
from ..Component import Component
2729
import abc
2830

@@ -85,6 +87,16 @@ def openPersistentLoop(self, writer, kernel):
8587
persistentLabel = Label(label="PersistentLoopStart", comment="")
8688
module.add(persistentLabel)
8789

90+
# Re-init sgprWaveIdx every persistent loop iteration: TDM init reads
91+
# s[sgprWaveIdx] but the same sgpr is later UNDEFed and reused as a temp,
92+
# so on the 2nd iteration the value would be stale.
93+
if kernel["enableTDMA"] or kernel["enableTDMB"]:
94+
wavelen = kernel["WavefrontSize"]
95+
with writer.allocTmpSgpr(1) as tmpSgprRes:
96+
module.add(VReadfirstlaneB32(sgpr(tmpSgprRes.idx), vgpr("Serial"), "first tId"))
97+
module.add(SLShiftRightB32(sgpr("WaveIdx"), ceil(log2(wavelen)), sgpr(tmpSgprRes.idx),
98+
"re-init WaveIdx for persistent loop iteration"))
99+
88100
# TODO remove?
89101
# kStr += inst("s_add_u32", sgpr("PersistentLoopIter"), sgpr("PersistentLoopIter"), hex(1), "Inc PersistentLoop Iter") # Back-up: not needed now
90102
#kStr += str(Code.WaitCnt(self.version, 0,0,"wait for outstanding stores"))

‎projects/hipblaslt/tensilelite/Tensile/Components/StreamK.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1345,9 +1345,10 @@ def partialsWriteBatch(self, writer, kernel, ss, batchIdx, applyAlpha, beta, edg
13451345
# for the top-left corner this thread will write. These are not changed
13461346
# across all the store loop iters.
13471347
if writer.db["ConservativeWaitCnt"] & 0x10:
1348-
module.add(SBarrier("debug"))
1348+
module.add(SBarrier(comment="debug"))
13491349
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
1350-
module.add(SBarrier("debug"))
1350+
module.add(SBarrier(comment="debug"))
1351+
13511352
if not edge and writer.db["ForceEdgeStores"]>=2:
13521353
module.add(writer.getBomb()) # should not get here
13531354
if edge and writer.db["AssertNoEdge"]:
@@ -1795,9 +1796,10 @@ def fixupBatch(self, writer, kernel, ss, batchIdx, edge, gwvw, \
17951796
# for the top-left corner this thread will write. These are not changed
17961797
# across all the store loop iters.
17971798
if writer.db["ConservativeWaitCnt"] & 0x10:
1798-
module.add(SBarrier("debug"))
1799+
module.add(SBarrier(comment="debug"))
17991800
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
1800-
module.add(SBarrier("debug"))
1801+
module.add(SBarrier(comment="debug"))
1802+
18011803
if not edge and writer.db["ForceEdgeStores"]>=2:
18021804
module.add(writer.getBomb()) # should not get here
18031805
if edge and writer.db["AssertNoEdge"]:
@@ -2136,9 +2138,10 @@ def fixupBatch(self, writer, kernel, ss, batchIdx, edge, gwvw, \
21362138
module.add(self.getEdgeMovInstType()(EXEC(), -1, "full mask -> exec"))
21372139

21382140
if writer.db["ConservativeWaitCnt"] & 0x40:
2139-
module.add(SBarrier("debug"))
2141+
module.add(SBarrier(comment="debug"))
21402142
module.add(SWaitCnt(vlcnt=0, vscnt=0, comment="ConservativeWaitCnt"))
2141-
module.add(SBarrier("debug"))
2143+
module.add(SBarrier(comment="debug"))
2144+
21422145
########################################
21432146
# End Not Atomic
21442147
########################################

‎projects/hipblaslt/tensilelite/Tensile/Components/TensorDataMover.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,9 @@ def calculateStartAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp
7575
mod.add(SMulI32(sgpr(tmpSgprIdx), tileStride, round(mt * bpe), f"stride * MT({mt}) * bpe({bpe})"))
7676
mod.add(SMulI32(sgpr(tmpSgprIdx), sgpr(tmpSgprIdx), sgpr(sgprWorkgroupName), "*= wgId)"))
7777
#add wave offset
78-
mod.add(VReadfirstlaneB32(sgpr(waveOffsetSgprIdx), vgpr(vgprThreadIdName), "first tId"))
79-
mod.add(SLShiftRightB32(sgpr(waveOffsetSgprIdx), ceil(log2(wavelen)), sgpr(waveOffsetSgprIdx), f"wId=fTid // {wavelen}"))
8078
if tp['isM']:
79+
mod.add(VReadfirstlaneB32(sgpr(waveOffsetSgprIdx), vgpr(vgprThreadIdName), "first tId"))
80+
mod.add(SLShiftRightB32(sgpr(waveOffsetSgprIdx), ceil(log2(wavelen)), sgpr(waveOffsetSgprIdx), f"wId=fTid // {wavelen}"))
8181
if not kernel["ProblemType"]["MetadataLayout"]:
8282
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), round(mt // numWaves), "woffset = wId * mt // numWaves"))
8383
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), sgpr("SizeL"), f"woffset *= stride (SizeL / 8 for metadata)"))
@@ -86,7 +86,7 @@ def calculateStartAddr(self, writer: "KernelWriterAssembly", kernel: Mapping, tp
8686
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), round(du * bpe // 2 // numWaves), "woffset = wId * du * bpe / 2 (sparse) // numWaves"))
8787
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), sgpr(sgprStrideName), f"woffset *= stride"))
8888
else:
89-
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), round(mt // numWaves * bpe // tdmSplit), "woffset = wId * mt // numWaves * bpe // tdmSplit"))
89+
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr("WaveIdx"), round(mt // numWaves * bpe // tdmSplit), "woffset = wId * mt // numWaves * bpe // tdmSplit"))
9090
mod.add(SMulI32(sgpr(waveOffsetSgprIdx), sgpr(waveOffsetSgprIdx), tdmSeparateStride, f"woffset *= stride"))
9191
mod.add(SAddU32(sgpr(tmpSgprIdx), sgpr(tmpSgprIdx), sgpr(waveOffsetSgprIdx), "+= woffset"))
9292
#add GSU offset
@@ -160,8 +160,7 @@ def calculateStartAddrWaveSeparated(self, writer: "KernelWriterAssembly", kernel
160160
mod.add(SMulI32(sgpr(tmpSgprIdx), tileStride, round(mt * bpe), f"tileStride * MT({mt}) * bpe({bpe})"))
161161
mod.add(SMulI32(sgpr(tmpSgprIdx), sgpr(tmpSgprIdx), sgpr(sgprWorkgroupName), "*= wgId)"))
162162
#add wave offset
163-
mod.add(VReadfirstlaneB32(sgpr(waveOffsetSgprIdx), vgpr(vgprThreadIdName), "first tId"))
164-
mod.add(SLShiftRightB32(sgpr(waveOffsetSgprIdx), ceil(log2(wavelen)) + 1, sgpr(waveOffsetSgprIdx), f"wCompId = fTid // wavelen({wavelen}) // 2"))
163+
mod.add(SLShiftRightB32(sgpr(waveOffsetSgprIdx), 1, sgpr("WaveIdx"), f"wCompId = fTid // wavelen({wavelen}) // 2)"))
165164
if ("MXS" in tc):
166165
mxDU = kernel["DepthU"] // kernel["ProblemType"][f"MXBlock{subTc}"]
167166
numMxKGroups = mxDU // mxUnit

0 commit comments

Comments
 (0)