Skip to content

Commit a2ce1ab

Browse files
AlexBrownAMDAlexBrownAMD
authored
Fix stream-k with mx scaling (#4388)
## Motivation Fix stream-k support for mx types with scaling. ## Technical Details Fixes global read address offset calculation for stream-k partial tiles. ## Test Plan Works in initial batch of local tests. More testing with other parameter combinations to be done, and a new test case to be added for CI.
1 parent fd621eb commit a2ce1ab

5 files changed

Lines changed: 247 additions & 14 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -307,9 +307,9 @@ def computeLoadSrd(self, writer, kernel, tP, stmp, tileStart):
307307
_DepthU = kernel["_DepthU%s" % tc]
308308
# swizzle
309309
if (tP["isSwizzled"] and tc == 'A'):
310-
_DepthU = (_DepthU * 16) # MI_M = 16
310+
_DepthU = (_DepthU * 16)
311311
elif (tP["isSwizzled"] and tc == 'B'):
312-
_DepthU = (_DepthU * 16) # MI_N = 16
312+
_DepthU = (_DepthU * 16)
313313

314314
gsucLabel = Label(label=writer.labels.getNameInc(f"GSUC_{tc}"), comment="")
315315
gsucLabelEnd = Label(label=writer.labels.getNameInc(f"GSUC_{tc}_End"), comment="")

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

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -176,15 +176,23 @@ def skExtraIters(self, writer, kernel, sSkExtraIters, sTmp):
176176
return module
177177

178178
@abc.abstractmethod
179-
def computeLoadSrd(self, writer, kernel, tc, sTmp):
179+
def computeLoadSrd(self, writer, kernel, tP, sTmp):
180180
pass
181181

182-
def computeLoadSrdCommon(self, writer, kernel, tc, sTmp):
182+
def computeLoadSrdCommon(self, writer, kernel, tP, sTmp):
183183
module = Module("StreamK Common computeLoadSrd")
184184

185+
tc = tP["tensorChar"]
186+
_DepthU = kernel["_DepthU%s" % tc]
187+
# swizzle
188+
if (tP["isSwizzled"] and tc == 'A'):
189+
_DepthU = (_DepthU * 16)
190+
elif (tP["isSwizzled"] and tc == 'B'):
191+
_DepthU = (_DepthU * 16)
192+
185193
tileStart = sTmp + 2
186194
# StreamK partial tile - offset to tile start index
187-
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("StreamKLocalStart"), src1=kernel["DepthU"], comment="StreamK tile start offset"))
195+
module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("StreamKLocalStart"), src1=_DepthU, comment="StreamK tile start offset"))
188196
strideL = writer.strideRef(tc, kernel["ProblemType"]["IndicesSummation"][0])
189197
module.add(writer.s_mul_u64_u32(sgpr(sTmp), sgpr(sTmp+1), sgpr(sTmp), strideL, comment="StreamK tile start offset"))
190198
# Overflow check removed
@@ -255,9 +263,16 @@ def graAddressesCommon(self, writer, kernel, tP, vTmp):
255263
module = Module("StreamK Common graAddresses")
256264

257265
tc = tP["tensorChar"]
266+
_DepthU = kernel["_DepthU%s" % tc]
267+
# swizzle
268+
if (tP["isSwizzled"] and tc == 'A'):
269+
_DepthU = (_DepthU * 16)
270+
elif (tP["isSwizzled"] and tc == 'B'):
271+
_DepthU = (_DepthU * 16)
272+
258273
# StreamK partial tile - offset to tile start index
259274
tmpOffset = writer.sgprPool.checkOut(2, "skStartOffset")
260-
module.add(SMulI32(dst=sgpr(tmpOffset), src0=sgpr("StreamKLocalStart"), src1=int(kernel["DepthU"] * tP["bpe"]), comment="StreamK tile start offset"))
275+
module.add(SMulI32(dst=sgpr(tmpOffset), src0=sgpr("StreamKLocalStart"), src1=int(_DepthU * tP["bpe"]), comment="StreamK tile start offset"))
261276
strideL = writer.strideRef(tc, kernel["ProblemType"]["IndicesSummation"][0])
262277
module.add(writer.s_mul_u64_u32(sgpr(tmpOffset), sgpr(tmpOffset+1), sgpr(tmpOffset), strideL, "StreamK tile start offset"))
263278
# Overflow check removed
@@ -1745,7 +1760,7 @@ def graWorkGroup(self, writer, kernel, tPA, tPB):
17451760
module = Module("StreamK Off graWorkGroup")
17461761
return module
17471762

