446 lines · plain
1// RUN: mlir-opt --allow-unregistered-dialect \2// RUN: --test-gpu-subgroup-reduce-lowering %s \3// RUN: | FileCheck %s --check-prefix=CHECK-SUB4 5// RUN: mlir-opt --allow-unregistered-dialect \6// RUN: --test-gpu-subgroup-reduce-lowering="expand-to-shuffles" %s \7// RUN: | FileCheck %s --check-prefix=CHECK-SHFL8 9// RUN: mlir-opt --allow-unregistered-dialect \10// RUN: --test-gpu-subgroup-reduce-lowering="expand-to-shuffles target=gfx942" %s \11// RUN: | FileCheck %s --check-prefixes=CHECK-GFX,CHECK-GFX912 13// RUN: mlir-opt --allow-unregistered-dialect \14// RUN: --test-gpu-subgroup-reduce-lowering="expand-to-shuffles target=gfx1030" %s \15// RUN: | FileCheck %s --check-prefixes=CHECK-GFX,CHECK-GFX1016 17// CHECK-SUB: gpu.module @kernels {18// CHECK-SHFL: gpu.module @kernels {19// CHECK-GFX9: gpu.module @kernels {20// CHECK-GFX10: gpu.module @kernels {21gpu.module @kernels {22 23 // CHECK-SUB-LABEL: gpu.func @kernel0(24 // CHECK-SUB-SAME: %[[ARG0:.+]]: vector<5xf16>)25 //26 // CHECK-SHFL-LABEL: gpu.func @kernel0(27 // CHECK-GFX-LABEL: gpu.func @kernel0(28 gpu.func @kernel0(%arg0: vector<5xf16>) kernel {29 // CHECK-SUB: %[[VZ:.+]] = arith.constant dense<0.0{{.*}}> : vector<5xf16>30 // CHECK-SUB: %[[E0:.+]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0], sizes = [2], strides = [1]} : vector<5xf16> to vector<2xf16>31 // CHECK-SUB: %[[R0:.+]] = gpu.subgroup_reduce add %[[E0]] : (vector<2xf16>) -> vector<2xf16>32 // CHECK-SUB: %[[V0:.+]] = vector.insert_strided_slice %[[R0]], %[[VZ]] {offsets = [0], strides = [1]} : vector<2xf16> into vector<5xf16>33 // CHECK-SUB: %[[E1:.+]] = vector.extract_strided_slice %[[ARG0]] {offsets = [2], sizes = [2], strides = [1]} : vector<5xf16> to vector<2xf16>34 // CHECK-SUB: %[[R1:.+]] = gpu.subgroup_reduce add %[[E1]] : (vector<2xf16>) -> vector<2xf16>35 // CHECK-SUB: %[[V1:.+]] = vector.insert_strided_slice %[[R1]], %[[V0]] {offsets = [2], strides = [1]} : vector<2xf16> into vector<5xf16>36 // CHECK-SUB: %[[E2:.+]] = vector.extract %[[ARG0]][4] : f16 from vector<5xf16>37 // CHECK-SUB: %[[R2:.+]] = gpu.subgroup_reduce add %[[E2]] : (f16) -> f1638 // CHECK-SUB: %[[V2:.+]] = vector.insert %[[R2]], %[[V1]] [4] : f16 into vector<5xf16>39 // CHECK-SUB: "test.consume"(%[[V2]]) : (vector<5xf16>) -> ()40 // CHECK-GFX9-COUNT-6: amdgpu.dpp41 // CHECK-GFX10-COUNT-4: amdgpu.dpp42 // CHECK-GFX10: rocdl.permlanex1643 // CHECK-GFX10-COUNT-2: rocdl.readlane44 %sum0 = gpu.subgroup_reduce add %arg0 : (vector<5xf16>) -> (vector<5xf16>)45 "test.consume"(%sum0) : (vector<5xf16>) -> ()46 47 // CHECK-SUB-COUNT-3: gpu.subgroup_reduce mul {{.+}} uniform48 // CHECK-SUB: "test.consume"49 // CHECK-GFX9-COUNT-6: amdgpu.dpp50 // CHECK-GFX10-COUNT-4: amdgpu.dpp51 // CHECK-GFX10: rocdl.permlanex1652 // CHECK-GFX10-COUNT-2: rocdl.readlane53 %sum1 = gpu.subgroup_reduce mul %arg0 uniform : (vector<5xf16>) -> (vector<5xf16>)54 "test.consume"(%sum1) : (vector<5xf16>) -> ()55 56 // CHECK-SUB-COUNT-3: gpu.subgroup_reduce mul {{.+}} cluster(size = 4)57 // CHECK-SUB: "test.consume"58 // CHECK-GFX-COUNT-2: amdgpu.dpp {{.+}}59 %sum2 = gpu.subgroup_reduce mul %arg0 cluster(size = 4) : (vector<5xf16>) -> (vector<5xf16>)60 "test.consume"(%sum2) : (vector<5xf16>) -> ()61 62 // CHECK-SUB-COUNT-3: gpu.subgroup_reduce mul {{.+}} uniform cluster(size = 4, stride = 2)63 // CHECK-SUB: "test.consume"64 %sum3 = gpu.subgroup_reduce mul %arg0 uniform cluster(size = 4, stride = 2) : (vector<5xf16>) -> (vector<5xf16>)65 "test.consume"(%sum3) : (vector<5xf16>) -> ()66 67 // CHECK-SUB: gpu.return68 gpu.return69 }70 71 // CHECK-SUB-LABEL: gpu.func @kernel1(72 // CHECK-SUB-SAME: %[[ARG0:.+]]: vector<1xf32>)73 //74 // CHECK-SHFL-LABEL: gpu.func @kernel1(75 // CHECK-GFX-LABEL: gpu.func @kernel1(76 gpu.func @kernel1(%arg0: vector<1xf32>) kernel {77 // CHECK-SUB: %[[E0:.+]] = vector.extract %[[ARG0]][0] : f32 from vector<1xf32>78 // CHECK-SUB: %[[R0:.+]] = gpu.subgroup_reduce add %[[E0]] : (f32) -> f3279 // CHECK-SUB: %[[V0:.+]] = vector.broadcast %[[R0]] : f32 to vector<1xf32>80 // CHECK-SUB: "test.consume"(%[[V0]]) : (vector<1xf32>) -> ()81 // CHECK-GFX9-COUNT-6: amdgpu.dpp82 // CHECK-GFX10-COUNT-4: amdgpu.dpp83 // CHECK-GFX10: rocdl.permlanex1684 // CHECK-GFX10-COUNT-2: rocdl.readlane85 %sum0 = gpu.subgroup_reduce add %arg0 : (vector<1xf32>) -> (vector<1xf32>)86 "test.consume"(%sum0) : (vector<1xf32>) -> ()87 88 // CHECK-SUB: gpu.subgroup_reduce add {{.+}} uniform : (f32) -> f3289 // CHECK-SUB: "test.consume"90 // CHECK-GFX9-COUNT-6: amdgpu.dpp91 // CHECK-GFX10-COUNT-4: amdgpu.dpp92 // CHECK-GFX10: rocdl.permlanex1693 // CHECK-GFX10-COUNT-2: rocdl.readlane94 %sum1 = gpu.subgroup_reduce add %arg0 uniform : (vector<1xf32>) -> (vector<1xf32>)95 "test.consume"(%sum1) : (vector<1xf32>) -> ()96 97 // Note stride is dropped because it is == 1.98 // CHECK-SUB: gpu.subgroup_reduce add {{.+}} cluster(size = 8) : (f32) -> f3299 // CHECK-SUB: "test.consume"100 // CHECK-GFX-COUNT-2: amdgpu.dpp {{.+}} quad_perm101 // CHECK-GFX: amdgpu.dpp {{.+}} row_half_mirror102 %sum2 = gpu.subgroup_reduce add %arg0 cluster(size = 8, stride = 1) : (vector<1xf32>) -> (vector<1xf32>)103 "test.consume"(%sum2) : (vector<1xf32>) -> ()104 105 // CHECK-SUB: gpu.subgroup_reduce add {{.+}} uniform cluster(size = 8, stride = 4) : (f32) -> f32106 // CHECK-SUB: "test.consume"107 // CHECK-GFX-NOT: amdgpu.dpp108 // CHECK-GFX10-NOT: rocdl.permlanex16109 %sum3 = gpu.subgroup_reduce add %arg0 uniform cluster(size = 8, stride = 4) : (vector<1xf32>) -> (vector<1xf32>)110 "test.consume"(%sum3) : (vector<1xf32>) -> ()111 112 // CHECK-SUB: gpu.return113 gpu.return114 }115 116 // These vectors fit the native shuffle size and should not be broken down.117 //118 // CHECK-SUB-LABEL: gpu.func @kernel2(119 // CHECK-SUB-SAME: %[[ARG0:.+]]: vector<3xi8>, %[[ARG1:.+]]: vector<4xi8>)120 //121 // CHECK-SHFL-LABEL: gpu.func @kernel2(122 //123 // CHECK-GFX-LABEL: gpu.func @kernel2(124 // CHECK-GFX-NOT: amdgpu.dpp125 gpu.func @kernel2(%arg0: vector<3xi8>, %arg1: vector<4xi8>) kernel {126 // CHECK-SUB: %[[R0:.+]] = gpu.subgroup_reduce add %[[ARG0]] : (vector<3xi8>) -> vector<3xi8>127 // CHECK-SUB: "test.consume"(%[[R0]]) : (vector<3xi8>) -> ()128 %sum0 = gpu.subgroup_reduce add %arg0 : (vector<3xi8>) -> (vector<3xi8>)129 "test.consume"(%sum0) : (vector<3xi8>) -> ()130 131 // CHECK-SUB: %[[R1:.+]] = gpu.subgroup_reduce add %[[ARG1]] : (vector<4xi8>) -> vector<4xi8>132 // CHECK-SUB: "test.consume"(%[[R1]]) : (vector<4xi8>) -> ()133 %sum1 = gpu.subgroup_reduce add %arg1 : (vector<4xi8>) -> (vector<4xi8>)134 "test.consume"(%sum1) : (vector<4xi8>) -> ()135 136 // CHECK-SUB: gpu.return137 gpu.return138 }139 140 // CHECK-SHFL-LABEL: gpu.func @kernel3(141 // CHECK-SHFL-SAME: %[[ARG0:.+]]: i32)142 // CHECK-GFX-LABEL: gpu.func @kernel3(143 gpu.func @kernel3(%arg0: i32) kernel {144 // CHECK-SHFL-DAG: %[[C1:.+]] = arith.constant 1 : i32145 // CHECK-SHFL-DAG: %[[C2:.+]] = arith.constant 2 : i32146 // CHECK-SHFL-DAG: %[[C4:.+]] = arith.constant 4 : i32147 // CHECK-SHFL-DAG: %[[C8:.+]] = arith.constant 8 : i32148 // CHECK-SHFL-DAG: %[[C16:.+]] = arith.constant 16 : i32149 // CHECK-SHFL-DAG: %[[C32:.+]] = arith.constant 32 : i32150 151 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[ARG0]], %[[C1]], %[[C32]] : i32152 // CHECK-SHFL: %[[A0:.+]] = arith.addi %[[ARG0]], %[[S0]] : i32153 // CHECK-SHFL: %[[S1:.+]], %{{.+}} = gpu.shuffle xor %[[A0]], %[[C2]], %[[C32]] : i32154 // CHECK-SHFL: %[[A1:.+]] = arith.addi %[[A0]], %[[S1]] : i32155 // CHECK-SHFL: %[[S2:.+]], %{{.+}} = gpu.shuffle xor %[[A1]], %[[C4]], %[[C32]] : i32156 // CHECK-SHFL: %[[A2:.+]] = arith.addi %[[A1]], %[[S2]] : i32157 // CHECK-SHFL: %[[S3:.+]], %{{.+}} = gpu.shuffle xor %[[A2]], %[[C8]], %[[C32]] : i32158 // CHECK-SHFL: %[[A3:.+]] = arith.addi %[[A2]], %[[S3]] : i32159 // CHECK-SHFL: %[[S4:.+]], %{{.+}} = gpu.shuffle xor %[[A3]], %[[C16]], %[[C32]] : i32160 // CHECK-SHFL: %[[A4:.+]] = arith.addi %[[A3]], %[[S4]] : i32161 // CHECK-SHFL: "test.consume"(%[[A4]]) : (i32) -> ()162 163 // CHECK-GFX9-COUNT-6: amdgpu.dpp164 165 // CHECK-GFX10-COUNT-4: amdgpu.dpp166 // CHECK-GFX10: rocdl.permlanex16167 // CHECK-GFX10-COUNT-2: rocdl.readlane168 %sum0 = gpu.subgroup_reduce add %arg0 : (i32) -> i32169 "test.consume"(%sum0) : (i32) -> ()170 171 // CHECK-SHFL: gpu.return172 gpu.return173 }174 175 // CHECK-SHFL-LABEL: gpu.func @kernel3_clustered(176 // CHECK-SHFL-SAME: %[[ARG0:.+]]: i32)177 //178 // CHECK-GFX-LABEL: gpu.func @kernel3_clustered(179 // CHECK-GFX-SAME: %[[ARG0:.+]]: i32)180 gpu.func @kernel3_clustered(%arg0: i32) kernel {181 // CHECK-SHFL-DAG: %[[C1:.+]] = arith.constant 1 : i32182 // CHECK-SHFL-DAG: %[[C2:.+]] = arith.constant 2 : i32183 // CHECK-SHFL-DAG: %[[C4:.+]] = arith.constant 4 : i32184 // CHECK-SHFL-DAG: %[[C32:.+]] = arith.constant 32 : i32185 186 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[ARG0]], %[[C1]], %[[C32]] : i32187 // CHECK-SHFL: %[[A0:.+]] = arith.addi %[[ARG0]], %[[S0]] : i32188 // CHECK-SHFL: %[[S1:.+]], %{{.+}} = gpu.shuffle xor %[[A0]], %[[C2]], %[[C32]] : i32189 // CHECK-SHFL: %[[A1:.+]] = arith.addi %[[A0]], %[[S1]] : i32190 // CHECK-SHFL: %[[S2:.+]], %{{.+}} = gpu.shuffle xor %[[A1]], %[[C4]], %[[C32]] : i32191 // CHECK-SHFL: %[[A2:.+]] = arith.addi %[[A1]], %[[S2]] : i32192 // CHECK-SHFL: "test.consume"(%[[A2]]) : (i32) -> ()193 194 // CHECK-GFX: %[[D0:.+]] = amdgpu.dpp %[[ARG0]] %[[ARG0]] quad_perm([1 : i32, 0 : i32, 3 : i32, 2 : i32]) {bound_ctrl = true} : i32195 // CHECK-GFX: %[[A0:.+]] = arith.addi %[[ARG0]], %[[D0]] : i32196 // CHECK-GFX: %[[D1:.+]] = amdgpu.dpp %[[A0]] %[[A0]] quad_perm([2 : i32, 3 : i32, 0 : i32, 1 : i32]) {bound_ctrl = true} : i32197 // CHECK-GFX: %[[A1:.+]] = arith.addi %[[A0]], %[[D1]] : i32198 // CHECK-GFX: %[[D2:.+]] = amdgpu.dpp %[[A1]] %[[A1]] row_half_mirror(unit) {bound_ctrl = true} : i32199 // CHECK-GFX: %[[A2:.+]] = arith.addi %[[A1]], %[[D2]] : i32200 201 // CHECK-GFX10: "test.consume"(%[[A2]]) : (i32) -> ()202 %sum0 = gpu.subgroup_reduce add %arg0 cluster(size = 8) : (i32) -> i32203 "test.consume"(%sum0) : (i32) -> ()204 205 // CHECK-SHFL: gpu.return206 gpu.return207 }208 209 // CHECK-SHFL-LABEL: gpu.func @kernel3_clustered_strided(210 // CHECK-SHFL-SAME: %[[ARG0:.+]]: i32)211 //212 // CHECK-GFX-LABEL: gpu.func @kernel3_clustered_strided(213 // CHECK-GFX-NOT: amdgpu.dpp214 gpu.func @kernel3_clustered_strided(%arg0: i32) kernel {215 // CHECK-SHFL-DAG: %[[C1:.+]] = arith.constant 4 : i32216 // CHECK-SHFL-DAG: %[[C2:.+]] = arith.constant 8 : i32217 // CHECK-SHFL-DAG: %[[C4:.+]] = arith.constant 16 : i32218 // CHECK-SHFL-DAG: %[[C32:.+]] = arith.constant 32 : i32219 220 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[ARG0]], %[[C1]], %[[C32]] : i32221 // CHECK-SHFL: %[[A0:.+]] = arith.addi %[[ARG0]], %[[S0]] : i32222 // CHECK-SHFL: %[[S1:.+]], %{{.+}} = gpu.shuffle xor %[[A0]], %[[C2]], %[[C32]] : i32223 // CHECK-SHFL: %[[A1:.+]] = arith.addi %[[A0]], %[[S1]] : i32224 // CHECK-SHFL: %[[S2:.+]], %{{.+}} = gpu.shuffle xor %[[A1]], %[[C4]], %[[C32]] : i32225 // CHECK-SHFL: %[[A2:.+]] = arith.addi %[[A1]], %[[S2]] : i32226 // CHECK-SHFL: "test.consume"(%[[A2]]) : (i32) -> ()227 %sum0 = gpu.subgroup_reduce add %arg0 cluster(size = 8, stride = 4) : (i32) -> i32228 "test.consume"(%sum0) : (i32) -> ()229 230 // CHECK-SHFL: gpu.return231 gpu.return232 }233 234 // CHECK-SHFL-LABEL: gpu.func @kernel4(235 // CHECK-SHFL-SAME: %[[ARG0:.+]]: vector<2xf16>)236 //237 // CHECK-GFX-LABEL: gpu.func @kernel4(238 // CHECK-GFX-NOT: amdgpu.dpp239 gpu.func @kernel4(%arg0: vector<2xf16>) kernel {240 // CHECK-SHFL-DAG: %[[C1:.+]] = arith.constant 1 : i32241 // CHECK-SHFL-DAG: %[[C2:.+]] = arith.constant 2 : i32242 // CHECK-SHFL-DAG: %[[C4:.+]] = arith.constant 4 : i32243 // CHECK-SHFL-DAG: %[[C8:.+]] = arith.constant 8 : i32244 // CHECK-SHFL-DAG: %[[C16:.+]] = arith.constant 16 : i32245 // CHECK-SHFL-DAG: %[[C32:.+]] = arith.constant 32 : i32246 247 // CHECK-SHFL: %[[V0:.+]] = vector.bitcast %[[ARG0]] : vector<2xf16> to vector<1xi32>248 // CHECK-SHFL: %[[I0:.+]] = vector.extract %[[V0]][0] : i32 from vector<1xi32>249 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[I0]], %[[C1]], %[[C32]] : i32250 // CHECK-SHFL: %[[BR0:.+]] = vector.broadcast %[[S0]] : i32 to vector<1xi32>251 // CHECK-SHFL: %[[BC0:.+]] = vector.bitcast %[[BR0]] : vector<1xi32> to vector<2xf16>252 // CHECK-SHFL: %[[ADD0:.+]] = arith.addf %[[ARG0]], %[[BC0]] : vector<2xf16>253 // CHECK-SHFL: %[[BC1:.+]] = vector.bitcast %[[ADD0]] : vector<2xf16> to vector<1xi32>254 // CHECK-SHFL: %[[I1:.+]] = vector.extract %[[BC1]][0] : i32 from vector<1xi32>255 // CHECK-SHFL: gpu.shuffle xor %[[I1]], %[[C2]], %[[C32]] : i32256 // CHECK-SHFL: arith.addf {{.+}} : vector<2xf16>257 // CHECK-SHFL: gpu.shuffle xor %{{.+}}, %[[C4]], %[[C32]] : i32258 // CHECK-SHFL: arith.addf {{.+}} : vector<2xf16>259 // CHECK-SHFL: gpu.shuffle xor %{{.+}}, %[[C8]], %[[C32]] : i32260 // CHECK-SHFL: arith.addf {{.+}} : vector<2xf16>261 // CHECK-SHFL: %[[SL:.+]], %{{.+}} = gpu.shuffle xor %{{.+}}, %[[C16]], %[[C32]] : i32262 // CHECK-SHFL: %[[BRL:.+]] = vector.broadcast %[[SL]] : i32 to vector<1xi32>263 // CHECK-SHFL: %[[BCL:.+]] = vector.bitcast %[[BRL]] : vector<1xi32> to vector<2xf16>264 // CHECK-SHFL: %[[ADDL:.+]] = arith.addf %{{.+}}, %[[BCL]] : vector<2xf16>265 // CHECK-SHFL: "test.consume"(%[[ADDL]]) : (vector<2xf16>) -> ()266 %sum0 = gpu.subgroup_reduce add %arg0 : (vector<2xf16>) -> (vector<2xf16>)267 "test.consume"(%sum0) : (vector<2xf16>) -> ()268 269 // CHECK-SHFL: gpu.return270 gpu.return271 }272 273 // CHECK-SHFL-LABEL: gpu.func @kernel4_clustered(274 // CHECK-SHFL-SAME: %[[ARG0:.+]]: vector<2xf16>)275 //276 // CHECK-GFX-LABEL: gpu.func @kernel4_clustered(277 // CHECK-GFX-NOT: amdgpu.dpp278 gpu.func @kernel4_clustered(%arg0: vector<2xf16>) kernel {279 // CHECK-SHFL-DAG: %[[C1:.+]] = arith.constant 1 : i32280 // CHECK-SHFL-DAG: %[[C2:.+]] = arith.constant 2 : i32281 // CHECK-SHFL-DAG: %[[C32:.+]] = arith.constant 32 : i32282 283 // CHECK-SHFL-COUNT-2: gpu.shuffle xor284 %sum0 = gpu.subgroup_reduce add %arg0 cluster(size = 4) : (vector<2xf16>) -> (vector<2xf16>)285 "test.consume"(%sum0) : (vector<2xf16>) -> ()286 287 // CHECK-SHFL: gpu.return288 gpu.return289 }290 291 // CHECK-SHFL-LABEL: gpu.func @kernel5(292 // CHECK-SHFL-SAME: %[[ARG0:.+]]: i16)293 //294 // CHECK-GFX-LABEL: gpu.func @kernel5(295 // CHECK-GFX-SAME: %[[ARG0:.+]]: i16)296 gpu.func @kernel5(%arg0: i16) kernel {297 // CHECK-SHFL: %[[E0:.+]] = arith.extui %[[ARG0]] : i16 to i32298 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[E0]], {{.+}} : i32299 // CHECK-SHFL: %[[T0:.+]] = arith.trunci %[[S0]] : i32 to i16300 // CHECK-SHFL: %[[A0:.+]] = arith.addi %[[ARG0]], %[[T0]] : i16301 // CHECK-SHFL: %[[E1:.+]] = arith.extui %[[A0]] : i16 to i32302 // CHECK-SHFL: %{{.+}}, %{{.+}} = gpu.shuffle xor %[[E1]], {{.+}} : i32303 // CHECK-SHFL-COUNT-3: gpu.shuffle xor304 // CHECK-SHFL: arith.trunci {{.+}} : i32 to i16305 // CHECK-SHFL: %[[AL:.+]] = arith.addi {{.+}} : i16306 // CHECK-SHFL: "test.consume"(%[[AL]]) : (i16) -> ()307 308 // CHECK-GFX9-COUNT-6: amdgpu.dpp309 310 // CHECK-GFX10: %[[D0:.+]] = amdgpu.dpp %[[ARG0]] %[[ARG0]] quad_perm([1 : i32, 0 : i32, 3 : i32, 2 : i32]) {bound_ctrl = true} : i16311 // CHECK-GFX10: %[[A0:.+]] = arith.addi %[[ARG0]], %[[D0]] : i16312 // CHECK-GFX10: %[[D1:.+]] = amdgpu.dpp %[[A0]] %[[A0]] quad_perm([2 : i32, 3 : i32, 0 : i32, 1 : i32]) {bound_ctrl = true} : i16313 // CHECK-GFX10: %[[A1:.+]] = arith.addi %[[A0]], %[[D1]] : i16314 // CHECK-GFX10: %[[D2:.+]] = amdgpu.dpp %[[A1]] %[[A1]] row_half_mirror(unit) {bound_ctrl = true} : i16315 // CHECK-GFX10: %[[A2:.+]] = arith.addi %[[A1]], %[[D2]] : i16316 // CHECK-GFX10: %[[D3:.+]] = amdgpu.dpp %[[A2]] %[[A2]] row_mirror(unit) {bound_ctrl = true} : i16317 // CHECK-GFX10: %[[A3:.+]] = arith.addi %[[A2]], %[[D3]] : i16318 // CHECK-GFX10: %[[P0:.+]] = rocdl.permlanex16 %[[A3]], %[[A3]], %c-1_i32, %c-1_i32, true, false : i16, i32319 // CHECK-GFX10: %[[A4:.+]] = arith.addi %[[A3]], %[[P0]] : i16320 // CHECK-GFX10: %[[R0:.+]] = rocdl.readlane %[[A4]], %{{.+}} : (i16, i32) -> i16321 // CHECK-GFX10: %[[R1:.+]] = rocdl.readlane %[[A4]], %{{.+}} : (i16, i32) -> i16322 // CHECK-GFX10: %[[A5:.+]] = arith.addi %[[R0]], %[[R1]] : i16323 // CHECK-GFX10: "test.consume"(%[[A5]]) : (i16) -> ()324 %sum0 = gpu.subgroup_reduce add %arg0 : (i16) -> i16325 "test.consume"(%sum0) : (i16) -> ()326 327 // CHECK-SHFL: gpu.return328 gpu.return329 }330 331 // CHECK-SHFL-LABEL: gpu.func @kernel5_clustered(332 // CHECK-SHFL-SAME: %[[ARG0:.+]]: i16)333 //334 // CHECK-GFX-LABEL: gpu.func @kernel5_clustered335 // CHECK-GFX-SAME: %[[ARG0:.+]]: i16)336 gpu.func @kernel5_clustered(%arg0: i16) kernel {337 // CHECK-SHFL: %[[E0:.+]] = arith.extui %[[ARG0]] : i16 to i32338 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[E0]], {{.+}} : i32339 // CHECK-SHFL: %[[T0:.+]] = arith.trunci %[[S0]] : i32 to i16340 // CHECK-SHFL: %[[A0:.+]] = arith.addi %[[ARG0]], %[[T0]] : i16341 // CHECK-SHFL: %[[E1:.+]] = arith.extui %[[A0]] : i16 to i32342 // CHECK-SHFL: %{{.+}}, %{{.+}} = gpu.shuffle xor %[[E1]], {{.+}} : i32343 // CHECK-SHFL-COUNT-2: gpu.shuffle xor344 // CHECK-SHFL: arith.trunci {{.+}} : i32 to i16345 // CHECK-SHFL: %[[AL:.+]] = arith.addi {{.+}} : i16346 // CHECK-SHFL: "test.consume"(%[[AL]]) : (i16) -> ()347 348 // CHECK-GFX: %[[VAR0:.+]] = amdgpu.dpp %[[ARG0]] %[[ARG0]] quad_perm([1 : i32, 0 : i32, 3 : i32, 2 : i32]) {bound_ctrl = true} : i16349 // CHECK-GFX: %[[VAR1:.+]] = arith.addi %[[ARG0]], %[[VAR0]] : i16350 // CHECK-GFX: %[[VAR2:.+]] = amdgpu.dpp %[[VAR1]] %[[VAR1]] quad_perm([2 : i32, 3 : i32, 0 : i32, 1 : i32]) {bound_ctrl = true} : i16351 // CHECK-GFX: %[[VAR3:.+]] = arith.addi %[[VAR1]], %[[VAR2]] : i16352 // CHECK-GFX: %[[VAR4:.+]] = amdgpu.dpp %[[VAR3]] %[[VAR3]] row_half_mirror(unit) {bound_ctrl = true} : i16353 // CHECK-GFX: %[[VAR5:.+]] = arith.addi %[[VAR3]], %[[VAR4]] : i16354 // CHECK-GFX: %[[VAR6:.+]] = amdgpu.dpp %[[VAR5]] %[[VAR5]] row_mirror(unit) {bound_ctrl = true} : i16355 // CHECK-GFX: %[[VAR7:.+]] = arith.addi %[[VAR5]], %[[VAR6]] : i16356 // CHECK-GFX: "test.consume"(%[[VAR7]]) : (i16) -> ()357 %sum0 = gpu.subgroup_reduce add %arg0 cluster(size = 16) : (i16) -> i16358 "test.consume"(%sum0) : (i16) -> ()359 360 // CHECK-SHFL: gpu.return361 gpu.return362 }363 364 // CHECK-SHFL-LABEL: gpu.func @kernel6(365 // CHECK-SHFL-SAME: %[[ARG0:.+]]: vector<3xi8>)366 //367 // CHECK-GFX-LABEL: gpu.func @kernel6(368 // CHECK-GFX-NOT: amdgpu.dpp369 gpu.func @kernel6(%arg0: vector<3xi8>) kernel {370 // CHECK-SHFL: %[[CZ:.+]] = arith.constant dense<0> : vector<4xi8>371 // CHECK-SHFL: %[[V0:.+]] = vector.insert_strided_slice %[[ARG0]], %[[CZ]] {offsets = [0], strides = [1]} : vector<3xi8> into vector<4xi8>372 // CHECK-SHFL: %[[BC0:.+]] = vector.bitcast %[[V0]] : vector<4xi8> to vector<1xi32>373 // CHECK-SHFL: %[[I0:.+]] = vector.extract %[[BC0]][0] : i32 from vector<1xi32>374 // CHECK-SHFL: %[[S0:.+]], %{{.+}} = gpu.shuffle xor %[[I0]], {{.+}} : i32375 // CHECK-SHFL: %[[BR0:.+]] = vector.broadcast %[[S0]] : i32 to vector<1xi32>376 // CHECK-SHFL: %[[BC1:.+]] = vector.bitcast %[[BR0]] : vector<1xi32> to vector<4xi8>377 // CHECK-SHFL: %[[ADD0:.+]] = arith.addi %[[V0]], %[[BC1]] : vector<4xi8>378 // CHECK-SHFL: %[[BC2:.+]] = vector.bitcast %[[ADD0]] : vector<4xi8> to vector<1xi32>379 // CHECK-SHFL: %[[I1:.+]] = vector.extract %[[BC2]][0] : i32 from vector<1xi32>380 // CHECK-SHFL-COUNT-4: gpu.shuffle xor381 // CHECK-SHFL: %[[ESS:.+]] = vector.extract_strided_slice %{{.+}} {offsets = [0], sizes = [3], strides = [1]} : vector<4xi8> to vector<3xi8>382 // CHECK-SHFL: "test.consume"(%[[ESS]]) : (vector<3xi8>) -> ()383 %sum0 = gpu.subgroup_reduce add %arg0 : (vector<3xi8>) -> (vector<3xi8>)384 "test.consume"(%sum0) : (vector<3xi8>) -> ()385 386 // CHECK-SHFL: gpu.return387 gpu.return388 }389 390 // CHECK-GFX-LABEL: gpu.func @kernel7(391 // CHECK-GFX-SAME: %[[ARG0:.+]]: f32)392 //393 // Checks, common to gfx942 and gfx1030, of394 // (1) quad_perm, followed by reduction resulting in reduction over 2 consecutive lanes,395 // (2) quad_perm, followed by reduction resulting in reduction over 4 consecutive lanes,396 // (3) row_half_mirror, followed by reduction resulting in reduction over 8 consecutive lanes, and397 // (4) row_mirror, followed by reduction resulting in reduction over 16 consecutive lanes.398 // CHECK-GFX: %[[D0:.+]] = amdgpu.dpp %[[ARG0]] %[[ARG0]] quad_perm([1 : i32, 0 : i32, 3 : i32, 2 : i32]) {bound_ctrl = true} : f32399 // CHECK-GFX: %[[A0:.+]] = arith.addf %[[ARG0]], %[[D0]] : f32400 // CHECK-GFX: %[[D1:.+]] = amdgpu.dpp %[[A0]] %[[A0]] quad_perm([2 : i32, 3 : i32, 0 : i32, 1 : i32]) {bound_ctrl = true} : f32401 // CHECK-GFX: %[[A1:.+]] = arith.addf %[[A0]], %[[D1]] : f32402 // CHECK-GFX: %[[D2:.+]] = amdgpu.dpp %[[A1]] %[[A1]] row_half_mirror(unit) {bound_ctrl = true} : f32403 // CHECK-GFX: %[[A2:.+]] = arith.addf %[[A1]], %[[D2]] : f32404 // CHECK-GFX: %[[D3:.+]] = amdgpu.dpp %[[A2]] %[[A2]] row_mirror(unit) {bound_ctrl = true} : f32405 // CHECK-GFX: %[[A3:.+]] = arith.addf %[[A2]], %[[D3]] : f32406 //407 // Now, on gfx942:408 // (1) Lane 15 gets broadcast to lanes [16, 32) and lane 31 gets broadcast to lanes [48, 64], after which409 // the reduction in lanes [16, 32) is over the full cluster of the first 32 lanes, and the reduction in lanes410 // [48, 64) is over the full cluster of the last 32 lanes.411 // (2) Update the reduction value in lanes [0, 16) and [32, 48) with the final reduction result from412 // lanes [16, 32) and [48, 64), respectively.413 // CHECK-GFX9: %[[BCAST15:.+]] = amdgpu.dpp %[[A3]] %[[A3]] row_bcast_15(unit) {row_mask = 10 : i32} : f32414 // CHECK-GFX9: %[[SUM:.+]] = arith.addf %[[A3]], %[[BCAST15]] : f32415 // CHECK-GFX9: %[[SWIZ:.+]] = amdgpu.swizzle_bitmode %[[SUM]] 0 31 0 : f32416 // CHECK-GFX9: "test.consume"(%[[SWIZ]]) : (f32) -> ()417 //418 // On gfx1030, the final step is to permute the lanes and perform final reduction:419 // CHECK-GFX10: rocdl.permlanex16420 // CHECK-GFX10: arith.addf421 // CHECK-GFX10: "test.consume"422 gpu.func @kernel7(%arg0: f32) kernel {423 %sum0 = gpu.subgroup_reduce add %arg0 cluster(size = 32) : (f32) -> (f32)424 "test.consume"(%sum0) : (f32) -> ()425 gpu.return426 }427 428 // CHECK-SHFL-LABEL: gpu.func @kernel_cluster_size_is_subgroup_size(429 // CHECK-SHFL-SAME: %[[ARG0:.+]]: vector<3xi8>)430 //431 // CHECK-GFX9-LABEL: gpu.func @kernel_cluster_size_is_subgroup_size(432 // CHECK-GFX9-NOT: amdgpu.dpp433 //434 // CHECK-GFX10-LABEL: gpu.func @kernel_cluster_size_is_subgroup_size(435 // CHECK-GFX10-NOT: amdgpu.dpp436 gpu.func @kernel_cluster_size_is_subgroup_size(%arg0: vector<3xi8>) kernel {437 // CHECK-SHFL-COUNT-5: gpu.shuffle xor438 %sum0 = gpu.subgroup_reduce add %arg0 cluster(size = 32) : (vector<3xi8>) -> (vector<3xi8>)439 "test.consume"(%sum0) : (vector<3xi8>) -> ()440 441 // CHECK-SHFL: gpu.return442 gpu.return443 }444}445 446