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

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<f16> {{tt.divisibility = 16 : i32}},
+                                                      %b_ptr: !tt.ptr<f16> {{tt.divisibility = 16 : i32}},
+                                                      %c_ptr: !tt.ptr<f32> {{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<f16>, !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<f16>, !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<f32>, !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)

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<f16>, %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<f16> -> tensor<32x32xf16, #dot>
+    tt.return
+  }
+}

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<f16> {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<f16> -> tensor<1x64x!tt.ptr<f16>, #dot1>
+    %3 = tt.addptr %2, %1 : tensor<1x64x!tt.ptr<f16>, #dot1>, tensor<1x64xi32, #dot1>
+    %4 = tt.broadcast %3 : tensor<1x64x!tt.ptr<f16>, #dot1> -> tensor<32x64x!tt.ptr<f16>, #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<f16>, #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

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<LinearLayout> computeTransposeShuffleMapping(
     RankedTensorType tensorType, const LinearLayout &regMapping,
-    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

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -971,11 +971,11 @@ struct BlockIOConversionBase : public LoadStoreConversionBase {
 
   static FailureOr<LinearLayout> computeTransposeShuffleMapping(
       RankedTensorType tensorType, const LinearLayout &regMapping,
-      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,

third_party/intel/lib/TritonIntelGPUTransforms/BlockIOUtils.cpp

@@ -496,43 +496,14 @@ DpasEncodingAttr getDpasLayout(RankedTensorType tensorTy) {
 
 FailureOr<LinearLayout> computeTransposeShuffleMapping(
     RankedTensorType tensorType, const LinearLayout &regMapping,
-    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<LinearLayout> computeTransposeShuffleMapping(
       return failure();
   }
 
-  if (numPackedVals > 1 && (widthToTranspose) != threadsPerWarp)
-    return failure();
+  Attribute encoding = tensorType.getEncoding();
+  std::optional<LinearLayout> llEncoding =
+      cast<DistributedEncodingTrait>(encoding).toLinearLayout(
+          tensorType.getShape());
+  assert(llEncoding.has_value() && "expected valid linear layout");
+  if (sizeInfo.transpose) {
+    auto dims = llvm::to_vector(llEncoding->getOutDimNames());
+    SmallVector<StringAttr> 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<std::vector<int32_t>> 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<std::vector<int32_t>> 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;
   }

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<ModuleOp>());
-      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;
       }