Stream-k initial support by AlexBrownAMD · Pull Request #1035 · ROCm/hipBLASLt

AlexBrownAMD · 2024-08-22T15:50:02Z

This PR introduces initial stream-k support into hipBLASLt. SGEMM, HGEMM, and BGEMM have been validated with a large number of parameter combinations and problem sizes, and this PR includes some new CI tests for these problem types. The stream-k algorithm helps deliver improved and stable performance with smaller library size. For full details, see the presentations or paper.

Most of the code changes for stream-k are isolated in a component class to help with code readability. Since stream-k is not compatible with GSU, all the GSU related code has also been moved from the KernelWriter classes into it's own component class, and an additional setting has been added to allow generating kernels with GSU disabled (GSU=0 disables it, GSU=1 is still the default, a new CI test for both settings is included). A component for persistent looping has also been added since stream-k operates on a loop. Descriptions of how to use the stream-k options and run-time environment variables are included in the comments in Common.py.

Full stream-k support in hipBLASLt is still a work in progress, and this PR is meant to introduce the initial core functionality. There are a number of known issues that I will work in in future PRs and there are tickets to track, including:

Support and validation for all data types
Consolidate some of the fixup and partials write procedures with the general global write module
Compare and consolidate the stream-k XCC mapping feature with the WorkGroupMappingXCC option
Group gemm support
Additional plans for code improvement and porting of more stream-k features from Tensile repo are described in tickets

Modularize GSU code so it is more contained, reduces branch points in KernelWriter code. Fix a temp register allocation bug. Fix alignment of some registers rather than rely on a specific order.

…ixes

#3) * Support for predicate arithmetic intensity added on top of streamK branch in Alex's repo

* Analytical grid size prediction model for Stream-K

Porting code gen changes for stream-k

Flag init

nakajee

First round review. I will spend more time later...

nakajee · 2024-08-22T16:13:54Z

    ##############################################################################
    class StoreConstConfig:
-        def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic):
+        def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic, isWorkspace=False):


Don't we need to change halfDataRegPerVI (line 105).
It refers DestDataType.

I think line 105 is OK. It is just checking if the write target is a half type (and excluding 8-bit types). The isWorkspace parameter makes bpeC use internal compute size instead of external write size so numVgprsPerDataPerVI is correct for the partials write case.

nakajee · 2024-08-22T16:26:47Z

@@ -0,0 +1,58 @@
+
+GlobalParameters:


I think we need skip here.
This should be only for gfx9xx

Good catch, just posted a fix for this this test as well as the other stream-k tests that were added

neoblizz · 2024-08-23T22:50:16Z

+             * `(n + d - 1)` would overflow.
+             */
+            template <typename N, typename D>
+            __device__ __host__ inline constexpr N safe_ceil_div(N n, D d)


Might be able to use ceilDivide: https://github.com/ROCm/hipBLASLt/blob/fecd96c2ecb4064aa3c11610207b565605061ecd/tensilelite/Tensile/TensileInstructions/Utils.py#L195C5-L195C15, safe_ceil_div is probably better version of the implementation.

nakajee · 2024-08-27T05:27:57Z

I tried HHS TN + StreamK: [0,1,2,3] but I see some issues.
(1) warnings
TensileInstructions::WARNING: RegisterPool::remove(0,95) pool0 still in use
TensileInstructions::WARNING: RegisterPool::remove(0,95) pool1 still in use
TensileInstructions::WARNING: RegisterPool::add(0,149) pool0 already available
TensileInstructions::WARNING: RegisterPool::add(0,149) pool1 already available
TensileInstructions::WARNING: RegisterPool::remove(0,95) pool0 still in use
TensileInstructions::WARNING: RegisterPool::remove(0,95) pool1 still in use
TensileInstructions::WARNING: RegisterPool::add(0,149) pool0 already available
TensileInstructions::WARNING: RegisterPool::add(0,149) pool1 already available
TensileInstructions::WARNING: RegisterPool::remove(0,95) pool0 still in use
TensileInstructions::WARNING: RegisterPool::remove(0,95) pool1 still in use
TensileInstructions::WARNING: RegisterPool::add(0,149) pool0 already available
TensileInstructions::WARNING: RegisterPool::add(0,149) pool1 already available

(2) error
File "tensilelite/Tensile/Components/StreamK.py", line 1792, in storeBranches
module.add(self.storeBranchesCommon(writer, kernel, skPartialsLabel, vectorWidths, element, tmpVgpr, cvtVgprStruct))
NameError: name 'element' is not defined

