[Mlir-commits] [mlir] [mlir][XeGPU] Add an SLM round-trip fallback for sg-to-lane convert_layout (PR #227884)

[Charitha Saumya]

================
@@ -1283,6 +1283,57 @@ gpu.func @convert_layout_repack_innermost_lane_data_3d() {
 }
 }
 
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL:     gpu.func @convert_layout_via_slm_lane_layout
+// CHECK-NOT:         xegpu.convert_layout
+// CHECK:             %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<16x16xf32> to vector<16x1xf32>
+// CHECK:             %[[SLM:.*]] = memref.alloca() : memref<4096xi8, 3>
+// CHECK:             %[[MD:.*]] = xegpu.create_mem_desc %[[SLM]] : memref<4096xi8, 3> -> !xegpu.mem_desc<64x16xf32>
+// CHECK:             %[[SGID:.*]] = gpu.subgroup_id : index
+// CHECK:             %[[ROWS:.*]] = arith.constant 16 : index
+// CHECK:             %[[BASE:.*]] = arith.muli %[[SGID]], %[[ROWS]] : index
+// CHECK:             gpu.lane_id
+// CHECK-COUNT-16:    xegpu.store_matrix %{{.*}}, %[[MD]]{{\[}}%{{.*}}, %{{.*}}] : vector<1x1xf32>, !xegpu.mem_desc<64x16xf32>, index, index
+// CHECK:             xegpu.fence memory_kind = slm, fence_scope = workgroup
+// CHECK-COUNT-16:    xegpu.load_matrix %[[MD]]{{\[}}%{{.*}}, %{{.*}}] : !xegpu.mem_desc<64x16xf32>, index, index -> vector<1x1xf32>
+// CHECK:             builtin.unrealized_conversion_cast %{{.*}} : vector<1x16xf32> to vector<16x16xf32>
+gpu.func @convert_layout_via_slm_lane_layout() kernel attributes {known_block_size = array<i32: 64, 1, 1>} {
+  %src = "test.some_op"() : () -> vector<16x16xf32>
+  %cvt = xegpu.convert_layout %src
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+      target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
----------------
charithaintc wrote:

what happens if lane_data spans multiple rows (non contagious ownership i.e. lane data = [2, 1]). This require emitting multiple llvm.loads. 

if this case is not handled consider adding a negative test and guard for it in the code. 

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