Skip to content

Add support for single wave stream-k kernels #1975

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
wants to merge 1 commit into
base: develop
Choose a base branch
from
+24 −5
Open

Add support for single wave stream-k kernels #1975

Show file tree
Hide file tree
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
15 changes: 12 additions & 3 deletions tensilelite/Tensile/Components/StreamK.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -756,6 +756,11 @@ def partialsWriteBatch(self, writer, kernel, ss, batchIdx, applyAlpha, beta, edg
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8Max), "0x47600000", comment="BF8 Max value 57344 as float32" ))
module.add(VMovB32(vgpr(cvtVgprStruct.vgprBF8Min), "0xc7600000", comment="BF8 Min value -57344 as float32" ))

if kernel["EnableMatrixInstruction"]:
WaveNum = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need to consider LSU to compute wave count here as well?

else:
WaveNum = kernel["NumThreads"] // kernel["WavefrontSize"]

storeCode = Module("Partials GroupLoadStore")
for elementIdx in range(len(batchElements)):
element = batchElements[elementIdx]
Expand All @@ -771,7 +776,7 @@ def partialsWriteBatch(self, writer, kernel, ss, batchIdx, applyAlpha, beta, edg
# kStr += inst("v_mul_lo_u32", , "Partials buffer address")
module.add(SMovB32(dst=sgpr(tmpS01), src=0, comment="Init sgpr offset"))
else:
increment = (kernel["WavefrontSize"] * 4) * storeWidth * writer.states.bpeCinternal
increment = (kernel["WavefrontSize"] * WaveNum) * storeWidth * writer.states.bpeCinternal
module.add(SAddU32(dst=sgpr(tmpS01), src0=sgpr(tmpS01), src1=increment, comment="Inc sgpr offset"))

# TODO StreamK need this packing code???
Expand Down Expand Up @@ -1089,7 +1094,6 @@ def fixupBatch(self, writer, kernel, ss, batchIdx, edge, gwvw, \
storesIssued = 0
tmpS01 = tmpSgpr # scratch sgprs

wavelen = kernel["WavefrontSize"]
# laneSGPRC = writer.states.laneSGPRCount
# always use gwvw for buffer load C for atomic_cmpswap
# bpm = self.bpeCexternal * atomicW
Expand Down Expand Up @@ -1126,6 +1130,11 @@ def fixupBatch(self, writer, kernel, ss, batchIdx, edge, gwvw, \
# accVgprRead = Code.Module("movaccVgpr")
# self.StoreCUnrollLoadCWaitComment = "waitcnt for LoadC" # this will be used later to identify waitcnt for loadC

if kernel["EnableMatrixInstruction"]:
WaveNum = kernel["MIWaveGroup"][0] * kernel["MIWaveGroup"][1]
else:
WaveNum = kernel["NumThreads"] // kernel["WavefrontSize"]

for elementIdx in range(0, len(batchElements)):
element = batchElements[elementIdx]
addrCVgpr = ss.elementAddr[elementIdx].addrCVgpr
Expand All @@ -1146,7 +1155,7 @@ def fixupBatch(self, writer, kernel, ss, batchIdx, edge, gwvw, \
# kStr += inst("v_mul_lo_u32", , "Partials buffer address")
module.add(SMovB32(dst=sgpr(tmpS01), src=0, comment="Init sgpr offset"))
else:
increment = (kernel["WavefrontSize"] * 4) * storeWidth * writer.states.bpeCinternal
increment = (kernel["WavefrontSize"] * WaveNum) * storeWidth * writer.states.bpeCinternal
module.add(SAddU32(dst=sgpr(tmpS01), src0=sgpr(tmpS01), src1=increment, comment="Inc sgpr offset"))

module.add(writer.readInput(kernel, ss, 'WS', kernel["ProblemType"]["ComputeDataType"], addrCalc, vc0, data, gwvw, addrCVgpr, sgpr(tmpS01)))
Expand Down
2 changes: 0 additions & 2 deletions tensilelite/Tensile/SolutionStructs/Solution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1089,8 +1089,6 @@ def assignDerivedParameters(
state["GlobalSplitUAlgorithm"] = "MultipleBuffer" # Set default Algorithm
if state["ProblemType"]["DataType"].isDouble():
reject(state, printRejectionReason, "Type {} for DataType not yet supported with StreamK".format(state["ProblemType"]["DataType"].toChar()))
if state["MIWaveGroup"][0] * state["MIWaveGroup"][1] != 4:
reject(state, printRejectionReason, "Stream-K requries MIWaveGroup0*MIWaveGroup1=4")
if not state["EnableMatrixInstruction"]:
reject(state, printRejectionReason, "Stream-K requires MatrixInstruction")
if isaInfoMap[isa].asmCaps["HasWMMA"]:
Expand Down
12 changes: 12 additions & 0 deletions tensilelite/Tensile/Tests/common/streamk/sk_hgemm_quick.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,9 @@ BenchmarkProblems:
- [16, 16, 16, 1, 1, 1,4, 2,2]
- [16, 16, 16, 1, 1, 1,2, 2,2]
- [16, 16, 16, 1, 1, 1,1, 2,2]
- [16, 16, 16, 1, 1, 1,1, 2,1]
- [16, 16, 16, 1, 1, 1,1, 1,2]
- [16, 16, 16, 1, 1, 1,1, 1,1]
- MIArchVgpr: [0]
- PrefetchGlobalRead: [2]
- PrefetchLocalRead: [1]
Expand Down Expand Up @@ -265,6 +268,9 @@ BenchmarkProblems:
- [16, 16, 16, 1, 1, 1,4, 2,2]
- [16, 16, 16, 1, 1, 1,2, 2,2]
- [16, 16, 16, 1, 1, 1,1, 2,2]
- [16, 16, 16, 1, 1, 1,1, 2,1]
- [16, 16, 16, 1, 1, 1,1, 1,2]
- [16, 16, 16, 1, 1, 1,1, 1,1]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will this change support this MFMA but with LSU=4 (e.g., - WorkGroup:[4,4,4]?
If so, pls add test with 2x1 or 1x2 waves with LSU=2 and 1x1 waves with LSU=2 and 4.

- MIArchVgpr: [0]
- PrefetchGlobalRead: [2]
- PrefetchLocalRead: [1]
Expand Down Expand Up @@ -461,6 +467,9 @@ BenchmarkProblems:
# - [16, 16, 16, 1, 1, 1,4, 2,2]
# - [16, 16, 16, 1, 1, 1,2, 2,2]
# - [16, 16, 16, 1, 1, 1,1, 2,2]
# - [16, 16, 16, 1, 1, 1,1, 2,1]
# - [16, 16, 16, 1, 1, 1,1, 1,2]
# - [16, 16, 16, 1, 1, 1,1, 1,1]
# - MIArchVgpr: [0]
# - PrefetchGlobalRead: [2]
# - PrefetchLocalRead: [1]
Expand Down Expand Up @@ -657,6 +666,9 @@ BenchmarkProblems:
# - [16, 16, 16, 1, 1, 1,4, 2,2]
# - [16, 16, 16, 1, 1, 1,2, 2,2]
# - [16, 16, 16, 1, 1, 1,1, 2,2]
# - [16, 16, 16, 1, 1, 1,1, 2,1]
# - [16, 16, 16, 1, 1, 1,1, 1,2]
# - [16, 16, 16, 1, 1, 1,1, 1,1]
# - MIArchVgpr: [0]
# - PrefetchGlobalRead: [2]
# - PrefetchLocalRead: [1]
Expand Down