[Mlir-commits] [mlir] [MLIR][XeGPU] Clean up the temporary layout usage in XeGPU test (PR #195739)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon May 4 13:52:08 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
This PR cleans up the XeGPU test to remove the temporary layout usage. All distribution and unrolling tests now don't use temporary layout from the operation and TensorDescriptor, since the recovery process won't honor the temporary layout and only depends on the anchor layout.
It also refactors the layout function implementation by removing recursive loops in getDistributeLayoutAttr(), and fixes two issues surfaced from the test clean up: adding layout recovery support for Extract/Insert op and tensor descriptor type.
---
Patch is 223.71 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/195739.diff
15 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h (+12)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp (+12-3)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+150-42)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+3)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+7-2)
- (modified) mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp (+5-5)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+3-3)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir (+2-2)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+2-2)
- (modified) mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir (+106-217)
- (modified) mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir (+84-95)
- (modified) mlir/test/Dialect/XeGPU/xegpu-blocking.mlir (+45-44)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir (+75-85)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir (+102-58)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir (+314-171)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 2dd8d9f610faf..bac8f413acd40 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -117,6 +117,18 @@ inferInsertStridedSliceSourceLayout(DistributeLayoutAttr resLayout,
ArrayRef<int64_t> resShape,
ArrayRef<int64_t> srcShape);
+/// Infers the source layout attribute for an insert operation.
+/// using same logic as inferInsertStridedSliceSourceLayout
+DistributeLayoutAttr inferInsertSourceLayout(DistributeLayoutAttr resLayout,
+ ArrayRef<int64_t> resShape,
+ ArrayRef<int64_t> srcShape);
+
+/// Infers the source layout attribute for an extract operation. Adds
+/// leading dimensions to the source layout to match the source shape size.
+DistributeLayoutAttr inferExtractSourceLayout(DistributeLayoutAttr resLayout,
+ ArrayRef<int64_t> resShape,
+ ArrayRef<int64_t> srcShape);
+
/// Infers the layout attribute for mask and offset operand for Chunked load
/// and store, given the anchor layout attribute for the value being load/store.
DistributeLayoutAttr
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 4fe15c625ea49..da48bff7b5048 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -419,11 +419,20 @@ LogicalResult LoadNdOp::verify() {
}
}
+ // Handle array_length. Two result shape conventions are accepted:
+ // * Legacy: leading array_length dimension prepended, e.g. descriptor
+ // 16x16 with array_length=2 -> [2, 16, 16].
+ // * Stacked 2D: array blocks stacked along the non-FCD (first) dimension,
+ // e.g. descriptor 16x16 with array_length=2 -> [32, 16].
auto array_len = tdescTy.getArrayLength();
- if (array_len > 1)
- tdescShape.insert(tdescShape.begin(), array_len);
+ SmallVector<int64_t> stackedShape(tdescShape);
+ SmallVector<int64_t> prependedShape(tdescShape);
+ if (array_len > 1 && !tdescShape.empty()) {
+ stackedShape[0] *= array_len;
+ prependedShape.insert(prependedShape.begin(), array_len);
+ }
- if (tdescShape != valueShape)
+ if (valueShape != stackedShape && valueShape != prependedShape)
return emitOpError() << "Result shape " << makeString(valueShape)
<< " is not consistent with tensor descriptor "
<< tdescTy;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d3925c40f9123..b35224f032d42 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -33,20 +33,6 @@
using namespace mlir;
-void xegpu::recoverTemporaryLayoutsDeprecated(Operation *op) {
- op->walk([&](Operation *nestOp) {
- for (OpOperand &opr : nestOp->getOpOperands()) {
- auto layout = getDistributeLayoutAttr(opr.get());
- setDistributeLayoutAttr(opr, layout);
- }
-
- for (OpResult result : nestOp->getOpResults()) {
- auto layout = getDistributeLayoutAttr(result);
- setDistributeLayoutAttr(result, layout);
- }
- });
-}
-
SmallVector<NamedAttribute>
xegpu::dropSgLayoutAndDataOnAttrs(ArrayRef<NamedAttribute> attrs) {
SmallVector<NamedAttribute> out;
@@ -83,6 +69,18 @@ xegpu::dropInstDataOnAttrs(ArrayRef<NamedAttribute> attrs) {
return out;
}
+// Sets the layout on a TensorDesc value by updating its type to include
+// the given layout, if the type does not already have a layout attached.
+static void setTensorDescLayout(Value val, xegpu::DistributeLayoutAttr layout) {
+ auto tensorDescTy = dyn_cast<xegpu::TensorDescType>(val.getType());
+ if (!tensorDescTy || tensorDescTy.getLayoutAttr())
+ return;
+ auto typeWithLayout = xegpu::TensorDescType::get(
+ tensorDescTy.getContext(), tensorDescTy.getShape(),
+ tensorDescTy.getElementType(), tensorDescTy.getEncoding(), layout);
+ val.setType(typeWithLayout);
+}
+
// the walkRegionBackward() is a recursive function
// the input rootOp is the function operation, which is also a region op.
// it recursively processes the region op in reverse topological order.
@@ -136,24 +134,23 @@ static void propagateResultsToRegularOperands(Operation *op) {
xegpu::DistributeLayoutAttr resLayout = getLayoutFromUsePoints(result);
Type resultType = result.getType();
- // recover layout for tensor Descriptor type, which is a special case since
- // its layout is not stored as an attribute but encoded in the type itself.
- // For vector type, we attach the layout as an attribute to op.
- if (auto tensorDescTy = dyn_cast<xegpu::TensorDescType>(resultType)) {
- auto layout = tensorDescTy.getLayoutAttr();
- if (!layout) {
- auto typeWithLayout = xegpu::TensorDescType::get(
- tensorDescTy.getContext(), tensorDescTy.getShape(),
- tensorDescTy.getElementType(), tensorDescTy.getEncoding(), resLayout);
- result.setType(typeWithLayout);
- }
- }
- if (isa<VectorType>(resultType) && resLayout)
+ if (!resLayout)
+ return;
+
+ // Recover layout for TensorDesc type results by updating the type to include
+ // the layout. For vector type
+ if (isa<xegpu::TensorDescType>(resultType))
+ setTensorDescLayout(result, resLayout);
+
+ // Recover layout for vector type results, or for multi-reduction ops which
+ // may reduce to a scalar that still needs a layout.
+ if (isa<VectorType>(resultType) || isa<vector::MultiDimReductionOp>(op))
xegpu::setTemporaryLayout(result, resLayout);
for (OpOperand &opr : op->getOpOperands()) {
xegpu::DistributeLayoutAttr operandLayout =
xegpu::inferSourceLayoutFromResult(opr, resLayout);
+ // Recover layout for vector operands
if (isa<VectorType>(opr.get().getType()) && operandLayout)
xegpu::setTemporaryLayout(opr, operandLayout);
}
@@ -189,9 +186,14 @@ static void propagateRegionResultsToYieldOperands(
if (successor.isParent()) {
// For parent successor, get layout from external use points of the
// parent op's results.
- layout = getLayoutFromUsePoints(regionBranchOp->getResult(i));
- if (layout)
- xegpu::setTemporaryLayout(regionBranchOp->getResult(i), layout);
+ auto regionResult = regionBranchOp->getResult(i);
+ layout = getLayoutFromUsePoints(regionResult);
+ if (layout) {
+ // set layout for the region op, like scf.loop
+ xegpu::setTemporaryLayout(regionResult, layout);
+ if (isa<xegpu::TensorDescType>(regionResult.getType()))
+ setTensorDescLayout(regionResult, layout);
+ }
} else {
// For region successor, get layout from the target region's block
// arg use points (e.g., "before/cond" region args for scf.while
@@ -200,7 +202,10 @@ static void propagateRegionResultsToYieldOperands(
}
if (!layout)
continue;
- if (isa<VectorType>(succOps[i].getType()))
+ auto operandType = succOps[i].getType();
+ if (isa<VectorType>(operandType) ||
+ dyn_cast<xegpu::TensorDescType>(operandType))
+ // recover layout for yield op operands
xegpu::setTemporaryLayout(yieldOp->getOpOperand(beginIdx + i), layout);
}
}
@@ -226,17 +231,10 @@ static void propagateRegionArgsToInits(mlir::RegionBranchOpInterface regionOp) {
continue;
// Recover layout for tensor_desc block args by updating the type.
- if (auto tensorDescTy =
- dyn_cast<xegpu::TensorDescType>(regionArg.getType())) {
- if (!tensorDescTy.getLayoutAttr()) {
- auto typeWithLayout = xegpu::TensorDescType::get(
- tensorDescTy.getContext(), tensorDescTy.getShape(),
- tensorDescTy.getElementType(), tensorDescTy.getEncoding(),
- layout);
- regionArg.setType(typeWithLayout);
- }
- }
+ if (isa<xegpu::TensorDescType>(regionArg.getType()))
+ setTensorDescLayout(regionArg, layout);
+ // Recover layout for region op operands, like scf.for's init operands.
// Find all predecessor values that flow into this block argument.
SmallVector<Value> predValues;
regionOp.getPredecessorValues(regionSuccessor, inputIdx, predValues);
@@ -293,7 +291,7 @@ bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
}
});
};
-
+ removeTemporaryLayoutAttrs(rootOp);
rootOp->walk([&](func::FuncOp func) {
processFunc(func.getBody(), func.getSymName());
});
@@ -484,6 +482,90 @@ xegpu::DistributeLayoutAttr xegpu::inferInsertStridedSliceSourceLayout(
return resLayout;
}
+/// Infers the source layout attribute for an insert operation
+/// given the result layout attribute, result shape, and source shape. Removes
+/// leading dimensions from the result layout to match the source shape size.
+// TODO: add propagation support for insert op
+xegpu::DistributeLayoutAttr
+xegpu::inferInsertSourceLayout(xegpu::DistributeLayoutAttr resLayout,
+ ArrayRef<int64_t> resShape,
+ ArrayRef<int64_t> srcShape) {
+
+ int srcShapeSize = srcShape.size();
+ int resShapeSize = resShape.size();
+ int dimDiff = resShapeSize - srcShapeSize;
+
+ if (dimDiff > 0) {
+ // assert that the leading dimensions being sliced off are not distributed
+ // (i.e. sg_layout and lane_layout for those dimensions are all 1)
+ auto resSgLayout = resLayout.getEffectiveSgLayoutAsInt();
+ auto resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
+ for (int i = 0; i < dimDiff; i++) {
+ assert((resSgLayout.size() == 0 || resSgLayout[i] == 1) &&
+ (resLaneLayout.size() == 0 || resLaneLayout[i] == 1) &&
+ "Leading dimensions being sliced off must not be distributed");
+ }
+ return resLayout.dropDims(llvm::to_vector(llvm::seq<int64_t>(0, dimDiff)));
+ }
+ return resLayout;
+}
+
+/// Infers the source layout attribute for extract operation
+/// given the result layout attribute, result shape, and source shape. Adds
+/// leading dimensions to the source layout to match the source shape size.
+// TODO: add layout attribute interface: expandDims() and use it here.
+// TODO: add propagation support for extract op
+xegpu::DistributeLayoutAttr
+xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
+ ArrayRef<int64_t> resShape,
+ ArrayRef<int64_t> srcShape) {
+
+ int srcShapeSize = srcShape.size();
+ int resShapeSize = resShape.size();
+ int dimDiff = srcShapeSize - resShapeSize;
+ auto context = resLayout.getContext();
+ // construct the source layout by adding unit dimensions to the front of
+ // result layout
+
+ SmallVector<int64_t> sgLayout(srcShapeSize, 1);
+ SmallVector<int64_t> sgData(srcShapeSize, 1);
+ SmallVector<int64_t> instData(srcShapeSize, 1);
+ SmallVector<int64_t> laneLayout(srcShapeSize, 1);
+ SmallVector<int64_t> laneData(srcShapeSize, 1);
+
+ if (dimDiff > 0) {
+ auto resSgLayout = resLayout.getEffectiveSgLayoutAsInt();
+ auto resSgData = resLayout.getEffectiveSgDataAsInt();
+ auto resInstData = resLayout.getEffectiveInstDataAsInt();
+ auto resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
+ auto resLaneData = resLayout.getEffectiveLaneDataAsInt();
+
+ for (int i = 0; i < resShapeSize; i++) {
+ sgLayout[dimDiff + i] = (resSgLayout.size() == 0) ? 1 : resSgLayout[i];
+ sgData[dimDiff + i] = (resSgData.size() == 0) ? 1 : resSgData[i];
+ instData[dimDiff + i] = (resInstData.size() == 0) ? 1 : resInstData[i];
+ laneLayout[dimDiff + i] =
+ (resLaneLayout.size() == 0) ? 1 : resLaneLayout[i];
+ laneData[dimDiff + i] = (resLaneData.size() == 0) ? 1 : resLaneData[i];
+ }
+
+ auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
+ if (v.empty())
+ return DenseI32ArrayAttr();
+ SmallVector<int32_t> v32(v.begin(), v.end());
+ return DenseI32ArrayAttr::get(context, v32);
+ };
+ auto srcLayout = xegpu::LayoutAttr::get(
+ context, resSgLayout.empty() ? nullptr : toAttr(sgLayout),
+ resSgData.empty() ? nullptr : toAttr(sgData),
+ resInstData.empty() ? nullptr : toAttr(instData),
+ resLaneLayout.empty() ? nullptr : toAttr(laneLayout),
+ resLaneData.empty() ? nullptr : toAttr(laneData), nullptr);
+ return srcLayout;
+ }
+ return resLayout;
+}
+
/// Infers the source layout attribute for a shape cast operation given the
/// result layout attribute, result shape, and source shape.
xegpu::DistributeLayoutAttr
@@ -1587,6 +1669,32 @@ xegpu::inferSourceLayoutFromResult(OpOperand &operand,
return resLayout;
}
+ // For vector::Insert Op, infer source layout from result layout using
+ // shapes.
+ if (auto insert = dyn_cast<vector::InsertOp>(op)) {
+ VectorType resVecTy = dyn_cast<VectorType>(insert.getResult().getType());
+ VectorType valueToStoreTy =
+ dyn_cast<VectorType>(insert.getValueToStore().getType());
+
+ if (idx == 0) {
+ return xegpu::inferInsertSourceLayout(resLayout, resVecTy.getShape(),
+ valueToStoreTy.getShape());
+ }
+ if (idx == 1)
+ return resLayout;
+ }
+
+ // For vector::Extract Op, infer source layout from result layout using
+ // shapes.
+ if (auto extract = dyn_cast<vector::ExtractOp>(op)) {
+ VectorType srcVecTy = dyn_cast<VectorType>(extract.getSource().getType());
+ VectorType resVecTy = dyn_cast<VectorType>(extract.getResult().getType());
+ if (!srcVecTy || !resVecTy)
+ return nullptr;
+ return xegpu::inferExtractSourceLayout(resLayout, resVecTy.getShape(),
+ srcVecTy.getShape());
+ }
+
// For vector::TransposeOp, infer source layout from result layout using
// permutation.
if (auto transpose = dyn_cast<vector::TransposeOp>(op)) {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index a5776ebce2e95..6613af2dfc164 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1445,6 +1445,9 @@ LogicalResult ResolveLayoutConflicts::run() {
return WalkResult::advance();
});
+ DBGS() << "IR after resolving layout conflicts:\n";
+ parentOp->dump();
+
return r.wasInterrupted() ? failure() : success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 8aa0758943cd1..af82effb9d379 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1562,7 +1562,6 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
target.addDynamicallyLegalOp<vector::ShapeCastOp, vector::StepOp,
vector::TransposeOp, vector::BroadcastOp,
- vector::MultiDimReductionOp,
vector::ConstantMaskOp, vector::CreateMaskOp>(
[=](Operation *op) -> bool {
// Check for either a SliceAttr or LayoutAttr on the result.
@@ -1570,7 +1569,13 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
xegpu::getTemporaryLayout(dyn_cast<OpResult>(op->getResult(0)));
return isLegal(layout);
});
-
+ target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
+ [=](Operation *op) -> bool {
+ // Check operand since the result maybe scalar not bearing layout..
+ auto layout =
+ xegpu::getTemporaryLayout(dyn_cast<vector::MultiDimReductionOp>(op)->getOpOperand(0));
+ return isLegal(layout);
+ });
target.addDynamicallyLegalOp<xegpu::LoadGatherOp>(
[=](xegpu::LoadGatherOp op) -> bool {
auto layout = op.getLayoutAttr();
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 2d1ce6eea17aa..12eb553b3bddc 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -136,10 +136,6 @@ xegpu::DistributeLayoutAttr xegpu::getDistributeLayoutAttr(const Value value) {
if (!value)
return nullptr;
- if (auto tdescTy =
- dyn_cast_if_present<xegpu::TensorDescType>(value.getType()))
- return tdescTy.getLayoutAttr();
-
if (auto result = dyn_cast<OpResult>(value)) {
Operation *defOp = result.getDefiningOp();
assert(defOp && "result must have a defining op");
@@ -162,10 +158,14 @@ xegpu::DistributeLayoutAttr xegpu::getDistributeLayoutAttr(const Value value) {
if (auto loop = dyn_cast_if_present<LoopLikeOpInterface>(parentOp)) {
OpOperand *tiedInit = loop.getTiedLoopInit(arg);
if (tiedInit)
- return getDistributeLayoutAttr(tiedInit->get());
+ return getTemporaryLayout(*tiedInit);
}
}
+ if (auto tdescTy =
+ dyn_cast_if_present<xegpu::TensorDescType>(value.getType()))
+ return tdescTy.getLayoutAttr();
+
return nullptr;
}
xegpu::DistributeLayoutAttr
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 6448db93c3f40..0d73985502e3f 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -248,8 +248,8 @@ func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
%cst_small = arith.constant dense<1> : vector<4x64xi8>
%cst_large = arith.constant dense<0> : vector<8x64xi8>
%insert = vector.insert_strided_slice %cst_small, %cst_large {offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
- %tdesc = xegpu.create_nd_tdesc %arg0 : memref<8x64xi8> -> !xegpu.tensor_desc<8x64xi8, #xegpu.layout<inst_data = [8, 64]>>
- xegpu.store_nd %insert, %tdesc[0, 0] <{layout = #xegpu.layout<inst_data = [8, 64]>}>: vector<8x64xi8>, !xegpu.tensor_desc<8x64xi8, #xegpu.layout<inst_data = [8, 64]>>
+ %tdesc = xegpu.create_nd_tdesc %arg0 : memref<8x64xi8> -> !xegpu.tensor_desc<8x64xi8>
+ xegpu.store_nd %insert, %tdesc[0, 0] <{layout = #xegpu.layout<inst_data = [8, 64]>}>: vector<8x64xi8>, !xegpu.tensor_desc<8x64xi8>
return
}
}
@@ -333,7 +333,7 @@ func.func @vector_shape_cast_expand_and_merge(%arg0: memref<256xf16>, %arg1: mem
%4 = vector.shape_cast %2 : vector<2x4x32xf16> to vector<1x256xf16>
%5 = vector.shape_cast %4 : vector<1x256xf16> to vector<256xf16>
- xegpu.store %5, %arg1[%0], %cst <{layout = #xegpu.layout<inst_data = [32] >}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
+ xegpu.store %5, %arg1[%0], %cst <{layout = #xegpu.layout<inst_data = [32]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
return
}
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index d4ad9087149c1..2c28073ee1c01 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -133,9 +133,9 @@ gpu.module @test {
%tdesc_src = xegpu.create_nd_tdesc %src : memref<32x64xf32> -> !xegpu.tensor_desc<32x64xf32>
%load = xegpu.load_nd %tdesc_src[0, 0] : !xegpu.tensor_desc<32x64xf32> -> vector<32x64xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [1] : vector<32x64xf32> to vector<32xf32>
- %tdesc_dst = xegpu.create_nd_tdesc %dst : memref<32xf32> -> !xegpu.tensor_desc<32xf32, #xegpu.layout<sg_layout = [32], sg_data = [1]>>
+ %tdesc_dst = xegpu.create_nd_tdesc %dst : memref<32xf32> -> !xegpu.tensor_desc<32xf32>
xegpu.store_nd %reduce, %tdesc_dst[0] <{layout = #xegpu.layout<sg_layout = [32], sg_data = [1]>}>
- : vector<32xf32>, !xegpu.tensor_desc<32xf32, #xegpu.layout<sg_layout = [32], sg_data = [1]>>
+ : vector<32xf32>, !xegpu.tensor_desc<32xf32>
gpu.r...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/195739
More information about the Mlir-commits
mailing list