[Mlir-commits] [mlir] [MLIR][XeGPU] Distribute `vector.step` with sliced layout (PR #182010)

[Jianhui Li]

================
@@ -1132,4 +1132,31 @@ gpu.func
     gpu.return
   }
 
+// CHECK-LABEL: gpu.func @vector_step_slice
+// CHECK:         (%[[SG_ID:[0-9a-zA-Z]+]]: index) {
+// CHECK-NEXT:    %[[SG_ID_IN_SLICED_DIM:.*]] = affine.apply #map()[%[[SG_ID]]]
+// CHECK-NEXT:    %[[SG_ID_IN_SLICED_DIM_VEC:.*]] = vector.broadcast %[[SG_ID_IN_SLICED_DIM]] : index to vector<1xindex>
----------------
Jianhui-Li wrote:

Why we use broadcast here? I think step just distributed to step_0 + lane_id. In case all lane shares one step value, all lanes map to lane 0 so get step_0. 
So : 
%5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>} : vector<**1**xindex>
=> 
%5 = vector.step : vector<1xindex>


%5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>} : vector<**16**xindex>
=> 
%5 = vector.step : vector<1xindex>
%6 = %5 + %lane_id


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