1748-
def computeLoadSrd(self, writer, kernel, tc, sTmp):
1763+
def computeLoadSrd(self, writer, kernel, tP, sTmp):
17491764
module = Module("StreamK Off computeLoadSrd")
17501765
return module
17511766

@@ -1861,9 +1876,9 @@ def graWorkGroup(self, writer, kernel, tPA, tPB):
18611876

18621877
return module
18631878

1864-
def computeLoadSrd(self, writer, kernel, tc, sTmp):
1879+
def computeLoadSrd(self, writer, kernel, tP, sTmp):
18651880
module = Module("StreamK Basic computeLoadSrd")
1866-
module.add(self.computeLoadSrdCommon(writer, kernel, tc, sTmp))
1881+
module.add(self.computeLoadSrdCommon(writer, kernel, tP, sTmp))
18671882
return module
18681883

18691884
def computeStoreSrdStart(self, writer, kernel):
@@ -1986,9 +2001,9 @@ def graWorkGroup(self, writer, kernel, tPA, tPB):
19862001

19872002
return module
19882003

1989-
def computeLoadSrd(self, writer, kernel, tc, sTmp):
2004+
def computeLoadSrd(self, writer, kernel, tP, sTmp):
19902005
module = Module("StreamK TwoTileOriginal computeLoadSrd")
1991-
module.add(self.computeLoadSrdCommon(writer, kernel, tc, sTmp))
2006+
module.add(self.computeLoadSrdCommon(writer, kernel, tP, sTmp))
19922007
return module
19932008

19942009
def computeStoreSrdStart(self, writer, kernel):
@@ -2256,9 +2271,9 @@ def graWorkGroup(self, writer, kernel, tPA, tPB):
22562271

22572272
return module
22582273

2259-
def computeLoadSrd(self, writer, kernel, tc, sTmp):
2274+
def computeLoadSrd(self, writer, kernel, tP, sTmp):
22602275
module = Module("StreamK TwoTileDPFirst computeLoadSrd")
2261-
module.add(self.computeLoadSrdCommon(writer, kernel, tc, sTmp))
2276+
module.add(self.computeLoadSrdCommon(writer, kernel, tP, sTmp))
22622277
return module
22632278

22642279
def computeStoreSrdStart(self, writer, kernel):

‎projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3620,7 +3620,7 @@ def computeLoadSrd(self, kernel, tP, tc, indices, bpe):
36203620
strideF, comment="tlu=0, scaled tile-offset by stride"))
36213621

36223622
skComponent = Component.StreamK.find(self)
3623-
module.add(skComponent.computeLoadSrd(self, kernel, tc, stmp))
3623+
module.add(skComponent.computeLoadSrd(self, kernel, tP, stmp))
36243624

36253625
gsuComponent = Component.GSU.find(self)
36263626
module.add(gsuComponent.computeLoadSrd(self, kernel, tP, stmp, tileStart))

