Support to load transpose B matrix of sub-group-size=32.#7488
Support to load transpose B matrix of sub-group-size=32.#7488chengjunlu with Copilot wants to merge 9 commits into
Conversation
|
@copilot Add a E2E functional test in test_block_load.py as well. |
Added |
|
@copilot Please fix CI failure. |
Fixed in the latest commit. Two issues were causing CI failures in
|
|
@copilot Fix CI failure. |
Fixed in the latest commit. The LIT test |
|
@copilot Fix CI failure. |
Fixed in the latest two commits. The broadcast check in |
PR #7487 replaced the width-comparison guard in
computeTransposeShuffleMappingwith a linear-layout comparison, unblocking column-major (transposed) B matrix 2D block loads whenthreadsPerWarp=32. Without that fix, anyf16/opsPerChan=2B-matrix column-major load with tpw=32 silently fell back to gather becausedpasInstShapeB()[1]=16 ≠ threadsPerWarp=32.Changes
Core fix (
BlockIOUtils.cpp,BlockIOUtils.h,LoadStoreOpToLLVM.cpp)computeTransposeShuffleMappingsignature changed from(numPackedVals, tileHeight)→(const BlockIOTileSizeInfo &sizeInfo).numPackedVals > 1 && widthToTranspose != threadsPerWarp→return failure().New check: builds
transPackedLayoutfrom lane/register bases and compares againstexpectedLoadUnpackLayoutderived 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.LIT tests
test/TritonIntelGPU/2d-block-load-to-llvm.mlir—block_load_dot_b_subgroup32_transpose:ttig.2d_block_load {column_major}with tpw=32/f16 lowers totriton_gen.2Dblockload ... transpose = true.test/TritonIntelGPU/LowerTo2DBlockLoad/pointer-load.mlir—pointer_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 andbase_height/base_width.Unit test (
BlockIOUtilsTest.cpp, new)Calls
validate2DBlockLoadTiledirectly:ColumnMajorB_Tpw32_F16_Accepted— tpw=32, f16/opsPerChan=2,memContiguousDim=0→true(wasfalsebefore).ColumnMajorB_Tpw16_F16_Accepted— regression: tpw=16 column-major B still accepted.