perf(tensilelite): reduce SGPR usage on gfx1250 StreamKForceDPOnly kernels#9730
Open
jaopaulolc wants to merge 18 commits into
Open
perf(tensilelite): reduce SGPR usage on gfx1250 StreamKForceDPOnly kernels#9730jaopaulolc wants to merge 18 commits into
jaopaulolc wants to merge 18 commits into
Conversation
Relax the PrefetchGL2 Stream-K guard to permit the DP-first two-tile DP-only case (StreamK==3 + StreamKForceDPOnly), which builds and passes numerical validation on gfx1250 (FFM). Add StreamK+PAP gfx1250 test configs covering ClusterDim and PrefetchGL2 on top of StreamKForceDPOnly.
…x1250 Broaden the PrefetchGL2 Stream-K guard from DP-only to all StreamK==3 (DP-first two-tile), including the non-DP-only fixup path. Both build and pass numerical validation on gfx1250 (FFM). Add non-DP-only StreamK+PAP gfx1250 test configs for ClusterDim and PrefetchGL2 (plus an A/B control).
…or on gfx1250 These 2-source VALU shifts (dst, shiftAmount, value) were declared under the 4-operand VOP3 format template. Because a partial .operand_fields override inherits all base-format fields and only patches the named ones, the VOP3 base left a phantom third source (src2) in each descriptor. The StinkyIRVerifier reads the descriptor's operand fields and reported the phantom src[2] as a missing operand, e.g. under PrefetchGL2 on gfx1250: [StinkyIRVerifier] Register width validation failed: Instruction 'v_lshrrev_b64' missing operand src[2] The emitted assembly was always correct (rocisa builds exactly two sources and the text emitter is bounded by the real operand list); this was a hardware-model metadata bug, not a codegen defect. Move both instructions to the existing VOP3_2SRC format (dst + src0 + src1), preserving the 64-bit (b64) and 16-bit (b16) size overrides, so the descriptor exposes exactly dst/src0/src1 with no phantom src2 and the spurious warning is gone. AIHPBLAS-4142
…comments (AIHPBLAS-4142) Strip explanatory/inline YAML comments from the 5 new gfx1250 MXFP4 Stream-K test configs (SPDX license header retained). Broaden coverage to match the existing MXFP4 configs (sk_mxf4gemm_tdm_pap.yaml / sk_mxf4gemm_tdm_cluster.yaml): expand the MatrixInstruction sweep to the full MT 32/64/128/256 spread ([16,16,128,..1,1], [..2,2], [..4,4], [..8,8]) and use comparable problem sizes, while preserving each file's distinguishing feature params (ClusterDim 1-D [2,1]/[4,1]/[8,1], PrefetchGL2, StreamKForceDPOnly, PrefetchAcrossPersistent). Cluster files use [2048,256,1,1024]/[2048,512,1,1056] so problem sizes divide evenly by cluster X-extents; PrefetchGL2 files keep GSU=0 + strided batch. Validated on FFM (validateAll, NumElementsToValidate:-1) for all 5 configs: every generated solution assembled and PASSED, zero StinkyIRVerifier / missing-operand warnings, clientExit=0.
…criptor tests (AIHPBLAS-4142) Guards the phantom-src2 fix that moved these 2-source VALU shifts from the 4-field VOP3 format to the 3-field VOP3_2SRC format. Adds descriptor-level assertions (HwInstDescTest) that each op exposes exactly dst/src0/src1 with no phantom src2, plus a contrast test proving a genuine 3-source VOP3 op (v_lshl_add_u32) still carries src2. Also adds verifier-level tests (RegisterWidthValidationTest) that a well-formed 2-source shift verifies clean with no "missing operand src[2]" / "Register width validation failed" error. Verified these tests fail on the pre-fix descriptor (reproducing the exact "missing operand src[2]" verifier error) and pass post-fix.
… (AIHPBLAS-4142) Cover the broadened PrefetchGL2 validator that now allows DP-first Stream-K (StreamK==3), both StreamKForceDPOnly variants, while still rejecting other non-zero Stream-K modes (1, 2, 4, 5) and leaving StreamK==0 allowed. Builds a real gfx1250 Solution end-to-end and asserts on the reject diagnostic.
…sterDim configs via fork params (AIHPBLAS-4142) Replace the 5 duplicated MXFP4 Stream-K gfx1250 test yamls (per StreamKForceDPOnly mode / per PrefetchGL2 control) with 2 self-contained fork-based configs that sweep the same coverage: - sk_mxf4gemm_pap_prefetchgl2.yaml: ForkParameters PrefetchGL2:[0,1] + StreamKForceDPOnly:[0,1] (the PGL2:[0] leg is the A/B control), keeping StreamK:[3], PrefetchAcrossPersistent:[1], PrefetchGlobalRead:[1,2], GSU=0, strided batch. Replaces sk_mxf4gemm_tdm_pap_prefetchgl2.yaml, sk_mxf4gemm_pap_prefetchgl2_nodponly.yaml and sk_mxf4gemm_pap_prefetchgl2_nodponly_control.yaml. - sk_mxf4gemm_pap_cluster.yaml: ForkParameters ClusterDim [[2,1],[4,1],[8,1]] (1-D, Y-extent 1) + StreamKForceDPOnly:[0,1], keeping StreamK:[3], PrefetchAcrossPersistent:[1]. Replaces sk_mxf4_force_dp_only_cluster_pap.yaml and sk_mxf4gemm_cluster_pap_nodponly.yaml. Both preserve the broadened MXFP4 MI/MT spread (MT 32/64/128/256, DepthU 256) and the previously-validated problem sizes. Validated on the FFM (gfx1250) functional simulator: file A 128/128 solutions PASSED (clientExit=0), file B 144/144 PASSED (clientExit=0), 0 StinkyIRVerifier warnings, all expected fork legs present (PGL0/PGL1 x SKFDPO0/SKFDPO1; CD2_1/CD4_1/CD8_1 x SKFDPO0/SKFDPO1).
…fork configs (AIHPBLAS-4142) Mirror the two consolidated MXFP4 Stream-K fork configs with MXFP8 coverage, as two new fork-based files: - sk_mxf8gemm_pap_prefetchgl2.yaml: ForkParameters PrefetchGL2:[0,1] (the [0] leg is the A/B control) x StreamKForceDPOnly:[0,1], with StreamK:[3], PrefetchAcrossPersistent:[1], GSU=0, strided batch. - sk_mxf8gemm_pap_cluster.yaml: ForkParameters 1-D ClusterDim [[2,1],[4,1],[8,1]] (Y-extent 1) x StreamKForceDPOnly:[0,1], with StreamK:[3], PrefetchAcrossPersistent:[1]. Both use MXFP8's own ProblemType (F8/s/s, HPA, TN, MXBlock 32) and its MatrixInstruction / DepthU / problem sizes taken from the existing sk_mxf8gemm_tdm_pap.yaml (and split_pap tail sizes) rather than MXFP4's values; K stays divisible by 32 for MXBlock32. The cluster file uses M=2048 so cluster M-tiles stay divisible by the 2/4/8 X-extents across all MIs. Validated on the gfx1250 FFM simulator: every generated solution assembles and passes validateAll (A: 96 PASSED/0 FAILED; B: 108 PASSED/0 FAILED), clientExit=0, 0 StinkyIRVerifier warnings; both PGL0/PGL1 and SKFDPO0/SKFDPO1 legs present in A, and CD2/CD4/CD8 with SKFDPO0/SKFDPO1 present in B.
…eDPOnly on gfx1250 JIRA ID : TBD Under StreamKForceDPOnly (SK3 DP-first, gfx1250) every workgroup processes complete tiles, so the StreamK partials/fixup workspace-reduction path is never reached: computeStoreSrdStartCommon and storeBranchesCommon early-return, computeWorkspaceSrd (the sole AddressWS->SrdWS init) and the tc=='WS' workspace store are never emitted, and the epilogue SrdWS borrow is already guarded by "SrdWS" in self.sgprs. The 4-aligned SrdWS SGPR (plus alignment padding) is therefore dead and is no longer allocated for DP-only kernels. AddressWS and AddressFlags are intentionally kept: they are kernarg-loaded with a fixed host layout (ContractionSolution.cpp appends ws/Flags for every streamK>0 && atomic==0 kernel regardless of DP-only) and AddressFlags still has live runtime readers under DP-only. Gated entirely on StreamKForceDPOnly; the non-DP-only path is unchanged. FFM-validated (gfx1250, MXFP4+MXFP8, PAP cluster + prefetchgl2, SKFDPO 0 and 1): all solutions assemble, validateAll all PASSED, clientExit=0.
…amKForceDPOnly JIRA ID : TBD Under StreamKForceDPOnly (SK3 DP-first, gfx1250) every workgroup spans complete tiles, so the per-tile local iteration bounds are compile-time constants: StreamKLocalStart == 0 and StreamKLocalEnd == ItersPerTile. These two persistent SGPRs are therefore no longer allocated for DP-only kernels; every reader is constant-folded or removed behind a StreamKForceDPOnly guard: - KernelWriter.py: drop StreamKLocalStart/End from the SK3 persistent SGPR list. - StreamK.py (shared Common + SK3): skTileIndex no longer sets them; computeLoadSrdCommon / graAddressesCommon / declareStaggerParmsCommon / tailLoopNumIterCommon early-return (offset 0, full tile); calculateLoopNumIterCommon uses loopCounter = ItersPerTile; the tailloop and graWorkGroup start/end checks are skipped (WG always starts and finishes its tile). - KernelWriterAssembly.py: the tailloopInNll skip is gated off; TDM wave-separated StreamK K-offset applier is a no-op (start==0) and the tail applier derives (ItersPerTile-1) from the ItersPerTile VGPR const instead of StreamKLocalEnd. - SubtileGREmit.py: the subtile StreamK K-offset applier is a no-op (start==0). StreamKForceDPOnly is only valid with StreamK==3 (Solution.py rejects otherwise), so StreamKDynamic/Hybrid keep these SGPRs and are unaffected. Gated entirely on StreamKForceDPOnly; the non-DP-only path is unchanged. FFM-validated (gfx1250, MXFP4+MXFP8, PAP cluster + prefetchgl2, SKFDPO 0 and 1): all solutions assemble, validateAll all PASSED, clientExit=0, 0 asserts.
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
… StreamKForceDPOnly on gfx1250 JIRA ID : AIHPBLAS-4145 Under StreamKForceDPOnly (SK3 DP-first, gfx1250) the reduction is always forced to the single-kernel tree-fixup path: the Synchronizer buffer is always non-null, so AddressFlags != 0 is an invariant. Every runtime `s_cmp_eq_u64 AddressFlags, 0` DP-vs-parallel-reduction branch therefore always resolves to the DP/not-parallel-reduction side. Fold those dead branches at codegen time so DP-only kernels emit no AddressFlags compare: - KernelWriterAssembly.py prefetchAcrossPersistent (PAP SK_SkipNllPAP): gate the AddressFlags skip compare behind `not StreamKForceDPOnly`, keeping the StreamKIter >= StreamKIterEnd last-tile check. - KernelWriter.py prefetchAcrossPersistentSubtile (Subtile PAP): same. - StreamK.py stridedBatchOrGeneralBatch (General-Batched SupportUserArgs): under DP-only emit an unconditional SBranch to the general-batched load target instead of the AddressFlags compare. - KernelWriterAssembly.py allocPostLoopSrd (initializeSrdAddressFlagsCheck call site): under DP-only emit an unconditional SBranch to the general-batched Srd-init target (the component method reads AddressFlags). - KernelWriterAssembly.py globalWriteElements GSU-split store branch: assert this dead AddressFlags reader is never reached under DP-only (SK3 sets GlobalSplitU==0, forcing noGSUBranch=True / gsuLimit==1). No kernarg-layout change: AddressWS/AddressFlags kernel args and their SGPR defines are unchanged, so this cannot alter assembly correctness; it only removes dead runtime branches. All other DP-only AddressFlags/AddressWS readers (computeStoreSrdStartCommon, storeBranchesCommon, writePartialsCommon, computeWorkspaceSrd, preLoop, graWorkGroup) are already dead via existing StreamKForceDPOnly early-returns; StreamKDynamic/Hybrid readers are never reached since StreamKForceDPOnly requires StreamK==3. FFM-validated (gfx1250, MXFP4+MXFP8, PAP cluster + prefetchgl2, SKFDPO 0/1): all solutions assemble, validateAll all PASSED, clientExit=0, 0 FAILED, 0 asserts, 0 StinkyIRVerifier. Generated DP-only .s has zero AddressFlags compares; non-DP-only .s retains all 10. Peak .amdhsa_next_free_sgpr unchanged vs prior branch (DP-only base 98, +PGL2 102, +ClusterDim 98).
…DPOnly on gfx1250 JIRA ID : AIHPBLAS-4145 Under StreamKForceDPOnly (SK3 DP-first, gfx1250) the reduction is always the single-kernel tree path (getSKReduction returns tree, Flags == Synchronizer, never parallel) and the workspace partials/fixup path is never reached. After the reader-fold refactor, AddressWS/AddressFlags have no live runtime reader in DP-only kernels, yet they were still allocated as persistent kernarg SGPRs and declared in the .kd metadata + host kernarg buffer. Because they are defined persistently (kept live for the late non-DP readers) they sit at the SGPR high-water mark, so the dead pair inflates the peak. Drop AddressWS/AddressFlags for DP-only kernels, atomically across all three positional-layout layers so the kernarg offsets stay consistent host<->device: - KernelWriter.py: gate the AddressWS/AddressFlags SGPR define on `not StreamKForceDPOnly` (and refresh the SrdWS-pool comment). - Components/Signature.py: gate the same two args out of the .kd metadata. - ContractionSolution.cpp (singleCallArgs): gate the ws/Flags kernarg append with `sizeMapping.streamKForceDPOnly == 0`. The SrdWS pool add (addSgprVarToPool) already no-ops on an undefined var. Measured impact (gfx1250, MT64x64x256 MI16x16x128, MXFP4 and MXFP8): DP-only peak .amdhsa_next_free_sgpr 98->94 (base / +ClusterDim) and 102->98 (+PrefetchGL2); .kernarg_segment_size 176->160 (-16B); zero AddressWS/AddressFlags symbol references in the DP-only assembly. Non-DP-only kernels are byte-identical to before (same .s md5, next_free_sgpr 102/106, kernarg 176). FFM-validated (gfx1250, MXFP4+MXFP8, PAP cluster + prefetchgl2, SKFDPO 0/1): all 4 fork configs assemble, validateAll all PASSED (476/476), clientExit=0, 0 FAILED, 0 DID_NOT_SATISFY_ASSERTS, 0 StinkyIRVerifier on both legs. The DP-only host<->device kernarg alignment (StridesD/Alpha/... offsets) is proven correct by the passing numerical validation after the -16B layout change.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #9730 +/- ##
===========================================
- Coverage 69.59% 68.56% -1.03%
===========================================
Files 2771 2764 -7
Lines 452699 459597 +6898
Branches 66658 71443 +4785
===========================================
+ Hits 315022 315092 +70
- Misses 117337 122197 +4860
- Partials 20340 22308 +1968
*This pull request uses carry forward flags. Click here to find out more.
🚀 New features to boost your workflow:
|
jaopaulolc
added
the
rocm:gemm algos
…mK fork configs (AIHPBLAS-4145)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
JIRA ID : AIHPBLAS-4145
Motivation
SK3 Stream-K kernels running in
StreamKForceDPOnlymode on gfx1250 allocate persistent SGPRs — and carry kernarg pointers — that the DP-only path can never use. In DP-only mode every workgroup processes complete tiles, so there is no partials/fixup/workspace-reduction path, the per-tile local iteration bounds are compile-time constants, and the workspace/synchronizer-flag pointers are dead. Trimming these dead/constant SGPRs and the dead workspace kernargs lowers SGPR pressure and shrinks the kernarg segment on the DP-only SK3 kernels (including +PrefetchGL2/+ClusterDim on gfx1250), with zero effect on the non-DP-only path.Technical Details
Four independent reductions, all gated on
kernel["StreamKForceDPOnly"](only ever set forStreamK == 3—Solution.pyrejects it otherwise, soStreamKDynamic/StreamKHybridare unaffected):Skip the StreamK workspace SRD SGPRs under DP-only: the 4-aligned
SrdWSworkspace SRD is only used by the partials/fixup reduction path, which never runs under DP-only (computeStoreSrdStartCommon/storeBranchesCommonearly-return,computeWorkspaceSrdand thetc=='WS'store are never emitted, and the epilogueSrdWSborrow is already guarded by"SrdWS" in self.sgprs). It is no longer allocated for DP-only kernels.Fold
StreamKLocalStart/StreamKLocalEndto constants under DP-only:StreamKLocalStart == 0andStreamKLocalEnd == ItersPerTile, so both persistent SGPRs are dropped from the SK3 allocation list and every reader is constant-folded or removed behind aStreamKForceDPOnlyguard —skTileIndexno longer sets them;computeLoadSrdCommon/graAddressesCommon/declareStaggerParmsCommon/tailLoopNumIterCommonearly-return (offset 0, full tile);calculateLoopNumIterCommonusesloopCounter = ItersPerTile; the tailloop/graWorkGroupstart-and-finish checks are skipped; the KWAtailloopInNllskip is gated off; the TDM wave-separated K-offset applier is a no-op (start==0) while the tail applier derives(ItersPerTile-1)from theItersPerTileVGPR const; and the subtile K-offset applier is a no-op.Make the StreamK workspace-flag readers dead under DP-only: the reduction is always the single-kernel tree path (the Synchronizer buffer is always non-null, so
AddressFlags != 0is an invariant — confirmed host-side bygetSKReductionreturningtreewithFlags == Synchronizer, neverparallel). Every runtimes_cmp_eq_u64 AddressFlags, 0DP-vs-parallel-reduction branch therefore always resolves to the DP side, so those dead branches are folded out at codegen: the PAP (prefetchAcrossPersistent) and Subtile-PAP (prefetchAcrossPersistentSubtile) skip compares are gated behindnot StreamKForceDPOnly(keeping theStreamKIter >= StreamKIterEndlast-tile check); the General-BatchedstridedBatchOrGeneralBatchandinitializeSrdAddressFlagsCheck(SupportUserArgs) paths emit an unconditional branch to the general-batched target; and theglobalWriteElementsGSU-split store branch (unreachable under DP-only since SK3 forcesnoGSUBranch) is asserted dead. This makes the AddressWS/AddressFlags kernargs fully dead in DP-only assembly without changing the kernarg layout.Drop the StreamK workspace kernargs under DP-only: with the readers dead,
AddressWS/AddressFlagshave no live reader yet were still allocated as persistent kernarg SGPRs at the SGPR high-water mark. They are now dropped for DP-only kernels, changed atomically across all three positional-layout layers so kernarg offsets stay consistent host↔device:KernelWriter.pygates the AddressWS/AddressFlags SGPR define onnot StreamKForceDPOnly;Components/Signature.pygates the same two args out of the.kdmetadata; andContractionSolution.cpp(singleCallArgs) gates the ws/Flags kernarg append withsizeMapping.streamKForceDPOnly == 0. (TheSrdWSpool add already no-ops on an undefined var.)The 6 SK tile-constants (
ItersPerTile/MagicNumberItersPerTile/MagicShiftItersPerTile/SKItersPerWG/skGrid/skTiles) are already VGPR-offloaded on gfx1250 SK3 (isStreamKConstantsToVgprEnabled), appearing in the DP-only.sonly as transient.set sgpr…aliases that areUNDEF'd immediately after areadfirstlanefrom their VGPR home, so they reuse freed slots below the high-water mark and do not contribute to peak SGPR; no change is made there.Test Plan
Functional validation on the gfx1250 FFM simulator of the four committed Stream-K regression configs under
Tensile/Tests/common/streamk/gfx1250/core/—sk_mxf4gemm_pap_cluster.yaml,sk_mxf4gemm_pap_prefetchgl2.yaml,sk_mxf8gemm_pap_cluster.yaml,sk_mxf8gemm_pap_prefetchgl2.yaml— each of which forksStreamKForceDPOnly: [0, 1](covering both the DP-only and non-DP-only paths) across MXFP4/MXFP8, PrefetchGL2, and ClusterDim. Per config: all solutions assemble,validateAllall PASSED,clientExit=0, 0 FAILED, 0DID_NOT_SATISFY_ASSERTS, 0 StinkyIRVerifier. The kernarg change touchesContractionSolution.cpp, so the host library + client were rebuilt and the full four-config validation re-run so the numerical results confirm the DP-only host↔device kernarg alignment. SGPR/kernarg counts were read from single-kernel generation (MI [16,16,128,1,1,2,2,2,2], MT64x64x256,KeepBuildTmp: True) via.amdhsa_next_free_sgprand.kernarg_segment_size.Test Result
All four FFM regression configs pass
validateAllon both the DP-only (SKFDPO1) and non-DP-only (SKFDPO0) legs:clientExit=0, 0 FAILED, 0 DID_NOT_SATISFY_ASSERTS, 0 StinkyIRVerifier. The generated DP-only.shas zeroAddressWS/AddressFlagssymbol references; the non-DP-only.sis byte-for-byte identical before/after this branch.Peak
.amdhsa_next_free_sgpr(identical for MXFP4 and MXFP8):The DP-only
.kernarg_segment_sizedrops from 176 → 160 bytes (-16B); downstream kernarg offsets (StridesD/Alpha/…) still align host↔device (proven by the passing numerical validation across the -16B layout change). The non-DP-only path and GSU baseline are unchanged.Submission Checklist
Risk level
Low–Medium. The workspace-SRD, iteration-constant, and dead-branch changes are pure DP-only codegen changes gated on
StreamKForceDPOnly(valid only forStreamK == 3); the non-DP-only path and all other Stream-K variants are untouched (SKFDPO0 and GSU SGPR counts unchanged, non-DP-only.sbyte-identical, non-DP-only FFM legs pass). The kernarg drop additionally changes the DP-only kernarg positional layout across device (SGPR define +.kdmetadata) and host (ContractionSolution.cpp) in lockstep; a mismatch would corrupt downstream args, so it is validated by the full four-config FFM numerical validation confirming correct host↔device alignment. Any missed DP-only reader of a removed SGPR fails loudly at kernel generation, and all targeted configs assemble and validate clean.Notes
This branch (
users/jolabega/reduce-sgpr-streamk-dponly-gfx1250) is currently stacked onusers/jolabega/test-gfx1250-features-with-streamk(PR #9713); until #9713 merges intodevelop, this PR's diff also includes #9713's commits. Rebasing ontodevelopafter #9713 merges reduces the diff to just the four commits described above.