diff --git a/python/test/unit/intel/test_block_load.py b/python/test/unit/intel/test_block_load.py index 337725f4fb..707b996be8 100644 --- a/python/test/unit/intel/test_block_load.py +++ b/python/test/unit/intel/test_block_load.py @@ -329,3 +329,92 @@ def test_block_tdesc_column_major_load(device, tmp_path: pathlib.Path): # C should equal A x B_stored.T (since B was loaded via column_major). expected = torch.mm(A.to(torch.float32), B_stored.to(torch.float32).T) torch.testing.assert_close(C, expected, atol=1e-1, rtol=1e-2) + + +# --------------------------------------------------------------------------- +# Column-major descriptor load — threadsPerWarp=32 +# --------------------------------------------------------------------------- +@pytest.mark.skipif(not is_xpu(), reason="Block descriptor tests are specific to the XPU backend") +@pytest.mark.xfail( + not (triton.runtime.driver.active.get_current_target().arch['has_subgroup_2d_block_io'] + and triton.runtime.driver.active.get_current_target().arch['has_subgroup_matrix_multiply_accumulate']), + reason="Block loads and/or DPAS not supported on this architecture", run=False) +def test_block_tdesc_column_major_load_subgroup32(device, tmp_path: pathlib.Path): + """Verify column_major descriptor load with ``threadsPerWarp = 32``, f16/opsPerChan=2. + + This exercises the fix in ``computeTransposeShuffleMapping`` that unblocks + B-matrix column-major 2D block loads when ``threadsPerWarp = 32``. Before + the fix, the old width-comparison guard rejected this configuration because + ``dpasInstShapeB()[1] = 16 != threadsPerWarp = 32``, causing a silent + fallback to gather instead of a hardware transpose load. + + The kernel is identical in structure to ``test_block_tdesc_column_major_load`` + (same M/K/N tile sizes) but uses ``threadsPerWarp = 32`` with a matching + DPAS config (``warpsPerCTA = [4, 2]``). The DPAS per-warp tile shapes + remain ``A = [8, 16], B = [16, 16], C = [8, 16]`` and the full tile fits + in two repetitions along each dimension (numRep_M = numRep_K = numRep_N = 2). + """ + M, K, N = 64, 32, 64 + + ir = f""" + #dpas = #ttig.dpas<{{repeatCount = 8, systolicDepth = 8, executionSize = 16, opsPerChan = 2, threadsPerWarp = 32, warpsPerCTA = [4, 2], repCluster = [1, 1], A = [8, 16], B = [16, 16], C = [8, 16]}}> + #dot0 = #ttg.dot_op<{{opIdx = 0, parent = #dpas, kWidth=1}}> + #dot1 = #ttg.dot_op<{{opIdx = 1, parent = #dpas, kWidth=2}}> + module attributes {{ttig.min_sg_size = 16 : i32, ttig.support_bfloat16_conversion, ttig.support_subgroup_matrix_multiply_accumulate, ttig.support_2d_block_io, "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.target = "xpu", "ttg.threads-per-warp" = 32 : i32}} {{ + tt.func public @column_major_load_subgroup32(%a_ptr: !tt.ptr {{tt.divisibility = 16 : i32}}, + %b_ptr: !tt.ptr {{tt.divisibility = 16 : i32}}, + %c_ptr: !tt.ptr {{tt.divisibility = 16 : i32}}) attributes {{noinline = false}} {{ + %c0_i32 = arith.constant 0 : i32 + %c1_i64 = arith.constant 1 : i64 + %cM_i32 = arith.constant {M} : i32 + %cK_i32 = arith.constant {K} : i32 + %cN_i32 = arith.constant {N} : i32 + %cK_i64 = arith.constant {K} : i64 + %cN_i64 = arith.constant {N} : i64 + + // A: [{M}, {K}] row-major. + %desc_a = tt.make_tensor_descriptor %a_ptr, [%cM_i32, %cK_i32], [%cK_i64, %c1_i64] + : !tt.ptr, !tt.tensordesc<{M}x{K}xf16> + %a = tt.descriptor_load %desc_a [%c0_i32, %c0_i32] {{ttig.block_io = "row_major"}} + : !tt.tensordesc<{M}x{K}xf16> -> tensor<{M}x{K}xf16, #dot0> + + // B stored as [{N}, {K}] row-major in memory (N rows, K contiguous cols). + // column_major load: the permuteDescDim logic swaps the last two descriptor + // dimensions before extracting surface parameters, so the result is + // tensor<{K}x{N}xf16, #dot1> (the transposed view of the [{N},{K}] storage). + %desc_b = tt.make_tensor_descriptor %b_ptr, [%cN_i32, %cK_i32], [%cK_i64, %c1_i64] + : !tt.ptr, !tt.tensordesc<{N}x{K}xf16> + %b = tt.descriptor_load %desc_b [%c0_i32, %c0_i32] {{ttig.block_io = "column_major"}} + : !tt.tensordesc<{N}x{K}xf16> -> tensor<{K}x{N}xf16, #dot1> + + // C = A x B (dimensions: [{M},{K}] x [{K},{N}] -> [{M},{N}]) + %C_init = arith.constant dense<0.000000e+00> : tensor<{M}x{N}xf32, #dpas> + %C = tt.dot %a, %b, %C_init, inputPrecision = tf32 + : tensor<{M}x{K}xf16, #dot0> * tensor<{K}x{N}xf16, #dot1> -> tensor<{M}x{N}xf32, #dpas> + + // Store C: the dpas encoding supports 2D block stores via a descriptor. + %desc_c = tt.make_tensor_descriptor %c_ptr, [%cM_i32, %cN_i32], [%cN_i64, %c1_i64] + : !tt.ptr, !tt.tensordesc<{M}x{N}xf32, #dpas> + tt.descriptor_store %desc_c [%c0_i32, %c0_i32], %C {{ttig.block_io = "row_major"}} + : !tt.tensordesc<{M}x{N}xf32, #dpas>, tensor<{M}x{N}xf32, #dpas> + + tt.return + }} + }} + """ + + torch.manual_seed(42) + A = torch.randn((M, K), dtype=torch.float16, device=device) + # B is stored as [N, K] in memory (the transposed representation of a [K, N] B matrix). + B_stored = torch.randn((N, K), dtype=torch.float16, device=device) + C = torch.empty((M, N), dtype=torch.float32, device=device) + + temp_file = tmp_path / "test_block_tdesc_column_major_load_subgroup32.ttgir" + temp_file.write_text(ir) + kernel = triton.compile(str(temp_file)) + + kernel[(1, 1, 1)](A, B_stored, C) + + # C should equal A x B_stored.T (since B was loaded via column_major). + expected = torch.mm(A.to(torch.float32), B_stored.to(torch.float32).T) + torch.testing.assert_close(C, expected, atol=1e-1, rtol=1e-2) diff --git a/test/TritonIntelGPU/2d-block-load-to-llvm.mlir b/test/TritonIntelGPU/2d-block-load-to-llvm.mlir index add58319d0..c8cb49be04 100644 --- a/test/TritonIntelGPU/2d-block-load-to-llvm.mlir +++ b/test/TritonIntelGPU/2d-block-load-to-llvm.mlir @@ -195,3 +195,20 @@ module attributes {"ttg.num-warps" = 8 : i32, "ttg.threads-per-warp" = 32 : i32, tt.return } } + +// ----- + +// COM: Test subgroup-size=32 DotOp-B column-major (transposed) block load with +// COM: f16/opsPerChan=2. The old width-comparison check rejected this because +// COM: dpasInstShapeB()[1]=16 != threadsPerWarp=32. The new linear-layout check +// COM: correctly accepts it. +#dpas = #ttig.dpas<{repeatCount = 8, systolicDepth = 8, executionSize = 16, opsPerChan = 2, threadsPerWarp = 32, warpsPerCTA = [2, 2], repCluster = [1, 1], A = [16, 16], B = [16, 16], C = [16, 16]}> +#dot = #ttg.dot_op<{opIdx = 1, parent = #dpas, kWidth = 2}> +module attributes {"ttg.num-warps" = 8 : i32, "ttg.threads-per-warp" = 32 : i32, "ttig.support_2d_block_io"} { + // CHECK-LABEL: @block_load_dot_b_subgroup32_transpose + tt.func public @block_load_dot_b_subgroup32_transpose(%arg0: !tt.ptr, %arg1: i32, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32) { + // CHECK: triton_gen.2Dblockload {{.*}} transpose = true + %0 = ttig.2d_block_load %arg0, %arg1, %arg2, %arg3[%arg4, %arg5] {column_major} : !tt.ptr -> tensor<32x32xf16, #dot> + tt.return + } +} diff --git a/test/TritonIntelGPU/LowerTo2DBlockLoad/pointer-load.mlir b/test/TritonIntelGPU/LowerTo2DBlockLoad/pointer-load.mlir index 341256c5d3..4efedc88a9 100644 --- a/test/TritonIntelGPU/LowerTo2DBlockLoad/pointer-load.mlir +++ b/test/TritonIntelGPU/LowerTo2DBlockLoad/pointer-load.mlir @@ -247,6 +247,28 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, "ttg.thr // ----- +// COM: Column-major pointer load with sub-group-size=32 and f16/opsPerChan=2. +// COM: The old width-comparison check rejected this because dpasInstShapeB()[1]=16 +// COM: != threadsPerWarp=32. The new linear-layout check correctly accepts it. +#dpas = #ttig.dpas<{repeatCount = 8, systolicDepth = 8, executionSize = 16, opsPerChan = 2, threadsPerWarp = 32, warpsPerCTA = [2, 2], repCluster = [1, 1], A = [16, 16], B = [16, 16], C = [16, 16]}> +#dot1 = #ttg.dot_op<{opIdx = 1, parent = #dpas, kWidth = 2}> +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, ttig.support_2d_block_io} { + // CHECK-LABEL: tt.func @pointer_load_column_major_subgroup32 + tt.func @pointer_load_column_major_subgroup32(%arg0: !tt.ptr {tt.divisibility = 16 : i32}) -> tensor<32x64xf16, #dot1> { + %0 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #dot1}>> + %1 = tt.expand_dims %0 {axis = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #dot1}>> -> tensor<1x64xi32, #dot1> + %2 = tt.splat %arg0 : !tt.ptr -> tensor<1x64x!tt.ptr, #dot1> + %3 = tt.addptr %2, %1 : tensor<1x64x!tt.ptr, #dot1>, tensor<1x64xi32, #dot1> + %4 = tt.broadcast %3 : tensor<1x64x!tt.ptr, #dot1> -> tensor<32x64x!tt.ptr, #dot1> + // CHECK: %[[P:.*]] = arith.constant 128 : i32 + // CHECK: ttig.2d_block_load_from_ptr %4, %[[P]] {column_major} {base_height = 1 : i32, base_width = 32 : i32} + %5 = tt.load %4 {ttig.block_io = "column_major"} : tensor<32x64x!tt.ptr, #dot1> + tt.return %5 : tensor<32x64xf16, #dot1> + } +} + +// ----- + // COM: Env var TRITON_INTEL_ONE_MATRIX_PER_LOAD_BT=1 forces the attribute on // COM: all loads, even those without it originally. // RUN: env TRITON_INTEL_ONE_MATRIX_PER_LOAD_BT=1 triton-opt %s -split-input-file --tritonintelgpu-lower-to-2d-block-load | FileCheck %s --check-prefix=ENV-CHECK diff --git a/third_party/intel/include/Dialect/TritonIntelGPU/Transforms/BlockIOUtils.h b/third_party/intel/include/Dialect/TritonIntelGPU/Transforms/BlockIOUtils.h index 99a015bc73..433d8dae90 100644 --- a/third_party/intel/include/Dialect/TritonIntelGPU/Transforms/BlockIOUtils.h +++ b/third_party/intel/include/Dialect/TritonIntelGPU/Transforms/BlockIOUtils.h @@ -90,7 +90,7 @@ DpasEncodingAttr getDpasLayout(RankedTensorType tensorTy); /// and the LLVM lowering (to produce the actual mapping). FailureOr computeTransposeShuffleMapping( RankedTensorType tensorType, const LinearLayout ®Mapping, - int64_t numElemsPerLoad, unsigned numPackedVals, unsigned tileHeight, + int64_t numElemsPerLoad, const BlockIOTileSizeInfo &sizeInfo, unsigned threadsPerWarp, bool hasDPASOperandType, MLIRContext *ctx); /// Check whether the packed element size and tile width satisfy the diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp index 87f64048dd..e062fd0845 100644 --- a/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp +++ b/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp @@ -971,11 +971,11 @@ struct BlockIOConversionBase : public LoadStoreConversionBase { static FailureOr computeTransposeShuffleMapping( RankedTensorType tensorType, const LinearLayout ®Mapping, - int64_t numElemsPerLoad, unsigned numPackedVals, unsigned tileHeight, + int64_t numElemsPerLoad, const BlockIOTileSizeInfo &sizeInfo, unsigned threadsPerWarp, bool hasDPASOperandType, MLIRContext *ctx) { return triton::gpu::intel::computeTransposeShuffleMapping( - tensorType, regMapping, numElemsPerLoad, numPackedVals, tileHeight, - threadsPerWarp, hasDPASOperandType, ctx); + tensorType, regMapping, numElemsPerLoad, sizeInfo, threadsPerWarp, + hasDPASOperandType, ctx); } /// Build a Block2DLoadConfig from a validated BlockIOTileSizeInfo. @@ -1050,8 +1050,8 @@ struct BlockIOConversionBase : public LoadStoreConversionBase { LinearLayout::identity1D(cfg.numElemsPerLoad, kRegister, kRegister); if (cfg.isTransposeRequired) { auto maybeShuffleMapping = computeTransposeShuffleMapping( - tensorType, cfg.regMapping, cfg.numElemsPerLoad, cfg.numPackedVals, - cfg.tileHeight, threadsPerWarp, !!cfg.packedDPASOperandType, ctx); + tensorType, cfg.regMapping, cfg.numElemsPerLoad, sizeInfo, + threadsPerWarp, !!cfg.packedDPASOperandType, ctx); assert(succeeded(maybeShuffleMapping) && "validate2DBlockLoadTile should have rejected this configuration"); cfg.shuffleMapping = *maybeShuffleMapping; @@ -3966,8 +3966,9 @@ static LogicalResult lowerBlockLoad2D( Value oldVal = b.extract_element(ret, b.i32_val(valueIndex)); Value newVal = oldVal; if (cfg.tileWidth < cfg.threadsPerWarp) { - assert(cfg.tileWidth * 2 == cfg.threadsPerWarp && - "Expecting tileWidth to be 2x threadsPerWarp"); + assert(cfg.threadsPerWarp % cfg.tileWidth == 0 && + "tileWidth must evenly divide threadsPerWarp for broadcast " + "row replication"); Value threadId = getThreadId(rewriter, loc); newVal = targetInfo.shuffleIdx( rewriter, loc, oldVal, diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/BlockIOUtils.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/BlockIOUtils.cpp index 4163cb7734..d1996905cf 100644 --- a/third_party/intel/lib/TritonIntelGPUTransforms/BlockIOUtils.cpp +++ b/third_party/intel/lib/TritonIntelGPUTransforms/BlockIOUtils.cpp @@ -496,43 +496,14 @@ DpasEncodingAttr getDpasLayout(RankedTensorType tensorTy) { FailureOr computeTransposeShuffleMapping( RankedTensorType tensorType, const LinearLayout ®Mapping, - int64_t numElemsPerLoad, unsigned numPackedVals, unsigned tileHeight, + int64_t numElemsPerLoad, const BlockIOTileSizeInfo &sizeInfo, unsigned threadsPerWarp, bool hasDPASOperandType, MLIRContext *ctx) { StringAttr kRegister = StringAttr::get(ctx, "register"); + StringAttr kLane = StringAttr::get(ctx, "lane"); LinearLayout shuffleMapping = LinearLayout::identity1D(numElemsPerLoad, kRegister, kRegister); - // Improve this. The current 2D block load only transposes the matrix at - // i32 granularity. We still need to perform an additional in-register - // transpose from i32 -> (N × ElemSizeInBits) tiles, using the tile width. - // At the moment, we can only achieve this using a bitcast operation, - // which implicitly uses the sub-group size as the transpose width. To - // optimize further, we should implement this with inline VISA - // instructions. - - // tileHeight becomes width after transposing. - unsigned widthToTranspose = tileHeight; if (hasDPASOperandType) { - // For the DPAS related layout, we will do the shuffle at first in the - // unpacking of the elements at the DPAS operands granularity. - // And then we will do the transposing. So the transposing width is DPAS - // op shapes. - DpasEncodingAttr::OpIdx opIdx = getOpIdx(tensorType); - DpasEncodingAttr dpasLayout = getDpasLayout(tensorType); - switch (opIdx) { - case DpasEncodingAttr::OpIdx::OperandA: { - widthToTranspose = dpasLayout.getDPASInstShapeA()[1]; - break; - } - case DpasEncodingAttr::OpIdx::OperandB: { - widthToTranspose = dpasLayout.getDPASInstShapeB()[1]; - break; - } - case DpasEncodingAttr::OpIdx::OperandC: { - widthToTranspose = dpasLayout.getDPASInstShapeC()[1]; - break; - } - } // For shuffle the transposed Dot operands matrix, we can shuffle the // loaded matrix in an reverse order. auto invertMapping = regMapping.invert(); @@ -555,8 +526,65 @@ FailureOr computeTransposeShuffleMapping( return failure(); } - if (numPackedVals > 1 && (widthToTranspose) != threadsPerWarp) - return failure(); + Attribute encoding = tensorType.getEncoding(); + std::optional llEncoding = + cast(encoding).toLinearLayout( + tensorType.getShape()); + assert(llEncoding.has_value() && "expected valid linear layout"); + if (sizeInfo.transpose) { + auto dims = llvm::to_vector(llEncoding->getOutDimNames()); + SmallVector loadDimName = {dims[sizeInfo.rowDim], + dims[sizeInfo.colDim]}; + LinearLayout blockLoadLayoutWithInSubgroup = + llEncoding->sublayout({kRegister, kLane}, loadDimName); + LinearLayout expectedLoadUnpackLayout = + blockLoadLayoutWithInSubgroup + .resizeOutDim(dims[sizeInfo.rowDim], + sizeInfo.tileWidth * sizeInfo.numElemPerPackedVal) + .resizeOutDim(dims[sizeInfo.colDim], sizeInfo.tileHeight) + .resizeInDim(kRegister, numElemsPerLoad); + // Compose the load return value layout of the transposed packed type. + std::vector> regBases; + int32_t colBase = 1, rowBase = 1; + for (int32_t i = 1; i < sizeInfo.numElemPerPackedVal; i *= 2) { + regBases.push_back({rowBase, 0}); + rowBase <<= 1; + } + std::vector> laneBases; + for (int32_t i = 1; i < threadsPerWarp; i *= 2) { + if (i < sizeInfo.tileHeight) { + laneBases.push_back({0, colBase}); + colBase <<= 1; + } else { + laneBases.push_back({rowBase, 0}); + rowBase <<= 1; + } + } + for (int32_t i = 1; i < numElemsPerLoad / sizeInfo.numElemPerPackedVal; + i *= 2) { + if (i < mlir::ceil(sizeInfo.tileHeight, (int32_t)threadsPerWarp)) { + regBases.push_back({0, colBase}); + colBase <<= 1; + } else { + regBases.push_back({rowBase, 0}); + rowBase <<= 1; + } + } + LinearLayout transPackedLayout = LinearLayout( + {{kRegister, regBases}, {kLane, laneBases}}, ArrayRef(loadDimName)); + if (transPackedLayout != expectedLoadUnpackLayout) { + // Improve this. The current 2D block load only transposes the matrix at + // i32 granularity. We still need to perform an additional in-register + // transpose from i32 -> (N × ElemSizeInBits) tiles, using the tile width. + // At the moment, we can only achieve this using a bitcast operation, + // which implicitly uses the sub-group size as the transpose width. To + // optimize further, we should implement this with inline VISA + // instructions. + // E.g. a case that loads fp8 B matrix with DPAS opsChan=2 layout under + // sub-group-size=32. + return failure(); + } + } return shuffleMapping; } @@ -618,8 +646,7 @@ bool validate2DBlockLoadTile(const LinearLayout &ll, unsigned memContiguousDim, } if (failed(computeTransposeShuffleMapping( - tensorType, regMapping, numElemsPerLoad, - sizeInfo.numElemPerPackedVal, sizeInfo.tileHeight, threadsPerWarp, + tensorType, regMapping, numElemsPerLoad, sizeInfo, threadsPerWarp, hasDPASOperandType, ctx))) return false; } diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/LowerTo2DBlockLoad.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/LowerTo2DBlockLoad.cpp index 1f49d7c474..5bd2681819 100644 --- a/third_party/intel/lib/TritonIntelGPUTransforms/LowerTo2DBlockLoad.cpp +++ b/third_party/intel/lib/TritonIntelGPUTransforms/LowerTo2DBlockLoad.cpp @@ -487,17 +487,16 @@ struct TritonIntelGPULowerTo2DBlockLoadPass } } - // For broadcast loads, the LLVM lowering's row replication - // requires tileWidth >= threadsPerWarp or tileWidth * 2 == threadsPerWarp. - // Reject unsupported configurations. + // For broadcast loads, the LLVM lowering's row replication uses + // shuffleIdx(threadId % tileWidth), which requires tileWidth to evenly + // divide threadsPerWarp. Reject unsupported configurations. if (isBroadcast && tileHeight > 1 && tileWidth > 0) { unsigned threadsPerWarp = ttg::TritonGPUDialect::getThreadsPerWarp( op->getParentOfType()); - if (tileWidth < (int)threadsPerWarp && - (unsigned)tileWidth * 2 != threadsPerWarp) { + if ((unsigned)tileWidth < threadsPerWarp && + threadsPerWarp % (unsigned)tileWidth != 0) { LDBG("Broadcast load tile width " << tileWidth - << " incompatible with " - "threadsPerWarp " + << " does not divide threadsPerWarp " << threadsPerWarp << " for: " << *op); return; } diff --git a/third_party/intel/unittest/Dialect/TritonIntelGPU/BlockIOUtilsTest.cpp b/third_party/intel/unittest/Dialect/TritonIntelGPU/BlockIOUtilsTest.cpp new file mode 100644 index 0000000000..09f5af36b0 --- /dev/null +++ b/third_party/intel/unittest/Dialect/TritonIntelGPU/BlockIOUtilsTest.cpp @@ -0,0 +1,76 @@ +// Tests for validate2DBlockLoadTile in BlockIOUtils. +// +// PR #7487 changed computeTransposeShuffleMapping from a simple width +// comparison to a linear-layout comparison, enabling column-major B matrix +// loads with sub-group-size=32. These tests verify: +// 1. tpw=32, f16/opsPerChan=2, column-major B is now accepted. +// 2. tpw=16, f16/opsPerChan=2, column-major B is still accepted (regression). + +#include "intel/include/Dialect/TritonIntelGPU/Transforms/BlockIOUtils.h" +#include "intel/include/Dialect/TritonIntelGPU/IR/Dialect.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "llvm/Support/Signals.h" +#include + +using namespace mlir; +using namespace mlir::triton::gpu; +using namespace mlir::triton::gpu::intel; + +namespace { + +class BlockIOUtilsTest : public ::testing::Test { +public: + void SetUp() override { + ctx.getOrLoadDialect(); + ctx.getOrLoadDialect(); + } + + DpasEncodingAttr makeDpas(ArrayRef warpsPerCTA, + unsigned threadsPerWarp, unsigned opsPerChan) { + return DpasEncodingAttr::get(&ctx, /*repeatCount=*/8, /*systolicDepth=*/8, + /*executionSize=*/16, opsPerChan, warpsPerCTA, + /*repCluster=*/{1, 1}, threadsPerWarp, + std::nullopt); + } + +protected: + MLIRContext ctx; +}; + +// Test that validate2DBlockLoadTile accepts a column-major B matrix load with +// sub-group-size=32 and f16/opsPerChan=2. +// +// Old code: computeTransposeShuffleMapping failed because +// numPackedVals(=2) > 1 && dpasInstShapeB()[1](=16) != threadsPerWarp(=32) +// New code: linear-layout comparison succeeds for this configuration. +TEST_F(BlockIOUtilsTest, ColumnMajorB_Tpw32_F16_Accepted) { + auto dpas = makeDpas(/*warpsPerCTA=*/{2, 2}, /*threadsPerWarp=*/32, + /*opsPerChan=*/2); + auto dot = DotOperandEncodingAttr::get(&ctx, /*opIdx=*/1, dpas, /*kWidth=*/2); + auto f16 = Float16Type::get(&ctx); + SmallVector shape = {32, 64}; // K=32, N=64 + auto tensorType = RankedTensorType::get(shape, f16, dot); + auto ll = cast(dot).toLinearLayout(shape); + // column_major: memContiguousDim = rank-2 = 0 (K direction) + EXPECT_TRUE(validate2DBlockLoadTile(ll, /*memContiguousDim=*/0, + /*elemSizeInBits=*/16, tensorType)); +} + +// Regression test: validate2DBlockLoadTile must still accept a column-major B +// matrix load with sub-group-size=16 and f16/opsPerChan=2 after PR #7487. +TEST_F(BlockIOUtilsTest, ColumnMajorB_Tpw16_F16_Accepted) { + auto dpas = makeDpas(/*warpsPerCTA=*/{4, 2}, /*threadsPerWarp=*/16, + /*opsPerChan=*/2); + auto dot = DotOperandEncodingAttr::get(&ctx, /*opIdx=*/1, dpas, /*kWidth=*/2); + auto f16 = Float16Type::get(&ctx); + SmallVector shape = {32, 64}; // K=32, N=64 + auto tensorType = RankedTensorType::get(shape, f16, dot); + auto ll = cast(dot).toLinearLayout(shape); + // column_major: memContiguousDim = rank-2 = 0 (K direction) + EXPECT_TRUE(validate2DBlockLoadTile(ll, /*memContiguousDim=*/0, + /*elemSizeInBits=*/16, tensorType)); +} + +} // namespace diff --git a/third_party/intel/unittest/Dialect/TritonIntelGPU/CMakeLists.txt b/third_party/intel/unittest/Dialect/TritonIntelGPU/CMakeLists.txt index 1a2b312f1f..374f107c92 100644 --- a/third_party/intel/unittest/Dialect/TritonIntelGPU/CMakeLists.txt +++ b/third_party/intel/unittest/Dialect/TritonIntelGPU/CMakeLists.txt @@ -18,3 +18,14 @@ add_triton_ut( TritonIntelGPUTransforms TritonNvidiaGPUTransforms ) +add_triton_ut( + NAME BlockIOUtils + SRCS BlockIOUtilsTest.cpp + LIBS + TritonIntelGPUIR + TritonGPUIR + TritonGPUTransforms + TritonIntelAnalysis + TritonIntelGPUTransforms + TritonNvidiaGPUTransforms +)