Add CHECKS

Author: nbpatel

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

@@ -338,8 +338,6 @@ struct WgToSgVectorBroadcastOp
                   ConversionPatternRewriter &rewriter) const override {
     VectorType resultType = op.getResult().getType();
     ArrayRef<int64_t> wgShape = resultType.getShape();
-    if (!resultType)
-      return failure();
 
     xegpu::LayoutAttr layout = xegpu::getLayoutAttr(op.getResult());
     if (!layout || !layout.getSgLayout())
@@ -348,17 +346,17 @@ struct WgToSgVectorBroadcastOp
     SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;
     VectorType newResultType =
         VectorType::get(sgShape, resultType.getElementType());
-    SmallVector<Value> newBroadcasts;
 
+    SmallVector<Value> newBroadcastOps;
     for (size_t i = 0; i < adaptor.getOperands().front().size(); ++i) {
       auto newBroadcast = rewriter.create<vector::BroadcastOp>(
           op.getLoc(), newResultType, adaptor.getOperands().front()[i]);
       xegpu::setLayoutAttr(newBroadcast->getResult(0),
                            layout.dropSgLayoutAndData());
-      newBroadcasts.push_back(newBroadcast.getResult());
+      newBroadcastOps.push_back(newBroadcast.getResult());
     }
 
-    rewriter.replaceOpWithMultiple(op, {newBroadcasts});
+    rewriter.replaceOpWithMultiple(op, {newBroadcastOps});
     return success();
   }
 };
@@ -556,11 +554,7 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
 
   target.addDynamicallyLegalOp<vector::BroadcastOp>(
       [=](vector::BroadcastOp op) -> bool {
-        auto resultType = dyn_cast<VectorType>(op.getResult().getType());
-        if (!resultType)
-          return true;
-        xegpu::LayoutAttr layout = xegpu::getLayoutAttr(op.getResult());
-        return isLegal(layout);
+        return isLegal(xegpu::getLayoutAttr(op.getResult()));
       });
 
   target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(

mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir

@@ -111,6 +111,10 @@ gpu.module @test_round_robin_assignment {
     %load =  xegpu.load_nd %tdesc
       : !xegpu.tensor_desc<24x1xf32, #xegpu.layout<sg_layout = [4, 1], sg_data = [2, 1], lane_layout = [2, 1], lane_data = [1, 1]>>
       -> vector<24x1xf32>
+    // CHECK-COUNT-3: vector.broadcast {{.*}}
+    // CHECK-SAME-COUNT-3: {layout_result_0 = #xegpu.layout<lane_layout = [2, 1], lane_data = [1, 1]>}
+    // CHECK-SAME-COUNT-3: : vector<2x1xf32> to vector<2x4xf32>
+    // CHECK-NOT: vector.broadcast
     %broadcast = vector.broadcast %load 
       {layout_result_0 = #xegpu.layout<sg_layout = [4, 1], sg_data = [2, 4], lane_layout = [2, 1], lane_data = [1, 1]>}
       : vector<24x1xf32> to vector<24x8xf32>

mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir

@@ -170,15 +170,16 @@ gpu.func @test_dpas_no_sg_data(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
     gpu.return
   }
 
-
-// CHECK-LABEL: test_broadcast
+  // CHECK-LABEL: test_broadcast
   // CHECK-SAME: %[[ARG_0:.*]]: memref<24x1xf32>
   gpu.func @test_broadcast(%src: memref<24x1xf32>) {
     %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x1xf32>
       -> !xegpu.tensor_desc<24x1xf32, #xegpu.layout<sg_layout = [2, 1], sg_data = [12, 1], lane_layout = [2, 1], lane_data = [1, 1]>>
     %load =  xegpu.load_nd %tdesc
       : !xegpu.tensor_desc<24x1xf32, #xegpu.layout<sg_layout = [2, 1], sg_data = [12, 1], lane_layout = [2, 1], lane_data = [1, 1]>>
       -> vector<24x1xf32>
+    // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 1], lane_data = [1, 1]>}
+    // CHECK-SAME: : vector<12x1xf32> to vector<12x8xf32>
     %broadcast = vector.broadcast %load 
       {layout_result_0 = #xegpu.layout<sg_layout = [2, 1], sg_data = [12, 8], lane_layout = [2, 1], lane_data = [1, 1]>}
       : vector<24x1xf32> to vector<24x8xf32>