1982 lines · plain
1// RUN: mlir-opt %s --allow-unregistered-dialect --split-input-file \2// RUN: --test-vector-warp-distribute=rewrite-warp-ops-to-scf-if | FileCheck %s --check-prefix=CHECK-SCF-IF3 4// RUN: mlir-opt %s --allow-unregistered-dialect --split-input-file \5// RUN: --test-vector-warp-distribute="hoist-uniform" | FileCheck --check-prefixes=CHECK-HOIST %s6 7// RUN: mlir-opt %s --allow-unregistered-dialect --split-input-file \8// RUN: --test-vector-warp-distribute="hoist-uniform distribute-transfer-write max-transfer-write-elements=4" \9// RUN: | FileCheck --check-prefixes=CHECK-D %s10 11// RUN: mlir-opt %s --allow-unregistered-dialect --split-input-file \12// RUN: --test-vector-warp-distribute=propagate-distribution --canonicalize \13// RUN: | FileCheck --check-prefixes=CHECK-PROP %s14 15// RUN: mlir-opt %s --allow-unregistered-dialect --split-input-file \16// RUN: --test-vector-warp-distribute="hoist-uniform distribute-transfer-write propagate-distribution" \17// RUN: --canonicalize | FileCheck --check-prefixes=CHECK-DIST-AND-PROP %s18 19// CHECK-SCF-IF-DAG: #[[$TIMES2:.*]] = affine_map<()[s0] -> (s0 * 2)>20// CHECK-SCF-IF-DAG: #[[$TIMES4:.*]] = affine_map<()[s0] -> (s0 * 4)>21// CHECK-SCF-IF-DAG: #[[$TIMES8:.*]] = affine_map<()[s0] -> (s0 * 8)>22// CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, 3>23// CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, 3>24// CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, 3>25// CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, 3>26 27// CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if(28// CHECK-SCF-IF-SAME: %[[laneid:.*]]: index,29// CHECK-SCF-IF-SAME: %[[v0:.*]]: vector<4xf32>, %[[v1:.*]]: vector<8xf32>)30func.func @rewrite_warp_op_to_scf_if(%laneid: index,31 %v0: vector<4xf32>, %v1: vector<8xf32>) {32// CHECK-SCF-IF-DAG: %[[c0:.*]] = arith.constant 0 : index33// CHECK-SCF-IF: %[[is_lane_0:.*]] = arith.cmpi eq, %[[laneid]], %[[c0]]34 35// CHECK-SCF-IF: %[[buffer_v0:.*]] = memref.get_global @__shared_128xf3236// CHECK-SCF-IF: %[[s0:.*]] = affine.apply #[[$TIMES4]]()[%[[laneid]]]37// CHECK-SCF-IF: vector.transfer_write %[[v0]], %[[buffer_v0]][%[[s0]]]38// CHECK-SCF-IF: %[[buffer_v1:.*]] = memref.get_global @__shared_256xf3239// CHECK-SCF-IF: %[[s1:.*]] = affine.apply #[[$TIMES8]]()[%[[laneid]]]40// CHECK-SCF-IF: vector.transfer_write %[[v1]], %[[buffer_v1]][%[[s1]]]41 42// CHECK-SCF-IF-DAG: gpu.barrier43// CHECK-SCF-IF-DAG: %[[buffer_def_0:.*]] = memref.get_global @__shared_32xf3244// CHECK-SCF-IF-DAG: %[[buffer_def_1:.*]] = memref.get_global @__shared_64xf3245 46// CHECK-SCF-IF: scf.if %[[is_lane_0]] {47 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]48 args(%v0, %v1 : vector<4xf32>, vector<8xf32>) -> (vector<1xf32>, vector<2xf32>) {49 ^bb0(%arg0: vector<128xf32>, %arg1: vector<256xf32>):50// CHECK-SCF-IF: %[[arg1:.*]] = vector.transfer_read %[[buffer_v1]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<256xf32, 3>, vector<256xf32>51// CHECK-SCF-IF: %[[arg0:.*]] = vector.transfer_read %[[buffer_v0]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<128xf32, 3>, vector<128xf32>52// CHECK-SCF-IF: %[[def_0:.*]] = "some_def"(%[[arg0]]) : (vector<128xf32>) -> vector<32xf32>53// CHECK-SCF-IF: %[[def_1:.*]] = "some_def"(%[[arg1]]) : (vector<256xf32>) -> vector<64xf32>54 %2 = "some_def"(%arg0) : (vector<128xf32>) -> vector<32xf32>55 %3 = "some_def"(%arg1) : (vector<256xf32>) -> vector<64xf32>56// CHECK-SCF-IF: vector.transfer_write %[[def_0]], %[[buffer_def_0]][%[[c0]]]57// CHECK-SCF-IF: vector.transfer_write %[[def_1]], %[[buffer_def_1]][%[[c0]]]58 gpu.yield %2, %3 : vector<32xf32>, vector<64xf32>59 }60// CHECK-SCF-IF: }61// CHECK-SCF-IF: gpu.barrier62// CHECK-SCF-IF: %[[o1:.*]] = affine.apply #[[$TIMES2]]()[%[[laneid]]]63// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, 3>, vector<2xf32>64// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, 3>, vector<1xf32>65// CHECK-SCF-IF: "some_use"(%[[r0]]) : (vector<1xf32>) -> ()66// CHECK-SCF-IF: "some_use"(%[[r1]]) : (vector<2xf32>) -> ()67 "some_use"(%r#0) : (vector<1xf32>) -> ()68 "some_use"(%r#1) : (vector<2xf32>) -> ()69 return70}71 72// -----73 74// CHECK-D-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 * 2 + 32)>75 76// CHECK-DIST-AND-PROP-LABEL: func @warp(77// CHECK-HOIST: memref.subview78// CHECK-HOIST: memref.subview79// CHECK-HOIST: memref.subview80// CHECK-HOIST: gpu.warp_execute_on_lane_081 82// CHECK-D: %[[R:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<2xf32>, vector<1xf32>) {83// CHECK-D: arith.addf {{.*}} : vector<32xf32>84// CHECK-D: arith.addf {{.*}} : vector<64xf32>85// CHECK-D: gpu.yield %{{.*}}, %{{.*}} : vector<64xf32>, vector<32xf32>86// CHECK-D-DAG: vector.transfer_write %[[R]]#1, %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<1xf32>, memref<128xf3287// CHECK-D-DAG: %[[ID1:.*]] = affine.apply #[[MAP1]]()[%{{.*}}]88// CHECK-D-DAG: vector.transfer_write %[[R]]#0, %{{.*}}[%[[ID1]]] {in_bounds = [true]} : vector<2xf32>, memref<128xf3289 90// CHECK-DIST-AND-PROP-NOT: gpu.warp_execute_on_lane_091// CHECK-DIST-AND-PROP: vector.transfer_read {{.*}} vector<1xf32>92// CHECK-DIST-AND-PROP: vector.transfer_read {{.*}} vector<1xf32>93// CHECK-DIST-AND-PROP: vector.transfer_read {{.*}} vector<2xf32>94// CHECK-DIST-AND-PROP: vector.transfer_read {{.*}} vector<2xf32>95// CHECK-DIST-AND-PROP: arith.addf {{.*}} : vector<1xf32>96// CHECK-DIST-AND-PROP: arith.addf {{.*}} : vector<2xf32>97// CHECK-DIST-AND-PROP: vector.transfer_write {{.*}} : vector<1xf32>98// CHECK-DIST-AND-PROP: vector.transfer_write {{.*}} : vector<2xf32>99 100func.func @warp(%laneid: index, %arg1: memref<1024xf32>, %arg2: memref<1024xf32>,101 %arg3: memref<1024xf32>, %gid : index) {102 gpu.warp_execute_on_lane_0(%laneid)[32] {103 %sa = memref.subview %arg1[%gid] [128] [1] : memref<1024xf32> to memref<128xf32, strided<[1], offset: ?>>104 %sb = memref.subview %arg2[%gid] [128] [1] : memref<1024xf32> to memref<128xf32, strided<[1], offset: ?>>105 %sc = memref.subview %arg3[%gid] [128] [1] : memref<1024xf32> to memref<128xf32, strided<[1], offset: ?>>106 %c0 = arith.constant 0 : index107 %c32 = arith.constant 32 : index108 %cst = arith.constant 0.000000e+00 : f32109 %2 = vector.transfer_read %sa[%c0], %cst : memref<128xf32, strided<[1], offset: ?>>, vector<32xf32>110 %3 = vector.transfer_read %sa[%c32], %cst : memref<128xf32, strided<[1], offset: ?>>, vector<32xf32>111 %4 = vector.transfer_read %sb[%c0], %cst : memref<128xf32, strided<[1], offset: ?>>, vector<64xf32>112 %5 = vector.transfer_read %sb[%c32], %cst : memref<128xf32, strided<[1], offset: ?>>, vector<64xf32>113 %6 = arith.addf %2, %3 : vector<32xf32>114 %7 = arith.addf %4, %5 : vector<64xf32>115 vector.transfer_write %6, %sc[%c0] : vector<32xf32>, memref<128xf32, strided<[1], offset: ?>>116 vector.transfer_write %7, %sc[%c32] : vector<64xf32>, memref<128xf32, strided<[1], offset: ?>>117 }118 return119}120 121// -----122 123// CHECK-D-LABEL: func @warp_extract(124// CHECK-D: %[[WARPOP:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>, vector<1x1xf32>)125// CHECK-D: "test.dummy_op"126// CHECK-D: "test.dummy_op"127// CHECK-D: gpu.yield %{{.*}}, %{{.*}} : vector<1xf32>, vector<1x1xf32>128// CHECK-D: }129// CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {130// CHECK-D: vector.transfer_write %[[WARPOP]]#1, %{{.*}}[%{{.*}}] {{.*}} : vector<1x1xf32>131// CHECK-D: }132// CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {133// CHECK-D: vector.transfer_write %[[WARPOP]]#0, %{{.*}}[%{{.*}}] {{.*}} : vector<1xf32>134// CHECK-D: }135 136func.func @warp_extract(%laneid: index, %arg1: memref<1024x1024xf32>, %gid : index) {137 gpu.warp_execute_on_lane_0(%laneid)[32] {138 %c0 = arith.constant 0 : index139 %v = "test.dummy_op"() : () -> (vector<1xf32>)140 %v1 = "test.dummy_op"() : () -> (vector<1x1xf32>)141 vector.transfer_write %v1, %arg1[%c0, %c0] : vector<1x1xf32>, memref<1024x1024xf32>142 vector.transfer_write %v, %arg1[%c0, %c0] : vector<1xf32>, memref<1024x1024xf32>143 }144 return145}146 147// -----148 149// Check that we can distribute writes of the maximum allowed number of elements.150 151// CHECK-D-LABEL: func @warp_extract_4_elems(152// CHECK-D: %[[WARPOP:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>, vector<4x1xf32>)153// CHECK-D: "test.dummy_op"154// CHECK-D: "test.dummy_op"155// CHECK-D: gpu.yield %{{.*}}, %{{.*}} : vector<4xf32>, vector<4x1xf32>156// CHECK-D: }157// CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {158// CHECK-D: vector.transfer_write %[[WARPOP]]#1, %{{.*}}[%{{.*}}] {{.*}} : vector<4x1xf32>159// CHECK-D: }160// CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {161// CHECK-D: vector.transfer_write %[[WARPOP]]#0, %{{.*}}[%{{.*}}] {{.*}} : vector<4xf32>162// CHECK-D: }163 164func.func @warp_extract_4_elems(%laneid: index, %arg1: memref<1024x1024xf32>, %gid : index) {165 gpu.warp_execute_on_lane_0(%laneid)[32] {166 %c0 = arith.constant 0 : index167 %v = "test.dummy_op"() : () -> (vector<4xf32>)168 %v1 = "test.dummy_op"() : () -> (vector<4x1xf32>)169 vector.transfer_write %v1, %arg1[%c0, %c0] : vector<4x1xf32>, memref<1024x1024xf32>170 vector.transfer_write %v, %arg1[%c0, %c0] : vector<4xf32>, memref<1024x1024xf32>171 }172 return173}174 175// -----176 177// Check that we do not distribute writes larger than the maximum allowed178// number of elements.179 180// CHECK-D-LABEL: func @warp_extract_5_elems(181// CHECK-D: arith.constant 0 : index182// CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {183// CHECK-D: %[[V:.+]] = "test.dummy_op"184// CHECK-D: %[[V1:.+]] = "test.dummy_op"185// CHECK-D: vector.transfer_write %[[V1]], %{{.*}}[%{{.*}}] {{.*}} : vector<5x1xf32>186// CHECK-D: vector.transfer_write %[[V]], %{{.*}}[%{{.*}}] {{.*}} : vector<5xf32>187// CHECK-D: }188 189func.func @warp_extract_5_elems(%laneid: index, %arg1: memref<1024x1024xf32>, %gid : index) {190 gpu.warp_execute_on_lane_0(%laneid)[32] {191 %c0 = arith.constant 0 : index192 %v = "test.dummy_op"() : () -> (vector<5xf32>)193 %v1 = "test.dummy_op"() : () -> (vector<5x1xf32>)194 vector.transfer_write %v1, %arg1[%c0, %c0] : vector<5x1xf32>, memref<1024x1024xf32>195 vector.transfer_write %v, %arg1[%c0, %c0] : vector<5xf32>, memref<1024x1024xf32>196 }197 return198}199 200// -----201 202// Check that we do not distribute writes larger than the maximum allowed203// number of elements, or multiples of the maximum number of elements.204 205// CHECK-D-LABEL: func @warp_extract_8_elems(206// CHECK-D: arith.constant 0 : index207// CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {208// CHECK-D: %[[V:.+]] = "test.dummy_op"209// CHECK-D: %[[V1:.+]] = "test.dummy_op"210// CHECK-D: vector.transfer_write %[[V1]], %{{.*}}[%{{.*}}] {{.*}} : vector<8x1xf32>211// CHECK-D: vector.transfer_write %[[V]], %{{.*}}[%{{.*}}] {{.*}} : vector<8xf32>212// CHECK-D: }213 214func.func @warp_extract_8_elems(%laneid: index, %arg1: memref<1024x1024xf32>, %gid : index) {215 gpu.warp_execute_on_lane_0(%laneid)[32] {216 %c0 = arith.constant 0 : index217 %v = "test.dummy_op"() : () -> (vector<8xf32>)218 %v1 = "test.dummy_op"() : () -> (vector<8x1xf32>)219 vector.transfer_write %v1, %arg1[%c0, %c0] : vector<8x1xf32>, memref<1024x1024xf32>220 vector.transfer_write %v, %arg1[%c0, %c0] : vector<8xf32>, memref<1024x1024xf32>221 }222 return223}224 225// -----226 227// CHECK-PROP-LABEL: func @warp_dead_result(228func.func @warp_dead_result(%laneid: index) -> (vector<1xf32>) {229 // CHECK-PROP: %[[R:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>)230 %r:3 = gpu.warp_execute_on_lane_0(%laneid)[32] ->231 (vector<1xf32>, vector<1xf32>, vector<1xf32>) {232 %2 = "some_def"() : () -> (vector<32xf32>)233 %3 = "some_def"() : () -> (vector<32xf32>)234 %4 = "some_def"() : () -> (vector<32xf32>)235 // CHECK-PROP: gpu.yield %{{.*}} : vector<32xf32>236 gpu.yield %2, %3, %4 : vector<32xf32>, vector<32xf32>, vector<32xf32>237 }238 // CHECK-PROP: return %[[R]] : vector<1xf32>239 return %r#1 : vector<1xf32>240}241 242// -----243 244// CHECK-PROP-LABEL: func @warp_propagate_operand(245// CHECK-PROP-SAME: %[[ID:.*]]: index, %[[V:.*]]: vector<4xf32>)246func.func @warp_propagate_operand(%laneid: index, %v0: vector<4xf32>)247 -> (vector<4xf32>) {248 %r = gpu.warp_execute_on_lane_0(%laneid)[32]249 args(%v0 : vector<4xf32>) -> (vector<4xf32>) {250 ^bb0(%arg0 : vector<128xf32>) :251 gpu.yield %arg0 : vector<128xf32>252 }253 // CHECK-PROP: return %[[V]] : vector<4xf32>254 return %r : vector<4xf32>255}256 257// -----258 259#map0 = affine_map<()[s0] -> (s0 * 2)>260 261// CHECK-PROP-LABEL: func @warp_propagate_elementwise(262func.func @warp_propagate_elementwise(%laneid: index, %dest: memref<1024xf32>) {263 %c0 = arith.constant 0 : index264 %c32 = arith.constant 0 : index265 %cst = arith.constant 0.000000e+00 : f32266 // CHECK-PROP: %[[R:.*]]:4 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>, vector<1xf32>, vector<2xf32>, vector<2xf32>)267 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32] ->268 (vector<1xf32>, vector<2xf32>) {269 // CHECK-PROP: %[[V0:.*]] = "some_def"() : () -> vector<32xf32>270 // CHECK-PROP: %[[V1:.*]] = "some_def"() : () -> vector<32xf32>271 // CHECK-PROP: %[[V2:.*]] = "some_def"() : () -> vector<64xf32>272 // CHECK-PROP: %[[V3:.*]] = "some_def"() : () -> vector<64xf32>273 // CHECK-PROP: gpu.yield %[[V0]], %[[V1]], %[[V2]], %[[V3]] : vector<32xf32>, vector<32xf32>, vector<64xf32>, vector<64xf32>274 %2 = "some_def"() : () -> (vector<32xf32>)275 %3 = "some_def"() : () -> (vector<32xf32>)276 %4 = "some_def"() : () -> (vector<64xf32>)277 %5 = "some_def"() : () -> (vector<64xf32>)278 %6 = arith.addf %2, %3 : vector<32xf32>279 %7 = arith.addf %4, %5 : vector<64xf32>280 gpu.yield %6, %7 : vector<32xf32>, vector<64xf32>281 }282 // CHECK-PROP: %[[A0:.*]] = arith.addf %[[R]]#2, %[[R]]#3 : vector<2xf32>283 // CHECK-PROP: %[[A1:.*]] = arith.addf %[[R]]#0, %[[R]]#1 : vector<1xf32>284 %id2 = affine.apply #map0()[%laneid]285 // CHECK-PROP: vector.transfer_write %[[A1]], {{.*}} : vector<1xf32>, memref<1024xf32>286 // CHECK-PROP: vector.transfer_write %[[A0]], {{.*}} : vector<2xf32>, memref<1024xf32>287 vector.transfer_write %r#0, %dest[%laneid] : vector<1xf32>, memref<1024xf32>288 vector.transfer_write %r#1, %dest[%id2] : vector<2xf32>, memref<1024xf32>289 return290}291 292// -----293 294// CHECK-PROP-LABEL: func @warp_propagate_scalar_arith(295// CHECK-PROP: %[[r:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} {296// CHECK-PROP: %[[some_def0:.*]] = "some_def"297// CHECK-PROP: %[[some_def1:.*]] = "some_def"298// CHECK-PROP: gpu.yield %[[some_def0]], %[[some_def1]]299// CHECK-PROP: }300// CHECK-PROP: arith.addf %[[r]]#0, %[[r]]#1 : f32301func.func @warp_propagate_scalar_arith(%laneid: index) {302 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {303 %0 = "some_def"() : () -> (f32)304 %1 = "some_def"() : () -> (f32)305 %2 = arith.addf %0, %1 : f32306 gpu.yield %2 : f32307 }308 vector.print %r : f32309 return310}311 312// -----313 314// CHECK-PROP-LABEL: func @warp_propagate_cast(315// CHECK-PROP-NOT: gpu.warp_execute_on_lane_0316// CHECK-PROP: %[[result:.*]] = arith.sitofp %{{.*}} : i32 to f32317// CHECK-PROP: return %[[result]]318func.func @warp_propagate_cast(%laneid : index, %i : i32) -> (f32) {319 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {320 %casted = arith.sitofp %i : i32 to f32321 gpu.yield %casted : f32322 }323 return %r : f32324}325 326// -----327 328#map0 = affine_map<()[s0] -> (s0 * 2)>329 330// CHECK-PROP-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 * 2)>331 332// CHECK-PROP: func @warp_propagate_read333// CHECK-PROP-SAME: (%[[ID:.*]]: index334func.func @warp_propagate_read(%laneid: index, %src: memref<1024xf32>, %dest: memref<1024xf32>) {335// CHECK-PROP-NOT: warp_execute_on_lane_0336// CHECK-PROP-DAG: %[[R0:.*]] = vector.transfer_read %arg1[%[[ID]]], %{{.*}} : memref<1024xf32>, vector<1xf32>337// CHECK-PROP-DAG: %[[ID2:.*]] = affine.apply #[[MAP0]]()[%[[ID]]]338// CHECK-PROP-DAG: %[[R1:.*]] = vector.transfer_read %arg1[%[[ID2]]], %{{.*}} : memref<1024xf32>, vector<2xf32>339// CHECK-PROP: vector.transfer_write %[[R0]], {{.*}} : vector<1xf32>, memref<1024xf32>340// CHECK-PROP: vector.transfer_write %[[R1]], {{.*}} : vector<2xf32>, memref<1024xf32>341 %c0 = arith.constant 0 : index342 %c32 = arith.constant 0 : index343 %cst = arith.constant 0.000000e+00 : f32344 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32] ->(vector<1xf32>, vector<2xf32>) {345 %2 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>, vector<32xf32>346 %3 = vector.transfer_read %src[%c32], %cst : memref<1024xf32>, vector<64xf32>347 gpu.yield %2, %3 : vector<32xf32>, vector<64xf32>348 }349 %id2 = affine.apply #map0()[%laneid]350 vector.transfer_write %r#0, %dest[%laneid] : vector<1xf32>, memref<1024xf32>351 vector.transfer_write %r#1, %dest[%id2] : vector<2xf32>, memref<1024xf32>352 return353}354 355// -----356 357// CHECK-PROP-LABEL: func @fold_vector_broadcast(358// CHECK-PROP: %[[r:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (vector<1xf32>)359// CHECK-PROP: %[[some_def:.*]] = "some_def"360// CHECK-PROP: gpu.yield %[[some_def]] : vector<1xf32>361// CHECK-PROP: vector.print %[[r]] : vector<1xf32>362func.func @fold_vector_broadcast(%laneid: index) {363 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {364 %0 = "some_def"() : () -> (vector<1xf32>)365 %1 = vector.broadcast %0 : vector<1xf32> to vector<32xf32>366 gpu.yield %1 : vector<32xf32>367 }368 vector.print %r : vector<1xf32>369 return370}371 372// -----373 374// CHECK-PROP-LABEL: func @extract_vector_broadcast(375// CHECK-PROP: %[[r:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (vector<1xf32>)376// CHECK-PROP: %[[some_def:.*]] = "some_def"377// CHECK-PROP: gpu.yield %[[some_def]] : vector<1xf32>378// CHECK-PROP: %[[broadcasted:.*]] = vector.broadcast %[[r]] : vector<1xf32> to vector<2xf32>379// CHECK-PROP: vector.print %[[broadcasted]] : vector<2xf32>380func.func @extract_vector_broadcast(%laneid: index) {381 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>) {382 %0 = "some_def"() : () -> (vector<1xf32>)383 %1 = vector.broadcast %0 : vector<1xf32> to vector<64xf32>384 gpu.yield %1 : vector<64xf32>385 }386 vector.print %r : vector<2xf32>387 return388}389 390// -----391 392// CHECK-PROP-LABEL: func @extract_scalar_vector_broadcast(393// CHECK-PROP: %[[r:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (f32)394// CHECK-PROP: %[[some_def:.*]] = "some_def"395// CHECK-PROP: gpu.yield %[[some_def]] : f32396// CHECK-PROP: %[[broadcasted:.*]] = vector.broadcast %[[r]] : f32 to vector<2xf32>397// CHECK-PROP: vector.print %[[broadcasted]] : vector<2xf32>398func.func @extract_scalar_vector_broadcast(%laneid: index) {399 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>) {400 %0 = "some_def"() : () -> (f32)401 %1 = vector.broadcast %0 : f32 to vector<64xf32>402 gpu.yield %1 : vector<64xf32>403 }404 vector.print %r : vector<2xf32>405 return406}407 408// -----409 410// CHECK-PROP-LABEL: func @warp_scf_for(411// CHECK-PROP: %[[INI:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>) {412// CHECK-PROP: %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>413// CHECK-PROP: gpu.yield %[[INI1]] : vector<128xf32>414// CHECK-PROP: }415// CHECK-PROP: %[[F:.*]] = scf.for %[[IT:.+]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[FARG:.*]] = %[[INI]]) -> (vector<4xf32>) {416// CHECK-PROP: %[[A:.*]] = arith.addi %[[IT]], %{{.*}} : index417// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%[[FARG]] : vector<4xf32>) -> (vector<4xf32>) {418// CHECK-PROP: ^bb0(%[[ARG:.*]]: vector<128xf32>):419// CHECK-PROP: %[[ACC:.*]] = "some_def"(%[[A]], %[[ARG]]) : (index, vector<128xf32>) -> vector<128xf32>420// CHECK-PROP: gpu.yield %[[ACC]] : vector<128xf32>421// CHECK-PROP: }422// CHECK-PROP: scf.yield %[[W]] : vector<4xf32>423// CHECK-PROP: }424// CHECK-PROP: "some_use"(%[[F]]) : (vector<4xf32>) -> ()425func.func @warp_scf_for(%arg0: index) {426 %c128 = arith.constant 128 : index427 %c1 = arith.constant 1 : index428 %c0 = arith.constant 0 : index429 %0 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>) {430 %ini = "some_def"() : () -> (vector<128xf32>)431 %3 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini) -> (vector<128xf32>) {432 %add = arith.addi %arg3, %c1 : index433 %acc = "some_def"(%add, %arg4) : (index, vector<128xf32>) -> (vector<128xf32>)434 scf.yield %acc : vector<128xf32>435 }436 gpu.yield %3 : vector<128xf32>437 }438 "some_use"(%0) : (vector<4xf32>) -> ()439 return440}441 442// -----443 444// CHECK-PROP-LABEL: func @warp_scf_for_use_from_above(445// CHECK-PROP: %[[INI:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>, vector<4xf32>) {446// CHECK-PROP: %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>447// CHECK-PROP: %[[USE:.*]] = "some_def_above"() : () -> vector<128xf32>448// CHECK-PROP: gpu.yield %[[INI1]], %[[USE]] : vector<128xf32>, vector<128xf32>449// CHECK-PROP: }450// CHECK-PROP: %[[F:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[FARG:.*]] = %[[INI]]#0) -> (vector<4xf32>) {451// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%[[FARG]], %[[INI]]#1 : vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) {452// CHECK-PROP: ^bb0(%[[ARG0:.*]]: vector<128xf32>, %[[ARG1:.*]]: vector<128xf32>):453// CHECK-PROP: %[[ACC:.*]] = "some_def"(%[[ARG0]], %[[ARG1]]) : (vector<128xf32>, vector<128xf32>) -> vector<128xf32>454// CHECK-PROP: gpu.yield %[[ACC]] : vector<128xf32>455// CHECK-PROP: }456// CHECK-PROP: scf.yield %[[W]] : vector<4xf32>457// CHECK-PROP: }458// CHECK-PROP: "some_use"(%[[F]]) : (vector<4xf32>) -> ()459func.func @warp_scf_for_use_from_above(%arg0: index) {460 %c128 = arith.constant 128 : index461 %c1 = arith.constant 1 : index462 %c0 = arith.constant 0 : index463 %0 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>) {464 %ini = "some_def"() : () -> (vector<128xf32>)465 %use_from_above = "some_def_above"() : () -> (vector<128xf32>)466 %3 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini) -> (vector<128xf32>) {467 %acc = "some_def"(%arg4, %use_from_above) : (vector<128xf32>, vector<128xf32>) -> (vector<128xf32>)468 scf.yield %acc : vector<128xf32>469 }470 gpu.yield %3 : vector<128xf32>471 }472 "some_use"(%0) : (vector<4xf32>) -> ()473 return474}475 476// -----477// CHECK-PROP-LABEL: func.func @warp_scf_for_local_loop_bounds478// CHECK-PROP: (%{{.*}}: index, %[[ARG1:[a-zA-Z0-9]+]]: index) {479// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%[[ARG1]] : index) -> (vector<4xf32>) {480// CHECK-PROP: ^bb0(%{{.*}}: index):481// CHECK-PROP: %[[T2:.*]] = "some_def"() : () -> vector<128xf32>482// CHECK-PROP: gpu.yield %[[T2]] : vector<128xf32>483// CHECK-PROP: }484// CHECK-PROP: %[[FOR:.*]] = scf.for %{{.*}} to %[[ARG1]] step %{{.*}} iter_args(%{{.*}}) -> (vector<4xf32>) {485// CHECK-PROP: %[[W2:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32]486// CHECK-PROP-SAME: args(%{{.*}} : vector<4xf32>) -> (vector<4xf32>) {487// CHECK-PROP: ^bb0(%{{.*}}: vector<128xf32>):488// CHECK-PROP: gpu.yield %{{.*}} : vector<128xf32>489// CHECK-PROP: }490// CHECK-PROP: scf.yield %[[W2]] : vector<4xf32>491// CHECK-PROP: }492// CHECK-PROP: "some_use"(%[[FOR]]) : (vector<4xf32>) -> ()493// CHECK-PROP: return494func.func @warp_scf_for_local_loop_bounds(%arg0: index, %bound: index) {495 %c1 = arith.constant 1 : index496 %c0 = arith.constant 0 : index497 %0 = gpu.warp_execute_on_lane_0(%arg0)[32]498 args(%bound : index) -> (vector<4xf32>) {499 ^bb0(%arg1: index):500 %ini = "some_def"() : () -> (vector<128xf32>)501 %3 = scf.for %arg3 = %c0 to %arg1 step %c1 iter_args(%arg4 = %ini) -> (vector<128xf32>) {502 %acc = "some_def"(%arg4) : (vector<128xf32>) -> (vector<128xf32>)503 scf.yield %acc : vector<128xf32>504 }505 gpu.yield %3 : vector<128xf32>506 }507 "some_use"(%0) : (vector<4xf32>) -> ()508 return509}510 511// -----512 513// CHECK-PROP-LABEL: func @warp_scf_for_swap(514// CHECK-PROP: %[[INI:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>, vector<4xf32>) {515// CHECK-PROP: %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>516// CHECK-PROP: %[[INI2:.*]] = "some_def"() : () -> vector<128xf32>517// CHECK-PROP: gpu.yield %[[INI1]], %[[INI2]] : vector<128xf32>, vector<128xf32>518// CHECK-PROP: }519// CHECK-PROP: %[[F:.*]]:2 = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[FARG1:.*]] = %[[INI]]#0, %[[FARG2:.*]] = %[[INI]]#1) -> (vector<4xf32>, vector<4xf32>) {520// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%[[FARG1]], %[[FARG2]] : vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) {521// CHECK-PROP: ^bb0(%[[ARG1:.*]]: vector<128xf32>, %[[ARG2:.*]]: vector<128xf32>):522// CHECK-PROP: %[[ACC1:.*]] = "some_def"(%[[ARG1]]) : (vector<128xf32>) -> vector<128xf32>523// CHECK-PROP: %[[ACC2:.*]] = "some_def"(%[[ARG2]]) : (vector<128xf32>) -> vector<128xf32>524// CHECK-PROP: gpu.yield %[[ACC2]], %[[ACC1]] : vector<128xf32>, vector<128xf32>525// CHECK-PROP: }526// CHECK-PROP: scf.yield %[[W]]#0, %[[W]]#1 : vector<4xf32>, vector<4xf32>527// CHECK-PROP: }528// CHECK-PROP: "some_use"(%[[F]]#0) : (vector<4xf32>) -> ()529// CHECK-PROP: "some_use"(%[[F]]#1) : (vector<4xf32>) -> ()530func.func @warp_scf_for_swap(%arg0: index) {531 %c128 = arith.constant 128 : index532 %c1 = arith.constant 1 : index533 %c0 = arith.constant 0 : index534 %0:2 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>, vector<4xf32>) {535 %ini1 = "some_def"() : () -> (vector<128xf32>)536 %ini2 = "some_def"() : () -> (vector<128xf32>)537 %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini1, %arg5 = %ini2) -> (vector<128xf32>, vector<128xf32>) {538 %acc1 = "some_def"(%arg4) : (vector<128xf32>) -> (vector<128xf32>)539 %acc2 = "some_def"(%arg5) : (vector<128xf32>) -> (vector<128xf32>)540 scf.yield %acc2, %acc1 : vector<128xf32>, vector<128xf32>541 }542 gpu.yield %3#0, %3#1 : vector<128xf32>, vector<128xf32>543 }544 "some_use"(%0#0) : (vector<4xf32>) -> ()545 "some_use"(%0#1) : (vector<4xf32>) -> ()546 return547}548 549// -----550 551// CHECK-PROP-LABEL: func @warp_scf_for_swap_no_yield(552// CHECK-PROP: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {553// CHECK-PROP-NEXT: gpu.warp_execute_on_lane_0(%{{.*}})[32] {554// CHECK-PROP-NEXT: "some_op"() : () -> ()555// CHECK-PROP-NEXT: }556// CHECK-PROP-NEXT: }557func.func @warp_scf_for_swap_no_yield(%arg0: index) {558 %c128 = arith.constant 128 : index559 %c1 = arith.constant 1 : index560 %c0 = arith.constant 0 : index561 gpu.warp_execute_on_lane_0(%arg0)[32] {562 scf.for %arg3 = %c0 to %c128 step %c1 {563 "some_op"() : () -> ()564 }565 }566 return567}568 569// -----570// scf.for result is not distributed in this case.571// CHECK-PROP-LABEL: func @warp_scf_for_broadcasted_result(572// CHECK-PROP: %[[W0:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {573// CHECK-PROP: %[[INI:.*]] = "some_def"() : () -> vector<1xf32>574// CHECK-PROP: gpu.yield %[[INI]] : vector<1xf32>575// CHECK-PROP: }576// CHECK-PROP: %[[F:.*]] = scf.for {{.*}} iter_args(%[[ARG2:.*]] = %[[W0]]) -> (vector<1xf32>) {577// CHECK-PROP: %[[W1:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%[[ARG2]] : vector<1xf32>) -> (vector<1xf32>) {578// CHECK-PROP: ^bb0(%{{.*}}: vector<1xf32>):579// CHECK-PROP: %[[T0:.*]] = "some_op"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>580// CHECK-PROP: gpu.yield %[[T0]] : vector<1xf32>581// CHECK-PROP: }582// CHECK-PROP: scf.yield %[[W1]] : vector<1xf32>583func.func @warp_scf_for_broadcasted_result(%arg0: index) -> vector<1xf32> {584 %c128 = arith.constant 128 : index585 %c1 = arith.constant 1 : index586 %c0 = arith.constant 0 : index587 %2 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<1xf32>) {588 %ini = "some_def"() : () -> (vector<1xf32>)589 %0 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini) -> (vector<1xf32>) {590 %1 = "some_op"(%arg4) : (vector<1xf32>) -> (vector<1xf32>)591 scf.yield %1 : vector<1xf32>592 }593 gpu.yield %0 : vector<1xf32>594 }595 return %2 : vector<1xf32>596}597 598// -----599 600#map = affine_map<()[s0] -> (s0 * 4)>601#map1 = affine_map<()[s0] -> (s0 * 128 + 128)>602#map2 = affine_map<()[s0] -> (s0 * 4 + 128)>603 604// CHECK-PROP-LABEL: func @warp_scf_for_multiple_yield(605// CHECK-PROP: gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {606// CHECK-PROP-NEXT: "some_def"() : () -> vector<32xf32>607// CHECK-PROP-NEXT: gpu.yield %{{.*}} : vector<32xf32>608// CHECK-PROP-NEXT: }609// CHECK-PROP-NOT: gpu.warp_execute_on_lane_0610// CHECK-PROP: vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>611// CHECK-PROP: vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>612// CHECK-PROP: %{{.*}}:2 = scf.for {{.*}} -> (vector<4xf32>, vector<4xf32>) {613// CHECK-PROP-NOT: gpu.warp_execute_on_lane_0614// CHECK-PROP: vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>615// CHECK-PROP: vector.transfer_read {{.*}} : memref<?xf32>, vector<4xf32>616// CHECK-PROP: arith.addf {{.*}} : vector<4xf32>617// CHECK-PROP: arith.addf {{.*}} : vector<4xf32>618// CHECK-PROP: scf.yield {{.*}} : vector<4xf32>, vector<4xf32>619// CHECK-PROP: }620func.func @warp_scf_for_multiple_yield(%arg0: index, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {621 %c256 = arith.constant 256 : index622 %c128 = arith.constant 128 : index623 %c1 = arith.constant 1 : index624 %c0 = arith.constant 0 : index625 %cst = arith.constant 0.000000e+00 : f32626 %0:3 = gpu.warp_execute_on_lane_0(%arg0)[32] ->627 (vector<1xf32>, vector<4xf32>, vector<4xf32>) {628 %def = "some_def"() : () -> (vector<32xf32>)629 %r1 = vector.transfer_read %arg2[%c0], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>630 %r2 = vector.transfer_read %arg2[%c128], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>631 %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %r1, %arg5 = %r2)632 -> (vector<128xf32>, vector<128xf32>) {633 %o1 = affine.apply #map1()[%arg3]634 %o2 = affine.apply #map2()[%arg3]635 %4 = vector.transfer_read %arg1[%o1], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>636 %5 = vector.transfer_read %arg1[%o2], %cst {in_bounds = [true]} : memref<?xf32>, vector<128xf32>637 %6 = arith.addf %4, %arg4 : vector<128xf32>638 %7 = arith.addf %5, %arg5 : vector<128xf32>639 scf.yield %6, %7 : vector<128xf32>, vector<128xf32>640 }641 gpu.yield %def, %3#0, %3#1 : vector<32xf32>, vector<128xf32>, vector<128xf32>642 }643 %1 = affine.apply #map()[%arg0]644 vector.transfer_write %0#1, %arg2[%1] {in_bounds = [true]} : vector<4xf32>, memref<?xf32>645 %2 = affine.apply #map2()[%arg0]646 vector.transfer_write %0#2, %arg2[%2] {in_bounds = [true]} : vector<4xf32>, memref<?xf32>647 "some_use"(%0#0) : (vector<1xf32>) -> ()648 return649}650 651// -----652// CHECK-PROP-LABEL: func.func @warp_scf_for_unused_for_result(653// CHECK-PROP: %[[W0:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<4xf32>, vector<4xf32>) {654// CHECK-PROP: %[[INI0:.*]] = "some_def"() : () -> vector<128xf32>655// CHECK-PROP: %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>656// CHECK-PROP: gpu.yield %[[INI0]], %[[INI1]] : vector<128xf32>, vector<128xf32>657// CHECK-PROP: }658// CHECK-PROP: %[[F:.*]]:2 = scf.for %{{.*}} iter_args(%{{.*}} = %[[W0]]#0, %{{.*}} = %[[W0]]#1) -> (vector<4xf32>, vector<4xf32>) {659// CHECK-PROP: %[[W1:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%{{.*}} : vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) {660// CHECK-PROP: %[[ACC0:.*]] = "some_def"(%{{.*}}) : (vector<128xf32>, index) -> vector<128xf32>661// CHECK-PROP: %[[ACC1:.*]] = "some_def"(%{{.*}}) : (index, vector<128xf32>, vector<128xf32>) -> vector<128xf32>662// CHECK-PROP: gpu.yield %[[ACC1]], %[[ACC0]] : vector<128xf32>, vector<128xf32>663// CHECK-PROP: }664// CHECK-PROP: scf.yield %[[W1]]#0, %[[W1]]#1 : vector<4xf32>, vector<4xf32>665// CHECK-PROP: }666// CHECK-PROP: "some_use"(%[[F]]#0) : (vector<4xf32>) -> ()667func.func @warp_scf_for_unused_for_result(%arg0: index) {668 %c128 = arith.constant 128 : index669 %c1 = arith.constant 1 : index670 %c0 = arith.constant 0 : index671 %0 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>) {672 %ini = "some_def"() : () -> (vector<128xf32>)673 %ini1 = "some_def"() : () -> (vector<128xf32>)674 %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini, %arg5 = %ini1) -> (vector<128xf32>, vector<128xf32>) {675 %add = arith.addi %arg3, %c1 : index676 %1 = "some_def"(%arg5, %add) : (vector<128xf32>, index) -> (vector<128xf32>)677 %acc = "some_def"(%add, %arg4, %1) : (index, vector<128xf32>, vector<128xf32>) -> (vector<128xf32>)678 scf.yield %acc, %1 : vector<128xf32>, vector<128xf32>679 }680 gpu.yield %3#0 : vector<128xf32>681 }682 "some_use"(%0) : (vector<4xf32>) -> ()683 return684}685 686// -----687// CHECK-PROP-LABEL: func.func @warp_scf_for_swapped_for_results(688// CHECK-PROP: %[[W0:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<8xf32>, vector<4xf32>, vector<4xf32>) {689// CHECK-PROP-NEXT: %[[INI0:.*]] = "some_def"() : () -> vector<256xf32>690// CHECK-PROP-NEXT: %[[INI1:.*]] = "some_def"() : () -> vector<128xf32>691// CHECK-PROP-NEXT: %[[INI2:.*]] = "some_def"() : () -> vector<128xf32>692// CHECK-PROP-NEXT: gpu.yield %[[INI0]], %[[INI1]], %[[INI2]] : vector<256xf32>, vector<128xf32>, vector<128xf32>693// CHECK-PROP-NEXT: }694// CHECK-PROP-NEXT: %[[F0:.*]]:3 = scf.for {{.*}} iter_args(%{{.*}} = %[[W0]]#0, %{{.*}} = %[[W0]]#1, %{{.*}} = %[[W0]]#2) -> (vector<8xf32>, vector<4xf32>, vector<4xf32>) {695// CHECK-PROP-NEXT: %[[W1:.*]]:3 = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%{{.*}} :696// CHECK-PROP-SAME: vector<8xf32>, vector<4xf32>, vector<4xf32>) -> (vector<8xf32>, vector<4xf32>, vector<4xf32>) {697// CHECK-PROP-NEXT: ^bb0(%{{.*}}: vector<256xf32>, %{{.*}}: vector<128xf32>, %{{.*}}: vector<128xf32>):698// CHECK-PROP-NEXT: %[[T3:.*]] = "some_def_1"(%{{.*}}) : (vector<256xf32>) -> vector<256xf32>699// CHECK-PROP-NEXT: %[[T4:.*]] = "some_def_2"(%{{.*}}) : (vector<128xf32>) -> vector<128xf32>700// CHECK-PROP-NEXT: %[[T5:.*]] = "some_def_3"(%{{.*}}) : (vector<128xf32>) -> vector<128xf32>701// CHECK-PROP-NEXT: gpu.yield %[[T3]], %[[T4]], %[[T5]] : vector<256xf32>, vector<128xf32>, vector<128xf32>702// CHECK-PROP-NEXT: }703// CHECK-PROP-NEXT: scf.yield %[[W1]]#0, %[[W1]]#1, %[[W1]]#2 : vector<8xf32>, vector<4xf32>, vector<4xf32>704// CHECK-PROP-NEXT: }705// CHECK-PROP-NEXT: "some_use_1"(%[[F0]]#2) : (vector<4xf32>) -> ()706// CHECK-PROP-NEXT: "some_use_2"(%[[F0]]#1) : (vector<4xf32>) -> ()707// CHECK-PROP-NEXT: "some_use_3"(%[[F0]]#0) : (vector<8xf32>) -> ()708func.func @warp_scf_for_swapped_for_results(%arg0: index) {709 %c128 = arith.constant 128 : index710 %c1 = arith.constant 1 : index711 %c0 = arith.constant 0 : index712 %0:3 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>, vector<4xf32>, vector<8xf32>) {713 %ini1 = "some_def"() : () -> (vector<256xf32>)714 %ini2 = "some_def"() : () -> (vector<128xf32>)715 %ini3 = "some_def"() : () -> (vector<128xf32>)716 %3:3 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini1, %arg5 = %ini2, %arg6 = %ini3) -> (vector<256xf32>, vector<128xf32>, vector<128xf32>) {717 %acc1 = "some_def_1"(%arg4) : (vector<256xf32>) -> (vector<256xf32>)718 %acc2 = "some_def_2"(%arg5) : (vector<128xf32>) -> (vector<128xf32>)719 %acc3 = "some_def_3"(%arg6) : (vector<128xf32>) -> (vector<128xf32>)720 scf.yield %acc1, %acc2, %acc3 : vector<256xf32>, vector<128xf32>, vector<128xf32>721 }722 gpu.yield %3#2, %3#1, %3#0 : vector<128xf32>, vector<128xf32>, vector<256xf32>723 }724 "some_use_1"(%0#0) : (vector<4xf32>) -> ()725 "some_use_2"(%0#1) : (vector<4xf32>) -> ()726 "some_use_3"(%0#2) : (vector<8xf32>) -> ()727 return728}729 730// -----731 732// CHECK-PROP-LABEL: func @vector_reduction(733// CHECK-PROP-SAME: %[[laneid:.*]]: index)734// CHECK-PROP-DAG: %[[c1:.*]] = arith.constant 1 : i32735// CHECK-PROP-DAG: %[[c2:.*]] = arith.constant 2 : i32736// CHECK-PROP-DAG: %[[c4:.*]] = arith.constant 4 : i32737// CHECK-PROP-DAG: %[[c8:.*]] = arith.constant 8 : i32738// CHECK-PROP-DAG: %[[c16:.*]] = arith.constant 16 : i32739// CHECK-PROP-DAG: %[[c32:.*]] = arith.constant 32 : i32740// CHECK-PROP: %[[warp_op:.*]] = gpu.warp_execute_on_lane_0(%[[laneid]])[32] -> (vector<1xf32>) {741// CHECK-PROP: gpu.yield %{{.*}} : vector<32xf32>742// CHECK-PROP: }743// CHECK-PROP: %[[a:.*]] = vector.extract %[[warp_op]][0] : f32 from vector<1xf32>744// CHECK-PROP: %[[r0:.*]], %{{.*}} = gpu.shuffle xor %[[a]], %[[c1]], %[[c32]]745// CHECK-PROP: %[[a0:.*]] = arith.addf %[[a]], %[[r0]]746// CHECK-PROP: %[[r1:.*]], %{{.*}} = gpu.shuffle xor %[[a0]], %[[c2]], %[[c32]]747// CHECK-PROP: %[[a1:.*]] = arith.addf %[[a0]], %[[r1]]748// CHECK-PROP: %[[r2:.*]], %{{.*}} = gpu.shuffle xor %[[a1]], %[[c4]], %[[c32]]749// CHECK-PROP: %[[a2:.*]] = arith.addf %[[a1]], %[[r2]]750// CHECK-PROP: %[[r3:.*]], %{{.*}} = gpu.shuffle xor %[[a2]], %[[c8]], %[[c32]]751// CHECK-PROP: %[[a3:.*]] = arith.addf %[[a2]], %[[r3]]752// CHECK-PROP: %[[r4:.*]], %{{.*}} = gpu.shuffle xor %[[a3]], %[[c16]], %[[c32]]753// CHECK-PROP: %[[a4:.*]] = arith.addf %[[a3]], %[[r4]]754// CHECK-PROP: return %[[a4]] : f32755func.func @vector_reduction(%laneid: index) -> (f32) {756 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {757 %0 = "some_def"() : () -> (vector<32xf32>)758 %1 = vector.reduction <add>, %0 : vector<32xf32> into f32759 gpu.yield %1 : f32760 }761 return %r : f32762}763 764// -----765 766// CHECK-PROP-LABEL: func @warp_distribute(767// CHECK-PROP-SAME: %[[ID:[a-zA-Z0-9]+]]768// CHECK-PROP-SAME: %[[SRC:[a-zA-Z0-9]+]]769// CHECK-PROP-SAME: %[[DEST:[a-zA-Z0-9]+]]770// CHECK-PROP: gpu.warp_execute_on_lane_0(%[[ID]])[32]771// CHECK-PROP-NEXT: "some_def"() : () -> vector<4096xf32>772// CHECK-PROP-NEXT: %{{.*}} = vector.reduction773// CHECK-PROP: %[[DEF:.*]] = arith.divf %{{.*}}, %{{.*}} : vector<1xf32>774// CHECK-PROP-NOT: gpu.warp_execute_on_lane_0775// CHECK-PROP: scf.for776// CHECK-PROP: %{{.*}} = arith.subf %{{.*}}, %[[DEF]] : vector<1xf32>777func.func @warp_distribute(%arg0: index, %src: memref<128xf32>, %dest: memref<128xf32>){778 %cst = arith.constant 0.000000e+00 : f32779 %c0 = arith.constant 0 : index780 %c1 = arith.constant 1 : index781 %c128 = arith.constant 128 : index782 %f0 = arith.constant 0.000000e+00 : f32783 gpu.warp_execute_on_lane_0(%arg0)[32]{784 %cst_1 = arith.constant dense<2.621440e+05> : vector<1xf32>785 %0 = "some_def"() : () -> (vector<4096xf32>)786 %1 = vector.reduction <add>, %0, %cst : vector<4096xf32> into f32787 %2 = vector.broadcast %1 : f32 to vector<1xf32>788 %3 = arith.divf %2, %cst_1 : vector<1xf32>789 scf.for %arg1 = %c0 to %c128 step %c1 {790 %4 = vector.transfer_read %src[%arg1], %f0 {in_bounds = [true]} : memref<128xf32>, vector<1xf32>791 %5 = arith.subf %4, %3 : vector<1xf32>792 vector.transfer_write %5, %dest[%arg1] : vector<1xf32>, memref<128xf32>793 }794 }795 return796}797 798// -----799 800func.func @vector_reduction(%laneid: index, %m0: memref<4x2x32xf32>, %m1: memref<f32>) {801 %c0 = arith.constant 0: index802 %f0 = arith.constant 0.0: f32803 // CHECK-D: %[[R:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<f32>) {804 // CHECK-D: gpu.warp_execute_on_lane_0(%{{.*}})[32] {805 // CHECK-D: vector.transfer_write %[[R]], %{{.*}}[] : vector<f32>, memref<f32>806 gpu.warp_execute_on_lane_0(%laneid)[32] {807 %0 = vector.transfer_read %m0[%c0, %c0, %c0], %f0 {in_bounds = [true]} : memref<4x2x32xf32>, vector<32xf32>808 %1 = vector.transfer_read %m1[], %f0 : memref<f32>, vector<f32>809 %2 = vector.extract %1[] : f32 from vector<f32>810 %3 = vector.reduction <add>, %0 : vector<32xf32> into f32811 %4 = arith.addf %3, %2 : f32812 %5 = vector.broadcast %4 : f32 to vector<f32>813 vector.transfer_write %5, %m1[] : vector<f32>, memref<f32>814 }815 return816}817 818// -----819 820// CHECK-PROP-LABEL: func @vector_reduction_large(821// CHECK-PROP-SAME: %[[laneid:.*]]: index)822// CHECK-PROP-DAG: %[[c1:.*]] = arith.constant 1 : i32823// CHECK-PROP-DAG: %[[c2:.*]] = arith.constant 2 : i32824// CHECK-PROP-DAG: %[[c4:.*]] = arith.constant 4 : i32825// CHECK-PROP-DAG: %[[c8:.*]] = arith.constant 8 : i32826// CHECK-PROP-DAG: %[[c16:.*]] = arith.constant 16 : i32827// CHECK-PROP-DAG: %[[c32:.*]] = arith.constant 32 : i32828// CHECK-PROP: %[[warp_op:.*]] = gpu.warp_execute_on_lane_0(%[[laneid]])[32] -> (vector<2xf32>) {829// CHECK-PROP: gpu.yield %{{.*}} : vector<64xf32>830// CHECK-PROP: }831// CHECK-PROP: %[[a:.*]] = vector.reduction <add>, %[[warp_op]] : vector<2xf32> into f32832// CHECK-PROP: %[[r0:.*]], %{{.*}} = gpu.shuffle xor %[[a]], %[[c1]], %[[c32]]833// CHECK-PROP: %[[a0:.*]] = arith.addf %[[a]], %[[r0]]834// CHECK-PROP: %[[r1:.*]], %{{.*}} = gpu.shuffle xor %[[a0]], %[[c2]], %[[c32]]835// CHECK-PROP: %[[a1:.*]] = arith.addf %[[a0]], %[[r1]]836// CHECK-PROP: %[[r2:.*]], %{{.*}} = gpu.shuffle xor %[[a1]], %[[c4]], %[[c32]]837// CHECK-PROP: %[[a2:.*]] = arith.addf %[[a1]], %[[r2]]838// CHECK-PROP: %[[r3:.*]], %{{.*}} = gpu.shuffle xor %[[a2]], %[[c8]], %[[c32]]839// CHECK-PROP: %[[a3:.*]] = arith.addf %[[a2]], %[[r3]]840// CHECK-PROP: %[[r4:.*]], %{{.*}} = gpu.shuffle xor %[[a3]], %[[c16]], %[[c32]]841// CHECK-PROP: %[[a4:.*]] = arith.addf %[[a3]], %[[r4]]842// CHECK-PROP: return %[[a4]] : f32843func.func @vector_reduction_large(%laneid: index) -> (f32) {844 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {845 %0 = "some_def"() : () -> (vector<64xf32>)846 %1 = vector.reduction <add>, %0 : vector<64xf32> into f32847 gpu.yield %1 : f32848 }849 return %r : f32850}851 852// -----853 854// CHECK-PROP-LABEL: func @vector_reduction_acc(855// CHECK-PROP-SAME: %[[laneid:.*]]: index)856// CHECK-PROP-DAG: %[[c1:.*]] = arith.constant 1 : i32857// CHECK-PROP-DAG: %[[c2:.*]] = arith.constant 2 : i32858// CHECK-PROP-DAG: %[[c4:.*]] = arith.constant 4 : i32859// CHECK-PROP-DAG: %[[c8:.*]] = arith.constant 8 : i32860// CHECK-PROP-DAG: %[[c16:.*]] = arith.constant 16 : i32861// CHECK-PROP-DAG: %[[c32:.*]] = arith.constant 32 : i32862// CHECK-PROP: %[[warp_op:.*]]:2 = gpu.warp_execute_on_lane_0(%[[laneid]])[32] -> (vector<2xf32>, f32) {863// CHECK-PROP: gpu.yield %{{.*}}, %{{.*}} : vector<64xf32>, f32864// CHECK-PROP: }865// CHECK-PROP: %[[a:.*]] = vector.reduction <add>, %[[warp_op]]#0 : vector<2xf32> into f32866// CHECK-PROP: %[[r0:.*]], %{{.*}} = gpu.shuffle xor %[[a]], %[[c1]], %[[c32]]867// CHECK-PROP: %[[a0:.*]] = arith.addf %[[a]], %[[r0]]868// CHECK-PROP: %[[r1:.*]], %{{.*}} = gpu.shuffle xor %[[a0]], %[[c2]], %[[c32]]869// CHECK-PROP: %[[a1:.*]] = arith.addf %[[a0]], %[[r1]]870// CHECK-PROP: %[[r2:.*]], %{{.*}} = gpu.shuffle xor %[[a1]], %[[c4]], %[[c32]]871// CHECK-PROP: %[[a2:.*]] = arith.addf %[[a1]], %[[r2]]872// CHECK-PROP: %[[r3:.*]], %{{.*}} = gpu.shuffle xor %[[a2]], %[[c8]], %[[c32]]873// CHECK-PROP: %[[a3:.*]] = arith.addf %[[a2]], %[[r3]]874// CHECK-PROP: %[[r4:.*]], %{{.*}} = gpu.shuffle xor %[[a3]], %[[c16]], %[[c32]]875// CHECK-PROP: %[[a4:.*]] = arith.addf %[[a3]], %[[r4]]876// CHECK-PROP: %[[a5:.*]] = arith.addf %[[a4]], %[[warp_op]]#1877// CHECK-PROP: return %[[a5]] : f32878func.func @vector_reduction_acc(%laneid: index) -> (f32) {879 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {880 %0 = "some_def"() : () -> (vector<64xf32>)881 %1 = "some_def"() : () -> (f32)882 %2 = vector.reduction <add>, %0, %1 : vector<64xf32> into f32883 gpu.yield %2 : f32884 }885 return %r : f32886}887 888// -----889 890// CHECK-PROP-LABEL: func @warp_duplicate_yield(891func.func @warp_duplicate_yield(%laneid: index) -> (vector<1xf32>, vector<1xf32>) {892 // CHECK-PROP: %{{.*}}:2 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>, vector<1xf32>)893 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>, vector<1xf32>) {894 %2 = "some_def"() : () -> (vector<32xf32>)895 %3 = "some_def"() : () -> (vector<32xf32>)896 %4 = arith.addf %2, %3 : vector<32xf32>897 %5 = arith.addf %2, %2 : vector<32xf32>898// CHECK-PROP-NOT: arith.addf899// CHECK-PROP: gpu.yield %{{.*}}, %{{.*}} : vector<32xf32>, vector<32xf32>900 gpu.yield %4, %5 : vector<32xf32>, vector<32xf32>901 }902 return %r#0, %r#1 : vector<1xf32>, vector<1xf32>903}904 905// -----906 907// CHECK-PROP-LABEL: func @warp_constant(908// CHECK-PROP: %[[C:.*]] = arith.constant dense<2.000000e+00> : vector<1xf32>909// CHECK-PROP: return %[[C]] : vector<1xf32>910func.func @warp_constant(%laneid: index) -> (vector<1xf32>) {911 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {912 %cst = arith.constant dense<2.0> : vector<32xf32>913 gpu.yield %cst : vector<32xf32>914 }915 return %r : vector<1xf32>916}917 918// -----919 920// TODO: We could use warp shuffles instead of broadcasting the entire vector.921 922// CHECK-PROP-LABEL: func.func @vector_extract_1d(923// CHECK-PROP-DAG: %[[C5_I32:.*]] = arith.constant 5 : i32924// CHECK-PROP: %[[R:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<2xf32>) {925// CHECK-PROP: %[[V:.*]] = "some_def"() : () -> vector<64xf32>926// CHECK-PROP: gpu.yield %[[V]] : vector<64xf32>927// CHECK-PROP: }928// CHECK-PROP: %[[E:.*]] = vector.extract %[[R]][1] : f32 from vector<2xf32>929// CHECK-PROP: %[[SHUFFLED:.*]], %{{.*}} = gpu.shuffle idx %[[E]], %[[C5_I32]]930// CHECK-PROP: return %[[SHUFFLED]] : f32931func.func @vector_extract_1d(%laneid: index) -> (f32) {932 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {933 %0 = "some_def"() : () -> (vector<64xf32>)934 %1 = vector.extract %0[9] : f32 from vector<64xf32>935 gpu.yield %1 : f32936 }937 return %r : f32938}939 940// -----941 942// CHECK-PROP-LABEL: func.func @vector_extract_2d(943// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<5x3xf32>) {944// CHECK-PROP: %[[V:.*]] = "some_def"945// CHECK-PROP: gpu.yield %[[V]] : vector<5x96xf32>946// CHECK-PROP: }947// CHECK-PROP: %[[E:.*]] = vector.extract %[[W]][2] : vector<3xf32> from vector<5x3xf32>948// CHECK-PROP: return %[[E]]949func.func @vector_extract_2d(%laneid: index) -> (vector<3xf32>) {950 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<3xf32>) {951 %0 = "some_def"() : () -> (vector<5x96xf32>)952 %1 = vector.extract %0[2] : vector<96xf32> from vector<5x96xf32>953 gpu.yield %1 : vector<96xf32>954 }955 return %r : vector<3xf32>956}957 958// -----959 960// CHECK-PROP-LABEL: func.func @vector_extract_2d_broadcast_scalar(961// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<5x96xf32>) {962// CHECK-PROP: %[[V:.*]] = "some_def"963// CHECK-PROP: gpu.yield %[[V]] : vector<5x96xf32>964// CHECK-PROP: }965// CHECK-PROP: %[[E:.*]] = vector.extract %[[W]][1, 2] : f32 from vector<5x96xf32>966// CHECK-PROP: return %[[E]]967func.func @vector_extract_2d_broadcast_scalar(%laneid: index) -> (f32) {968 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {969 %0 = "some_def"() : () -> (vector<5x96xf32>)970 %1 = vector.extract %0[1, 2] : f32 from vector<5x96xf32>971 gpu.yield %1 : f32972 }973 return %r : f32974}975 976// -----977 978// CHECK-PROP-LABEL: func.func @vector_extract_2d_broadcast(979// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<5x96xf32>) {980// CHECK-PROP: %[[V:.*]] = "some_def"981// CHECK-PROP: gpu.yield %[[V]] : vector<5x96xf32>982// CHECK-PROP: }983// CHECK-PROP: %[[E:.*]] = vector.extract %[[W]][2] : vector<96xf32> from vector<5x96xf32>984// CHECK-PROP: return %[[E]]985func.func @vector_extract_2d_broadcast(%laneid: index) -> (vector<96xf32>) {986 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<96xf32>) {987 %0 = "some_def"() : () -> (vector<5x96xf32>)988 %1 = vector.extract %0[2] : vector<96xf32> from vector<5x96xf32>989 gpu.yield %1 : vector<96xf32>990 }991 return %r : vector<96xf32>992}993 994// -----995 996// CHECK-PROP-LABEL: func.func @vector_extract_3d(997// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<8x4x96xf32>) {998// CHECK-PROP: %[[V:.*]] = "some_def"999// CHECK-PROP: gpu.yield %[[V]] : vector<8x128x96xf32>1000// CHECK-PROP: }1001// CHECK-PROP: %[[E:.*]] = vector.extract %[[W]][2] : vector<4x96xf32> from vector<8x4x96xf32>1002// CHECK-PROP: return %[[E]]1003func.func @vector_extract_3d(%laneid: index) -> (vector<4x96xf32>) {1004 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<4x96xf32>) {1005 %0 = "some_def"() : () -> (vector<8x128x96xf32>)1006 %1 = vector.extract %0[2] : vector<128x96xf32> from vector<8x128x96xf32>1007 gpu.yield %1 : vector<128x96xf32>1008 }1009 return %r : vector<4x96xf32>1010}1011 1012// -----1013 1014// CHECK-PROP-LABEL: func.func @vector_extract_0d(1015// CHECK-PROP: %[[R:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<f32>) {1016// CHECK-PROP: %[[V:.*]] = "some_def"() : () -> vector<f32>1017// CHECK-PROP: gpu.yield %[[V]] : vector<f32>1018// CHECK-PROP: }1019// CHECK-PROP: %[[E:.*]] = vector.extract %[[R]][] : f32 from vector<f32>1020// CHECK-PROP: return %[[E]] : f321021func.func @vector_extract_0d(%laneid: index) -> (f32) {1022 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {1023 %0 = "some_def"() : () -> (vector<f32>)1024 %1 = vector.extract %0[] : f32 from vector<f32>1025 gpu.yield %1 : f321026 }1027 return %r : f321028}1029 1030// -----1031 1032// CHECK-PROP-LABEL: func.func @vector_extract_1element(1033// CHECK-PROP: %[[R:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {1034// CHECK-PROP: %[[V:.*]] = "some_def"() : () -> vector<1xf32>1035// CHECK-PROP: gpu.yield %[[V]] : vector<1xf32>1036// CHECK-PROP: }1037// CHECK-PROP: %[[E:.*]] = vector.extract %[[R]][0] : f32 from vector<1xf32>1038// CHECK-PROP: return %[[E]] : f321039func.func @vector_extract_1element(%laneid: index) -> (f32) {1040 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {1041 %0 = "some_def"() : () -> (vector<1xf32>)1042 %c0 = arith.constant 0 : index1043 %1 = vector.extract %0[%c0] : f32 from vector<1xf32>1044 gpu.yield %1 : f321045 }1046 return %r : f321047}1048 1049// -----1050 1051// CHECK-PROP: #[[$map:.*]] = affine_map<()[s0] -> (s0 ceildiv 3)>1052// CHECK-PROP: #[[$map1:.*]] = affine_map<()[s0] -> (s0 mod 3)>1053// CHECK-PROP-LABEL: func.func @vector_extract_1d(1054// CHECK-PROP-SAME: %[[LANEID:.*]]: index, %[[POS:.*]]: index1055// CHECK-PROP-DAG: %[[C32:.*]] = arith.constant 32 : i321056// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<3xf32>) {1057// CHECK-PROP: %[[V:.*]] = "some_def"1058// CHECK-PROP: gpu.yield %[[V]] : vector<96xf32>1059// CHECK-PROP: }1060// CHECK-PROP: %[[FROM_LANE:.*]] = affine.apply #[[$map]]()[%[[POS]]]1061// CHECK-PROP: %[[DISTR_POS:.*]] = affine.apply #[[$map1]]()[%[[POS]]]1062// CHECK-PROP: %[[EXTRACTED:.*]] = vector.extract %[[W]][%[[DISTR_POS]]] : f32 from vector<3xf32>1063// CHECK-PROP: %[[FROM_LANE_I32:.*]] = arith.index_cast %[[FROM_LANE]] : index to i321064// CHECK-PROP: %[[SHUFFLED:.*]], %{{.*}} = gpu.shuffle idx %[[EXTRACTED]], %[[FROM_LANE_I32]], %[[C32]] : f321065// CHECK-PROP: return %[[SHUFFLED]]1066func.func @vector_extract_1d(%laneid: index, %pos: index) -> (f32) {1067 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {1068 %0 = "some_def"() : () -> (vector<96xf32>)1069 %1 = vector.extract %0[%pos] : f32 from vector<96xf32>1070 gpu.yield %1 : f321071 }1072 return %r : f321073}1074 1075// -----1076 1077// Index-typed values cannot be shuffled at the moment.1078 1079// CHECK-PROP-LABEL: func.func @vector_extract_1d_index(1080// CHECK-PROP: gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (index) {1081// CHECK-PROP: "some_def"1082// CHECK-PROP: vector.extract1083// CHECK-PROP: gpu.yield {{.*}} : index1084// CHECK-PROP: }1085func.func @vector_extract_1d_index(%laneid: index, %pos: index) -> (index) {1086 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (index) {1087 %0 = "some_def"() : () -> (vector<96xindex>)1088 %1 = vector.extract %0[%pos] : index from vector<96xindex>1089 gpu.yield %1 : index1090 }1091 return %r : index1092}1093 1094// -----1095 1096// CHECK-PROP: func @lane_dependent_warp_propagate_read1097// CHECK-PROP-SAME: %[[ID:.*]]: index1098func.func @lane_dependent_warp_propagate_read(1099 %laneid: index, %src: memref<1x1024xf32>, %dest: memref<1x1024xf32>) {1100 // CHECK-PROP-DAG: %[[C0:.*]] = arith.constant 0 : index1101 // CHECK-PROP-NOT: gpu.warp_execute_on_lane_01102 // CHECK-PROP-DAG: %[[R0:.*]] = vector.transfer_read %arg1[%[[C0]], %[[ID]]], %{{.*}} : memref<1x1024xf32>, vector<1x1xf32>1103 // CHECK-PROP: vector.transfer_write %[[R0]], {{.*}} : vector<1x1xf32>, memref<1x1024xf32>1104 %c0 = arith.constant 0 : index1105 %cst = arith.constant 0.000000e+00 : f321106 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x1xf32>) {1107 %2 = vector.transfer_read %src[%c0, %c0], %cst : memref<1x1024xf32>, vector<1x32xf32>1108 gpu.yield %2 : vector<1x32xf32>1109 }1110 vector.transfer_write %r, %dest[%c0, %laneid] : vector<1x1xf32>, memref<1x1024xf32>1111 return1112}1113 1114// -----1115 1116func.func @warp_propagate_read_3d(%laneid: index, %src: memref<32x4x32xf32>) -> vector<1x1x4xf32> {1117 %c0 = arith.constant 0 : index1118 %cst = arith.constant 0.000000e+00 : f321119 %r = gpu.warp_execute_on_lane_0(%laneid)[1024] -> (vector<1x1x4xf32>) {1120 %2 = vector.transfer_read %src[%c0, %c0, %c0], %cst : memref<32x4x32xf32>, vector<32x4x32xf32>1121 gpu.yield %2 : vector<32x4x32xf32>1122 }1123 return %r : vector<1x1x4xf32>1124}1125 1126// CHECK-PROP-DAG: #[[$ID0MAP:.+]] = affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 8) * 32)>1127// CHECK-PROP-DAG: #[[$ID1MAP:.+]] = affine_map<()[s0] -> ((s0 floordiv 8) mod 4)>1128// CHECK-PROP-DAG: #[[$ID2MAP:.+]] = affine_map<()[s0] -> ((s0 floordiv 8) floordiv 32)>1129// CHECK-PROP-LABEL: func.func @warp_propagate_read_3d1130// CHECK-PROP-SAME: (%[[LANE:.+]]: index, %[[SRC:.+]]: memref<32x4x32xf32>)1131// CHECK-PROP-DAG: %[[ID0:.+]] = affine.apply #[[$ID0MAP]]()[%[[LANE]]]1132// CHECK-PROP-DAG: %[[ID1:.+]] = affine.apply #[[$ID1MAP]]()[%[[LANE]]]1133// CHECK-PROP-DAG: %[[ID2:.+]] = affine.apply #[[$ID2MAP]]()[%[[LANE]]]1134// CHECK-PROP: %[[READ:.+]] = vector.transfer_read %[[SRC]][%[[ID2]], %[[ID1]], %[[ID0]]], %{{.+}} : memref<32x4x32xf32>, vector<1x1x4xf32>1135// CHECK-PROP: return %[[READ]] : vector<1x1x4xf32>1136 1137// -----1138 1139func.func @warp_propagate_read_broadcast(%laneid: index, %src: memref<32x1xf32>) -> vector<1x4xf32> {1140 %c0 = arith.constant 0 : index1141 %cst = arith.constant 0.000000e+00 : f321142 %r = gpu.warp_execute_on_lane_0(%laneid)[512] -> (vector<1x4xf32>) {1143 %2 = vector.transfer_read %src[%c0, %c0], %cst {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d0, 0)>} : memref<32x1xf32>, vector<32x64xf32>1144 gpu.yield %2 : vector<32x64xf32>1145 }1146 return %r : vector<1x4xf32>1147}1148 1149// CHECK-PROP-DAG: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 16)>1150// CHECK-PROP-DAG: #[[$READMAP:.+]] = affine_map<(d0, d1) -> (d0, 0)>1151// CHECK-PROP-LABEL: func.func @warp_propagate_read_broadcast1152// CHECK-PROP-SAME: (%[[LANE:.+]]: index, %[[SRC:.+]]: memref<32x1xf32>)1153// CHECK-PROP: %[[C0:.+]] = arith.constant 0 : index1154// CHECK-PROP: %[[ID:.+]] = affine.apply #[[$MAP]]()[%[[LANE]]]1155// CHECK-PROP: %[[READ:.+]] = vector.transfer_read %[[SRC]][%[[ID]], %[[C0]]], %{{.+}} {in_bounds = [true, true], permutation_map = #[[$READMAP]]} : memref<32x1xf32>, vector<1x4xf32>1156// CHECK-PROP: return %[[READ]] : vector<1x4xf32>1157 1158// -----1159 1160// CHECK-PROP: func @dont_duplicate_read1161func.func @dont_duplicate_read(1162 %laneid: index, %src: memref<1024xf32>) -> vector<1xf32> {1163 %c0 = arith.constant 0 : index1164 %cst = arith.constant 0.000000e+00 : f321165// CHECK-PROP: gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {1166// CHECK-PROP-NEXT: vector.transfer_read1167// CHECK-PROP-NEXT: "blocking_use"1168// CHECK-PROP-NEXT: gpu.yield1169 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {1170 %2 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>, vector<32xf32>1171 "blocking_use"(%2) : (vector<32xf32>) -> ()1172 gpu.yield %2 : vector<32xf32>1173 }1174 return %r : vector<1xf32>1175}1176 1177// -----1178 1179// CHECK-PROP: func @dedup1180func.func @dedup(%laneid: index, %v0: vector<4xf32>, %v1: vector<4xf32>)1181 -> (vector<1xf32>, vector<1xf32>) {1182 1183 // CHECK-PROP: %[[SINGLE_RES:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (vector<1xf32>) {1184 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]1185 args(%v0, %v1 : vector<4xf32>, vector<4xf32>) -> (vector<1xf32>, vector<1xf32>) {1186 ^bb0(%arg0: vector<128xf32>, %arg1: vector<128xf32>):1187 1188 // CHECK-PROP: %[[SINGLE_VAL:.*]] = "some_def"(%{{.*}}) : (vector<128xf32>) -> vector<32xf32>1189 %2 = "some_def"(%arg0) : (vector<128xf32>) -> vector<32xf32>1190 1191 // CHECK-PROP: gpu.yield %[[SINGLE_VAL]] : vector<32xf32>1192 gpu.yield %2, %2 : vector<32xf32>, vector<32xf32>1193 }1194 1195 // CHECK-PROP: return %[[SINGLE_RES]], %[[SINGLE_RES]] : vector<1xf32>, vector<1xf32>1196 return %r#0, %r#1 : vector<1xf32>, vector<1xf32>1197}1198 1199// -----1200 1201// CHECK-SCF-IF: func @warp_execute_has_broadcast_semantics1202func.func @warp_execute_has_broadcast_semantics(%laneid: index, %s0: f32, %v0: vector<f32>, %v1: vector<1xf32>, %v2: vector<1x1xf32>)1203 -> (f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) {1204 // CHECK-SCF-IF-DAG: %[[C0:.*]] = arith.constant 0 : index1205 1206 // CHECK-SCF-IF: scf.if{{.*}}{1207 %r:4 = gpu.warp_execute_on_lane_0(%laneid)[32]1208 args(%s0, %v0, %v1, %v2 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) -> (f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) {1209 ^bb0(%bs0: f32, %bv0: vector<f32>, %bv1: vector<1xf32>, %bv2: vector<1x1xf32>):1210 1211 // CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, 3>, vector<1x1xf32>1212 // CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, 3>, vector<1xf32>1213 // CHECK-SCF-IF: vector.transfer_read {{.*}}[]{{.*}} : memref<f32, 3>, vector<f32>1214 // CHECK-SCF-IF: memref.load {{.*}}[%[[C0]]] : memref<1xf32, 3>1215 // CHECK-SCF-IF: "some_def_0"(%{{.*}}) : (f32) -> f321216 // CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<f32>) -> vector<f32>1217 // CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>1218 // CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<1x1xf32>) -> vector<1x1xf32>1219 // CHECK-SCF-IF: memref.store {{.*}}[%[[C0]]] : memref<1xf32, 3>1220 // CHECK-SCF-IF: vector.transfer_write {{.*}}[] : vector<f32>, memref<f32, 3>1221 // CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]]] {in_bounds = [true]} : vector<1xf32>, memref<1xf32, 3>1222 // CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x1xf32>, memref<1x1xf32, 3>1223 1224 %rs0 = "some_def_0"(%bs0) : (f32) -> f321225 %rv0 = "some_def_1"(%bv0) : (vector<f32>) -> vector<f32>1226 %rv1 = "some_def_1"(%bv1) : (vector<1xf32>) -> vector<1xf32>1227 %rv2 = "some_def_1"(%bv2) : (vector<1x1xf32>) -> vector<1x1xf32>1228 1229 // CHECK-SCF-IF-NOT: gpu.yield1230 gpu.yield %rs0, %rv0, %rv1, %rv2 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>1231 }1232 1233 // CHECK-SCF-IF: gpu.barrier1234 // CHECK-SCF-IF: %[[RV2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, 3>, vector<1x1xf32>1235 // CHECK-SCF-IF: %[[RV1:.*]] = vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, 3>, vector<1xf32>1236 // CHECK-SCF-IF: %[[RV0:.*]] = vector.transfer_read {{.*}}[]{{.*}} : memref<f32, 3>, vector<f32>1237 // CHECK-SCF-IF: %[[RS0:.*]] = memref.load {{.*}}[%[[C0]]] : memref<1xf32, 3>1238 // CHECK-SCF-IF: return %[[RS0]], %[[RV0]], %[[RV1]], %[[RV2]] : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>1239 return %r#0, %r#1, %r#2, %r#3 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>1240}1241 1242// -----1243 1244// CHECK-SCF-IF-DAG: #[[$TIMES2:.*]] = affine_map<()[s0] -> (s0 * 2)>1245 1246// CHECK-SCF-IF: func @warp_execute_nd_distribute1247// CHECK-SCF-IF-SAME: (%[[LANEID:.*]]: index1248func.func @warp_execute_nd_distribute(%laneid: index, %v0: vector<1x64x1xf32>, %v1: vector<1x2x128xf32>)1249 -> (vector<1x64x1xf32>, vector<1x2x128xf32>) {1250 // CHECK-SCF-IF-DAG: %[[C0:.*]] = arith.constant 0 : index1251 1252 // CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[LANEID]], %c0, %c0] {in_bounds = [true, true, true]} : vector<1x64x1xf32>, memref<32x64x1xf32, 3>1253 // CHECK-SCF-IF: %[[RID:.*]] = affine.apply #[[$TIMES2]]()[%[[LANEID]]]1254 // CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[C0]], %[[RID]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x2x128xf32>, memref<1x64x128xf32, 3>1255 // CHECK-SCF-IF: gpu.barrier1256 1257 // CHECK-SCF-IF: scf.if{{.*}}{1258 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]1259 args(%v0, %v1 : vector<1x64x1xf32>, vector<1x2x128xf32>) -> (vector<1x64x1xf32>, vector<1x2x128xf32>) {1260 ^bb0(%arg0: vector<32x64x1xf32>, %arg1: vector<1x64x128xf32>):1261 1262 // CHECK-SCF-IF-DAG: %[[SR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<32x64x1xf32, 3>, vector<32x64x1xf32>1263 // CHECK-SCF-IF-DAG: %[[SR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<1x64x128xf32, 3>, vector<1x64x128xf32>1264 // CHECK-SCF-IF: %[[W0:.*]] = "some_def_0"(%[[SR0]]) : (vector<32x64x1xf32>) -> vector<32x64x1xf32>1265 // CHECK-SCF-IF: %[[W1:.*]] = "some_def_1"(%[[SR1]]) : (vector<1x64x128xf32>) -> vector<1x64x128xf32>1266 // CHECK-SCF-IF-DAG: vector.transfer_write %[[W0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<32x64x1xf32>, memref<32x64x1xf32, 3>1267 // CHECK-SCF-IF-DAG: vector.transfer_write %[[W1]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x64x128xf32>, memref<1x64x128xf32, 3>1268 1269 %r0 = "some_def_0"(%arg0) : (vector<32x64x1xf32>) -> vector<32x64x1xf32>1270 %r1 = "some_def_1"(%arg1) : (vector<1x64x128xf32>) -> vector<1x64x128xf32>1271 1272 // CHECK-SCF-IF-NOT: gpu.yield1273 gpu.yield %r0, %r1 : vector<32x64x1xf32>, vector<1x64x128xf32>1274 }1275 1276 // CHECK-SCF-IF: gpu.barrier1277 // CHECK-SCF-IF: %[[WID:.*]] = affine.apply #[[$TIMES2]]()[%[[LANEID]]]1278 // CHECK-SCF-IF-DAG: %[[R0:.*]] = vector.transfer_read %{{.*}}[%[[LANEID]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<32x64x1xf32, 3>, vector<1x64x1xf32>1279 // CHECK-SCF-IF-DAG: %[[R1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[WID]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<1x64x128xf32, 3>, vector<1x2x128xf32>1280 // CHECK-SCF-IF: return %[[R0]], %[[R1]] : vector<1x64x1xf32>, vector<1x2x128xf32>1281 return %r#0, %r#1 : vector<1x64x1xf32>, vector<1x2x128xf32>1282}1283 1284// -----1285 1286// CHECK-PROP: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 ceildiv 3)>1287// CHECK-PROP: #[[$MAP1:.*]] = affine_map<()[s0] -> (s0 mod 3)>1288// CHECK-PROP-LABEL: func @vector_insert_1d(1289// CHECK-PROP-SAME: %[[LANEID:.*]]: index, %[[POS:.*]]: index1290// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<3xf32>, f32)1291// CHECK-PROP: %[[INSERTING_LANE:.*]] = affine.apply #[[$MAP]]()[%[[POS]]]1292// CHECK-PROP: %[[INSERTING_POS:.*]] = affine.apply #[[$MAP1]]()[%[[POS]]]1293// CHECK-PROP: %[[SHOULD_INSERT:.*]] = arith.cmpi eq, %[[LANEID]], %[[INSERTING_LANE]] : index1294// CHECK-PROP: %[[R:.*]] = scf.if %[[SHOULD_INSERT]] -> (vector<3xf32>) {1295// CHECK-PROP: %[[INSERT:.*]] = vector.insert %[[W]]#1, %[[W]]#0 [%[[INSERTING_POS]]]1296// CHECK-PROP: scf.yield %[[INSERT]]1297// CHECK-PROP: } else {1298// CHECK-PROP: scf.yield %[[W]]#01299// CHECK-PROP: }1300// CHECK-PROP: return %[[R]]1301func.func @vector_insert_1d(%laneid: index, %pos: index) -> (vector<3xf32>) {1302 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<3xf32>) {1303 %0 = "some_def"() : () -> (vector<96xf32>)1304 %f = "another_def"() : () -> (f32)1305 %1 = vector.insert %f, %0[%pos] : f32 into vector<96xf32>1306 gpu.yield %1 : vector<96xf32>1307 }1308 return %r : vector<3xf32>1309}1310 1311// -----1312 1313// CHECK-PROP-LABEL: func @vector_insert_1d_broadcast(1314// CHECK-PROP-SAME: %[[LANEID:.*]]: index, %[[POS:.*]]: index1315// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<96xf32>, f32)1316// CHECK-PROP: %[[VEC:.*]] = "some_def"1317// CHECK-PROP: %[[VAL:.*]] = "another_def"1318// CHECK-PROP: gpu.yield %[[VEC]], %[[VAL]]1319// CHECK-PROP: vector.insert %[[W]]#1, %[[W]]#0 [%[[POS]]] : f32 into vector<96xf32>1320func.func @vector_insert_1d_broadcast(%laneid: index, %pos: index) -> (vector<96xf32>) {1321 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<96xf32>) {1322 %0 = "some_def"() : () -> (vector<96xf32>)1323 %f = "another_def"() : () -> (f32)1324 %1 = vector.insert %f, %0[%pos] : f32 into vector<96xf32>1325 gpu.yield %1 : vector<96xf32>1326 }1327 return %r : vector<96xf32>1328}1329 1330// -----1331 1332// CHECK-PROP-LABEL: func @vector_insert_0d(1333// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<f32>, f32)1334// CHECK-PROP: %[[VEC:.*]] = "some_def"1335// CHECK-PROP: %[[VAL:.*]] = "another_def"1336// CHECK-PROP: gpu.yield %[[VEC]], %[[VAL]]1337// CHECK-PROP: vector.broadcast %[[W]]#1 : f32 to vector<f32>1338func.func @vector_insert_0d(%laneid: index) -> (vector<f32>) {1339 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<f32>) {1340 %0 = "some_def"() : () -> (vector<f32>)1341 %f = "another_def"() : () -> (f32)1342 %1 = vector.insert %f, %0[] : f32 into vector<f32>1343 gpu.yield %1 : vector<f32>1344 }1345 return %r : vector<f32>1346}1347 1348// -----1349 1350// CHECK-PROP-LABEL: func @vector_insert_1d(1351// CHECK-PROP-SAME: %[[LANEID:.*]]: index1352// CHECK-PROP-DAG: %[[C26:.*]] = arith.constant 26 : index1353// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<3xf32>, f32)1354// CHECK-PROP: %[[VEC:.*]] = "some_def"1355// CHECK-PROP: %[[VAL:.*]] = "another_def"1356// CHECK-PROP: gpu.yield %[[VEC]], %[[VAL]]1357// CHECK-PROP: %[[SHOULD_INSERT:.*]] = arith.cmpi eq, %[[LANEID]], %[[C26]]1358// CHECK-PROP: %[[R:.*]] = scf.if %[[SHOULD_INSERT]] -> (vector<3xf32>) {1359// CHECK-PROP: %[[INSERT:.*]] = vector.insert %[[W]]#1, %[[W]]#0 [1]1360// CHECK-PROP: scf.yield %[[INSERT]]1361// CHECK-PROP: } else {1362// CHECK-PROP: scf.yield %[[W]]#01363// CHECK-PROP: }1364// CHECK-PROP: return %[[R]]1365func.func @vector_insert_1d(%laneid: index) -> (vector<3xf32>) {1366 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<3xf32>) {1367 %0 = "some_def"() : () -> (vector<96xf32>)1368 %f = "another_def"() : () -> (f32)1369 %1 = vector.insert %f, %0[76] : f32 into vector<96xf32>1370 gpu.yield %1 : vector<96xf32>1371 }1372 return %r : vector<3xf32>1373}1374 1375// -----1376 1377// CHECK-PROP-LABEL: func @vector_insert_2d_distr_src(1378// CHECK-PROP-SAME: %[[LANEID:.*]]: index1379// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<3xf32>, vector<4x3xf32>)1380// CHECK-PROP: %[[VEC:.*]] = "some_def"1381// CHECK-PROP: %[[VAL:.*]] = "another_def"1382// CHECK-PROP: gpu.yield %[[VAL]], %[[VEC]]1383// CHECK-PROP: %[[INSERT:.*]] = vector.insert %[[W]]#0, %[[W]]#1 [2] : vector<3xf32> into vector<4x3xf32>1384// CHECK-PROP: return %[[INSERT]]1385func.func @vector_insert_2d_distr_src(%laneid: index) -> (vector<4x3xf32>) {1386 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<4x3xf32>) {1387 %0 = "some_def"() : () -> (vector<4x96xf32>)1388 %s = "another_def"() : () -> (vector<96xf32>)1389 %1 = vector.insert %s, %0[2] : vector<96xf32> into vector<4x96xf32>1390 gpu.yield %1 : vector<4x96xf32>1391 }1392 return %r : vector<4x3xf32>1393}1394 1395// -----1396 1397// CHECK-PROP-LABEL: func @vector_insert_2d_distr_pos(1398// CHECK-PROP-SAME: %[[LANEID:.*]]: index1399// CHECK-PROP: %[[C19:.*]] = arith.constant 19 : index1400// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<96xf32>, vector<4x96xf32>)1401// CHECK-PROP: %[[VEC:.*]] = "some_def"1402// CHECK-PROP: %[[VAL:.*]] = "another_def"1403// CHECK-PROP: gpu.yield %[[VAL]], %[[VEC]]1404// CHECK-PROP: %[[SHOULD_INSERT:.*]] = arith.cmpi eq, %[[LANEID]], %[[C19]]1405// CHECK-PROP: %[[R:.*]] = scf.if %[[SHOULD_INSERT]] -> (vector<4x96xf32>) {1406// CHECK-PROP: %[[INSERT:.*]] = vector.insert %[[W]]#0, %[[W]]#1 [3] : vector<96xf32> into vector<4x96xf32>1407// CHECK-PROP: scf.yield %[[INSERT]]1408// CHECK-PROP: } else {1409// CHECK-PROP: scf.yield %[[W]]#11410// CHECK-PROP: }1411// CHECK-PROP: return %[[R]]1412func.func @vector_insert_2d_distr_pos(%laneid: index) -> (vector<4x96xf32>) {1413 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<4x96xf32>) {1414 %0 = "some_def"() : () -> (vector<128x96xf32>)1415 %s = "another_def"() : () -> (vector<96xf32>)1416 %1 = vector.insert %s, %0[79] : vector<96xf32> into vector<128x96xf32>1417 gpu.yield %1 : vector<128x96xf32>1418 }1419 return %r : vector<4x96xf32>1420}1421 1422// -----1423 1424// CHECK-PROP-LABEL: func @vector_insert_2d_broadcast(1425// CHECK-PROP-SAME: %[[LANEID:.*]]: index1426// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<96xf32>, vector<4x96xf32>)1427// CHECK-PROP: %[[VEC:.*]] = "some_def"1428// CHECK-PROP: %[[VAL:.*]] = "another_def"1429// CHECK-PROP: gpu.yield %[[VAL]], %[[VEC]]1430// CHECK-PROP: %[[INSERT:.*]] = vector.insert %[[W]]#0, %[[W]]#1 [2] : vector<96xf32> into vector<4x96xf32>1431// CHECK-PROP: return %[[INSERT]]1432func.func @vector_insert_2d_broadcast(%laneid: index) -> (vector<4x96xf32>) {1433 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<4x96xf32>) {1434 %0 = "some_def"() : () -> (vector<4x96xf32>)1435 %s = "another_def"() : () -> (vector<96xf32>)1436 %1 = vector.insert %s, %0[2] : vector<96xf32> into vector<4x96xf32>1437 gpu.yield %1 : vector<4x96xf32>1438 }1439 return %r : vector<4x96xf32>1440}1441 1442// -----1443// CHECK-PROP-LABEL: func.func @vector_extract_strided_slice_2d_distr_inner(1444// CHECK-RPOP-SAME: %[[LANEID:.*]]: index1445// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (vector<64x1xf32>) {1446// CHECK-PROP: %[[VEC:.*]] = "some_def"() : () -> vector<64x32xf32>1447// CHECK-PROP: gpu.yield %[[VEC]] : vector<64x32xf32>1448// CHECK-PROP: %[[EXTRACT:.*]] = vector.extract_strided_slice %[[W]]1449// CHECK-PROP-SAME: {offsets = [8], sizes = [24], strides = [1]} : vector<64x1xf32> to vector<24x1xf32>1450// CHECK-PROP: return %[[EXTRACT]] : vector<24x1xf32>1451func.func @vector_extract_strided_slice_2d_distr_inner(%laneid: index) -> (vector<24x1xf32>) {1452 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<24x1xf32>) {1453 %0 = "some_def"() : () -> (vector<64x32xf32>)1454 %1 = vector.extract_strided_slice %0 { offsets = [8], sizes = [24], strides = [1]}1455 : vector<64x32xf32> to vector<24x32xf32>1456 gpu.yield %1 : vector<24x32xf32>1457 }1458 return %r : vector<24x1xf32>1459}1460 1461// -----1462// CHECK-PROP-LABEL: func.func @vector_extract_strided_slice_2d_distr_outer(1463// CHECK-PROP-SAME: %[[LANEID:.*]]: index1464// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (vector<1x64xf32>) {1465// CHECK-PROP: %[[VEC:.*]] = "some_def"() : () -> vector<32x64xf32>1466// CHECK-PROP: gpu.yield %[[VEC]] : vector<32x64xf32>1467// CHECK-PROP: %[[EXTRACT:.*]] = vector.extract_strided_slice %[[W]]1468// CHECK-PROP-SAME: {offsets = [0, 12], sizes = [1, 8], strides = [1, 1]} : vector<1x64xf32> to vector<1x8xf32>1469// CHECK-PROP: return %[[EXTRACT]] : vector<1x8xf32>1470func.func @vector_extract_strided_slice_2d_distr_outer(%laneid: index) -> (vector<1x8xf32>) {1471 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x8xf32>) {1472 %0 = "some_def"() : () -> (vector<32x64xf32>)1473 %1 = vector.extract_strided_slice %0 { offsets = [0, 12], sizes = [32, 8], strides = [1, 1]}1474 : vector<32x64xf32> to vector<32x8xf32>1475 gpu.yield %1 : vector<32x8xf32>1476 }1477 return %r : vector<1x8xf32>1478}1479 1480// -----1481// CHECK-PROP-LABEL: func.func @vector_insert_strided_slice_1d_to_2d(1482// CHECK-PROP-SAME: %[[LANEID:.*]]: index)1483// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0({{.*}} -> (vector<1xf32>, vector<64x1xf32>) {1484// CHECK-PROP: %[[SRC:.*]] = "some_def"() : () -> vector<32xf32>1485// CHECK-PROP: %[[DEST:.*]] = "some_def"() : () -> vector<64x32xf32>1486// CHECK-PROP: gpu.yield %[[SRC]], %[[DEST]] : vector<32xf32>, vector<64x32xf32>1487// CHECK-PROP: %[[INSERT:.*]] = vector.insert_strided_slice %[[W]]#0, %[[W]]#11488// CHECK-PROP-SAME: {offsets = [18, 0], strides = [1]} : vector<1xf32> into vector<64x1xf32>1489// CHECK-PROP: return %[[INSERT]] : vector<64x1xf32>1490func.func @vector_insert_strided_slice_1d_to_2d(%laneid: index) -> (vector<64x1xf32>) {1491 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<64x1xf32>) {1492 %0 = "some_def"() : () -> (vector<32xf32>)1493 %1 = "some_def"() : () -> (vector<64x32xf32>)1494 %2 = vector.insert_strided_slice %0, %1 { offsets = [18, 0], strides = [1]}1495 : vector<32xf32> into vector<64x32xf32>1496 gpu.yield %2 : vector<64x32xf32>1497 }1498 return %r : vector<64x1xf32>1499}1500 1501// -----1502// CHECK-PROP-LABEL: func.func @vector_insert_strided_slice_2d_to_2d(1503// CHECK-PROP-SAME: %[[LANEID:.*]]: index)1504// CHECK-PROP: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0{{.*}} -> (vector<16x1xf32>, vector<64x1xf32>) {1505// CHECK-PROP: %[[SRC:.*]] = "some_def"() : () -> vector<16x32xf32>1506// CHECK-PROP: %[[DEST:.*]] = "some_def"() : () -> vector<64x32xf32>1507// CHECK-PROP: gpu.yield %[[SRC]], %[[DEST]] : vector<16x32xf32>, vector<64x32xf32>1508// CHECK-PROP: %[[INSERT:.*]] = vector.insert_strided_slice %[[W]]#0, %[[W]]#1 {offsets = [36, 0], strides = [1, 1]} :1509// CHECK-PROP-SAME: vector<16x1xf32> into vector<64x1xf32>1510// CHECK-PROP: return %[[INSERT]] : vector<64x1xf32>1511func.func @vector_insert_strided_slice_2d_to_2d(%laneid: index) -> (vector<64x1xf32>) {1512 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<64x1xf32>) {1513 %0 = "some_def"() : () -> (vector<16x32xf32>)1514 %1 = "some_def"() : () -> (vector<64x32xf32>)1515 %2 = vector.insert_strided_slice %0, %1 { offsets = [36, 0], strides = [1, 1]}1516 : vector<16x32xf32> into vector<64x32xf32>1517 gpu.yield %2 : vector<64x32xf32>1518 }1519 return %r : vector<64x1xf32>1520}1521 1522// -----1523 1524// Make sure that all operands of the transfer_read op are properly propagated.1525// The vector.extract op cannot be propagated because index-typed1526// shuffles are not supported at the moment.1527 1528// CHECK-PROP: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 * 2)>1529// CHECK-PROP-LABEL: func @transfer_read_prop_operands(1530// CHECK-PROP-SAME: %[[IN2:[^ :]*]]: vector<1x2xindex>,1531// CHECK-PROP-SAME: %[[AR1:[^ :]*]]: memref<1x4x2xi32>,1532// CHECK-PROP-SAME: %[[AR2:[^ :]*]]: memref<1x4x1024xf32>)1533// CHECK-PROP-DAG: %[[C0:.*]] = arith.constant 0 : index1534// CHECK-PROP-DAG: %[[THREADID:.*]] = gpu.thread_id x1535// CHECK-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%[[THREADID]])[32] args(%[[IN2]]1536// CHECK-PROP: %[[GATHER:.*]] = vector.gather %[[AR1]][{{.*}}]1537// CHECK-PROP: %[[EXTRACT:.*]] = vector.extract %[[GATHER]][0] : vector<64xi32> from vector<1x64xi32>1538// CHECK-PROP: %[[CAST:.*]] = arith.index_cast %[[EXTRACT]] : vector<64xi32> to vector<64xindex>1539// CHECK-PROP: %[[EXTRACTELT:.*]] = vector.extract %[[CAST]][{{.*}}] : index from vector<64xindex>1540// CHECK-PROP: gpu.yield %[[EXTRACTELT]] : index1541// CHECK-PROP: %[[APPLY:.*]] = affine.apply #[[$MAP]]()[%[[THREADID]]]1542// CHECK-PROP: %[[TRANSFERREAD:.*]] = vector.transfer_read %[[AR2]][%[[C0]], %[[W]], %[[APPLY]]],1543// CHECK-PROP: return %[[TRANSFERREAD]]1544func.func @transfer_read_prop_operands(%in2: vector<1x2xindex>, %ar1 : memref<1x4x2xi32>, %ar2 : memref<1x4x1024xf32>)-> vector<2xf32> {1545 %0 = gpu.thread_id x1546 %c0_i32 = arith.constant 0 : index1547 %c0 = arith.constant 0 : index1548 %cst = arith.constant dense<0> : vector<1x64xi32>1549 %cst_0 = arith.constant dense<true> : vector<1x64xi1>1550 %cst_1 = arith.constant dense<3> : vector<64xindex>1551 %cst_2 = arith.constant dense<0> : vector<64xindex>1552 %cst_6 = arith.constant 0.000000e+00 : f321553 1554 %18 = gpu.warp_execute_on_lane_0(%0)[32] args(%in2 : vector<1x2xindex>) -> (vector<2xf32>) {1555 ^bb0(%arg4: vector<1x64xindex>):1556 %28 = vector.gather %ar1[%c0, %c0, %c0] [%arg4], %cst_0, %cst : memref<1x4x2xi32>, vector<1x64xindex>, vector<1x64xi1>, vector<1x64xi32> into vector<1x64xi32>1557 %29 = vector.extract %28[0] : vector<64xi32> from vector<1x64xi32>1558 %30 = arith.index_cast %29 : vector<64xi32> to vector<64xindex>1559 %36 = vector.extract %30[%c0_i32] : index from vector<64xindex>1560 %37 = vector.transfer_read %ar2[%c0, %36, %c0], %cst_6 {in_bounds = [true]} : memref<1x4x1024xf32>, vector<64xf32>1561 gpu.yield %37 : vector<64xf32>1562 }1563 return %18 : vector<2xf32>1564}1565 1566// -----1567 1568// Check that we don't fold vector.broadcast when each thread doesn't get the1569// same value.1570 1571// CHECK-PROP-LABEL: func @dont_fold_vector_broadcast(1572// CHECK-PROP: %[[r:.*]] = gpu.warp_execute_on_lane_0{{.*}} -> (vector<1x2xf32>)1573// CHECK-PROP: %[[some_def:.*]] = "some_def"1574// CHECK-PROP: %[[broadcast:.*]] = vector.broadcast %[[some_def]] : vector<64xf32> to vector<1x64xf32>1575// CHECK-PROP: gpu.yield %[[broadcast]] : vector<1x64xf32>1576// CHECK-PROP: vector.print %[[r]] : vector<1x2xf32>1577func.func @dont_fold_vector_broadcast(%laneid: index) {1578 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x2xf32>) {1579 %0 = "some_def"() : () -> (vector<64xf32>)1580 %1 = vector.broadcast %0 : vector<64xf32> to vector<1x64xf32>1581 gpu.yield %1 : vector<1x64xf32>1582 }1583 vector.print %r : vector<1x2xf32>1584 return1585}1586 1587// -----1588 1589func.func @warp_propagate_shape_cast(%laneid: index, %src: memref<32x4x32xf32>) -> vector<4xf32> {1590 %c0 = arith.constant 0 : index1591 %cst = arith.constant 0.000000e+00 : f321592 %r = gpu.warp_execute_on_lane_0(%laneid)[1024] -> (vector<4xf32>) {1593 %2 = vector.transfer_read %src[%c0, %c0, %c0], %cst : memref<32x4x32xf32>, vector<32x4x32xf32>1594 %3 = vector.shape_cast %2 : vector<32x4x32xf32> to vector<4096xf32>1595 gpu.yield %3 : vector<4096xf32>1596 }1597 return %r : vector<4xf32>1598}1599 1600// CHECK-PROP-LABEL: func.func @warp_propagate_shape_cast1601// CHECK-PROP: %[[READ:.+]] = vector.transfer_read {{.+}} : memref<32x4x32xf32>, vector<1x1x4xf32>1602// CHECK-PROP: %[[CAST:.+]] = vector.shape_cast %[[READ]] : vector<1x1x4xf32> to vector<4xf32>1603// CHECK-PROP: return %[[CAST]] : vector<4xf32>1604 1605// -----1606 1607func.func @warp_propagate_uniform_transfer_read(%laneid: index, %src: memref<4096xf32>, %index: index) -> vector<1xf32> {1608 %f0 = arith.constant 0.000000e+00 : f321609 %r = gpu.warp_execute_on_lane_0(%laneid)[64] -> (vector<1xf32>) {1610 %1 = vector.transfer_read %src[%index], %f0 {in_bounds = [true]} : memref<4096xf32>, vector<1xf32>1611 gpu.yield %1 : vector<1xf32>1612 }1613 return %r : vector<1xf32>1614}1615 1616// CHECK-PROP-LABEL: func.func @warp_propagate_uniform_transfer_read1617// CHECK-PROP-SAME: (%{{.+}}: index, %[[SRC:.+]]: memref<4096xf32>, %[[INDEX:.+]]: index)1618// CHECK-PROP: %[[READ:.+]] = vector.transfer_read %[[SRC]][%[[INDEX]]], %cst {in_bounds = [true]} : memref<4096xf32>, vector<1xf32>1619// CHECK-PROP: return %[[READ]] : vector<1xf32>1620 1621// -----1622 1623func.func @warp_propagate_multi_transfer_read(%laneid: index, %src: memref<4096xf32>, %index: index, %index1: index) -> (vector<1xf32>, vector<1xf32>) {1624 %f0 = arith.constant 0.000000e+00 : f321625 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[64] -> (vector<1xf32>, vector<1xf32>) {1626 %0 = vector.transfer_read %src[%index], %f0 {in_bounds = [true]} : memref<4096xf32>, vector<1xf32>1627 "some_use"(%0) : (vector<1xf32>) -> ()1628 %1 = vector.transfer_read %src[%index1], %f0 {in_bounds = [true]} : memref<4096xf32>, vector<1xf32>1629 gpu.yield %0, %1 : vector<1xf32>, vector<1xf32>1630 }1631 return %r#0, %r#1 : vector<1xf32>, vector<1xf32>1632}1633 1634// CHECK-PROP-LABEL: func.func @warp_propagate_multi_transfer_read1635// CHECK-PROP: gpu.warp_execute_on_lane_0{{.*}} -> (vector<1xf32>)1636// CHECK-PROP: %[[INNER_READ:.+]] = vector.transfer_read1637// CHECK-PROP: "some_use"(%[[INNER_READ]])1638// CHECK-PROP: gpu.yield %[[INNER_READ]] : vector<1xf32>1639// CHECK-PROP: vector.transfer_read1640 1641// -----1642 1643func.func @warp_propagate_dead_user_multi_read(%laneid: index, %src: memref<4096xf32>, %index: index, %index1: index) -> (vector<1xf32>) {1644 %f0 = arith.constant 0.000000e+00 : f321645 %r = gpu.warp_execute_on_lane_0(%laneid)[64] -> (vector<1xf32>) {1646 %0 = vector.transfer_read %src[%index], %f0 {in_bounds = [true]} : memref<4096xf32>, vector<64xf32>1647 %1 = vector.transfer_read %src[%index1], %f0 {in_bounds = [true]} : memref<4096xf32>, vector<64xf32>1648 %max = arith.maximumf %0, %1 : vector<64xf32>1649 gpu.yield %max : vector<64xf32>1650 }1651 return %r : vector<1xf32>1652}1653 1654// CHECK-PROP-LABEL: func.func @warp_propagate_dead_user_multi_read1655// CHECK-PROP-COUNT-2: vector.transfer_read {{.*}} vector<1xf32>1656// CHECK-PROP: arith.maximumf {{.*}} : vector<1xf32>1657 1658// -----1659 1660func.func @warp_propagate_masked_write(%laneid: index, %dest: memref<4096xf32>) {1661 %c0 = arith.constant 0 : index1662 gpu.warp_execute_on_lane_0(%laneid)[32] -> () {1663 %mask = "mask_def_0"() : () -> (vector<4096xi1>)1664 %mask2 = "mask_def_1"() : () -> (vector<32xi1>)1665 %0 = "some_def_0"() : () -> (vector<4096xf32>)1666 %1 = "some_def_1"() : () -> (vector<32xf32>)1667 vector.transfer_write %0, %dest[%c0], %mask : vector<4096xf32>, memref<4096xf32>1668 vector.transfer_write %1, %dest[%c0], %mask2 : vector<32xf32>, memref<4096xf32>1669 gpu.yield1670 }1671 return1672}1673 1674// CHECK-DIST-AND-PROP-LABEL: func.func @warp_propagate_masked_write(1675// CHECK-DIST-AND-PROP: %[[W:.*]]:4 = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>, vector<1xi1>, vector<128xf32>, vector<128xi1>) {1676// CHECK-DIST-AND-PROP: %[[M0:.*]] = "mask_def_0"1677// CHECK-DIST-AND-PROP: %[[M1:.*]] = "mask_def_1"1678// CHECK-DIST-AND-PROP: %[[V0:.*]] = "some_def_0"1679// CHECK-DIST-AND-PROP: %[[V1:.*]] = "some_def_1"1680// CHECK-DIST-AND-PROP: gpu.yield %[[V1]], %[[M1]], %[[V0]], %[[M0]]1681// CHECK-DIST-AND-PROP-SAME: vector<32xf32>, vector<32xi1>, vector<4096xf32>, vector<4096xi1>1682// CHECK-DIST-AND-PROP: }1683// CHECK-DIST-AND-PROP: vector.transfer_write %[[W]]#2, {{.*}}, %[[W]]#3 {in_bounds = [true]} : vector<128xf32>, memref<4096xf32>1684// CHECK-DIST-AND-PROP: vector.transfer_write %[[W]]#0, {{.*}}, %[[W]]#1 {in_bounds = [true]} : vector<1xf32>, memref<4096xf32>1685 1686// -----1687 1688func.func @warp_propagate_masked_transfer_read(%laneid: index, %src: memref<4096x4096xf32>, %index: index) -> (vector<2xf32>, vector<2x2xf32>) {1689 %f0 = arith.constant 0.000000e+00 : f321690 %c0 = arith.constant 0 : index1691 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[64] -> (vector<2xf32>, vector<2x2xf32>) {1692 %mask = "mask_def_0"() : () -> (vector<128xi1>)1693 %0 = vector.transfer_read %src[%c0, %index], %f0, %mask {in_bounds = [true]} : memref<4096x4096xf32>, vector<128xf32>1694 %mask2 = "mask_def_1"() : () -> (vector<128x2xi1>)1695 %1 = vector.transfer_read %src[%c0, %index], %f0, %mask2 {in_bounds = [true, true]} : memref<4096x4096xf32>, vector<128x2xf32>1696 gpu.yield %0, %1 : vector<128xf32>, vector<128x2xf32>1697 }1698 return %r#0, %r#1 : vector<2xf32>, vector<2x2xf32>1699}1700 1701// CHECK-PROP-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 * 2)>1702// CHECK-PROP-DAG: #[[$MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>1703// CHECK-PROP-LABEL: func.func @warp_propagate_masked_transfer_read1704// CHECK-PROP-SAME: %[[ARG0:.+]]: index, {{.*}}, %[[ARG2:.+]]: index1705// CHECK-PROP: %[[C0:.*]] = arith.constant 0 : index1706// CHECK-PROP: %[[R:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[64] -> (vector<2xi1>, vector<2x2xi1>) {1707// CHECK-PROP: %[[M0:.*]] = "mask_def_0"1708// CHECK-PROP: %[[M1:.*]] = "mask_def_1"1709// CHECK-PROP: gpu.yield %[[M0]], %[[M1]] : vector<128xi1>, vector<128x2xi1>1710// CHECK-PROP: }1711// CHECK-PROP: %[[DIST_READ_IDX0:.+]] = affine.apply #[[$MAP0]]()[%[[ARG0]]]1712// CHECK-PROP: vector.transfer_read {{.*}}[%[[DIST_READ_IDX0]], %[[ARG2]]], {{.*}}, %[[R]]#1 {{.*}} vector<2x2xf32>1713// CHECK-PROP: %[[DIST_READ_IDX1:.+]] = affine.apply #[[$MAP1]]()[%[[ARG2]], %[[ARG0]]]1714// CHECK-PROP: vector.transfer_read {{.*}}[%[[C0]], %[[DIST_READ_IDX1]]], {{.*}}, %[[R]]#0 {{.*}} vector<2xf32>1715 1716// -----1717 1718func.func @warp_propagate_nontrivial_map_masked_transfer_read(%laneid: index, %src: memref<4096x4096xf32>, %index: index) -> vector<2xf32> {1719 %f0 = arith.constant 0.000000e+00 : f321720 %c0 = arith.constant 0 : index1721 %r = gpu.warp_execute_on_lane_0(%laneid)[64] -> (vector<2xf32>) {1722 %mask = "mask_def_0"() : () -> (vector<128xi1>)1723 %0 = vector.transfer_read %src[%index, %c0], %f0, %mask {in_bounds = [true], permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<4096x4096xf32>, vector<128xf32>1724 gpu.yield %0 : vector<128xf32>1725 }1726 return %r : vector<2xf32>1727}1728 1729// CHECK-PROP-DAG: #[[$MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>1730// CHECK-PROP-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (d0)>1731// CHECK-PROP-LABEL: func.func @warp_propagate_nontrivial_map_masked_transfer_read1732// CHECK-PROP-SAME: %[[ARG0:.+]]: index, {{.*}}, %[[ARG2:.+]]: index1733// CHECK-PROP: %[[C0:.*]] = arith.constant 0 : index1734// CHECK-PROP: %[[R:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[64] -> (vector<2xi1>) {1735// CHECK-PROP: %[[M0:.*]] = "mask_def_0"1736// CHECK-PROP: gpu.yield %[[M0]] : vector<128xi1>1737// CHECK-PROP: }1738// CHECK-PROP: %[[DIST_READ_IDX0:.+]] = affine.apply #[[$MAP0]]()[%[[ARG2]], %[[ARG0]]]1739// CHECK-PROP: vector.transfer_read {{.*}}[%[[DIST_READ_IDX0]], %[[C0]]], {{.*}}, %[[R]]1740// CHECK-PROP-SAME: permutation_map = #[[$MAP1]]} {{.*}} vector<2xf32>1741 1742// -----1743 1744func.func @warp_propagate_masked_transfer_read_shared_mask(%laneid: index, %src: memref<4096x4096xf32>, %index: index, %index2: index, %mask_ub: index) -> (vector<2xf32>, vector<2xf32>) {1745 %f0 = arith.constant 0.000000e+00 : f321746 %c0 = arith.constant 0 : index1747 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[64] -> (vector<2xf32>, vector<2xf32>) {1748 %mask = vector.create_mask %mask_ub: vector<128xi1>1749 %0 = vector.transfer_read %src[%c0, %index], %f0, %mask {in_bounds = [true]} : memref<4096x4096xf32>, vector<128xf32>1750 %1 = vector.transfer_read %src[%c0, %index2], %f0, %mask {in_bounds = [true]} : memref<4096x4096xf32>, vector<128xf32>1751 gpu.yield %0, %1 : vector<128xf32>, vector<128xf32>1752 }1753 return %r#0, %r#1 : vector<2xf32>, vector<2xf32>1754}1755 1756// CHECK-PROP-LABEL: func.func @warp_propagate_masked_transfer_read_shared_mask1757// CHECK-PROP: vector.create_mask %{{.*}} : vector<2xi1>1758// CHECK-PROP: vector.transfer_read %{{.*}} : memref<4096x4096xf32>, vector<2xf32>1759// CHECK-PROP: vector.create_mask %{{.*}} : vector<2xi1>1760// CHECK-PROP: vector.transfer_read %{{.*}} : memref<4096x4096xf32>, vector<2xf32>1761 1762// -----1763 1764func.func @warp_propagate_unconnected_read_write(%laneid: index, %buffer: memref<128xf32>, %f1: f32) -> (vector<2xf32>, vector<4xf32>) {1765 %f0 = arith.constant 0.000000e+00 : f321766 %c0 = arith.constant 0 : index1767 %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>, vector<4xf32>) {1768 %cst = arith.constant dense<2.0> : vector<128xf32>1769 %0 = vector.transfer_read %buffer[%c0], %f0 {in_bounds = [true]} : memref<128xf32>, vector<128xf32>1770 vector.transfer_write %cst, %buffer[%c0] : vector<128xf32>, memref<128xf32>1771 %1 = vector.broadcast %f1 : f32 to vector<64xf32>1772 gpu.yield %1, %0 : vector<64xf32>, vector<128xf32>1773 }1774 return %r#0, %r#1 : vector<2xf32>, vector<4xf32>1775}1776 1777// Verify that the write comes after the read1778// CHECK-DIST-AND-PROP-LABEL: func.func @warp_propagate_unconnected_read_write(1779// CHECK-DIST-AND-PROP: %[[CST:.+]] = arith.constant dense<2.000000e+00> : vector<4xf32>1780// CHECK-DIST-AND-PROP: vector.transfer_read {{.*}} : memref<128xf32>, vector<4xf32>1781// CHECK-DIST-AND-PROP: vector.transfer_write %[[CST]], {{.*}} : vector<4xf32>, memref<128xf32>1782 1783// -----1784 1785func.func @warp_propagate_create_mask(%laneid: index, %m0: index) -> vector<1xi1> {1786 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xi1>) {1787 %1 = vector.create_mask %m0 : vector<32xi1>1788 gpu.yield %1 : vector<32xi1>1789 }1790 return %r : vector<1xi1>1791}1792 1793// CHECK-PROP-DAG: #[[$SUB:.*]] = affine_map<()[s0, s1] -> (-s0 + s1)>1794// CHECK-PROP-LABEL: func @warp_propagate_create_mask1795// CHECK-PROP-SAME: %[[LANEID:.+]]: index, %[[M0:.+]]: index1796// CHECK-PROP: %[[MDIST:.+]] = affine.apply #[[$SUB]]()[%[[LANEID]], %[[M0]]]1797// CHECK-PROP: vector.create_mask %[[MDIST]] : vector<1xi1>1798 1799// -----1800 1801func.func @warp_propagate_multi_dim_create_mask(%laneid: index, %m0: index, %m1: index, %m2: index) -> vector<1x2x4xi1> {1802 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x2x4xi1>) {1803 %1 = vector.create_mask %m0, %m1, %m2 : vector<16x4x4xi1>1804 gpu.yield %1 : vector<16x4x4xi1>1805 }1806 return %r : vector<1x2x4xi1>1807}1808 1809// CHECK-PROP-DAG: #[[$SUBM0:.*]] = affine_map<()[s0, s1] -> (s0 - s1 floordiv 2)>1810// CHECK-PROP-DAG: #[[$SUBM1:.*]] = affine_map<()[s0, s1] -> (s0 - s1 * 2 + (s1 floordiv 2) * 4)>1811// CHECK-PROP-LABEL: func @warp_propagate_multi_dim_create_mask1812// CHECK-PROP-SAME: %[[LANEID:.+]]: index, %[[M0:.+]]: index, %[[M1:.+]]: index, %[[M2:.+]]: index1813// CHECK-PROP: %[[DISTM0:.+]] = affine.apply #[[$SUBM0]]()[%[[M0]], %[[LANEID]]]1814// CHECK-PROP: %[[DISTM1:.+]] = affine.apply #[[$SUBM1]]()[%[[M1]], %[[LANEID]]]1815// CHECK-PROP: vector.create_mask %[[DISTM0]], %[[DISTM1]], %[[M2]] : vector<1x2x4xi1>1816 1817// -----1818 1819func.func @warp_propagate_nd_write(%laneid: index, %dest: memref<4x1024xf32>) {1820 %c0 = arith.constant 0 : index1821 gpu.warp_execute_on_lane_0(%laneid)[32] -> () {1822 %0 = "some_def"() : () -> (vector<4x1024xf32>)1823 vector.transfer_write %0, %dest[%c0, %c0] : vector<4x1024xf32>, memref<4x1024xf32>1824 gpu.yield1825 }1826 return1827}1828 1829// CHECK-DIST-AND-PROP: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 128)>1830 1831// CHECK-DIST-AND-PROP-LABEL: func.func @warp_propagate_nd_write(1832// CHECK-DIST-AND-PROP: %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1x128xf32>) {1833// CHECK-DIST-AND-PROP: %[[V0:.*]] = "some_def"1834// CHECK-DIST-AND-PROP: gpu.yield %[[V0]]1835// CHECK-DIST-AND-PROP-SAME: vector<4x1024xf32>1836// CHECK-DIST-AND-PROP: }1837 1838// CHECK-DIST-AND-PROP: %[[IDS:.+]]:2 = affine.delinearize_index %{{.*}} into (4, 8) : index, index1839// CHECK-DIST-AND-PROP: %[[INNER_ID:.+]] = affine.apply #map()[%[[IDS]]#1]1840// CHECK-DIST-AND-PROP: vector.transfer_write %[[W]], %{{.*}}[%[[IDS]]#0, %[[INNER_ID]]] {{.*}} : vector<1x128xf32>1841 1842// -----1843func.func @warp_propagate_duplicated_operands_in_yield(%laneid: index) {1844 %r:3 = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>, vector<1xf32>, vector<1xf32>) {1845 %0 = "some_def"() : () -> (vector<32xf32>)1846 %1 = "some_other_def"() : () -> (vector<32xf32>)1847 %2 = math.exp %1 : vector<32xf32>1848 gpu.yield %2, %0, %0 : vector<32xf32>, vector<32xf32>, vector<32xf32>1849 }1850 "some_use"(%r#0) : (vector<1xf32>) -> ()1851 return1852}1853 1854// CHECK-PROP-LABEL : func.func @warp_propagate_duplicated_operands_in_yield(1855// CHECK-PROP : %[[W:.*]] = gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {1856// CHECK-PROP : %{{.*}} = "some_def"() : () -> vector<32xf32>1857// CHECK-PROP : %[[T3:.*]] = "some_other_def"() : () -> vector<32xf32>1858// CHECK-PROP : gpu.yield %[[T3]] : vector<32xf32>1859// CHECK-PROP : }1860// CHECK-PROP : %[T1:.*] = math.exp %[[W]] : vector<1xf32>1861// CHECK-PROP : "some_use"(%[[T1]]) : (vector<1xf32>) -> ()1862 1863// -----1864 1865func.func @warp_step_distribute(%buffer: memref<128xindex>) {1866 %laneid = gpu.lane_id1867 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xindex>) {1868 %seq = vector.step : vector<32xindex>1869 gpu.yield %seq : vector<32xindex>1870 }1871 vector.transfer_write %r, %buffer[%laneid] : vector<1xindex>, memref<128xindex>1872 return1873}1874 1875// CHECK-PROP-LABEL: func.func @warp_step_distribute(1876// CHECK-PROP: %[[LANE_ID:.*]] = gpu.lane_id1877// CHECK-PROP: %[[LANE_ID_VEC:.*]] = vector.broadcast %[[LANE_ID]] : index to vector<1xindex>1878// CHECK-PROP: vector.transfer_write %[[LANE_ID_VEC]], %{{.*}} : vector<1xindex>, memref<128xindex>1879 1880// -----1881 1882func.func @negative_warp_step_more_than_warp_size(%laneid: index, %buffer: memref<128xindex>) {1883 %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xindex>) {1884 %seq = vector.step : vector<64xindex>1885 gpu.yield %seq : vector<64xindex>1886 }1887 vector.transfer_write %r, %buffer[%laneid] : vector<2xindex>, memref<128xindex>1888 return1889}1890 1891// CHECK-PROP-LABEL: @negative_warp_step_more_than_warp_size1892// CHECK-PROP-NOT: vector.broadcast1893// CHECK-PROP: vector.step : vector<64xindex>1894 1895// -----1896 1897func.func @warp_scf_if_no_yield_distribute(%buffer: memref<128xindex>, %pred : i1) {1898 %laneid = gpu.lane_id1899 %c0 = arith.constant 0 : index1900 1901 gpu.warp_execute_on_lane_0(%laneid)[32] {1902 %seq = vector.step : vector<32xindex>1903 scf.if %pred {1904 vector.store %seq, %buffer[%c0] : memref<128xindex>, vector<32xindex>1905 }1906 gpu.yield1907 }1908 return1909}1910 1911// CHECK-PROP-LABEL: func.func @warp_scf_if_no_yield_distribute(1912// CHECK-PROP-SAME: %[[ARG0:.+]]: memref<128xindex>, %[[ARG1:.+]]: i11913// CHECK-PROP: scf.if %[[ARG1]] {1914// CHECK-PROP: gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%{{.*}} : vector<1xindex>) {1915// CHECK-PROP: ^bb0(%[[ARG2:.+]]: vector<32xindex>):1916// CHECK-PROP: vector.store %[[ARG2]], %[[ARG0]][%{{.*}}] : memref<128xindex>, vector<32xindex>1917 1918// -----1919 1920func.func @warp_scf_if_distribute(%pred : i1) {1921 %laneid = gpu.lane_id1922 %c0 = arith.constant 0 : index1923 1924 %0 = gpu.warp_execute_on_lane_0(%laneid)[32] -> vector<1xf32> {1925 %seq1 = vector.step : vector<32xindex>1926 %seq2 = arith.constant dense<2> : vector<32xindex>1927 %0 = scf.if %pred -> (vector<32xf32>) {1928 %1 = "some_op"(%seq1) : (vector<32xindex>) -> (vector<32xf32>)1929 scf.yield %1 : vector<32xf32>1930 } else {1931 %2 = "other_op"(%seq2) : (vector<32xindex>) -> (vector<32xf32>)1932 scf.yield %2 : vector<32xf32>1933 }1934 gpu.yield %0 : vector<32xf32>1935 }1936 "some_use"(%0) : (vector<1xf32>) -> ()1937 1938 return1939}1940 1941// CHECK-PROP-LABEL: func.func @warp_scf_if_distribute(1942// CHECK-PROP-SAME: %[[ARG0:.+]]: i11943// CHECK-PROP: %[[SEQ2:.+]] = arith.constant dense<2> : vector<32xindex>1944// CHECK-PROP: %[[LANE_ID:.+]] = gpu.lane_id1945// CHECK-PROP: %[[SEQ1:.+]] = vector.broadcast %[[LANE_ID]] : index to vector<1xindex>1946// CHECK-PROP: %[[IF_YIELD_DIST:.+]] = scf.if %[[ARG0]] -> (vector<1xf32>) {1947// CHECK-PROP: %[[THEN_DIST:.+]] = gpu.warp_execute_on_lane_0(%[[LANE_ID]])[32] args(%[[SEQ1]] : vector<1xindex>) -> (vector<1xf32>) {1948// CHECK-PROP: ^bb0(%[[ARG1:.+]]: vector<32xindex>):1949// CHECK-PROP: %{{.*}} = "some_op"(%[[ARG1]]) : (vector<32xindex>) -> vector<32xf32>1950// CHECK-PROP: gpu.yield %{{.*}} : vector<32xf32>1951// CHECK-PROP: }1952// CHECK-PROP: scf.yield %[[THEN_DIST]] : vector<1xf32>1953// CHECK-PROP: } else {1954// CHECK-PROP: %[[ELSE_DIST:.+]] = gpu.warp_execute_on_lane_0(%[[LANE_ID]])[32] -> (vector<1xf32>) {1955// CHECK-PROP: %{{.*}} = "other_op"(%[[SEQ2]]) : (vector<32xindex>) -> vector<32xf32>1956// CHECK-PROP: gpu.yield %{{.*}} : vector<32xf32>1957// CHECK-PROP: }1958// CHECK-PROP: scf.yield %[[ELSE_DIST]] : vector<1xf32>1959// CHECK-PROP: }1960// CHECK-PROP: "some_use"(%[[IF_YIELD_DIST]]) : (vector<1xf32>) -> ()1961// CHECK-PROP: return1962// CHECK-PROP: }1963 1964// -----1965func.func @dedup_unused_result(%laneid : index) -> (vector<1xf32>) {1966 %r:3 = gpu.warp_execute_on_lane_0(%laneid)[32] ->1967 (vector<1xf32>, vector<2xf32>, vector<1xf32>) {1968 %2 = "some_def"() : () -> (vector<32xf32>)1969 %3 = "some_def"() : () -> (vector<64xf32>)1970 gpu.yield %2, %3, %2 : vector<32xf32>, vector<64xf32>, vector<32xf32>1971 }1972 %r0 = "some_use"(%r#2, %r#2) : (vector<1xf32>, vector<1xf32>) -> (vector<1xf32>)1973 return %r0 : vector<1xf32>1974}1975 1976// CHECK-PROP: func @dedup_unused_result1977// CHECK-PROP: %[[R:.*]] = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<1xf32>)1978// CHECK-PROP: %[[Y0:.*]] = "some_def"() : () -> vector<32xf32>1979// CHECK-PROP: %[[Y1:.*]] = "some_def"() : () -> vector<64xf32>1980// CHECK-PROP: gpu.yield %[[Y0]] : vector<32xf32>1981// CHECK-PROP: "some_use"(%[[R]], %[[R]]) : (vector<1xf32>, vector<1xf32>) -> vector<1xf32>1982