Add check for output layout

Author: nbpatel

mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

@@ -34,35 +34,6 @@ using namespace mlir;
 
 namespace {
 
-bool isDistributable(ArrayRef<int64_t> sgLayout, ArrayRef<int64_t> sgData,
-                     ArrayRef<int64_t> wgShape) {
-  // Check rank consistency
-  if (sgLayout.size() != sgData.size() || sgLayout.size() != wgShape.size())
-    return false;
-
-  for (size_t i = 0; i < sgLayout.size(); ++i) {
-    int64_t subgroupCount = sgLayout[i];
-    int64_t subgroupData = sgData[i];
-    int64_t shape = wgShape[i];
-
-    // Each subgroup must have positive data size
-    if (subgroupData <= 0 || subgroupCount <= 0 || shape <= 0)
-      return false;
-
-    // Total data assigned to all subgroups in this dimension
-    int64_t totalSubgroupData = subgroupCount * subgroupData;
-
-    // Subgroups must not collectively exceed the shape
-    if (totalSubgroupData > shape)
-      return false;
-
-    // Each subgroup's data must evenly divide the shape
-    if (shape % subgroupData != 0)
-      return false;
-  }
-  return true;
-}
-
 static std::pair<SmallVector<int64_t>, int>
 getSgShapeAndCount(ArrayRef<int64_t> shape, xegpu::LayoutAttr layout) {
   int count = 1;
@@ -393,7 +364,7 @@ struct WgToSgVectorBroadcastOp
     else
       return failure();
 
-    if (!isDistributable(sgLayout, sgShape, wgShape))
+    if (!xegpu::XeGPUDialect::isEvenlyDistributable(wgShape, layout))
       return failure();
 
     // Check if the srcShape has unit dim in dimensions being broadcasted,