Address comments

Author: nbpatel

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

@@ -148,15 +148,17 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
     if (!layout.getEffectiveInstDataAsInt().empty()) {
       SmallVector<int64_t> instData = layout.getEffectiveInstDataAsInt();
       // Remove leading unit dimensions from inst_data
+      // For example, if the inst_data is [1, 1, 32]
+      // it will pass [32] as the unroll/blocking size.
       // Skip it for xegpu nd ops since it will be 2D
       Operation *definingOp = value.getDefiningOp();
       bool skipLeadingUnitDimRemoval =
           definingOp &&
           (isa<xegpu::CreateNdDescOp, xegpu::LoadNdOp, xegpu::DpasOp,
                xegpu::StoreNdOp, xegpu::PrefetchNdOp>(definingOp));
       if (!skipLeadingUnitDimRemoval) {
-        while (!instData.empty() && instData.front() == 1)
-          instData.erase(instData.begin());
+        auto it = llvm::find_if(instData, [](auto val) { return val != 1; });
+        instData.erase(instData.begin(), it);
       }
       return instData;
     }

mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp

@@ -255,8 +255,6 @@ xegpu::extractVectorsWithShapeFromValue(OpBuilder &builder, Location loc,
             1);
   std::copy(shape.begin(), shape.end(), adjustedTargetShape.begin() + rankDiff);
 
-  int64_t adjustedTargetShapeRank = adjustedTargetShape.size();
-
   SmallVector<Value> result;
   for (SmallVector<int64_t> offsets :
        StaticTileOffsetRange(srcShape, adjustedTargetShape)) {
@@ -265,7 +263,7 @@ xegpu::extractVectorsWithShapeFromValue(OpBuilder &builder, Location loc,
         builder, loc, value, offsets, adjustedTargetShape, staticStrides);
 
     // Reshape to remove leading unit dims if needed
-    if (adjustedTargetShapeRank > targetShapeRank) {
+    if (srcShapeRank > targetShapeRank) {
       auto targetTy = VectorType::get(shape, vecTy.getElementType());
       slice = builder.create<vector::ShapeCastOp>(loc, targetTy, slice);
     }

mlir/test/Dialect/XeGPU/xegpu-blocking.mlir

@@ -685,7 +685,7 @@ gpu.module @test_kernel {
 
 // -----
 gpu.module @test_kernel {
-  // CHECK-LABEL: load_gather
+  // CHECK-LABEL: remove_unit_dim_inst_data
   // CHECK-SAME: [[arg0:%.+]]: ui64
   // CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<1x1x32xf32>
   // CHECK: [[cst_0:%.+]] = arith.constant dense<true> : vector<16xi1>
@@ -695,7 +695,7 @@ gpu.module @test_kernel {
   // CHECK: [[ld_1:%.+]] = xegpu.load [[arg0]][[[cst_2]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
   // CHECK: [[ins_0:%.+]] = vector.insert_strided_slice [[ld_0]], [[cst]] {offsets = [0, 0, 0], strides = [1]} : vector<16xf32> into vector<1x1x32xf32>
   // CHECK: [[ins_1:%.+]] = vector.insert_strided_slice [[ld_1]], [[ins_0]] {offsets = [0, 0, 16], strides = [1]} : vector<16xf32> into vector<1x1x32xf32>
-  gpu.func @load_gather(%src: ui64) -> vector<1x1x32xf32> {
+  gpu.func @remove_unit_dim_inst_data(%src: ui64) -> vector<1x1x32xf32> {
       %cst = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 16]>} dense<[[
       [0,   8,  16,  24,  32,  40,  48,  56,
       64,  72,  80,  88,  96, 104, 112, 120,