[Mlir-commits] [mlir] [MLIR][XeGPU] Improve workgroup to subgroup distribution pattern for mulit-reduction op (PR #182178)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Feb 18 15:08:47 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
This PR improves the multi-reduction op pattern in wg distribution to save the partially reduced tile to slm using the same rank as the tile before reduced, instead of forcing them to using 2d slm tile. This simplifies the layout assignment at later stage.
---
Patch is 29.20 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/182178.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+68-90)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir (+76-79)
``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 66eb7fc97aa1a..a4172bdc9d136 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1304,10 +1304,10 @@ static Value createAccumulator(ConversionPatternRewriter &rewriter,
/// This gives us a unique linear index for each combination of subgroup
/// positions in the specified dimensions, which is used for SLM row/column
/// addressing.
-static Value linearizeSubgroupIndices(ConversionPatternRewriter &rewriter,
- Location loc, ArrayRef<Value> sgIds,
- ArrayRef<int64_t> dims,
- ArrayRef<int64_t> sgLayout) {
+[[maybe_unused]] static Value
+linearizeSubgroupIndices(ConversionPatternRewriter &rewriter, Location loc,
+ ArrayRef<Value> sgIds, ArrayRef<int64_t> dims,
+ ArrayRef<int64_t> sgLayout) {
Value linearizedOffset = arith::ConstantIndexOp::create(rewriter, loc, 0);
int64_t stride = 1;
@@ -1370,6 +1370,9 @@ struct WgToSgMultiDimReductionOp
return failure();
auto originalSrcShape = srcType.getShape();
+ auto originalDstShape = dstType.getShape();
+ int srcVecRank = originalSrcShape.size();
+
xegpu::DistributeLayoutAttr layout =
xegpu::getTemporaryLayout(dyn_cast<OpResult>(op.getResult()));
if (!layout || !layout.isForWorkgroup())
@@ -1391,10 +1394,15 @@ struct WgToSgMultiDimReductionOp
// Step 1: perform local subgroup reductions with ZERO accumulator
SmallVector<Value> localReductions;
- SmallVector<int64_t> sgShape =
- getSgShapeAndCount(originalSrcShape, layout).first;
- VectorType newDstType = VectorType::get(sgShape, elemTy);
- for (auto sgSrc : adaptor.getSource()) {
+ SmallVector<int64_t> sgDstShape =
+ getSgShapeAndCount(originalDstShape, layout).first;
+ auto sgSrcs = adaptor.getSource();
+ auto sgSrcType = dyn_cast<VectorType>(sgSrcs.front().getType());
+ SmallVector<int64_t> sgSrcShape(sgSrcType.getShape().begin(),
+ sgSrcType.getShape().end());
+
+ VectorType newDstType = VectorType::get(sgDstShape, elemTy);
+ for (auto sgSrc : sgSrcs) {
// Create ZERO accumulator for local reduction
auto neutralLocalAcc =
createAccumulator(rewriter, loc, newDstType, op.getKind());
@@ -1430,44 +1438,38 @@ struct WgToSgMultiDimReductionOp
}
// Step 2: cross-subgroup reduction using SLM
-
- // Calculate total elements in local result
- int64_t localElements = computeProduct(sgShape);
-
- // Shape cast for SLM storage - store as [1, localElements]
- SmallVector<int64_t> storeShape2D = {1, localElements};
- VectorType storeType2D = VectorType::get(storeShape2D, elemTy);
- auto storeShapeCast = vector::ShapeCastOp::create(
- rewriter, loc, storeType2D, localReductions[0]);
- Value storeData = storeShapeCast.getResult();
-
- // Calculate SLM shape - rows for sg's in reduction dims, cols for total
- // result elements across all subgroups in non-reduction dimensions
- int64_t totalReductionSubgroups = 1;
- for (int64_t dim : crossSgReductionDims) {
- totalReductionSubgroups *= sgLayout[dim];
- }
-
- // Total result elements across all subgroups in non-reduction dimensions
- int64_t totalResultElements =
- localElements * computeProduct(sgLayout) / totalReductionSubgroups;
-
- SmallVector<int64_t> slmShape2D = {totalReductionSubgroups,
- totalResultElements};
+ auto slmStoreDataShape = sgSrcShape;
+ for (int64_t dim : reductionDims)
+ slmStoreDataShape[dim] = 1;
+ VectorType slmStoreDataType = VectorType::get(slmStoreDataShape, elemTy);
+ Value slmStoreData = vector::ShapeCastOp::create(
+ rewriter, loc, slmStoreDataType, localReductions[0]);
+
+ SmallVector<int64_t> slmShape(originalSrcShape.begin(),
+ originalSrcShape.end());
+ // for reduction dimension, SLM stores partial results from each subgroup
+ for (int64_t dim : reductionDims)
+ slmShape[dim] = originalSrcShape[dim] / sgSrcShape[dim];
// Allocate SLM
auto bitWidth = elemTy.getIntOrFloatBitWidth();
auto bytesPerElement = bitWidth / 8;
- int64_t slmElements = slmShape2D[0] * slmShape2D[1];
- auto slmSize = slmElements * bytesPerElement;
+ auto slmSize = computeProduct(slmShape) * bytesPerElement;
auto slmTy = MemRefType::get({slmSize}, rewriter.getI8Type(), {}, 3);
auto slm = memref::AllocaOp::create(rewriter, loc, slmTy);
- auto memDescType = xegpu::MemDescType::get(rewriter.getContext(),
- slmShape2D, elemTy, nullptr);
+ auto memDescType = xegpu::MemDescType::get(rewriter.getContext(), slmShape,
+ elemTy, nullptr);
auto memDesc =
xegpu::CreateMemDescOp::create(rewriter, loc, memDescType, slm);
+ // if localReductions have more than 1 result, not support
+ if (localReductions.size() > 1) {
+ return rewriter.notifyMatchFailure(
+ op,
+ "Multiple local reductions not supported in current implementation.");
+ }
+
// Step 4: Store local results to SLM
auto sgId = gpu::SubgroupIdOp::create(rewriter, loc,
rewriter.getIndexType(), nullptr);
@@ -1484,78 +1486,54 @@ struct WgToSgMultiDimReductionOp
return failure();
SmallVector<Value> sgIds = *sgIdsResult;
- // Row offset: linearize reduction dimension indices
- Value rowOffsetStore = linearizeSubgroupIndices(
- rewriter, loc, sgIds, crossSgReductionDims, sgLayout);
-
- // Column offset: linearize non-reduction dimension indices
- SmallVector<int64_t> nonReductionDims;
- for (size_t i = 0; i < sgLayout.size(); ++i) {
- if (!llvm::is_contained(reductionDims, static_cast<int64_t>(i))) {
- nonReductionDims.push_back(static_cast<int64_t>(i));
- }
+ SmallVector<OpFoldResult> slmStoreOffsets;
+ for (int i = 0; i < srcVecRank; ++i) {
+ Value dimVal = sgIds[i];
+ int64_t stride =
+ (llvm::is_contained(reductionDims, i)) ? 1 : sgSrcShape[i];
+ Value strideVal = arith::ConstantIndexOp::create(rewriter, loc, stride);
+ Value offsetVal = arith::MulIOp::create(rewriter, loc, dimVal, strideVal);
+ slmStoreOffsets.push_back(offsetVal);
}
- Value colOffset = linearizeSubgroupIndices(rewriter, loc, sgIds,
- nonReductionDims, sgLayout);
-
- Value localElementsVal =
- arith::ConstantIndexOp::create(rewriter, loc, localElements);
- colOffset =
- arith::MulIOp::create(rewriter, loc, colOffset, localElementsVal);
-
- SmallVector<OpFoldResult> storeOffsets2D = {rowOffsetStore, colOffset};
-
- xegpu::StoreMatrixOp::create(rewriter, loc, storeData, memDesc.getResult(),
- storeOffsets2D, /*layout=*/nullptr);
+ xegpu::StoreMatrixOp::create(rewriter, loc, slmStoreData,
+ memDesc.getResult(), slmStoreOffsets,
+ /*layout=*/nullptr);
gpu::BarrierOp::create(rewriter, loc);
// Step 5: Load from SLM for final reduction
- SmallVector<int64_t> loadShape2D = {totalReductionSubgroups, localElements};
- VectorType loadType2D = VectorType::get(loadShape2D, elemTy);
-
- // Load offsets - each subgroup loads its column based on non-reduction
- // position
- Value rowOffsetLoad = arith::ConstantIndexOp::create(rewriter, loc, 0);
-
- SmallVector<OpFoldResult> loadOffsets2D = {rowOffsetLoad, colOffset};
+ SmallVector<int64_t> slmLoadDataShape(sgSrcShape.begin(), sgSrcShape.end());
+ for (int64_t dim : reductionDims)
+ slmLoadDataShape[dim] = slmShape[dim];
+
+ SmallVector<OpFoldResult> slmLoadOffsets;
+ for (int i = 0; i < srcVecRank; ++i) {
+ Value dimVal = sgIds[i];
+ int64_t stride =
+ (llvm::is_contained(reductionDims, i)) ? 0 : sgSrcShape[i];
+ Value strideVal = arith::ConstantIndexOp::create(rewriter, loc, stride);
+ Value offsetVal = arith::MulIOp::create(rewriter, loc, dimVal, strideVal);
+ slmLoadOffsets.push_back(offsetVal);
+ }
- auto loadOp = xegpu::LoadMatrixOp::create(
- rewriter, loc, loadType2D, memDesc.getResult(), loadOffsets2D,
+ VectorType slmLoadType = VectorType::get(slmLoadDataShape, elemTy);
+ auto slmLoadOp = xegpu::LoadMatrixOp::create(
+ rewriter, loc, slmLoadType, memDesc.getResult(), slmLoadOffsets,
/*layout=*/nullptr);
// Step 6: Perform final reduction with ZERO accumulator
- SmallVector<int64_t> finalReductionDims = {0};
- SmallVector<int64_t> finalResultShape = {localElements};
- VectorType finalResultType = VectorType::get(finalResultShape, elemTy);
-
auto neutralFinalAcc =
- createAccumulator(rewriter, loc, finalResultType, op.getKind());
+ createAccumulator(rewriter, loc, newDstType, op.getKind());
auto finalReduce = vector::MultiDimReductionOp::create(
- rewriter, loc, finalResultType, op.getKind(), loadOp.getResult(),
- neutralFinalAcc, finalReductionDims);
+ rewriter, loc, newDstType, op.getKind(), slmLoadOp.getResult(),
+ neutralFinalAcc, reductionDims);
// Step 7: Add the original accumulator at the end
Value originalAcc = adaptor.getAcc()[0];
Value accToAdd = originalAcc;
- // Handle shape mismatch by shape casting
- if (originalAcc.getType() != finalReduce.getResult().getType()) {
- auto originalAccType = cast<VectorType>(originalAcc.getType());
- auto finalResultType =
- cast<VectorType>(finalReduce.getResult().getType());
-
- // If they have the same number of elements, just shape cast
- if (originalAccType.getNumElements() ==
- finalResultType.getNumElements()) {
- auto shapeCast = vector::ShapeCastOp::create(
- rewriter, loc, finalResultType, originalAcc);
- accToAdd = shapeCast.getResult();
- }
- }
-
auto finalResult = vector::makeArithReduction(
rewriter, loc, op.getKind(), finalReduce.getResult(), accToAdd);
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
index e2e94c5f0300f..9407f7f2357a2 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
@@ -646,30 +646,28 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<1x32xf32>
// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : vector<1x1x32xindex>
// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<1x1x32xi1>
- // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ // CHECK-DAG: %[[LOAD:.*]] = xegpu.load %[[ARG0:.*]][%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<1x32xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [1] : vector<1x1x32xf32> to vector<1x32xf32>
- // CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x32xf32> to vector<1x32xf32>
+ // CHECK-DAG: %[[CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x32xf32> to vector<1x1x32xf32>
// CHECK-DAG: %[[ALLOCA:.*]] = memref.alloca() : memref<4096xi8, 3>
- // CHECK-DAG: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<4096xi8, 3> -> !xegpu.mem_desc<32x32xf32>
+ // CHECK-DAG: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<4096xi8, 3> -> !xegpu.mem_desc<1x32x32xf32>
// CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index
- // CHECK-DAG: %[[AFFINE1:.*]] = affine.apply #map()[%[[SGID]]]
- // CHECK-DAG: %[[AFFINE2:.*]] = affine.apply #map1()[%[[SGID]]]
- // CHECK-DAG: %[[AFFINE3:.*]] = affine.apply #map2()[%[[SGID]]]
- // CHECK-DAG: %[[MUL1:.*]] = arith.muli {{.*}}, %[[C1:.*]] : index
- // CHECK-DAG: %[[ROW_OFFSET:.*]] = arith.addi %[[C0:.*]], %[[MUL1]] : index
- // CHECK-DAG: %[[MUL2:.*]] = arith.muli %[[AFFINE1]], %[[C1:.*]] : index
- // CHECK-DAG: %[[ADD1:.*]] = arith.addi %[[C0:.*]], %[[MUL2]] : index
- // CHECK-DAG: %[[MUL3:.*]] = arith.muli %[[AFFINE3]], %[[C1:.*]] : index
- // CHECK-DAG: %[[ADD2:.*]] = arith.addi %[[ADD1]], %[[MUL3]] : index
- // CHECK-DAG: %[[COL_OFFSET:.*]] = arith.muli %[[ADD2]], %[[C32:.*]] : index
- // CHECK-DAG: xegpu.store_matrix %[[SHAPE_CAST]], %[[MEM_DESC]][%[[ROW_OFFSET]], %[[COL_OFFSET]]] : vector<1x32xf32>, !xegpu.mem_desc<32x32xf32>, index, index
+ // CHECK-DAG: %[[AFF0:.*]] = affine.apply #map()[%[[SGID]]]
+ // CHECK-DAG: %[[AFF1:.*]] = affine.apply #map1()[%[[SGID]]]
+ // CHECK-DAG: %[[AFF2:.*]] = affine.apply #map2()[%[[SGID]]]
+ // CHECK-DAG: %[[ROW:.*]] = arith.muli %[[AFF0]], %[[C1A:.*]] : index
+ // CHECK-DAG: %[[COL0:.*]] = arith.muli %[[AFF1:.*]], %[[C1B:.*]] : index
+ // CHECK-DAG: %[[COL1:.*]] = arith.muli %[[AFF2]], %[[C32A:.*]] : index
+ // CHECK-DAG: xegpu.store_matrix %[[CAST]], %[[MEM_DESC]][%[[ROW]], %[[COL0]], %[[COL1]]] : vector<1x1x32xf32>, !xegpu.mem_desc<1x32x32xf32>, index, index, index
// CHECK-DAG: gpu.barrier
- // CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MEM_DESC]][%[[C0:.*]], %[[COL_OFFSET]]] : !xegpu.mem_desc<32x32xf32>, index, index -> vector<32x32xf32>
- // CHECK-DAG: %[[CST_3:.*]] = arith.constant dense<0.000000e+00> : vector<32xf32>
- // CHECK-DAG: %[[FINAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD_SLM]], %[[CST_3]] [0] : vector<32x32xf32> to vector<32xf32>
- // CHECK-DAG: %[[SHAPE_CAST_FINAL:.*]] = vector.shape_cast %[[CST]] : vector<1x32xf32> to vector<32xf32>
- // CHECK-DAG: %{{.*}} = arith.addf %[[FINAL_REDUCE]], %[[SHAPE_CAST_FINAL]] : vector<32xf32>
+ // CHECK-DAG: %[[ROW_L:.*]] = arith.muli %[[AFF0]], %[[C1C:.*]] : index
+ // CHECK-DAG: %[[COL0_L:.*]] = arith.muli %[[AFF1]], %[[C0:.*]] : index
+ // CHECK-DAG: %[[COL1_L:.*]] = arith.muli %[[AFF2]], %[[C32B:.*]] : index
+ // CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MEM_DESC]][%[[ROW_L]], %[[COL0_L]], %[[COL1_L]]] : !xegpu.mem_desc<1x32x32xf32>, index, index, index -> vector<1x32x32xf32>
+ // CHECK-DAG: %[[CST_3:.*]] = arith.constant dense<0.000000e+00> : vector<1x32xf32>
+ // CHECK-DAG: %[[FINAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD_SLM]], %[[CST_3]] [1] : vector<1x32x32xf32> to vector<1x32xf32>
+ // CHECK-DAG: %[[ADD:.*]] = arith.addf %[[FINAL_REDUCE]], %[[CST]] : vector<1x32xf32>
// CHECK-DAG: gpu.return
%cst_3 = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>, dims = [1]>} dense<1.0> : vector<1x32xf32>
%offset = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>} dense<0> : vector<1x32x32xindex>
@@ -688,7 +686,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[DIV1:.*]] = arith.divui %[[SGID]], %[[C4:.*]] : index
// CHECK-DAG: %[[REM2:.*]] = arith.remui %[[DIV1]], %[[C8:.*]] : index
// CHECK-DAG: %[[MUL1:.*]] = arith.muli %[[REM2]], %[[C32:.*]] : index
- // CHECK-DAG: %[[MUL2:.*]] = arith.muli %[[REM1]], %[[C32:.*]] : index
+ // CHECK-DAG: %[[MUL2:.*]] = arith.muli %[[REM1]], %[[C32_0:.*]] : index
// CHECK-DAG: %[[REM3:.*]] = arith.remui %[[MUL1]], %[[C256:.*]] : index
// CHECK-DAG: %[[REM4:.*]] = arith.remui %[[MUL2]], %[[C128:.*]] : index
// CHECK-DAG: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]][%[[REM3]], %[[REM4]]] : memref<256x128xf32> -> !xegpu.tensor_desc<32x32xf32>
@@ -699,19 +697,18 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[ALLOCA:.*]] = memref.alloca() : memref<4096xi8, 3>
// CHECK-DAG: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<4096xi8, 3> -> !xegpu.mem_desc<8x128xf32>
// CHECK-DAG: %[[SGID2:.*]] = gpu.subgroup_id : index
- // CHECK-DAG: %[[AFFINE1:.*]] = affine.apply #map3()[%[[SGID2]]]
- // CHECK-DAG: %[[AFFINE2:.*]] = affine.apply #map4()[%[[SGID2]]]
- // CHECK-DAG: %[[MUL3:.*]] = arith.muli {{.*}}, %[[C1:.*]] : index
- // CHECK-DAG: %[[ROW_OFFSET:.*]] = arith.addi %[[C0:.*]], %[[MUL3]] : index
- // CHECK-DAG: %[[MUL4:.*]] = arith.muli {{.*}}, %[[C1:.*]] : index
- // CHECK-DAG: %[[ADD1:.*]] = arith.addi %[[C0:.*]], %[[MUL4]] : index
- // CHECK-DAG: %[[COL_OFFSET:.*]] = arith.muli %[[ADD1]], %[[C32:.*]] : index
+ // CHECK-DAG: %[[AFFINE1:.*]] = affine.apply {{#map[[:alnum:]_]*}}()[%[[SGID2]]]
+ // CHECK-DAG: %[[AFFINE2:.*]] = affine.apply {{#map[[:alnum:]_]*}}()[%[[SGID2]]]
+ // CHECK-DAG: %[[ROW_OFFSET:.*]] = arith.muli %[[AFFINE1]], %[[C1:.*]] : index
+ // CHECK-DAG: %[[COL_OFFSET:.*]] = arith.muli %[[AFFINE2]], %[[C32_1:.*]] : index
// CHECK-DAG: xegpu.store_matrix %[[SHAPE_CAST]], %[[MEM_DESC]][%[[ROW_OFFSET]], %[[COL_OFFSET]]] : vector<1x32xf32>, !xegpu.mem_desc<8x128xf32>, index, index
// CHECK-DAG: gpu.barrier
- // CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MEM_DESC]][%[[C0:.*]], %[[COL_OFFSET]]] : !xegpu.mem_desc<8x128xf32>, index, index -> vector<8x32xf32>
+ // CHECK-DAG: %[[ZERO_ROW:.*]] = arith.muli %[[AFFINE1]], %[[C0:.*]] : index
+ // CHECK-DAG: %[[COL_OFFSET2:.*]] = arith.muli %[[AFFINE2]], %[[C32_2:.*]] : index
+ // CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MEM_DESC]][%[[ZERO_ROW]], %[[COL_OFFSET2]]] : !xegpu.mem_desc<8x128xf32>, index, index -> vector<8x32xf32>
// CHECK-DAG: %[[CST_CROSS_SG_1:.*]] = arith.constant dense<0.000000e+00> : vector<32xf32>
// CHECK-DAG: %[[FINAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD_SLM]], %[[CST_CROSS_SG_1]] [0] : vector<8x32xf32> to vector<32xf32>
- // CHECK-DAG: arith.addf %[[FINAL_REDUCE]], %[[CST:.*]] : vector<32xf32>
+ // CHECK-DAG: arith.addf %[[FINAL_REDUCE]], %[[CST]] : vector<32xf32>
%cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>, dims = [0]>} dense<0.0> : vector<128xf32>
%tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
@@ -732,36 +729,38 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[LOAD:.*]] = xegpu.load %{{.*}}[%[[CST_0]]], %[[CST_1]] <{chunk_size = 1 : i64}> : memref<?xf32>, vector<1x1x32x32xindex>, vector<1x1x32x32xi1> -> vector<1x1x32x32xf32>
// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<0.000000e+00> : vector<1x1xf32>
// CHECK-DAG: %[[LOCAL_REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_2]] [2, 3] : vector<1x1x32x32xf32> to vector<1x1xf32>
- // CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x1xf32> to vector<1x1xf32>
+ // CHECK-DAG: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[LOCAL_REDUCE]] : vector<1x1xf32> to vector<1x1x1x1xf32>
// CHECK-DAG: %[[ALLOCA:.*]] = memref.alloca() : memref<256xi8, 3>
- // CHECK-DAG: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<256xi8, 3> -> !xegpu.mem_desc<16x4xf32>
+ // CHECK-DAG: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<256xi8, 3> -> !xegpu.mem_desc<2x2x4x4xf32>
// CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index
- // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
- // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
- // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
- // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index
- // CHECK-DAG: %[[AFFINE1:.*]] = affine.apply #map()[%[[SGID]]]
- // CHECK-DAG: %[[AFFINE2:.*]] = affine.apply #map1()[%[[SGID]]]
- // CHECK-DAG: %[[AFFINE3:.*]] = affine.apply #map5()[%[[SGID]]]
- // CHECK-DAG: %[[AFFINE4:.*]] = affine.apply #map6()[%[[SGID]]]
- // CHECK-DAG: %[[AFFINE5:.*]] = affine.apply #m...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/182178
More information about the Mlir-commits
mailing list