This is the folkparameters I tried.
- 1LDSBuffer: [1]
- DepthU: [64]
- MatrixInstruction:
- [16, 16, 16, 1, 1, 4, 8, 4, 1] # MT 256x128
- PrefetchGlobalRead: [2]
- PrefetchLocalRead: [1]
- SourceSwap: [1]
- StaggerU: [0]
- StreamK: [0,1,2,3]

nakajee · 2024-08-27T05:28:28Z

Please make sure all CI tests pass on your local node

tensilelite tox
hipblaslt-test

AlexBrownAMD · 2024-08-28T23:06:42Z

Thanks for the notes. I fixed issue #2 as well as a bug in the xccmap. Still working on debugging issue #1.

nakajee · 2024-09-06T15:42:04Z

+            module.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr+2), src=vgpr("Serial"), comment="Wave 0 updates flags"))
+            module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=0, comment="Check for wave 0"))
+            module.add(SCBranchSCC0(labelName=skipFlagReset.getLabelName(), comment="Skip flag reset"))
+            module.add(SMovB32(dst=sgpr(tmpSgpr+2), src=0, comment="flag data"))


Small optimization.
tmpSgpr+2 is already 0 here based on the above condition.
We can remove SMovB32.

Good catch, just posted an update

nakajee · 2024-09-06T15:43:28Z

+            skipFlagReset = Label(label=writer.labels.getNameInc("SK_SkipFlagReset"), comment="")
+            module.add(VReadfirstlaneB32(dst=sgpr(tmpSgpr+2), src=vgpr("Serial"), comment="Wave 0 updates flags"))
+            module.add(SCmpEQU32(src0=sgpr(tmpSgpr+2), src1=0, comment="Check for wave 0"))
+            module.add(SCBranchSCC0(labelName=skipFlagReset.getLabelName(), comment="Skip flag reset"))


Quick question.
Why we need this code only for tensilelite?
I did not see it in Tensile.

In Tensile the synchronization flags are stored in the workspace, which is supplied by the user and could contain any values, so there is a pre-kernel that initializes the flags to 0.
In TensileLite the flags are stored in the Synchronizer buffer which is kept in the handle and initialized to 0. In this case we also have to reset the flags to 0 after it they have been read so that they're ready for the next kernel to run. The first version on this PR had all 4 wavefronts in a workgroup doing the s_store to the flag, which was creating a race condition between the WG doing the partial tile setting it to 1 and the WG doing the fixup resetting it to 0. The kernel would sometimes end with some flags set to 1 instead of 0 causing kernels that run afterwards to fail.
This update changes it so only wavefront 0 updates the flag value (and wavefronts in a WG are still synchronized with a barrier to ensure all partial data is fully written before setting the flag)

I just added one more test file based on the test that most easily reproduced the race condition

…current vgpr allocations

nakajee · 2024-09-06T21:14:50Z

Thanks for your update.
Since "noCI" label is added, please run all CI tests (tensilelite tox + hipblaslt-test) on your local node and paste the result here.

bnemanich · 2024-09-06T21:18:54Z

    std::shared_ptr<hipDeviceProp_t>                                                 deviceProp;

    auto adapter = get_library_and_adapter(&library, &deviceProp, handle->device);
+<<<<<<< HEAD


Merge issue.

Thanks, got it fixed

bnemanich · 2024-09-06T21:19:06Z

    std::shared_ptr<hipDeviceProp_t>                                                 deviceProp;

    auto adapter = get_library_and_adapter(&library, &deviceProp, handle->device);
+<<<<<<< HEAD


Merge issue.

AlexBrownAMD · 2024-09-07T00:51:34Z

Finished running tensilelite tox tests, working on hipblaslt-test

=================================================================================== 75 passed, 31 skipped, 44 warnings in 7744.02s (2:09:04) ====================================================================================
py310: exit 0 (7744.30 seconds) /home/stebrown/hipBLASLt2/tensilelite> py.test -v --basetemp=/home/stebrown/hipBLASLt2/tensilelite/.tox/py310/tmp --junit-xml=/home/stebrown/hipBLASLt2/tensilelite/python_tests.xml --junit-prefix=py310 --color=yes -n 4 --prebuilt-client=/home/stebrown/hipBLASLt2/tensilelite/.tox/py310/client/0_Build/client/tensile_client ./Tensile/Tests/ -m common pid=1141334
py310: OK (7888.08=setup[14.73]+cmd[1.80,127.24,7744.30] seconds)
congratulations :) (7888.16 seconds)

nakajee

Thanks for your update.
Looks good.

AlexBrownAMD · 2024-09-07T04:04:26Z

All hipblaslt-test test cases passing as well:

[----------] Global test environment tear-down
[==========] 48206 tests from 13 test suites ran. (2278421 ms total)
[ PASSED ] 48206 tests.

