[Mlir-commits] [mlir] [MLIR][XeGPU] Sg to Lane distribute convert_layout redistributing broadcasted data (PR #215645)

[Jianhui Li]

================
@@ -1578,4 +1578,75 @@ gpu.func @convert_layout_partial_subgroup() {
   } : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
   gpu.return
 }
+
+// A convert_layout from a layout that broadcasts the value over two groups of
+// eight lanes to a layout that gives each of the first eight lanes a row of
+// it. Lane `i` extracts row `i % 8` of the copy it holds, so lanes [0, 8) end
+// up with the first column and lanes [8, 16) with the second one, and the
+// shuffles hand lane `i` the two values of its row.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_to_lane_distributed
+// CHECK:         %[[SRC:.*]] = arith.constant dense<1.000000e+00> : vector<8x1xf8E8M0FNU>
+// CHECK:         %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x1xf8E8M0FNU> to vector<8xf8E8M0FNU>
+// CHECK:         %[[BITS:.*]] = vector.bitcast %[[FLAT]] : vector<8xf8E8M0FNU> to vector<8xi8>
+// CHECK:         %[[LANE:.*]] = gpu.lane_id
+// CHECK:         %[[C8:.*]] = arith.constant 8 : index
+// CHECK:         %[[ROW:.*]] = arith.remui %[[LANE]], %[[C8]] : index
+// CHECK:         %[[ZERO:.*]] = arith.constant dense<0> : vector<2xi8>
+// CHECK:         %[[ELEM:.*]] = vector.extract %[[BITS]][%[[ROW]]] : i8 from vector<8xi8>
+// CHECK:         %[[WIDTH:.*]] = arith.constant 16 : i32
+// CHECK:         %[[ROW_I32:.*]] = arith.index_cast %[[ROW]] : index to i32
+// CHECK:         %[[SHUF0:.*]], %{{.*}} = gpu.shuffle idx %[[ELEM]], %[[ROW_I32]], %[[WIDTH]] : i8
----------------
Jianhui-Li wrote:

I don't think we need to shuffle since each lane already has whole data. At least from this particular test case. 

https://github.com/llvm/llvm-project/pull/215645