Use min shape for dist_unit

nbpatel · nbpatel · commit 4612f646b6f8 · 2025-05-19T20:43:02.000Z
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -155,7 +155,7 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
     SmallVector<int64_t> distUnitShape(sgLayout.size());
     SmallVector<Value> localOffset(sgLayout.size());
     for (size_t i = 0; i < sgLayout.size(); i++) {
-      distUnitShape[i] = sgLayout[i] * sgShape[i];
+      distUnitShape[i] = std::min(sgLayout[i] * sgShape[i], wgShape[i]);
       localOffset[i] =
           rewriter.createOrFold<index::MulOp>(loc, sgIds[i], sgDataDim[i]);
     }

Original file line number	Diff line number	Diff line change
`@@ -155,7 +155,7 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {`
`155`	`155`	`SmallVector<int64_t> distUnitShape(sgLayout.size());`
`156`	`156`	`SmallVector<Value> localOffset(sgLayout.size());`
`157`	`157`	`for (size_t i = 0; i < sgLayout.size(); i++) {`
`158`		`- distUnitShape[i] = sgLayout[i] * sgShape[i];`
	`158`	`+ distUnitShape[i] = std::min(sgLayout[i] * sgShape[i], wgShape[i]);`
`159`	`159`	`localOffset[i] =`
`160`	`160`	`rewriter.createOrFold<index::MulOp>(loc, sgIds[i], sgDataDim[i]);`
`161`	`161`	`}`