Skip to content

align scale datatype with cublaslt for I8I8I #490

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

Closed
xiaohuguo2023 wants to merge 7 commits into from
Closed

align scale datatype with cublaslt for I8I8I #490

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
8f2eea7
set I8I8I scale datatype to fp32
xiaohuguo2023 Jun 27, 2025
12c56f3
enable loading fp32 scale datatype to LDS for I8I8I
xiaohuguo2023 Jun 28, 2025
8c3f397
Merge branch 'ROCm:develop' into users/xiaohuguo2023/add_I8I8I_scaled…
xiaohuguo2023 Jun 30, 2025
731aee2
add I8I8I scale scalar/vector test
xiaohuguo2023 Jul 2, 2025
d9bca2c
enable UseScaleAB = Scalar for I8I8I with scale data type fp32
xiaohuguo2023 Jul 2, 2025
c1a567a
Merge branch 'ROCm:develop' into users/xiaohuguo2023/add_I8I8I_scaled…
xiaohuguo2023 Jul 10, 2025
7e3d0de
Merge branch 'ROCm:develop' into users/xiaohuguo2023/add_I8I8I_scaled…
xiaohuguo2023 Aug 3, 2025
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
74 changes: 62 additions & 12 deletions projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
DPPModifiers, vgpr, sgpr, accvgpr, mgpr, ContinuousRegister, \
HWRegContainer
from rocisa.instruction import SGetPositivePCOffset, SLongBranchPositive, SCLongBranchScc0, SCLongBranchScc1, \
SMulInt64to32, VCvtBF16toFP32
SMulInt64to32, VCvtBF16toFP32, VCvtF32toI32, VCvtI32toF32
from rocisa.functions import vectorStaticDivide, vectorStaticRemainder, vectorUInt32CeilDivideAndRemainder, \
vectorStaticDivideAndRemainder, scalarStaticDivideAndRemainder, scalarStaticCeilDivide, \
scalarStaticRemainder, scalarUInt32DivideAndRemainder, sMagicDiv, vectorStaticMultiply, \
Expand Down Expand Up @@ -9177,8 +9177,22 @@ def localWriteBody(tP):
localWriteCVTCode.add(VCvtF16toF32(dst=vgpr(vgprTmp), src=paramList[0], sdwa=SDWAModifiers(src0_sel=src_sel), comment="convert to F32"))

# ScaleA/B
if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and kernel["ProblemType"]["DataType%s"%tc].numRegisters() > kernel["ProblemType"]["DataType"].numRegisters():
localWriteCVTCode.add(VMulF32(dst=vgpr(vgprTmp), src0=vgpr(vgprTmp), src1=sgpr("Scale%s"%tc), comment="Input *= scale %s"%tc))
if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (
kernel["ProblemType"]["DataType%s"%tc].numRegisters() > kernel["ProblemType"]["DataType"].numRegisters() or
# check if it is I8I8I
(
kernel["ProblemType"]["DataType%s"%tc].isInt8() and kernel["ProblemType"]["DestDataType%s"%tc].isInt8() and
kernel["ProblemType"]["ComputeDataType"].isInt32()
)
):
if kernel["ProblemType"]["ComputeDataType"].isInt32():
vgprTmpF32 = self.vgprPool.checkout(1)
localWriteCVTCode.add(VCvtI32toF32(dst=vgpr(vgprTmpF32), src=vgpr(vgprTmp), comment="Convert int32 to fp32 for scalar scaling"))
localWriteCVTCode.add(VMulF32(dst=vgpr(vgprTmpF32), src0=vgpr(vgprTmpF32), src1=sgpr("Scale%s"%tc), comment="Input *= scale %s"%tc))
localWriteCVTCode.add(VCvtF32toI32(dst=vgpr(vgprTmp), src=vgpr(vgprTmpF32), comment="Convert fp32 back to int32 after apply scalar scaling"))
self.vgprPool.checkin(vgprTmpF32)
else:
localWriteCVTCode.add(VMulF32(dst=vgpr(vgprTmp), src0=vgpr(vgprTmp), src1=sgpr("Scale%s"%tc), comment="Input *= scale %s"%tc))