‎projects/hipblaslt/tensilelite/Tensile/Tests/common/streamk/sk_mx32f4_quick.yaml‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
TestParameters:
2+
marks: [skip-gfx900, skip-gfx906, skip-gfx908, skip-gfx90a, skip-gfx940, skip-gfx941, skip-gfx942, skip-gfx1010, skip-gfx1011, skip-gfx1012, skip-gfx1030, skip-gfx1100, skip-gfx1101, skip-gfx1102, skip-gfx1200, skip-gfx1201]
3+
4+
GlobalParameters:
5+
NumElementsToValidate: -1
6+
MinimumRequiredVersion: 5.0.0
7+
PrintLevel: 1
8+
# PrintSolutionRejectionReason: True
9+
Device: 0
10+
CMakeBuildType: Debug
11+
# MergeFiles: False
12+
KernelTime: True
13+
MaxWorkspaceSize: 13421772800
14+
DataInitTypeA: 3
15+
DataInitTypeB: 3
16+
DataInitTypeC: 0
17+
DataInitTypeMXSA: 3
18+
DataInitTypeMXSB: 3
19+
DataInitTypeAlpha: 1
20+
DataInitTypeBeta: 0
21+
NumElementsToValidate: -1
22+
BoundsCheck: 0
23+
KeepBuildTmp: True
24+
MaxFileName: 128
25+
DeviceLDS: 163840
26+
MaxLDS: 163840
27+
28+
BenchmarkProblems:
29+
########################################
30+
# FP4SS
31+
########################################
32+
-
33+
- # ProblemType
34+
OperationType: GEMM
35+
DataType: F4
36+
DestDataType: S
37+
ComputeDataType: S
38+
HighPrecisionAccumulate: True
39+
MXBlockA: 32
40+
MXBlockB: 32
41+
TransposeA: 1
42+
TransposeB: 0
43+
UseBeta: True
44+
Batched: True
45+
Activation: True
46+
ActivationType: hipblaslt_all
47+
# UseScaleAB: "Scalar"
48+
# UseScaleCD: True
49+
UseScaleAlphaVec: 1
50+
UseBias: 1
51+
BiasDataTypeList: [s]
52+
53+
- # BenchmarkProblemSizeGroup - Standard
54+
InitialSolutionParameters:
55+
BenchmarkCommonParameters:
56+
- KernelLanguage: ["Assembly"]
57+
ForkParameters:
58+
- MatrixInstruction:
59+
# - [16, 16, 128, 1, 1, 1,1, 1,1]
60+
# - [16, 16, 128, 1, 1, 4,2, 2,2]
61+
- [16, 16, 128, 1, 1, 2,4, 2,2] # 64x128
62+
# - [16, 16, 128, 1, 1, 8,8, 2,2]
63+
# - [32, 32, 64, 1, 1, 1,1, 1,1]
64+
- [32, 32, 64, 1, 1, 2,2, 2,2] # 128x128
65+
# - [32, 32, 64, 1, 1, 4,2, 2,2]
66+
# - [32, 32, 64, 1, 1, 2,4, 2,2]
67+
- ForceDisableShadowInit: [True, False]
68+
# - UseSgprForGRO: [0,1]
69+
- UseSgprForGRO: [0]
70+
# - DepthU: [64, 128]
71+
- DepthU: [128]
72+
- AssertFree0ElementMultiple: [1]
73+
- AssertFree1ElementMultiple: [1]
74+
- AssertSummationElementMultiple: [64]
75+
- LocalReadVectorWidth: [16]
76+
# - PrefetchGlobalRead: [0,1,2]
77+
- PrefetchGlobalRead: [2]
78+
# - PrefetchLocalRead: [0,1]
79+
- PrefetchLocalRead: [1]
80+
# - PreloadKernArgs: [0,1]
81+
- PreloadKernArgs: [1]
82+
# - ClusterLocalRead: [0,1]
83+
- ClusterLocalRead: [1]
84+
- VectorWidthA: [1]
85+
- VectorWidthB: [1]
86+
- GlobalReadVectorWidthA: [16]
87+
- GlobalReadVectorWidthB: [16]
88+
- ScheduleIterAlg: [3]
89+
- InnerUnroll: [1]
90+
- TransposeLDS: [1]
91+
- WaveSeparateGlobalReadA: [0]
92+
- WaveSeparateGlobalReadB: [0]
93+
- 1LDSBuffer: [0]
94+
- GlobalReadPerMfma: [1]
95+
- LocalWritePerMfma: [-1]
96+
- StoreVectorWidth: [-1]
97+
- SourceSwap: [1]
98+
- StreamK: [3]
99+
BenchmarkJoinParameters:
100+
BenchmarkFinalParameters:
101+
- ProblemSizes:
102+
- Exact: [1025, 513, 1, 2048]
103+
- Exact: [1044, 532, 1, 2048]
104+
- Exact: [127, 127, 1, 640] #special cleanup case
105+
- Exact: [128, 128, 1, 128]
106+
- Exact: [129, 129, 1, 640]
107+
- BiasTypeArgs: ['s']
108+
- ActivationArgs:
109+
- [Enum: none]

