shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[MLIR][XeGPU] Add support for convert layout with scalar in Sg to WI distribution (#189721)

Browse Source
This commit is contained in:
Nishant Patel 2026-04-01 12:05:32 -07:00 committed by GitHub
parent 35ceec6a78
commit 150aa6f2d3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
View File

@ -1497,14 +1497,21 @@ struct SgToWiConvertLayout
ConversionPatternRewriter &rewriter) const override {
auto inputLayout = op.getInputLayoutAttr();
auto targetLayout = op.getTargetLayoutAttr();
auto resShape = cast<VectorType>(op.getResult().getType()).getShape();
SmallVector<int64_t> resShapeVec(resShape.begin(), resShape.end());
Type valType = op.getResult().getType();
if (valType.isIntOrFloat()) {
rewriter.replaceOp(op, op.getSource());
return success();
}
auto resShape = cast<VectorType>(valType).getShape();
SmallVector<int64_t> resShapeVec(resShape.begin(), resShape.end());
if (!inputLayout.isCompatibleWith(targetLayout, resShapeVec,
xegpu::LayoutKind::Lane)) {
return rewriter.notifyMatchFailure(
op, "lowering incompatible convert_layout not yet supported");
}
rewriter.replaceOp(op, adaptor.getSource());
return success();
}

15
mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
View File

@ -858,6 +858,21 @@ gpu.func @convert_layout_removed_when_compatible() {
}
}
// -----
gpu.module @xevm_module {
// CHECK-LABEL: gpu.func @convert_layout_scalar
// CHECK-NOT: xegpu.convert_layout
gpu.func @convert_layout_scalar() {
%0 = "some_op"() : () -> f32
%1 = xegpu.convert_layout %0
<{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16], lane_data = [1]>, dims = [0]>,
target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16], lane_data = [1]>, dims = [0]>}>
: f32
"some_use"(%1) : (f32) -> ()
gpu.return
}
}
// -----
// load_matrix and store_matrix with coordinate computation (offsets [0,0])
gpu.module @xevm_module {