924 lines · plain
1// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=pvc' -test-xegpu-sg-distribute \2// RUN: -allow-unregistered-dialect -canonicalize -cse %s | FileCheck %s3gpu.module @xevm_module{4// CHECK-LABEL: gpu.func @store_nd_1d5// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {6// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16]7// CHECK-SAME: -> (vector<1xf32>, !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>, index) {8// CHECK: gpu.yield %{{.*}} : vector<16xf32>,9// CHECK-SAME: !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>, index10// CHECK-NEXT: }11// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#1 : !xegpu.tensor_desc<16xf32,12// CHECK-SAME: #xegpu.layout<lane_layout = [16], lane_data = [1]>> to !xegpu.tensor_desc<16xf32> {resolve_simt_type_mismatch}13// CHECK-NEXT: xegpu.store_nd %[[W]]#0, %[[T1]][%[[W]]#2] : vector<1xf32>, !xegpu.tensor_desc<16xf32>14gpu.func @store_nd_1d(%laneid: index) {15 %c0 = arith.constant 0 : index16 gpu.warp_execute_on_lane_0(%laneid)[16] {17 %0 = "some_op"() : () -> !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>18 %cst = "some_op"() : () -> vector<16xf32>19 xegpu.store_nd %cst, %0 [%c0] {layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}20 : vector<16xf32>, !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>21 }22 gpu.return23}24 25// CHECK-LABEL: gpu.func @store_nd_2d26// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {27// CHECK: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16]28// CHECK-SAME: -> (vector<16x1xf16>, !xegpu.tensor_desc<16x16xf16,29// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index) {30// CHECK: gpu.yield %{{.*}} : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16,31// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index32// CHECK-NEXT: }33// CHECK-NEXT: %[[CAST:.*]] = vector.shape_cast %[[W]]#0 : vector<16x1xf16> to vector<16xf16>34// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#1 : !xegpu.tensor_desc<16x16xf16,35// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> to !xegpu.tensor_desc<16x16xf16> {resolve_simt_type_mismatch}36// CHECK-NEXT: xegpu.store_nd %[[CAST]], %[[T1]][%[[W]]#2, %[[W]]#3] : vector<16xf16>, !xegpu.tensor_desc<16x16xf16>37gpu.func @store_nd_2d(%laneid : index) {38 %c0 = arith.constant 0 : index39 gpu.warp_execute_on_lane_0(%laneid)[16] {40 %0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>41 %cst = "some_op"() : () -> vector<16x16xf16>42 xegpu.store_nd %cst, %0 [%c0, %c0] {layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}43 : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>44 }45 gpu.return46}47 48 49// CHECK-LABEL: gpu.func @load_nd_1d50// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {51// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16] -> (vector<1xf32>,52// CHECK-SAME: !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>, index) {53// CHECK: gpu.yield %{{.*}} : vector<16xf32>, !xegpu.tensor_desc<16xf32,54// CHECK-SAME: #xegpu.layout<lane_layout = [16], lane_data = [1]>>, index55// CHECK-NEXT: }56// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#1 : !xegpu.tensor_desc<16xf32,57// CHECK-SAME: #xegpu.layout<lane_layout = [16], lane_data = [1]>> to !xegpu.tensor_desc<16xf32> {resolve_simt_type_mismatch}58// CHECK-NEXT: xegpu.load_nd %[[T1]][%[[W]]#2] : !xegpu.tensor_desc<16xf32> -> vector<1xf32>59gpu.func @load_nd_1d(%laneid: index) {60 %c0 = arith.constant 0 : index61 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<1xf32>) {62 %0 = "some_op"() : () -> !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>63 %1 = xegpu.load_nd %0 [%c0] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} :64 !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>> -> vector<16xf32>65 gpu.yield %1 : vector<16xf32>66 }67 "some_user_op"(%r) : (vector<1xf32>) -> ()68 gpu.return69}70 71 72// CHECK-LABEL: gpu.func @load_nd_2d73// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {74// CHECK: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16] -> (vector<16x1xf16>, !xegpu.tensor_desc<16x16xf16,75// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index) {76// CHECK: gpu.yield %{{.*}} : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16,77// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index78// CHECK-NEXT: }79// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#1 : !xegpu.tensor_desc<16x16xf16,80// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> to !xegpu.tensor_desc<16x16xf16> {resolve_simt_type_mismatch}81// CHECK-NEXT: %[[T2:.*]] = xegpu.load_nd %[[T1]][%[[W]]#2, %[[W]]#3] : !xegpu.tensor_desc<16x16xf16> -> vector<16xf16>82// CHECK: vector.shape_cast %[[T2]] : vector<16xf16> to vector<16x1xf16>83gpu.func @load_nd_2d(%laneid: index) {84 %c0 = arith.constant 0 : index85 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<16x1xf16>) {86 %0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>87 %1 = xegpu.load_nd %0[%c0, %c0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}88 : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf16>89 gpu.yield %1 : vector<16x16xf16>90 }91 "some_user_op"(%r) : (vector<16x1xf16>) -> ()92 gpu.return93}94 95 96// CHECK-LABEL: gpu.func @load_nd_array_length97// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {98// CHECK: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16] -> (vector<2x16x1xf16>,99// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>,100// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index) {101// CHECK: gpu.yield %{{.*}} : vector<2x16x16xf16>, !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<102// CHECK-SAME: array_length = 2 : i64>, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index103// CHECK-NEXT: }104// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#1 : !xegpu.tensor_desc<16x16xf16,105// CHECK-SAME: #xegpu.block_tdesc_attr<array_length = 2 : i64>, #xegpu.layout<lane_layout = [1, 16],106// CHECK-SAME: lane_data = [1, 1]>> to !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>>107// CHECK-NEXT: %[[T2:.*]] = xegpu.load_nd %[[T1]][%[[W]]#2, %[[W]]#3] : !xegpu.tensor_desc<16x16xf16,108// CHECK-SAME: #xegpu.block_tdesc_attr<array_length = 2 : i64>> -> vector<32xf16>109// CHECK-NEXT: vector.shape_cast %[[T2]] : vector<32xf16> to vector<2x16x1xf16>110gpu.func @load_nd_array_length(%laneid: index) {111 %c0 = arith.constant 0 : index112 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2x16x1xf16>) {113 %0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>,114 #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>115 %1 = xegpu.load_nd %0[%c0, %c0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}116 : !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>,117 #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16x16xf16>118 gpu.yield %1 : vector<2x16x16xf16>119 }120 "some_user_op"(%r) : (vector<2x16x1xf16>) -> ()121 gpu.return122}123 124 125// CHECK-LABEL: gpu.func @dpas126// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {127// CHECK: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16] ->128// CHECK-SAME: (vector<8x1xf32>, vector<8x1xf16>, vector<16x1xf16>, vector<8x1xf32>) {129// CHECK: gpu.yield %{{.*}} : vector<8x16xf32>, vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32>130// CHECK-NEXT: }131// CHECK-DAG: %[[T1:.*]] = vector.shape_cast %[[W]]#1 : vector<8x1xf16> to vector<8xf16>132// CHECK-DAG: %[[T2:.*]] = vector.shape_cast %[[W]]#2 : vector<16x1xf16> to vector<16xf16>133// CHECK-DAG: %[[T3:.*]] = vector.shape_cast %[[W]]#3 : vector<8x1xf32> to vector<8xf32>134// CHECK-NEXT: %[[T4:.*]] = xegpu.dpas %[[T1]], %[[T2]], %[[T3]] : vector<8xf16>, vector<16xf16>, vector<8xf32> -> vector<8xf32>135// CHECK-NEXT: vector.shape_cast %[[T4]] : vector<8xf32> to vector<8x1xf32>136gpu.func @dpas(%laneid: index) {137 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<8x1xf32>) {138 %0 = "some_op"() : () -> vector<8x16xf16>139 %1 = "some_op"() : () -> vector<16x16xf16>140 %2 = "some_op"() : () -> vector<8x16xf32>141 %3 = xegpu.dpas %0, %1, %2142 {143 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,144 layout_operand_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,145 layout_operand_2 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,146 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>147 }148 : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>149 gpu.yield %3 : vector<8x16xf32>150 }151 "some_user_op"(%r) : (vector<8x1xf32>) -> ()152 gpu.return153}154 155 156 157// CHECK-LABEL: gpu.func @create_nd_tdesc_non_memref158// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: ui64, %[[ARG1:[0-9a-zA-Z]+]]: index) {159// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%[[ARG1]])[16] -> (!xegpu.tensor_desc<16x16xf16,160// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, ui64) {161// CHECK: gpu.yield %{{.*}} : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, ui64162// CHECK-NEXT: }163// CHECK-NEXT: %[[T1:.*]] = xegpu.create_nd_tdesc %[[W]]#1, shape : [64, 128], strides : [128, 1] : ui64 -> !xegpu.tensor_desc<16x16xf16>164// CHECK-NEXT: builtin.unrealized_conversion_cast %[[T1]] : !xegpu.tensor_desc<16x16xf16> to !xegpu.tensor_desc<16x16xf16,165// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> {resolve_simt_type_mismatch}166gpu.func @create_nd_tdesc_non_memref(%arg0: ui64, %laneid: index) {167 %c0 = arith.constant 0 : index168 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (!xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) {169 %0 = xegpu.create_nd_tdesc %arg0, shape:[64, 128], strides:[128, 1] : ui64 ->170 !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>171 gpu.yield %0 : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>172 }173 "some_user_op"(%r)174 : (!xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) -> ()175 gpu.return176}177 178 179// CHECK-LABEL: gpu.func @prefetch_2d180// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {181// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16] -> (!xegpu.tensor_desc<16x16xf16,182// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>, index, index) {183// CHECK: gpu.yield %{{.*}} : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>184// CHECK-SAME: , index, index185// CHECK-NEXT: }186// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#0 : !xegpu.tensor_desc<16x16xf16,187// CHECK-SAME: #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> to !xegpu.tensor_desc<16x16xf16> {resolve_simt_type_mismatch}188// CHECK-NEXT: xegpu.prefetch_nd %[[T1]][%[[W]]#1, %[[W]]#2]189// CHECK-SAME: <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}> : !xegpu.tensor_desc<16x16xf16>190gpu.func @prefetch_2d(%laneid: index) {191 %c0 = arith.constant 0 : index192 gpu.warp_execute_on_lane_0(%laneid)[16] {193 %0 = "some_op"() : ()194 -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>195 xegpu.prefetch_nd %0[%c0, %c0]196 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>197 : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>198 }199 gpu.return200}201 202 203// CHECK-LABEL: gpu.func @prefetch_1d204// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: index) {205// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%[[ARG0]])[16] -> (!xegpu.tensor_desc<16xf16,206// CHECK-SAME: #xegpu.layout<lane_layout = [16], lane_data = [1]>>, index) {207// CHECK: gpu.yield %{{.*}} : !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>, index208// CHECK-NEXT: }209// CHECK-NEXT: %[[T1:.*]] = builtin.unrealized_conversion_cast %[[W]]#0 : !xegpu.tensor_desc<16xf16,210// CHECK-SAME: #xegpu.layout<lane_layout = [16], lane_data = [1]>> to !xegpu.tensor_desc<16xf16> {resolve_simt_type_mismatch}211// CHECK-NEXT: xegpu.prefetch_nd %[[T1]][%[[W]]#1] <{l1_hint = #xegpu.cache_hint<cached>,212// CHECK-SAME: l2_hint = #xegpu.cache_hint<uncached>}> : !xegpu.tensor_desc<16xf16>213gpu.func @prefetch_1d(%laneid: index) {214 %c0 = arith.constant 0 : index215 gpu.warp_execute_on_lane_0(%laneid)[16] {216 %0 = "some_op"() : ()217 -> !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>218 xegpu.prefetch_nd %0[%c0]219 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>220 : !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>221 }222 gpu.return223}224 225 226// CHECK-LABEL: gpu.func @gpu_barrier({{.*}}) {227// CHECK: gpu.warp_execute_on_lane_0(%{{.*}})[16] -> ({{.*}}) {228// CHECK: gpu.yield %{{.*}}229// CHECK: }230// CHECK: %{{.*}} = xegpu.load_nd %{{.*}} : !xegpu.tensor_desc<16xf16> -> vector<1xf16>231// CHECK: gpu.barrier232gpu.func @gpu_barrier(%laneid: index) {233 %c0 = arith.constant 0 : index234 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<1xf16>) {235 %0 = "some_op"() : () -> !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>236 %1 = xegpu.load_nd %0[%c0]237 {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}238 : !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>> -> vector<16xf16>239 gpu.barrier240 gpu.yield %1 : vector<16xf16>241 }242 "some_user_op"(%r) : (vector<1xf16>) -> ()243 gpu.return244}245 246 247// CHECK-LABEL: gpu.func @vector_multi_reduction_dim1_distributed_dim0_reduction248// CHECK: %[[ACC:.*]] = arith.constant {{.*}} dense<0.000000e+00> : vector<32xf32>249// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16]250// CHECK-SAME: -> (vector<2xf32>, vector<16x2xf32>, vector<2xf32>) {251// CHECK: %[[SRC:.*]] = "some_def"() {{.*}} : () -> vector<16x32xf32>252// CHECK: gpu.yield %{{.*}}, %[[SRC]], %[[ACC]] : vector<32xf32>, vector<16x32xf32>, vector<32xf32>253// CHECK-NEXT: }254// CHECK: %[[T1:.*]] = vector.extract_strided_slice %[[W]]#1255// CHECK-SAME: {offsets = [0, 0], sizes = [16, 1], strides = [1, 1]} : vector<16x2xf32> to vector<16x1xf32>256// CHECK: %[[T2:.*]] = vector.shape_cast %[[T1]] : vector<16x1xf32> to vector<16xf32>257// CHECK: %[[T3:.*]] = vector.extract %[[W]]#2[0] : f32 from vector<2xf32>258// CHECK: %[[T4:.*]] = vector.reduction <add>, %[[T2]], %[[T3]] : vector<16xf32> into f32259// CHECK: %[[T5:.*]] = vector.extract_strided_slice %[[W]]#1260// CHECK-SAME: {offsets = [0, 1], sizes = [16, 1], strides = [1, 1]} : vector<16x2xf32> to vector<16x1xf32>261// CHECK: %[[T6:.*]] = vector.shape_cast %[[T5]] : vector<16x1xf32> to vector<16xf32>262// CHECK: %[[T7:.*]] = vector.extract %[[W]]#2[1] : f32 from vector<2xf32>263// CHECK: %[[T8:.*]] = vector.reduction <add>, %[[T6]], %[[T7]] : vector<16xf32> into f32264// CHECK: %[[T9:.*]] = vector.from_elements %[[T4]], %[[T8]] : vector<2xf32>265gpu.func @vector_multi_reduction_dim1_distributed_dim0_reduction(%laneid: index) {266 %c0 = arith.constant 0 : index267 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {268 %src = "some_def"()269 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}270 : () -> (vector<16x32xf32>)271 %acc = arith.constant272 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}273 dense<0.0> : vector<32xf32>274 %1 = vector.multi_reduction <add>, %src, %acc275 {276 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,277 layout_operand_1 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,278 layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>279 } [0]280 : vector<16x32xf32> to vector<32xf32>281 gpu.yield %1 : vector<32xf32>282 }283 "some_user_op"(%r) : (vector<2xf32>) -> ()284 gpu.return285}286 287 288// CHECK-LABEL: gpu.func @vector_multi_reduction_dim1_distributed_dim1_reduction289// CHECK: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2xf32>) {290// CHECK-NEXT: %[[SRC:.*]] = "some_def"() {{.*}} : () -> vector<2x16xf32>291// CHECK-NEXT: %[[T2:.*]] = vector.extract %[[SRC]][0] : vector<16xf32> from vector<2x16xf32>292// CHECK-NEXT: %[[T3:.*]] = vector.reduction <add>, %[[T2]], %cst : vector<16xf32> into f32293// CHECK-NEXT: %[[T4:.*]] = vector.extract %[[SRC]][1] : vector<16xf32> from vector<2x16xf32>294// CHECK-NEXT: %[[T5:.*]] = vector.reduction <add>, %[[T4]], %cst : vector<16xf32> into f32295// CHECK-NEXT: %[[T6:.*]] = vector.from_elements %[[T3]], %[[T5]] : vector<2xf32>296// CHECK-NEXT: gpu.yield %[[T6]] : vector<2xf32>297// CHECK-NEXT: }298gpu.func @vector_multi_reduction_dim1_distributed_dim1_reduction(%laneid: index) {299 %c0 = arith.constant 0 : index300 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {301 %src = "some_def"()302 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}303 : () -> (vector<2x16xf32>)304 %acc = arith.constant305 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>}306 dense<0.0> : vector<2xf32>307 %1 = vector.multi_reduction <add>, %src, %acc308 {309 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,310 layout_operand_1 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>,311 layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>312 }313 [1] : vector<2x16xf32> to vector<2xf32>314 gpu.yield %1 : vector<2xf32>315 }316 "some_user_op"(%r) : (vector<2xf32>) -> ()317 gpu.return318}319 320 321// CHECK-LABEL: gpu.func @vector_multi_reduction_dim0_distributed_dim1_reduction322// CHECK: %[[ACC:.*]] = arith.constant {{.*}} dense<0.000000e+00> : vector<32xf32>323// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2xf32>, vector<2x16xf32>, vector<2xf32>) {324// CHECK: %[[SRC:.*]] = "some_def"() {{.*}} : () -> vector<32x16xf32>325// CHECK: gpu.yield %9, %[[SRC]], %[[ACC]] : vector<32xf32>, vector<32x16xf32>, vector<32xf32>326// CHECK: }327// CHECK: %[[T1:.*]] = vector.extract %[[W]]#1[0] : vector<16xf32> from vector<2x16xf32>328// CHECK: %[[T2:.*]] = vector.extract %[[W]]#2[0] : f32 from vector<2xf32>329// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T1]], %[[T2]] : vector<16xf32> into f32330// CHECK: %[[T4:.*]] = vector.extract %[[W]]#1[1] : vector<16xf32> from vector<2x16xf32>331// CHECK: %[[T5:.*]] = vector.extract %[[W]]#2[1] : f32 from vector<2xf32>332// CHECK: %[[T6:.*]] = vector.reduction <add>, %[[T4]], %[[T5]] : vector<16xf32> into f32333// CHECK: %[[T7:.*]] = vector.from_elements %[[T3]], %[[T6]] : vector<2xf32>334gpu.func @vector_multi_reduction_dim0_distributed_dim1_reduction(%laneid: index) {335 %c0 = arith.constant 0 : index336 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {337 %src = "some_def"()338 {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}339 : () -> (vector<32x16xf32>)340 %acc = arith.constant341 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>}342 dense<0.0> : vector<32xf32>343 %1 = vector.multi_reduction <add>, %src, %acc344 {345 layout_operand_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,346 layout_operand_1 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>,347 layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>348 }349 [1] : vector<32x16xf32> to vector<32xf32>350 gpu.yield %1 : vector<32xf32>351 }352 "some_user_op"(%r) : (vector<2xf32>) -> ()353 gpu.return354}355 356 357// CHECK-LABEL: gpu.func @vector_multi_reduction_dim0_distributed_dim0_reduction358// CHECK: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2xf32>) {359// CHECK: %[[SRC:.*]] = "some_def"() {{.*}} : () -> vector<16x2xf32>360// CHECK: %[[T1:.*]] = vector.extract_strided_slice %[[SRC]]361// CHECK-SAME: {offsets = [0, 0], sizes = [16, 1], strides = [1, 1]} : vector<16x2xf32> to vector<16x1xf32>362// CHECK: %[[T2:.*]] = vector.shape_cast %[[T1]] {{.*}} : vector<16x1xf32> to vector<16xf32>363// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T2]], %{{.*}} : vector<16xf32> into f32364// CHECK: %[[T4:.*]] = vector.extract_strided_slice %[[SRC]]365// CHECK-SAME: {offsets = [0, 1], sizes = [16, 1], strides = [1, 1]} : vector<16x2xf32> to vector<16x1xf32>366// CHECK: %[[T5:.*]] = vector.shape_cast %[[T4]] {{.*}} : vector<16x1xf32> to vector<16xf32>367// CHECK: %[[T6:.*]] = vector.reduction <add>, %[[T5]], %{{.*}} : vector<16xf32> into f32368// CHECK: %[[T7:.*]] = vector.from_elements %[[T3]], %[[T6]] : vector<2xf32>369// CHECK: gpu.yield %[[T7]] : vector<2xf32>370// CHECK: }371gpu.func @vector_multi_reduction_dim0_distributed_dim0_reduction(%laneid: index) {372 %c0 = arith.constant 0 : index373 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {374 %src = "some_def"()375 {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}376 : () -> (vector<16x2xf32>)377 %acc = arith.constant378 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [0]>}379 dense<0.0> : vector<2xf32>380 %1 = vector.multi_reduction <add>, %src, %acc381 {382 layout_operand_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,383 layout_operand_1 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [0]>,384 layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [0]>385 }386 [0] : vector<16x2xf32> to vector<2xf32>387 gpu.yield %1 : vector<2xf32>388 }389 "some_user_op"(%r) : (vector<2xf32>) -> ()390 gpu.return391}392 393 394// CHECK-LABEL: gpu.func @scatter_ops_chunksize({{.*}}) {395// CHECK: %[[OFFSETS:.*]] = arith.constant {{.*}} dense<12> : vector<16xindex>396// CHECK: %[[MASKS:.*]] = arith.constant {{.*}} dense<true> : vector<16xi1>397// CHECK: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%{{.*}})[16]398// CHECK-SAME: -> (vector<1x8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>) {399// CHECK: gpu.yield %{{.*}}, %{{.*}}, %[[OFFSETS]], %[[MASKS]] :400// CHECK-SAME: vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>401// CHECK-NEXT: }402// CHECK-NEXT: %[[T1:.*]] = xegpu.load %[[W]]#1[%[[W]]#2], %[[W]]#3 <{chunk_size = 8 : i64}>403// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>404// CHECK-NEXT: xegpu.store %[[T1]], %[[W]]#1[%[[W]]#2], %[[W]]#3 <{chunk_size = 8 : i64}>405// CHECK-SAME: : vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>406gpu.func @scatter_ops_chunksize(%laneid: index, %src: memref<256xf16>) {407 gpu.warp_execute_on_lane_0(%laneid)[16] {408 %1 = arith.constant409 {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}410 dense<1>: vector<16xi1>411 %offset = arith.constant412 {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}413 dense<12> : vector<16xindex>414 %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}>415 {416 layout_operand_1 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,417 layout_operand_2 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,418 layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>419 }420 : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>421 xegpu.store %3, %src[%offset], %1 <{chunk_size=8}>422 {423 layout_operand_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>,424 layout_operand_2 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,425 layout_operand_3 = #xegpu.layout<lane_layout = [16], lane_data = [1]>426 }427 : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>428 }429 gpu.return430}431 432 433// CHECK-LABEL: gpu.func @scatter_ops({{.*}}) {434// CHECK: %[[OFFSETS:.*]] = arith.constant {{.*}} dense<12> : vector<16xindex>435// CHECK: %[[MASKS:.*]] = arith.constant {{.*}} dense<true> : vector<16xi1>436// CHECK: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%{{.*}})[16]437// CHECK-SAME: -> (vector<1xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>) {438// CHECK: gpu.yield %{{.*}}, %{{.*}}, %[[OFFSETS]], %[[MASKS]]439// CHECK-SAME: : vector<16xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>440// CHECK-NEXT: }441// CHECK-NEXT: %[[T1:.*]] = xegpu.load %[[W]]#1[%[[W]]#2], %[[W]]#3442// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<1xf16>443// CHECK-NEXT: xegpu.store %[[T1]], %[[W]]#1[%[[W]]#2], %[[W]]#3444// CHECK-SAME: : vector<1xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>445gpu.func @scatter_ops(%src: memref<256xf16>, %laneid: index) {446 gpu.warp_execute_on_lane_0(%laneid)[16] {447 %1 = arith.constant448 {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}449 dense<1> : vector<16xi1>450 %offset = arith.constant451 {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}452 dense<12> : vector<16xindex>453 %3 = xegpu.load %src[%offset], %1454 {455 layout_operand_1 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,456 layout_operand_2 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,457 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>458 } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16xf16>459 xegpu.store %3, %src[%offset], %1460 {461 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,462 layout_operand_2 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,463 layout_operand_3 = #xegpu.layout<lane_layout = [16], lane_data = [1]>464 }465 : vector<16xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>466 }467 gpu.return468}469 470 471// CHECK-LABEL: gpu.func @memref_extract_aligned_pointer_as_index(472// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (index, memref<256x256xf16>) {473// CHECK: gpu.yield %{{.*}}, %{{.*}} : index, memref<256x256xf16>474// CHECK-NEXT: }475// CHECK-NEXT: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[W]]#1 : memref<256x256xf16> -> index476// CHECK-NEXT: arith.index_cast %[[INTPTR]] : index to i64477gpu.func @memref_extract_aligned_pointer_as_index(%arg0 : memref<256x256xf16>, %laneid: index) {478 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (index) {479 %ptr = memref.extract_aligned_pointer_as_index %arg0 : memref<256x256xf16> -> index480 gpu.yield %ptr : index481 }482 %ptr_i64 = arith.index_cast %r : index to i64483 "some_user_op"(%ptr_i64) : (i64) -> ()484 gpu.return485}486 487 488 489// CHECK-LABEL: gpu.func @vector_transpose(490// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2x1xf32>, vector<1x2xf32>) {491// CHECK: %[[SRC:.*]] = "some_op"() {{.*}} : () -> vector<16x2xf32>492// CHECK: gpu.yield %{{.*}}, %[[SRC]] : vector<2x16xf32>, vector<16x2xf32>493// CHECK-NEXT: }494// CHECK-NEXT: %[[T1:.*]] = vector.transpose %[[W]]#1, [1, 0] : vector<1x2xf32> to vector<2x1xf32>495gpu.func @vector_transpose(%laneid: index) {496 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2x1xf32>) {497 %cst = "some_op"()498 {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}499 : () -> (vector<16x2xf32>)500 %transpose = vector.transpose %cst, [1, 0]501 {502 layout_operand_0 = #xegpu.layout<lane_layout = [16 , 1], lane_data = [1, 1]>,503 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>504 }505 : vector<16x2xf32> to vector<2x16xf32>506 gpu.yield %transpose : vector<2x16xf32>507 }508 "some_user_op"(%r) : (vector<2x1xf32>) -> ()509 gpu.return510}511 512 513// CHECK-LABEL: gpu.func @vector_bitcast(514// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<4x1xi16>, vector<4x2xi8>) {515// CHECK: %[[SRC:.*]] = "some_op"() {{.*}} : () -> vector<4x32xi8>516// CHECK: gpu.yield %{{.*}}, %[[SRC]] : vector<4x16xi16>, vector<4x32xi8>517// CHECK: }518// CHECK: vector.bitcast %[[W]]#1 : vector<4x2xi8> to vector<4x1xi16>519gpu.func @vector_bitcast(%laneid: index) {520 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<4x1xi16>) {521 %cst = "some_op"()522 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}523 : () -> (vector<4x32xi8>)524 %bitcast = vector.bitcast %cst525 {526 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>,527 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>528 }529 : vector<4x32xi8> to vector<4x16xi16>530 gpu.yield %bitcast : vector<4x16xi16>531 }532 "some_user_op"(%r) : (vector<4x1xi16>) -> ()533 gpu.return534}535 536 537// CHECK-LABEL: gpu.func @vector_shapecast_rank_increasing538// CHECK: %{{.*}}:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<1x1xf32>, vector<1xf32>) {539// CHECK: gpu.yield %{{.*}} : vector<1x16xf32>, vector<16xf32>540// CHECK: }541// CHECK: %{{.*}} = vector.shape_cast %{{.*}}#1 : vector<1xf32> to vector<1x1xf32>542gpu.func @vector_shapecast_rank_increasing(%laneid: index) {543 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<1x1xf32>) {544 %cst = "some_op"()545 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}546 : () -> (vector<16xf32>)547 %cast = vector.shape_cast %cst548 {549 layout_operand_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,550 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>551 }552 : vector<16xf32> to vector<1x16xf32>553 gpu.yield %cast : vector<1x16xf32>554 }555 "some_user_op"(%r) : (vector<1x1xf32>) -> ()556 gpu.return557}558 559 560// CHECK-LABEL: gpu.func @vector_shapecast_rank_reducing(561// CHECK: %{{.*}}:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<1xf32>, vector<1x1xf32>) {562// CHECK: gpu.yield %{{.*}} : vector<16xf32>, vector<1x16xf32>563// CHECK: }564// CHECK: %{{.*}} = vector.shape_cast %{{.*}}#1 : vector<1x1xf32> to vector<1xf32>565gpu.func @vector_shapecast_rank_reducing(%laneid: index) {566 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<1xf32>) {567 %cst = "some_op"()568 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}569 : () -> (vector<1x16xf32>)570 %cast = vector.shape_cast %cst571 {572 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,573 layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>574 }575 : vector<1x16xf32> to vector<16xf32>576 gpu.yield %cast : vector<16xf32>577 }578 "some_user_op"(%r) : (vector<1xf32>) -> ()579 gpu.return580}581 582 583// NOTE: Layouts are still valid, but distribution still requires a slice layout for the operand.584//585// CHECK-LABEL: gpu.func @vector_shapecast_unsupported586// CHECK: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<1x1xf32>) {587// CHECK: %[[T1:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<1x16xf32>588// CHECK: gpu.yield %[[T1]] : vector<1x16xf32>589// CHECK: }590// CHECK: "some_user_op"(%[[W]]) : (vector<1x1xf32>) -> ()591// CHECK: gpu.return592gpu.func @vector_shapecast_unsupported(%laneid: index) {593 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<1x1xf32>) {594 %cst = "some_op"()595 {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]> }596 : () -> (vector<16xf32>)597 %cast = vector.shape_cast %cst598 {599 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,600 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>601 }602 : vector<16xf32> to vector<1x16xf32>603 gpu.yield %cast : vector<1x16xf32>604 }605 "some_user_op"(%r) : (vector<1x1xf32>) -> ()606 gpu.return607}608 609 610// CHECK-LABEL: gpu.func @vector_extract_strided_slice_distributed_dim_fully_extracted611// CHECK-NEXT: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<8x1xf32>, vector<24x1xf32>) {612// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<24x16xf32>613// CHECK: gpu.yield %{{.*}}, %[[S]] : vector<8x16xf32>, vector<24x16xf32>614// CHECK-NEXT: }615// CHECK-NEXT: %[[T1:.*]] = vector.extract_strided_slice %[[W]]#1616// CHECK-SAME: {offsets = [8, 0], sizes = [8, 1], strides = [1, 1]} : vector<24x1xf32> to vector<8x1xf32>617// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<8x1xf32>) -> ()618gpu.func @vector_extract_strided_slice_distributed_dim_fully_extracted(%laneid: index) {619 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<8x1xf32>) {620 %0 = "some_def"() : () -> (vector<24x16xf32>)621 %1 = vector.extract_strided_slice %0 { offsets = [8, 0], sizes = [8, 16], strides = [1, 1],622 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,623 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>624 }625 : vector<24x16xf32> to vector<8x16xf32>626 gpu.yield %1 : vector<8x16xf32>627 }628 "some_use"(%r) : (vector<8x1xf32>) -> ()629 gpu.return630}631 632// CHECK-LABEL: gpu.func @vector_extract_strided_slice_non_distributed633// CHECK-NEXT: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<8x1xf32>, vector<24x1xf32>) {634// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<24x1xf32>635// CHECK: gpu.yield %{{.*}}, %[[S]] : vector<8x1xf32>, vector<24x1xf32>636// CHECK-NEXT: }637// CHECK-NEXT: %[[T1:.*]] = vector.extract_strided_slice %[[W]]#1638// CHECK-SAME: {offsets = [8, 0], sizes = [8, 1], strides = [1, 1]} : vector<24x1xf32> to vector<8x1xf32>639// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<8x1xf32>) -> ()640gpu.func @vector_extract_strided_slice_non_distributed(%laneid: index) {641 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<8x1xf32>) {642 %0 = "some_def"() : () -> (vector<24x1xf32>)643 %1 = vector.extract_strided_slice %0 { offsets = [8, 0], sizes = [8, 1], strides = [1, 1],644 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,645 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>646 }647 : vector<24x1xf32> to vector<8x1xf32>648 gpu.yield %1 : vector<8x1xf32>649 }650 "some_use"(%r) : (vector<8x1xf32>) -> ()651 gpu.return652}653 654// CHECK-LABEL: gpu.func @vector_extract_strided_slice_inner_distributed655// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<8x1xf32>, vector<24x4xf32>) {656// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<24x64xf32>657// CHECK: gpu.yield %{{.*}}, %[[S]] : vector<8x16xf32>, vector<24x64xf32>658// CHECK-NEXT: }659// CHECK-NEXT: %[[T1:.*]] = vector.extract_strided_slice %[[W]]#1660// CHECK-SAME: {offsets = [8, 3], sizes = [8, 1], strides = [1, 1]} : vector<24x4xf32> to vector<8x1xf32>661// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<8x1xf32>) -> ()662gpu.func @vector_extract_strided_slice_inner_distributed(%laneid: index) {663 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<8x1xf32>) {664 %0 = "some_def"() : () -> (vector<24x64xf32>)665 %1 = vector.extract_strided_slice %0 { offsets = [8, 48], sizes = [8, 16], strides = [1, 1],666 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,667 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>668 }669 : vector<24x64xf32> to vector<8x16xf32>670 gpu.yield %1 : vector<8x16xf32>671 }672 "some_use"(%r) : (vector<8x1xf32>) -> ()673 gpu.return674}675 676// CHECK-LABEL: gpu.func @vector_extract_strided_slice_outer_distributed677// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<1x16xf32>, vector<2x16xf32>) {678// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<32x16xf32>679// CHECK: gpu.yield %{{.*}}, %[[S]] : vector<16x16xf32>, vector<32x16xf32>680// CHECK: }681// CHECK-NEXT: %[[T1:.*]] = vector.extract %[[W]]#1[1] : vector<16xf32> from vector<2x16xf32>682// CHECK-NEXT: %[[T2:.*]] = vector.shape_cast %[[T1]] : vector<16xf32> to vector<1x16xf32>683// CHECK-NEXT: "some_use"(%[[T2]]) : (vector<1x16xf32>) -> ()684gpu.func @vector_extract_strided_slice_outer_distributed(%laneid: index) {685 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<1x16xf32>) {686 %0 = "some_def"() : () -> (vector<32x16xf32>)687 %1 = vector.extract_strided_slice %0 { offsets = [16], sizes = [16], strides = [1],688 layout_operand_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,689 layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>690 }691 : vector<32x16xf32> to vector<16x16xf32>692 gpu.yield %1 : vector<16x16xf32>693 }694 "some_use"(%r) : (vector<1x16xf32>) -> ()695 gpu.return696}697 698// CHECK-LABEL: gpu.func @vector_extract_strided_slice_1d699// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2xf32>, vector<4xf32>) {700// CHECK: %[[S:.*]] = "some_def"() : () -> vector<64xf32>701// CHECK: gpu.yield %{{.*}}, %[[S]] : vector<32xf32>, vector<64xf32>702// CHECK-NEXT: }703// CHECK-NEXT: %[[T1:.*]] = vector.extract_strided_slice %[[W]]#1704// CHECK-SAME: {offsets = [1], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>705// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<2xf32>) -> ()706gpu.func @vector_extract_strided_slice_1d(%laneid: index) {707 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {708 %0 = "some_def"() : () -> (vector<64xf32>)709 %1 = vector.extract_strided_slice %0 { offsets = [16], sizes = [32], strides = [1],710 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,711 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>712 }713 : vector<64xf32> to vector<32xf32>714 gpu.yield %1 : vector<32xf32>715 }716 "some_use"(%r) : (vector<2xf32>) -> ()717 gpu.return718}719 720// CHECK-LABEL: gpu.func @vector_extract_strided_slice_unsopported_offset721// CHECK: %{{.*}} = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2xf32>) {722// CHECK: }723// CHECK-NOT: %{{.*}} = vector.extract_strided_slice724gpu.func @vector_extract_strided_slice_unsopported_offset(%laneid: index) {725 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {726 %0 = "some_def"() : () -> (vector<64xf32>)727 %1 = vector.extract_strided_slice %0 { offsets = [3], sizes = [32], strides = [1],728 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,729 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>730 }731 : vector<64xf32> to vector<32xf32>732 gpu.yield %1 : vector<32xf32>733 }734 "some_use"(%r) : (vector<2xf32>) -> ()735 gpu.return736}737 738// CHECK-LABEL: gpu.func @vector_extract_strided_slice_unsopported_source739// CHECK: %{{.*}} = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<2xf32>) {740// CHECK: }741// CHECK-NOT: %{{.*}} = vector.extract_strided_slice742gpu.func @vector_extract_strided_slice_unsopported_source(%laneid: index) {743 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<2xf32>) {744 %0 = "some_def"() : () -> (vector<54xf32>)745 %1 = vector.extract_strided_slice %0 { offsets = [0], sizes = [32], strides = [1],746 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,747 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>748 }749 : vector<54xf32> to vector<32xf32>750 gpu.yield %1 : vector<32xf32>751 }752 "some_use"(%r) : (vector<2xf32>) -> ()753 gpu.return754}755 756 757// CHECK-LABEL: gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted758// CHECK-NEXT: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<64x1xf32>, vector<16x1xf32>, vector<64x1xf32>) {759// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<16x16xf32>760// CHECK-NEXT: %[[D:.*]] = "some_def"() : () -> vector<64x16xf32>761// CHECK: gpu.yield %{{.*}}, %[[S]], %[[D]] : vector<64x16xf32>, vector<16x16xf32>, vector<64x16xf32>762// CHECK-NEXT: }763// CHECK-NEXT: %[[T1:.*]] = vector.insert_strided_slice %[[W]]#1, %[[W]]#2764// CHECK-SAME: {offsets = [24, 0], strides = [1, 1]} : vector<16x1xf32> into vector<64x1xf32>765// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<64x1xf32>) -> ()766gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted(%laneid: index) {767 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<64x1xf32>) {768 %0 = "some_def"() : () -> (vector<16x16xf32>)769 %1 = "some_def"() : () -> (vector<64x16xf32>)770 %2 = vector.insert_strided_slice %0, %1 { offsets = [24, 0], strides = [1, 1],771 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,772 layout_operand_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,773 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>774 }775 : vector<16x16xf32> into vector<64x16xf32>776 gpu.yield %2 : vector<64x16xf32>777 }778 "some_use"(%r) : (vector<64x1xf32>) -> ()779 gpu.return780}781 782 783// CHECK-LABEL: gpu.func @vector_insert_strided_slice_non_distributed784// CHECK-NEXT: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<64x1xf32>, vector<16x1xf32>, vector<64x1xf32>) {785// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<16x1xf32>786// CHECK-NEXT: %[[D:.*]] = "some_def"() : () -> vector<64x1xf32>787// CHECK: gpu.yield %{{.*}}, %[[S]], %[[D]] : vector<64x1xf32>, vector<16x1xf32>, vector<64x1xf32>788// CHECK-NEXT: }789// CHECK-NEXT: %[[T1:.*]] = vector.insert_strided_slice %[[W]]#1, %[[W]]#2790// CHECK-SAME: {offsets = [24, 0], strides = [1, 1]} : vector<16x1xf32> into vector<64x1xf32>791// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<64x1xf32>) -> ()792gpu.func @vector_insert_strided_slice_non_distributed(%laneid: index) {793 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<64x1xf32>) {794 %0 = "some_def"() : () -> (vector<16x1xf32>)795 %1 = "some_def"() : () -> (vector<64x1xf32>)796 %2 = vector.insert_strided_slice %0, %1 { offsets = [24, 0], strides = [1, 1],797 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,798 layout_operand_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,799 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>800 }801 : vector<16x1xf32> into vector<64x1xf32>802 gpu.yield %2 : vector<64x1xf32>803 }804 "some_use"(%r) : (vector<64x1xf32>) -> ()805 gpu.return806}807 808// CHECK-LABEL: gpu.func @vector_insert_strided_slice_inner_distributed809// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<64x2xf32>, vector<16x1xf32>, vector<64x2xf32>) {810// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<16x16xf32>811// CHECK-NEXT: %[[D:.*]] = "some_def"() : () -> vector<64x32xf32>812// CHECK: gpu.yield %{{.*}}, %[[S]], %[[D]] : vector<64x32xf32>, vector<16x16xf32>, vector<64x32xf32>813// CHECK-NEXT: }814// CHECK-NEXT: %[[T1:.*]] = vector.insert_strided_slice %[[W]]#1, %[[W]]#2815// CHECK-SAME: {offsets = [24, 1], strides = [1, 1]} : vector<16x1xf32> into vector<64x2xf32>816// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<64x2xf32>) -> ()817gpu.func @vector_insert_strided_slice_inner_distributed(%laneid: index) {818 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<64x2xf32>) {819 %0 = "some_def"() : () -> (vector<16x16xf32>)820 %1 = "some_def"() : () -> (vector<64x32xf32>)821 %2 = vector.insert_strided_slice %0, %1 { offsets = [24, 16], strides = [1, 1],822 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,823 layout_operand_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,824 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>825 }826 : vector<16x16xf32> into vector<64x32xf32>827 gpu.yield %2 : vector<64x32xf32>828 }829 "some_use"(%r) : (vector<64x2xf32>) -> ()830 gpu.return831}832 833// CHECK-LABEL: gpu.func @vector_insert_strided_slice_outer_distributed834// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<3x32xf32>, vector<1x16xf32>, vector<3x32xf32>) {835// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<16x16xf32>836// CHECK-NEXT: %[[D:.*]] = "some_def"() : () -> vector<48x32xf32>837// CHECK: gpu.yield %{{.*}}, %[[S]], %[[D]] : vector<48x32xf32>, vector<16x16xf32>, vector<48x32xf32>838// CHECK-NEXT: }839// CHECK-NEXT: %[[T1:.*]] = vector.insert_strided_slice %[[W]]#1, %[[W]]#2840// CHECK-SAME: {offsets = [2, 4], strides = [1, 1]} : vector<1x16xf32> into vector<3x32xf32>841// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<3x32xf32>) -> ()842gpu.func @vector_insert_strided_slice_outer_distributed(%laneid: index) {843 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<3x32xf32>) {844 %0 = "some_def"() : () -> (vector<16x16xf32>)845 %1 = "some_def"() : () -> (vector<48x32xf32>)846 %2 = vector.insert_strided_slice %0, %1 { offsets = [32, 4], strides = [1, 1],847 layout_operand_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,848 layout_operand_1 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,849 layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>850 }851 : vector<16x16xf32> into vector<48x32xf32>852 gpu.yield %2 : vector<48x32xf32>853 }854 "some_use"(%r) : (vector<3x32xf32>) -> ()855 gpu.return856}857 858// CHECK-LABEL: gpu.func @vector_insert_strided_slice_1d859// CHECK: %[[W:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<3xf32>, vector<1xf32>, vector<3xf32>) {860// CHECK-NEXT: %[[S:.*]] = "some_def"() : () -> vector<16xf32>861// CHECK-NEXT: %[[D:.*]] = "some_def"() : () -> vector<48xf32>862// CHECK: gpu.yield %{{.*}}, %[[S]], %[[D]] : vector<48xf32>, vector<16xf32>, vector<48xf32>863// CHECK-NEXT: }864// CHECK-NEXT: %[[T1:.*]] = vector.insert_strided_slice %[[W]]#1, %[[W]]#2865// CHECK-SAME: {offsets = [1], strides = [1]} : vector<1xf32> into vector<3xf32>866// CHECK-NEXT: "some_use"(%[[T1]]) : (vector<3xf32>) -> ()867gpu.func @vector_insert_strided_slice_1d(%laneid: index) {868 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<3xf32>) {869 %0 = "some_def"() : () -> (vector<16xf32>)870 %1 = "some_def"() : () -> (vector<48xf32>)871 %2 = vector.insert_strided_slice %0, %1 { offsets = [16], strides = [1],872 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,873 layout_operand_1 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,874 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>875 }876 : vector<16xf32> into vector<48xf32>877 gpu.yield %2 : vector<48xf32>878 }879 "some_use"(%r) : (vector<3xf32>) -> ()880 gpu.return881}882 883// CHECK-LABEL: gpu.func @vector_insert_strided_slice_unsupported_source884// CHECK: %{{.*}} = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<3xf32>) {885// CHECK: }886// CHECK-NOT: %{{.*}} = vector.insert_strided_slice887gpu.func @vector_insert_strided_slice_unsupported_source(%laneid: index) {888 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<3xf32>) {889 %0 = "some_def"() : () -> (vector<8xf32>)890 %1 = "some_def"() : () -> (vector<48xf32>)891 %2 = vector.insert_strided_slice %0, %1 { offsets = [16], strides = [1],892 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,893 layout_operand_1 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,894 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>895 }896 : vector<8xf32> into vector<48xf32>897 gpu.yield %2 : vector<48xf32>898 }899 "some_use"(%r) : (vector<3xf32>) -> ()900 gpu.return901}902 903// CHECK-LABEL: gpu.func @vector_insert_strided_slice_unsupported_offset904// CHECK: %{{.*}} = gpu.warp_execute_on_lane_0(%{{.*}})[16] -> (vector<3xf32>) {905// CHECK: }906// CHECK-NOT: %{{.*}} = vector.insert_strided_slice907gpu.func @vector_insert_strided_slice_unsupported_offset(%laneid: index) {908 %r = gpu.warp_execute_on_lane_0(%laneid)[16] -> (vector<3xf32>) {909 %0 = "some_def"() : () -> (vector<16xf32>)910 %1 = "some_def"() : () -> (vector<48xf32>)911 %2 = vector.insert_strided_slice %0, %1 { offsets = [3], strides = [1],912 layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,913 layout_operand_1 = #xegpu.layout<lane_layout = [16], lane_data = [1]>,914 layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>915 }916 : vector<16xf32> into vector<48xf32>917 gpu.yield %2 : vector<48xf32>918 }919 "some_use"(%r) : (vector<3xf32>) -> ()920 gpu.return921}922 923}924