‎projects/hipblaslt/tensilelite/Tensile/Tests/common/streamk/sk_mx32f8_quick.yaml‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
TestParameters:
2+
marks: [skip-gfx900, skip-gfx906, skip-gfx908, skip-gfx90a, skip-gfx940, skip-gfx941, skip-gfx942, skip-gfx1010, skip-gfx1011, skip-gfx1012, skip-gfx1030, skip-gfx1100, skip-gfx1101, skip-gfx1102, skip-gfx1200, skip-gfx1201]
3+
4+
GlobalParameters:
5+
NumElementsToValidate: -1
6+
MinimumRequiredVersion: 5.0.0
7+
PrintLevel: 1
8+
# PrintSolutionRejectionReason: True
9+
Device: 0
10+
CMakeBuildType: Debug
11+
# MergeFiles: False
12+
KernelTime: True
13+
MaxWorkspaceSize: 13421772800
14+
DataInitTypeA: 3
15+
DataInitTypeB: 3
16+
DataInitTypeC: 0
17+
DataInitTypeMXSA: 3
18+
DataInitTypeMXSB: 3
19+
DataInitTypeAlpha: 1
20+
DataInitTypeBeta: 0
21+
NumElementsToValidate: -1
22+
BoundsCheck: 0
23+
KeepBuildTmp: True
24+
MaxFileName: 128
25+
DeviceLDS: 163840
26+
MaxLDS: 163840
27+
28+
BenchmarkProblems:
29+
########################################
30+
# FP8SS
31+
########################################
32+
-
33+
- # ProblemType
34+
OperationType: GEMM
35+
DataType: F8
36+
DestDataType: S
37+
ComputeDataType: S
38+
HighPrecisionAccumulate: True
39+
MXBlockA: 32
40+
MXBlockB: 32
41+
TransposeA: 1
42+
TransposeB: 0
43+
UseBeta: True
44+
Batched: True
45+
Activation: True
46+
ActivationType: hipblaslt_all
47+
# UseScaleAB: "Scalar"
48+
# UseScaleCD: True
49+
UseScaleAlphaVec: 1
50+
UseBias: 1
51+
BiasDataTypeList: [s]
52+
53+
- # BenchmarkProblemSizeGroup - Standard
54+
InitialSolutionParameters:
55+
BenchmarkCommonParameters:
56+
- KernelLanguage: ["Assembly"]
57+
ForkParameters:
58+
- MatrixInstruction:
59+
# - [16, 16, 128, 1, 1, 1,1, 1,1]
60+
# - [16, 16, 128, 1, 1, 4,2, 2,2]
61+
- [16, 16, 128, 1, 1, 2,4, 2,2] # 64x128
62+
# - [16, 16, 128, 1, 1, 8,8, 2,2]
63+
# - [32, 32, 64, 1, 1, 1,1, 1,1]
64+
- [32, 32, 64, 1, 1, 2,2, 2,2] # 128x128
65+
# - [32, 32, 64, 1, 1, 4,2, 2,2]
66+
# - [32, 32, 64, 1, 1, 2,4, 2,2]
67+
- ForceDisableShadowInit: [True, False]
68+
# - UseSgprForGRO: [0,1]
69+
- UseSgprForGRO: [0]
70+
# - DepthU: [64, 128]
71+
- DepthU: [128]
72+
- AssertFree0ElementMultiple: [1]
73+
- AssertFree1ElementMultiple: [1]
74+
- AssertSummationElementMultiple: [64]
75+
- LocalReadVectorWidth: [16]
76+
# - PrefetchGlobalRead: [0,1,2]
77+
- PrefetchGlobalRead: [2]
78+
# - PrefetchLocalRead: [0,1]
79+
- PrefetchLocalRead: [1]
80+
# - PreloadKernArgs: [0,1]
81+
- PreloadKernArgs: [1]
82+
# - ClusterLocalRead: [0,1]
83+
- ClusterLocalRead: [1]
84+
- VectorWidthA: [1]
85+
- VectorWidthB: [1]
86+
- GlobalReadVectorWidthA: [16]
87+
- GlobalReadVectorWidthB: [16]
88+
- ScheduleIterAlg: [3]
89+
- InnerUnroll: [1]
90+
- TransposeLDS: [1]
91+
- WaveSeparateGlobalReadA: [0]
92+
- WaveSeparateGlobalReadB: [0]
93+
- 1LDSBuffer: [0]
94+
- GlobalReadPerMfma: [1]
95+
- LocalWritePerMfma: [-1]
96+
- StoreVectorWidth: [-1]
97+
- SourceSwap: [1]
98+
- StreamK: [3]
99+
BenchmarkJoinParameters:
100+
BenchmarkFinalParameters:
101+
- ProblemSizes:
102+
- Exact: [1025, 513, 1, 2048]
103+
- Exact: [1044, 532, 1, 2048]
104+
- Exact: [127, 127, 1, 640] #special cleanup case
105+
- Exact: [128, 128, 1, 128]
106+
- Exact: [129, 129, 1, 640]
107+
- BiasTypeArgs: ['s']
108+
- ActivationArgs:
109+
- [Enum: none]

0 commit comments

Comments
 (0)