if kernel["ProblemType"]["StochasticRounding"]:
vRand = vgprTmp+1 #seed
Expand Down Expand Up @@ -10986,8 +11000,12 @@ def globalWriteElements(self, kernel, tPA, tPB, vectorWidths_2, vectorWidths_1,
if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and \
isSingleKernel and \
((kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters()) or \
(kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters())):
assert(kernel["ProblemType"]["ComputeDataType"].isSingle())
(kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters()) or
(
kernel["ProblemType"]["ComputeDataType"].isInt32() and
kernel["ProblemType"]["DataType"].isInt8() and kernel["ProblemType"]["DestDataType"].isInt8()
)):
assert(kernel["ProblemType"]["ComputeDataType"].isSingle() or kernel["ProblemType"]["ComputeDataType"].isInt32())
sgprScaleA = self.sgprPool.checkOut(1, preventOverflow=False)
sgprScaleB = self.sgprPool.checkOut(1, preventOverflow=False)
for i,name in enumerate(['A','B']):
Expand Down Expand Up @@ -11049,8 +11067,15 @@ def globalWriteElements(self, kernel, tPA, tPB, vectorWidths_2, vectorWidths_1,
module.add(allocPostLoopSrdSuppress("ScaleB", labelStrB, sgprLength=sgpr("SizeJ")))
module.add(SMulI32(dst=sgpr("SrdScaleA+2"), src0=hex(self.states.bpeCinternal), src1=sgpr("SrdScaleA+2"), comment="ScaleAVec scaled by BPE"))# scaled by BPE
module.add(SMulI32(dst=sgpr("SrdScaleB+2"), src0=hex(self.states.bpeCinternal), src1=sgpr("SrdScaleB+2"), comment="ScaleBVec scaled by BPE"))# scaled by BPE
vectorDataTypes.scaleA.dataType = kernel["ProblemType"]["ComputeDataType"]
vectorDataTypes.scaleB.dataType = kernel["ProblemType"]["ComputeDataType"]
inputType = kernel["ProblemType"]["DataType"]
outputType = kernel["ProblemType"]["DestDataType"]
computeType = kernel["ProblemType"]["ComputeDataType"]
if inputType.isInt8() and outputType.isInt8() and computeType.isInt32():
vectorDataTypes.scaleA.dataType = DataType('s')
vectorDataTypes.scaleB.dataType = DataType('s')
else:
vectorDataTypes.scaleA.dataType = kernel["ProblemType"]["ComputeDataType"]
vectorDataTypes.scaleB.dataType = kernel["ProblemType"]["ComputeDataType"]

factorDim0Label = Label(self.labels.getNameInc("Load_FactorDim_0"), "")
factorDim1Label = Label(self.labels.getNameInc("Load_FactorDim_1"), "")
Expand Down Expand Up @@ -11239,15 +11264,34 @@ def globalWriteElements(self, kernel, tPA, tPB, vectorWidths_2, vectorWidths_1,

if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and (((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0) or kernel["_GlobalAccumulation"] == 'MultipleBufferSingleKernel') and \
((kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters()) or \
(kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters())):
assert(kernel["ProblemType"]["ComputeDataType"].isSingle())
(kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters()) or \
(
kernel["ProblemType"]["ComputeDataType"].isInt32() and
kernel["ProblemType"]["DataType"].isInt8() and kernel["ProblemType"]["DestDataType"].isInt8()
)):
assert(kernel["ProblemType"]["ComputeDataType"].isSingle() or kernel["ProblemType"]["ComputeDataType"].isInt32())
newAlphaVgpr = self.vgprPool.checkOut(1)
module.add(VMovB32(dst=vgpr(newAlphaVgpr), src=sgpr("Alpha")))
module.add(SWaitCnt(lgkmcnt=0, kmcnt=0, comment="wait for scaleAB load"))
if kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters():
module.add(VMulF32(dst=vgpr(newAlphaVgpr), src0=vgpr(newAlphaVgpr), src1=sgpr(sgprScaleA)))
else:
if kernel["ProblemType"]["ComputeDataType"].isInt32():
tmpAlphaF32 = self.vgprPool.checkout(1)
module.add(VCvtI32toF32(dst=vgpr(tmpAlphaF32), src=vgpr(newAlphaVgpr), comment="Convert int32 to fp32 for scalar scaling"))
module.add(VMulF32(dst=vgpr(tmpAlphaF32), src0=vgpr(tmpAlphaF32), src1=sgpr(sgprScaleA), comment="Input *= scale %s"%tc))
module.add(VCvtF32toI32(dst=vgpr(newAlphaVgpr), src=vgpr(tmpAlphaF32), comment="Convert fp32 back to int32 after apply scalar scaling"))
self.vgprPool.checkin(tmpAlphaF32)
if kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters():
module.add(VMulF32(dst=vgpr(newAlphaVgpr), src0=vgpr(newAlphaVgpr), src1=sgpr(sgprScaleB)))
else:
if kernel["ProblemType"]["ComputeDataType"].isInt32():
tmpAlphaF32 = self.vgprPool.checkout(1)
module.add(VCvtI32toF32(dst=vgpr(tmpAlphaF32), src=vgpr(newAlphaVgpr), comment="Convert int32 to fp32 for scalar scaling"))
module.add(VMulF32(dst=vgpr(tmpAlphaF32), src0=vgpr(tmpAlphaF32), src1=sgpr(sgprScaleB), comment="Input *= scale %s"%tc))
module.add(VCvtF32toI32(dst=vgpr(newAlphaVgpr), src=vgpr(tmpAlphaF32), comment="Convert fp32 back to int32 after apply scalar scaling"))
self.vgprPool.checkin(tmpAlphaF32)

module.add(SNop(waitState=0, comment="1 wait states"))
if kernel["StreamK"] > 0:
oldAlpha = self.sgprPool.checkOut(1)
Expand Down Expand Up @@ -11605,8 +11649,12 @@ def globalWriteElements(self, kernel, tPA, tPB, vectorWidths_2, vectorWidths_1,

if kernel["ProblemType"]["UseScaleAB"] == "Scalar" and kernel["StreamK"] > 0 and \
((kernel["ProblemType"]["DataTypeA"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters()) or \
(kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters())):
assert(kernel["ProblemType"]["ComputeDataType"].isSingle())
(kernel["ProblemType"]["DataTypeB"].numRegisters() <= kernel["ProblemType"]["DataType"].numRegisters()) or \
(
kernel["ProblemType"]["ComputeDataType"].isInt32() and
kernel["ProblemType"]["DataType"].isInt8() and kernel["ProblemType"]["DestDataType"].isInt8()
)):
assert(kernel["ProblemType"]["ComputeDataType"].isSingle() or kernel["ProblemType"]["ComputeDataType"].isInt32())
module.add(SMovB32(dst=sgpr("Alpha"), src=sgpr(oldAlpha), comment="Restore alpha value"))
self.sgprPool.checkIn(oldAlpha)

Expand Down Expand Up @@ -12604,7 +12652,9 @@ def addVectorLocalStore(self, kernel, addressStr: str, offsetVgpr, shiftOffset,
else:
printExit("[Compute fp32] Unrecognized data type %s."%str(dataType))
elif kernel["ProblemType"]["ComputeDataType"].isInt32():
if dataType == kernel["ProblemType"]["ComputeDataType"]:
if kernel["ProblemType"]["DataType"].isInt8():
pass # this is to deal with I8I8I scaledatatype = fp32
elif dataType == kernel["ProblemType"]["ComputeDataType"]:
pass # Same, no need to convert
else:
printExit("[Compute int32] Unrecognized data type %s."%str(dataType))
Expand Down
139 changes: 139 additions & 0 deletions projects/hipblaslt/tensilelite/Tensile/Tests/common/streamk/sk_I8gemm_scale_test.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
TestParameters:
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] # not supported by arch

GlobalParameters:
NumElementsToValidate: -1
BoundsCheck: False
KernelTime: False
DataInitTypeA: 21
DataInitTypeB: 21
DataInitTypeC: 21
DataInitTypeAlpha: 1
DataInitTypeBeta: 1
DataInitTypeBias: 21
DataInitTypeScaleA: 2
DataInitTypeScaleB: 2
DataInitTypeScaleC: 1
DataInitTypeScaleD: 2
DataInitTypeScaleAlphaVec: 21
MaxWorkspaceSize: 134217728
PrintSolutionRejectionReason: True

BenchmarkProblems:
- # I8I8I NN
-
OperationType: GEMM
DataType: I8
DestDataType: I8
ComputeDataType: I
HighPrecisionAccumulate: True
TransposeA: 0
TransposeB: 0
UseBeta: True
Batched: True
Activation: False
UseScaleAB: "Scalar"

@jichangjichang jichangjichang Jul 4, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Did you try to verify "Vector" or "Scalar"?

@xiaohuguo2023 xiaohuguo2023 Jul 4, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

both

UseScaleAlphaVec: 1
-
InitialSolutionParameters:
BenchmarkCommonParameters:
- KernelLanguage: ["Assembly"]
ForkParameters:
- 1LDSBuffer: [-1]
- ClusterLocalRead: [1]
- DepthU: [ 64 ]
- ExpandPointerSwap: [1]
- GlobalReadPerMfma: [1]
- GlobalReadVectorWidthA: [4]
- GlobalReadVectorWidthB: [4]
- InnerUnroll: [1]
- MatrixInstruction:
- [16,16,32, 1, 1, 2,2, 2,2] # 16x16
# - [16,16,64, 1, 1, 1,1, 1,1]
# - [32,32,32, 1, 1, 1,1, 1,1]
- LdsBlockSizePerPadA: [0]
- LdsBlockSizePerPadB: [0]
- LdsPadA: [0]
- LdsPadB: [0]
- LocalWritePerMfma: [-1]
- NumElementsPerBatchStore: [0]
- PrefetchGlobalRead: [2]
- PrefetchLocalRead: [1]
- ScheduleIterAlg: [3]
- SourceSwap: [1]
- StorePriorityOpt: [0]
- StoreVectorWidth: [-1]
- StreamK: [3]
- VectorWidthA: [1]
- VectorWidthB: [1]
- WaveSeparateGlobalReadB: [1]
BenchmarkJoinParameters:
BenchmarkFinalParameters:
- ProblemSizes:
- Exact: [128, 128, 1, 128]
- Exact: [1024, 1024, 1, 1024]
- Exact: [4111, 4111, 1, 4111]
- ActivationArgs:
- [Enum: relu]

- # I8BS NT ScaleAB=Vector ScaleCD ScaleAlpha Bias
-
OperationType: GEMM
DataType: I8
DestDataType: I8
ComputeDataType: I
HighPrecisionAccumulate: True
TransposeA: 0
TransposeB: 1
UseBeta: True
Batched: True
Activation: True
#UseScaleAB: "Scalar"
UseScaleAB: "Vector"
UseScaleAlphaVec: 1
# UseBias: 0
# BiasDataTypeList: [I]
UseE: True
-
InitialSolutionParameters:
BenchmarkCommonParameters:
- KernelLanguage: ["Assembly"]
ForkParameters:
- 1LDSBuffer: [-1]
- ClusterLocalRead: [1]
- DepthU: [ 64 ]
- ExpandPointerSwap: [1]
- GlobalReadPerMfma: [1]
- GlobalReadVectorWidthA: [4]
- GlobalReadVectorWidthB: [4]
- InnerUnroll: [1]
- MatrixInstruction:
- [16,16,32, 1, 1, 2,2, 2,2] # 16x16
# - [16,16,64, 1, 1, 1,1, 1,1]
# - [32,32,32, 1, 1, 1,1, 1,1]
- LdsBlockSizePerPadA: [0]
- LdsBlockSizePerPadB: [0]
- LdsPadA: [0]
- LdsPadB: [0]
- LocalWritePerMfma: [-1]
- NumElementsPerBatchStore: [0]
- PrefetchGlobalRead: [2]
- PrefetchLocalRead: [1]
- ScheduleIterAlg: [3]
- SourceSwap: [1]
- StorePriorityOpt: [0]
- StoreVectorWidth: [-1]
- StreamK: [3]
- VectorWidthA: [1]
- VectorWidthB: [1]
- WaveSeparateGlobalReadB: [1]
BenchmarkJoinParameters:
BenchmarkFinalParameters:
- ProblemSizes:
- Exact: [128, 128, 1, 128]
- Exact: [1024, 1024, 1, 1024]
- Exact: [4111, 4111, 1, 4111]
# - BiasTypeArgs: ['I']
- ActivationArgs:
- [Enum: relu]