Skip to content

Support to load transpose B matrix of sub-group-size=32.#7488

Draft
chengjunlu with Copilot wants to merge 9 commits into
mainfrom
copilot/add-lit-and-unit-tests
Draft

Support to load transpose B matrix of sub-group-size=32.#7488
chengjunlu with Copilot wants to merge 9 commits into
mainfrom
copilot/add-lit-and-unit-tests

Conversation

Copilot AI commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

PR #7487 replaced the width-comparison guard in computeTransposeShuffleMapping with a linear-layout comparison, unblocking column-major (transposed) B matrix 2D block loads when threadsPerWarp=32. Without that fix, any f16/opsPerChan=2 B-matrix column-major load with tpw=32 silently fell back to gather because dpasInstShapeB()[1]=16 ≠ threadsPerWarp=32.

Changes

Core fix (BlockIOUtils.cpp, BlockIOUtils.h, LoadStoreOpToLLVM.cpp)

  • computeTransposeShuffleMapping signature changed from (numPackedVals, tileHeight)(const BlockIOTileSizeInfo &sizeInfo).
  • Old check: numPackedVals > 1 && widthToTranspose != threadsPerWarpreturn failure().
    New check: builds transPackedLayout from lane/register bases and compares against expectedLoadUnpackLayout derived from the encoding's linear layout. Correctly accepts f16/opsPerChan=2 at tpw=32; still rejects e.g. fp8/opsPerChan=2 at tpw=32 where layouts don't match.
  • Minor grammar fix in an existing code comment.

LIT tests

  • test/TritonIntelGPU/2d-block-load-to-llvm.mlirblock_load_dot_b_subgroup32_transpose: ttig.2d_block_load {column_major} with tpw=32/f16 lowers to triton_gen.2Dblockload ... transpose = true.
  • test/TritonIntelGPU/LowerTo2DBlockLoad/pointer-load.mlirpointer_load_column_major_subgroup32: pointer-tensor path (tt.load → ttig.2d_block_load_from_ptr) works for tpw=32 column-major B, producing the correct pitch and base_height/base_width.

Unit test (BlockIOUtilsTest.cpp, new)

Calls validate2DBlockLoadTile directly:

  • ColumnMajorB_Tpw32_F16_Accepted — tpw=32, f16/opsPerChan=2, memContiguousDim=0true (was false before).
  • ColumnMajorB_Tpw16_F16_Accepted — regression: tpw=16 column-major B still accepted.

Copilot AI changed the title Add LIT and unit tests for PR #7487 (tpw=32 column-major B matrix transpose) Add LIT and unit tests for tpw=32 column-major B matrix transpose (PR #7487) Jul 16, 2026
Copilot AI requested a review from chengjunlu July 16, 2026 02:28
@chengjunlu

Copy link
Copy Markdown
Contributor

@copilot Add a E2E functional test in test_block_load.py as well.

Copilot AI commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

@copilot Add a E2E functional test in test_block_load.py as well.

Added test_block_tdesc_column_major_load_subgroup32 to test_block_load.py in commit f00c2b3. The test mirrors the existing test_block_tdesc_column_major_load (M=64, K=32, N=64 single-tile matmul with column-major B load) but uses threadsPerWarp = 32 with warpsPerCTA = [4, 2], exercising the computeTransposeShuffleMapping fix directly and verifying correctness against torch.mm(A, B_stored.T).

@chengjunlu chengjunlu changed the title Add LIT and unit tests for tpw=32 column-major B matrix transpose (PR #7487) Support to load transpose B matrix of sub-group-size=32. Jul 16, 2026
@chengjunlu
chengjunlu marked this pull request as ready for review July 16, 2026 03:15
@chengjunlu

Copy link
Copy Markdown
Contributor

@copilot Please fix CI failure.

Copilot AI commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

@copilot Please fix CI failure.

Fixed in the latest commit. Two issues were causing CI failures in BlockIOUtilsTest.cpp:

  1. clang-format: off-by-one space in makeDpas parameter alignment (29→28 spaces)
  2. Build error: LinearLayout unqualified — it lives in mlir::triton which was not using-ed; replaced both LinearLayout ll = ... with auto ll = ...

@chengjunlu

Copy link
Copy Markdown
Contributor

@copilot Fix CI failure.

Copilot AI commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

@copilot Fix CI failure.

Fixed in the latest commit. The LIT test pointer_load_column_major_subgroup32 had "ttg.num-warps" = 8 but warpsPerCTA = [2, 2] (4 warps total) — tt.make_range validates the slice layout against the module's num-warps and rejected the mismatch. Changed to "ttg.num-warps" = 4 to match.

@chengjunlu

Copy link
Copy Markdown
Contributor

@copilot Fix CI failure.

Copilot AI commented Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

@copilot Fix CI failure.

Fixed in the latest two commits. The broadcast check in LowerTo2DBlockLoad.cpp was using tileWidth * 2 == threadsPerWarp, which rejected tpw=32/tileWidth=8 (factor-4 replication). Relaxed to threadsPerWarp % tileWidth == 0 to allow any replication factor. Updated the corresponding assert in LoadStoreOpToLLVM.cpp to match — the shuffleIdx(threadId % tileWidth) row-replication logic is correct for any tileWidth that evenly divides threadsPerWarp.

@chengjunlu
chengjunlu removed the request for review from whitneywhtsang July 17, 2026 04:43
@chengjunlu
chengjunlu marked this pull request as draft July 17, 2026 04:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants