[DRAFT] Add Lds transpose load ds read tr16 #2043
Draft
+772
−30
Conversation
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
- Implemented rock.lds_transpose_load TD definition supporting f16 and bf16. - Added verifier to ensure source memref is in workgroup (LDS) memory and indices match rank. - Implemented lowering pattern to amdgpu.transpose_load. - Created MLIR tests covering FP16 and BF16 loads with FileCheck patterns.
This commit introduces the full implementation of LDS transpose load handling used in threadwise read and single-buffering and double buffering pipelines. It adds logic for computing per-lane base offsets, generating LdsTransposeLoadOp instructions, and managing vectorized fragment loading for MFMA operations. The implementation supports multiple layout kinds (e.g., L16x16, L32x16, L32x8) and dynamically expands offsets for multi-K fused cases.This enables more flexible data movement between LDS and registers for MFMA input tiles.
f16 and bf16 data types, with multiple K-dimension configurations and schedule versions. Add CFG file to restrict execution to gfx950 architecture only, ensuring tests run exclusively on supported hardware. All test cases have passed validation under gfx950.
buffering pipeline.
1 task
for double buffering. Added global decision context to propagate hwtranspose::Decision from BlockwiseGemm to ThreadwiseReadIntoOp. Updated ThreadwiseReadIntoOp to attach hwtranspose attributes when a valid decision is available. Fixed double buffering handling to ensure correct LDS access.
decision handling.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Motivation
close: https://github.com/ROCm/rocMLIR-internal/issues/1858
Technical Details
This change adds full support for LDS transpose load integration within both single-buffering and double-buffering pipelines.
The implementation enables transpose-aware LDS loading for operands A and B, provided that both matrices use compatible memory layouts.
Currently, the logic performs iterations over the K dimension, while iteration over M and N dimensions is still under development and will be refined in the next update.
Future work will focus on performance evaluation and optimization of bank conflict patterns during LDS access.
Test Plan
Basic functionality was verified using existing MFMA pipeline tests for both single and double buffering.
Next, I will extend the tests to cover various matrix layout configurations and measure execution performance.
A detailed performance table and LDS bank conflict statistics will be added in comment later to quantify the improvements.
Test Result
Submission Checklist