AlexBrownAMD and others added 30 commits June 4, 2024 20:00

Add tuning parameters with updated doc comments

848b6f6

SK Port: (Host) Debug settings and Grid Selector (#1)

6ff49e1

Merge branch 'ROCm:develop' into StreamK

c85586d

Stream-K kernel arguments, metadata, and launch grid code

f25ba0b

Increase default max workspace for tensile tuning client

1e19915

Increase default workspace for bench client (#2)

223aedc

Merge branch 'ROCm:develop' into StreamK

e1e38a7

Merge remote-tracking branch 'origin' into StreamK

d0d4bc3

Merge remote-tracking branch 'origin/StreamK' into StreamK

e9a56b0

Update gitignore

7809986

Option to disable GSU in kernels.

cbf4ef3

Modularize GSU code so it is more contained, reduces branch points in KernelWriter code. Fix a temp register allocation bug. Fix alignment of some registers rather than rely on a specific order.

UpdatedSolutionStructs, reject cases and golbal accumulation code

b92f313

Stream-k code gen work in progress

ca5221d

Some bug fixes and more work on persistent loop

8f5a915

Store branches and close persistent loop

0d5cec7

Add code for workspace SRD, working on partials write, some updates/f…

2296748

…ixes

Support for predicate arithmetic intensity added on top of streamK br… (

ac84fa9

#3) * Support for predicate arithmetic intensity added on top of streamK branch in Alex's repo

Write partials and some bug fixes

259f6b4

Grid prediction (#4)

a100b1f

* Analytical grid size prediction model for Stream-K

Fixup/partials, flag init, bug fixes

5ce3017

Enable sk grid and fix tile calculation

6a50f3d

Merge branch 'StreamK' into StreamKGen

4df66e5

Merge pull request #6 from AlexBrownAMD/StreamKGen

fc0f23d

Porting code gen changes for stream-k

Some build fixes

16fa505

Fix race condition on flag reset

9e8fa21

Flag point to input synchronizer

5642e5c

point flag to synchronizer

2f882d6

remove space allocated for flags

0c5159d

Some bug fixes, general case sgemm tests passing

445e1f1

Merge pull request #7 from adityalj/flag_init

8367400

Flag init

nakajee reviewed Aug 22, 2024

View reviewed changes

nakajee added the gfx94x Run CI on gfx94x label Aug 22, 2024

eidenyoshida added the noCI Disable testing on supported CI systems: math libraries CI has this feature enabled.. label Aug 22, 2024

AlexBrownAMD added 3 commits August 23, 2024 17:16

Update targets for new tests

cfdaf49

Remove duplicate code and addreject for unsupported config

44f2d35

If stream-k enabled, force GSU=0

e7d7d57

neoblizz reviewed Aug 23, 2024

View reviewed changes

neoblizz previously approved these changes Aug 23, 2024

View reviewed changes

Bug fixes for SKXCCM and SK=2

bfbfde1

AlexBrownAMD dismissed neoblizz’s stale review via bfbfde1 August 28, 2024 17:39

Fix race condition when resetting synchronizer flags

e6de95f

AlexBrownAMD requested a review from solaslin as a code owner September 6, 2024 15:23

nakajee reviewed Sep 6, 2024

View reviewed changes

AlexBrownAMD added 4 commits September 6, 2024 16:04

Disable LRA backup, not needed for stream-k, and not compatible with …

fc61d87

…current vgpr allocations

Small optimization from PR review

e304cf7

Add new test case based on race condition

256da9b

Merge remote-tracking branch 'upstream/develop' into StreamKMerge2

3bbeff9

bnemanich reviewed Sep 6, 2024

View reviewed changes

Fix merge issue

21abe72

nakajee approved these changes Sep 7, 2024

View reviewed changes

KKyang approved these changes Sep 7, 2024

View reviewed changes

AlexBrownAMD merged commit af5dd3f into ROCm:develop Sep 7, 2024

alex391a mentioned this pull request Sep 8, 2024

Make WorkGroupMappingXCC and WorkGroupMappingXCCGroup as GPU kernel arguments #1058

Merged

Conversation

AlexBrownAMD commented Aug 22, 2024

Uh oh!

nakajee left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nakajee commented Aug 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nakajee commented Aug 27, 2024

Uh oh!

AlexBrownAMD commented Aug 28, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nakajee commented Sep 6, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AlexBrownAMD commented Sep 7, 2024

Uh oh!

nakajee left a comment

Choose a reason for hiding this comment

Uh oh!

AlexBrownAMD commented Sep 7, 2024

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

nakajee commented Aug 27, 2024 •

edited

Loading