839 lines · plain
1// RUN: mlir-opt --transform-interpreter --split-input-file -canonicalize -cse %s | FileCheck %s2 3!type = memref<2 x 32 x f32>4!type1d = memref<32 x f32>5 6// CHECK-LABEL: func.func @blocks_3d(7// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>8// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>9// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>10func.func @blocks_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {11 %c9 = arith.constant 9 : index12 %c7 = arith.constant 7 : index13 %one = arith.constant 1 : index14// CHECK: gpu.launch15// CHECK: %[[BLKX:.*]] = gpu.block_id x16// CHECK: %[[BLKY:.*]] = gpu.block_id y17// CHECK: memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]]]18// CHECK: memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]]]19 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)20 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)21 {22 scf.forall (%i, %j) in (%c7, %c9) {23 %4 = memref.load %x[%i, %j] : !type24 %5 = memref.load %y[%i, %j] : !type25 %6 = math.fma %alpha, %4, %5 : f3226 memref.store %6, %y[%i, %j] : !type27 } { mapping = [#gpu.block<x>, #gpu.block<y>]}28 gpu.terminator29 }30 return %y : !type31}32 33module attributes {transform.with_named_sequence} {34 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {35 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op36 transform.gpu.map_forall_to_blocks %funcop grid_dims = [12, 9, 1] : (!transform.any_op) -> !transform.any_op37 transform.yield38 }39}40 41// -----42 43!type = memref<2 x 32 x f32>44!type1d = memref<32 x f32>45 46// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 floordiv 128)>47 48// CHECK-LABEL: func.func @warpgroup_3d(49// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>50// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>51// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>52func.func @warpgroup_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {53 %c1 = arith.constant 1 : index54 %c3 = arith.constant 3 : index55 %one = arith.constant 1 : index56 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index57 // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index58 // CHECK-DAG: %[[C384:.*]] = arith.constant 384 : index59 // CHECK-DAG: %[[C512:.*]] = arith.constant 512 : index60 61// CHECK: gpu.launch62// CHECK: %[[TIDX:.*]] = gpu.thread_id x63// CHECK: %[[TIDY:.*]] = gpu.thread_id y64// CHECK-DAG: %[[WG:.*]] = affine.apply #[[$MAP]]()[%[[TIDX]]]65// CHECK-DAG: %[[CMPX:.*]] = arith.cmpi ult, %[[TIDX]], %[[C384]] : index66// CHECK-DAG: %[[CMPY:.*]] = arith.cmpi ult, %[[TIDY]], %[[C1]] : index67// CHECK: %[[COND:.*]] = arith.andi %[[CMPX]], %[[CMPY]] : i168// CHECK: scf.if %[[COND]]69// CHECK: memref.load %[[ARGX]][%[[WG]], %[[TIDY]]]70// CHECK: memref.load %[[ARGY]][%[[WG]], %[[TIDY]]]71 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)72 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)73 {74 scf.forall (%i, %j) in (%c3, %c1) {75 %4 = memref.load %x[%i, %j] : !type76 %5 = memref.load %y[%i, %j] : !type77 %6 = math.fma %alpha, %4, %5 : f3278 memref.store %6, %y[%i, %j] : !type79 } { mapping = [#gpu.warpgroup<x>, #gpu.warpgroup<y>]}80 gpu.terminator81 }82 return %y : !type83}84 85module attributes {transform.with_named_sequence} {86 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {87 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op88 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [512, 2, 1] : (!transform.any_op) -> !transform.any_op89 transform.yield90 }91}92 93// -----94 95!type = memref<2 x 32 x f32>96!type1d = memref<32 x f32>97 98// CHECK-DAG: #map = affine_map<()[s0] -> (s0 floordiv 16)>99 100// CHECK-LABEL: func.func @warp_3d(101// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>102// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>103// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>104func.func @warp_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {105 %c2 = arith.constant 2 : index106 %c3 = arith.constant 3 : index107 %one = arith.constant 1 : index108 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index109 // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index110 // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index111 // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index112 // CHECK-DAG: %[[c64:.*]] = arith.constant 64 : index113 114// CHECK: gpu.launch115// CHECK: %[[TIDX:.*]] = gpu.thread_id x116// CHECK: %[[TIDY:.*]] = gpu.thread_id y117// CHECK-DAG: %[[W:.*]] = affine.apply #[[$MAP]]()[%[[TIDX]]]118// CHECK-DAG: %[[CMPX:.*]] = arith.cmpi ult, %[[TIDX]], %[[C32]] : index119// CHECK-DAG: %[[CMPY:.*]] = arith.cmpi ult, %[[TIDY]], %[[C3]] : index120// CHECK: %[[COND:.*]] = arith.andi %[[CMPX]], %[[CMPY]] : i1121// CHECK: scf.if %[[COND]]122// CHECK: memref.load %[[ARGX]][%[[W]], %[[TIDY]]]123// CHECK: memref.load %[[ARGY]][%[[W]], %[[TIDY]]]124 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)125 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)126 {127 scf.forall (%i, %j, %k) in (%c2, %c3, %c3) {128 %4 = memref.load %x[%i, %j] : !type129 %5 = memref.load %y[%i, %j] : !type130 %6 = math.fma %alpha, %4, %5 : f32131 memref.store %6, %y[%i, %j] : !type132 } { mapping = [#gpu.warp<x>, #gpu.warp<y>, #gpu.warp<z>]}133 gpu.terminator134 }135 return %y : !type136}137 138module attributes {transform.with_named_sequence} {139 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {140 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op141 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [64, 4, 3] warp_size = 16: (!transform.any_op) -> !transform.any_op142 transform.yield143 }144}145 146// -----147 148!type = memref<2 x 32 x f32>149!type1d = memref<32 x f32>150 151// CHECK-LABEL: func.func @threads_3d(152// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>153// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>154// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>155func.func @threads_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {156 %one = arith.constant 1 : index157 %c12 = arith.constant 12 : index158 %c9 = arith.constant 9 : index159 %c7 = arith.constant 7 : index160// CHECK: %[[C1:.*]] = arith.constant 1 : index161// CHECK: %[[C12:.*]] = arith.constant 12 : index162// CHECK: %[[C9:.*]] = arith.constant 9 : index163// CHECK: %[[C7:.*]] = arith.constant 7 : index164// CHECK: gpu.launch async [%{{.*}}] blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C1]], %{{.*}} = %[[C1]], %{{.*}} = %[[C1]]) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C12]], %{{.*}} = %[[C9]], %{{.*}} = %[[C1]])165// CHECK: %[[TIDX:.*]] = gpu.thread_id x166// CHECK: %[[TIDY:.*]] = gpu.thread_id y167// CHECK: arith.cmpi ult, %[[TIDX]], %[[C9]] : index168// CHECK: arith.cmpi ult, %[[TIDY]], %[[C7]] : index169// CHECK: memref.load %[[ARGX]][%[[TIDY]], %[[TIDX]]]170// CHECK: memref.load %[[ARGY]][%[[TIDY]], %[[TIDX]]]171// CHECK: gpu.barrier172// CHECK: arith.cmpi ult, %[[TIDY]], %[[C1]] : index173// CHECK: memref.load %[[ARGT]][%[[TIDX]]]174// CHECK: gpu.barrier175 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)176 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)177 {178 scf.forall (%i, %j) in (%c7, %c9) {179 %4 = memref.load %x[%i, %j] : !type180 %5 = memref.load %y[%i, %j] : !type181 %6 = math.fma %alpha, %4, %5 : f32182 memref.store %6, %y[%i, %j] : !type183 } { mapping = [#gpu.thread<y>, #gpu.thread<x>]}184 scf.forall (%i) in (%c12) {185 %7 = memref.load %t[%i] : !type1d186 %8 = arith.addf %alpha, %7 : f32187 memref.store %8, %t[%i] : !type1d188 } {mapping = [#gpu.thread<x>] }189 gpu.terminator190 }191 return %y : !type192}193 194module attributes {transform.with_named_sequence} {195 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {196 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op197 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] : (!transform.any_op) -> !transform.any_op198 transform.yield199 }200}201 202// -----203 204!type4d = memref<32x64x4x32xf32>205 206// CHECK-LABEL: func.func @saxpy4d(207// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<32x64x4x32xf32>208// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<32x64x4x32xf32>209func.func @saxpy4d(%x: !type4d, %y: !type4d, %alpha : f32) -> !type4d {210 %c32 = arith.constant 32 : index211 %c64 = arith.constant 64 : index212 %c4 = arith.constant 4 : index213// CHECK: %[[C32:.*]] = arith.constant 32 : index214// CHECK: %[[C64:.*]] = arith.constant 64 : index215// CHECK: %[[C4:.*]] = arith.constant 4 : index216// CHECK: %[[C1:.*]] = arith.constant 1 : index217// CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C64]], %{{.*}} = %[[C1]]) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C4]], %{{.*}} = %[[C1]])218// CHECK: %[[BLKX:.*]] = gpu.block_id x219// CHECK: %[[BLKY:.*]] = gpu.block_id y220// CHECK: %[[TIDX:.*]] = gpu.thread_id x221// CHECK: %[[TIDY:.*]] = gpu.thread_id y222// CHECK: memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]]223// CHECK: memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]]224 scf.forall (%i, %j) in (%c32, %c64) {225 scf.forall (%k, %l) in (%c4, %c32) {226 %4 = memref.load %x[%i, %j, %k, %l] : !type4d227 %5 = memref.load %y[%i, %j, %k, %l] : !type4d228 %6 = math.fma %alpha, %4, %5 : f32229 memref.store %6, %y[%i, %j, %k, %l] : !type4d230 } { mapping = [#gpu.thread<y>, #gpu.thread<x>] }231 } { mapping = [#gpu.block<x>, #gpu.block<y>] }232 return %y : !type4d233}234 235module attributes {transform.with_named_sequence} {236 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {237 %funcop = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op238 %gpuLaunch = transform.gpu.map_forall_to_blocks %funcop { generate_gpu_launch } : (!transform.any_op) -> !transform.any_op239 transform.gpu.map_nested_forall_to_threads %gpuLaunch block_dims = [32, 4, 1] : (!transform.any_op) -> !transform.any_op240 transform.yield241 }242}243 244// -----245 246!type = memref<2 x 32 x f32>247!type1d = memref<32 x f32>248 249// CHECK-LABEL: func.func @saxpy2d_no_barrier(250func.func @saxpy2d_no_barrier(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {251 %one = arith.constant 1 : index252 %c12 = arith.constant 12 : index253 %c9 = arith.constant 9 : index254 %c7 = arith.constant 7 : index255// CHECK-NOT: gpu.barrier256// CHECK: return257 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)258 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)259 {260 scf.forall (%i, %j) in (%c7, %c9) {261 %4 = memref.load %x[%i, %j] : !type262 %5 = memref.load %y[%i, %j] : !type263 %6 = math.fma %alpha, %4, %5 : f32264 memref.store %6, %y[%i, %j] : !type265 } { mapping = [#gpu.thread<y>, #gpu.thread<x>] }266 gpu.terminator267 }268 return %y : !type269}270 271module attributes {transform.with_named_sequence} {272 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {273 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op274 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] sync_after_distribute = false : (!transform.any_op) -> !transform.any_op275 transform.yield276 }277}278 279// -----280 281!type = memref<32x32xf32>282// CHECK-LABEL: func.func @saxpy2d_singleloop(283// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<32x32xf32>284// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<32x32xf32>285func.func @saxpy2d_singleloop(%x: !type, %y: !type, %stream : !gpu.async.token) -> !type {286 %c32 = arith.constant 32 : index287 %one = arith.constant 1 : index288 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)289 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)290 {291// CHECK: %[[TIDX:.*]] = gpu.thread_id x292// CHECK: memref.load %[[ARGX]][%[[TIDX]], %[[TIDX]]]293// CHECK: memref.load %[[ARGY]][%[[TIDX]], %[[TIDX]]]294 scf.forall (%i) in (%c32) {295 %4 = memref.load %x[%i, %i] : !type296 %5 = memref.load %y[%i, %i] : !type297 %6 = arith.mulf %4, %5 : f32298 memref.store %6, %y[%i, %i] : !type299 } { mapping = [#gpu.thread<x>] }300 gpu.terminator301 }302 return %y : !type303}304 305module attributes {transform.with_named_sequence} {306 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {307 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op308 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [32, 1, 1] : (!transform.any_op) -> !transform.any_op309 transform.yield310 }311}312 313// -----314 315!type = memref<3 x 2 x 32 x f32>316!type1d = memref<32 x f32>317 318// CHECK-LABEL: func.func @saxpy3d_fold_id_z(319func.func @saxpy3d_fold_id_z(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {320 %one = arith.constant 1 : index321 %c12 = arith.constant 12 : index322 %c9 = arith.constant 9 : index323 %c7 = arith.constant 7 : index324// CHECK: %[[C0:.+]] = arith.constant 0 : index325// CHECK-NOT: gpu.thread_id z326 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)327 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)328 {329 scf.forall (%i, %j, %k) in (%one, %c7, %c9) {330// CHECK: memref.load %{{.*}}[%[[C0]],331// CHECK: memref.load %{{.*}}[%[[C0]],332 %4 = memref.load %x[%i, %j, %k] : !type333 %5 = memref.load %y[%i, %j, %k] : !type334 %6 = math.fma %alpha, %4, %5 : f32335// CHECK: memref.store %{{.*}}, %{{.*}}[%[[C0]]336 memref.store %6, %y[%i, %j, %k] : !type337 } { mapping = [#gpu.thread<z>, #gpu.thread<y>, #gpu.thread<x>] }338 gpu.terminator339 }340 return %y : !type341}342 343module attributes {transform.with_named_sequence} {344 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {345 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op346 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] sync_after_distribute = false : (!transform.any_op) -> !transform.any_op347 transform.yield348 }349}350 351 352// -----353 354!type = memref<2 x 32 x f32>355!type1d = memref<32 x f32>356 357// CHECK-DAG: #[[$MAPWGLIN:.*]] = affine_map<()[s0, s1, s2] -> (s0 + s1 * 32 + s2 * 256)>358// CHECK-DAG: #[[$MAPWGX:.*]] = affine_map<()[s0, s1] -> (((s0 + s1 * 32) floordiv 128) mod 2)>359// CHECK-DAG: #[[$MAPWGY:.*]] = affine_map<()[s0, s1, s2] -> (s2 + ((s0 + s1 * 32) floordiv 128) floordiv 2)>360 361// CHECK-LABEL: func.func @warpgroup_linear(362// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>363// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>364// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>365func.func @warpgroup_linear(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {366 %c2 = arith.constant 2 : index367 %c3 = arith.constant 3 : index368 %one = arith.constant 1 : index369 370// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index371// CHECK-DAG: %[[C768:.*]] = arith.constant 768 : index372// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index373// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index374// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index375 376// CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id x377// CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id y378// CHECK-DAG: %[[TIDZ:.*]] = gpu.thread_id z379// CHECK-DAG: %[[WIDLIN:.*]] = affine.apply #[[$MAPWGLIN]]()[%[[TIDX]], %[[TIDY]], %[[TIDZ]]]380// CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWGX]]()[%[[TIDX]], %[[TIDY]]]381// CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWGY]]()[%[[TIDX]], %[[TIDY]], %[[TIDZ]]]382// CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[WIDLIN]], %[[C768]] : index383// CHECK: scf.if %[[CMPLIN]]384// CHECK: memref.load %[[ARGX]][%[[WIDX]], %[[WIDY]]]385// CHECK: memref.load %[[ARGY]][%[[WIDX]], %[[WIDY]]]386 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)387 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)388 {389 scf.forall (%i, %j) in (%c2, %c3) {390 %4 = memref.load %x[%i, %j] : !type391 %5 = memref.load %y[%i, %j] : !type392 %6 = math.fma %alpha, %4, %5 : f32393 memref.store %6, %y[%i, %j] : !type394 } { mapping = [#gpu.warpgroup<linear_dim_0>, #gpu.warpgroup<linear_dim_1>]}395 gpu.terminator396 }397 return %y : !type398}399 400module attributes {transform.with_named_sequence} {401 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {402 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op403 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [32, 8, 4] : (!transform.any_op) -> !transform.any_op404 transform.yield405 }406}407 408// -----409 410!type = memref<2 x 32 x f32>411!type1d = memref<32 x f32>412 413// CHECK-DAG: #[[$MAPWLIN:.*]] = affine_map<()[s0, s1, s2] -> (s0 + s1 * 32 + s2 * 256)>414// CHECK-DAG: #[[$MAPWX:.*]] = affine_map<()[s0, s1, s2] -> ((s1 + s2 * 8 + s0 floordiv 32) mod 2)>415// CHECK-DAG: #[[$MAPWY:.*]] = affine_map<()[s0, s1, s2] -> ((s1 + s2 * 8 + s0 floordiv 32) floordiv 2)>416 417// CHECK-LABEL: func.func @warp_linear(418// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>419// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>420// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>421func.func @warp_linear(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {422 %c2 = arith.constant 2 : index423 %c3 = arith.constant 3 : index424 %one = arith.constant 1 : index425 426// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index427// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index428// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index429// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index430// CHECK-DAG: %[[C192:.*]] = arith.constant 192 : index431 432// CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id x433// CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id y434// CHECK-DAG: %[[TIDZ:.*]] = gpu.thread_id z435// CHECK-DAG: %[[WIDLIN:.*]] = affine.apply #[[$MAPWLIN]]()[%[[TIDX]], %[[TIDY]], %[[TIDZ]]]436// CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWX]]()[%[[TIDX]], %[[TIDY]], %[[TIDZ]]]437// CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWY]]()[%[[TIDX]], %[[TIDY]], %[[TIDZ]]]438// CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[WIDLIN]], %[[C192]] : index439// CHECK: scf.if %[[CMPLIN]]440// CHECK: memref.load %[[ARGX]][%[[WIDX]], %[[WIDY]]]441// CHECK: memref.load %[[ARGY]][%[[WIDX]], %[[WIDY]]]442 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)443 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)444 {445 scf.forall (%i, %j) in (%c2, %c3) {446 %4 = memref.load %x[%i, %j] : !type447 %5 = memref.load %y[%i, %j] : !type448 %6 = math.fma %alpha, %4, %5 : f32449 memref.store %6, %y[%i, %j] : !type450 } { mapping = [#gpu.warp<linear_dim_0>, #gpu.warp<linear_dim_1>]}451 gpu.terminator452 }453 return %y : !type454}455 456module attributes {transform.with_named_sequence} {457 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {458 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op459 transform.gpu.map_nested_forall_to_threads %funcop block_dims = [32, 8, 4] : (!transform.any_op) -> !transform.any_op460 transform.yield461 }462}463 464// -----465 466!type = memref<2 x 32 x f32>467!type1d = memref<32 x f32>468 469// CHECK-DAG: #[[$MAPWX:.*]] = affine_map<()[s0, s1] -> (((s0 + s1 * 18) floordiv 32) mod 3)>470// CHECK-DAG: #[[$MAPWY:.*]] = affine_map<()[s0, s1] -> ((((s0 + s1 * 18) floordiv 32) mod 6) floordiv 3)>471 472// CHECK-DAG: #[[$MAPLIN:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 18)>473// CHECK-DAG: #[[$MAPLX:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 18) mod 10)>474// CHECK-DAG: #[[$MAPLY:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 18) floordiv 10)>475 476// CHECK-LABEL: func.func @map_multi_level_linear(477func.func @map_multi_level_linear(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {478 %one = arith.constant 1 : index479 %c10 = arith.constant 10 : index480 %c9 = arith.constant 9 : index481 %c7 = arith.constant 7 : index482 %c1 = arith.constant 1 : index483 %c2 = arith.constant 2 : index484 %c3 = arith.constant 3 : index485 486 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index487 // CHECK-DAG: %[[C11:.*]] = arith.constant 11 : index488 // CHECK-DAG: %[[C18:.*]] = arith.constant 18 : index489 // CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index490 // CHECK-DAG: %[[C192:.*]] = arith.constant 192 : index491 492 // check that both the thread level and the warp level got distributed.493 // CHECK-NOT: #gpu.thread494 // CHECK-NOT: #gpu.warp495 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)496 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)497 {498 // CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id x499 // CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id y500 scf.forall (%i, %j) in (%c7, %c9) {501 %4 = memref.load %x[%i, %j] : !type502 %5 = memref.load %y[%i, %j] : !type503 %6 = math.fma %alpha, %4, %5 : f32504 memref.store %6, %y[%i, %j] : !type505 } { mapping = [#gpu.thread<y>, #gpu.thread<x>]}506 507 // CHECK-DAG: %[[LIN:.*]] = affine.apply #[[$MAPLIN]]()[%[[TIDX]], %[[TIDY]]]508 // CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWX]]()[%[[TIDX]], %[[TIDY]]]509 // CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWY]]()[%[[TIDX]], %[[TIDY]]]510 // CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[LIN]], %[[C192]] : index511 // CHECK: scf.if %[[CMPLIN]]512 scf.forall (%i, %j, %k) in (%c3, %c2, %c1) {513 %7 = memref.load %x[%i, %j] : !type514 %8 = arith.addf %alpha, %7 : f32515 memref.store %8, %y[%i, %j] : !type516 } {mapping = [#gpu.warp<linear_dim_0>, #gpu.warp<linear_dim_1>, #gpu.warp<linear_dim_2>] }517 518 // CHECK-DAG: %[[LIDX:.*]] = affine.apply #[[$MAPLX]]()[%[[TIDX]], %[[TIDY]]]519 // CHECK-DAG: %[[LIDY:.*]] = affine.apply #[[$MAPLY]]()[%[[TIDX]], %[[TIDY]]]520 // CHECK-DAG: %[[COND:.*]] = arith.cmpi ult, %[[LIN]], %[[C20]] : index521 // CHECK: scf.if %[[COND]]522 // CHECK: memref.load %{{.*}}[%[[LIDX]]] : memref<32xf32>523 // CHECK: memref.store %{{.*}}[%[[LIDY]]] : memref<32xf32>524 scf.forall (%i, %j) in (%c10, %c2) {525 %7 = memref.load %t[%i] : !type1d526 %8 = arith.addf %alpha, %7 : f32527 memref.store %8, %t[%j] : !type1d528 } {mapping = [#gpu.thread<linear_dim_0>, #gpu.thread<linear_dim_1>] }529 gpu.terminator530 }531 return %y : !type532}533 534module attributes {transform.with_named_sequence} {535 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {536 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op537 transform.gpu.map_nested_forall_to_threads %funcop538 block_dims = [18, 11, 1] : (!transform.any_op) -> !transform.any_op539 transform.yield540 }541}542 543// -----544 545!type = memref<2 x 32 x f32>546!type1d = memref<32 x f32>547 548// CHECK-DAG: #[[$MAPBLIN:.*]] = affine_map<()[s0, s1, s2] -> (s0 + s1 * 12 + s2 * 108)> 549// CHECK-DAG: #[[$MAPBX:.*]] = affine_map<()[s0, s1, s2] -> ((s0 + s1 * 12 + s2 * 108) mod 7)>550// CHECK-DAG: #[[$MAPBY:.*]] = affine_map<()[s0, s1, s2] -> ((s0 + s1 * 12 + s2 * 108) floordiv 7)>551 552// CHECK-LABEL: func.func @block_linear_existing_launch(553// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>554// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>555// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>556func.func @block_linear_existing_launch(557 %x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {558 %c9 = arith.constant 9 : index559 %c7 = arith.constant 7 : index560 %one = arith.constant 1 : index561 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index562 // CHECK-DAG: %[[C9:.*]] = arith.constant 9 : index563 // CHECK-DAG: %[[C12:.*]] = arith.constant 12 : index564 // CHECK-DAG: %[[C63:.*]] = arith.constant 63 : index565// CHECK: gpu.launch async [{{.*}}] blocks({{.*}}) in (%{{.*}} = %[[C12]], %{{.*}} = %[[C9]], %{{.*}} = %[[C1]]) threads566// CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x567// CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y568// CHECK-DAG: %[[BIDZ:.*]] = gpu.block_id z569// CHECK-DAG: %[[BIDLIN:.*]] = affine.apply #[[$MAPBLIN]]()[%[[BIDX]], %[[BIDY]], %[[BIDZ]]]570// CHECK-DAG: %[[BLX:.*]] = affine.apply #[[$MAPBX]]()[%[[BIDX]], %[[BIDY]], %[[BIDZ]]]571// CHECK-DAG: %[[BLY:.*]] = affine.apply #[[$MAPBY]]()[%[[BIDX]], %[[BIDY]], %[[BIDZ]]]572// CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[BIDLIN]], %[[C63]] : index573// CHECK: scf.if %[[CMPLIN]]574// CHECK: memref.load %[[ARGX]][%[[BLX]], %[[BLY]]]575// CHECK: memref.load %[[ARGY]][%[[BLX]], %[[BLY]]]576 %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)577 threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)578 {579 scf.forall (%i, %j) in (%c7, %c9) {580 %4 = memref.load %x[%i, %j] : !type581 %5 = memref.load %y[%i, %j] : !type582 %6 = math.fma %alpha, %4, %5 : f32583 memref.store %6, %y[%i, %j] : !type584 } { mapping = [#gpu.block<linear_dim_0>, #gpu.block<linear_dim_1>]}585 gpu.terminator586 }587 return %y : !type588}589 590module attributes {transform.with_named_sequence} {591 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {592 %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op593 transform.gpu.map_forall_to_blocks %funcop grid_dims = [12, 9, 1] : (!transform.any_op) -> !transform.any_op594 transform.yield595 }596}597 598// -----599 600!type = memref<2 x 32 x f32>601!type1d = memref<32 x f32>602 603// CHECK-DAG: #[[$MAPBX:.*]] = affine_map<()[s0] -> (s0 mod 7)>604// CHECK-DAG: #[[$MAPBY:.*]] = affine_map<()[s0, s1, s2] -> (s1 + s2 * 9 + s0 floordiv 7)>605 606// CHECK-LABEL: func.func @block_linear_generate_launch(607// CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>608// CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>609// CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32>610func.func @block_linear_generate_launch(611 %x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {612 %c9 = arith.constant 9 : index613 %c7 = arith.constant 7 : index614 %one = arith.constant 1 : index615 616 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index617 // CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index618 // CHECK-DAG: %[[C9:.*]] = arith.constant 9 : index619// CHECK: gpu.launch blocks({{.*}}) in (%{{.*}} = %[[C7]], %{{.*}} = %[[C9]], %{{.*}} = %[[C1]]) threads620// CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x621// CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y622// CHECK-DAG: %[[BIDZ:.*]] = gpu.block_id z623// CHECK-DAG: %[[BLX:.*]] = affine.apply #[[$MAPBX]]()[%[[BIDX]]]624// CHECK-DAG: %[[BLY:.*]] = affine.apply #[[$MAPBY]]()[%[[BIDX]], %[[BIDY]], %[[BIDZ]]]625// CHECK: memref.load %[[ARGX]][%[[BLX]], %[[BLY]]]626// CHECK: memref.load %[[ARGY]][%[[BLX]], %[[BLY]]]627 scf.forall (%i, %j) in (%c7, %c9) {628 %4 = memref.load %x[%i, %j] : !type629 %5 = memref.load %y[%i, %j] : !type630 %6 = math.fma %alpha, %4, %5 : f32631 memref.store %6, %y[%i, %j] : !type632 } { mapping = [#gpu.block<linear_dim_0>, #gpu.block<linear_dim_1>]}633 634 return %y : !type635}636 637module attributes {transform.with_named_sequence} {638 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {639 %funcop = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op640 transform.gpu.map_forall_to_blocks %funcop generate_gpu_launch : (!transform.any_op) -> !transform.any_op641 transform.yield642 }643}644 645// -----646 647#map = affine_map<(d0) -> (d0 * 128)>648#map1 = affine_map<(d0) -> (d0 * 32)>649 650// CHECK-DAG: #[[$MAPB:.*]] = affine_map<()[s0] -> (s0 * 128)>651// CHECK-DAG: #[[$MAPW:.*]] = affine_map<()[s0, s1, s2] -> (s2 * 32 + ((s0 + s1 * 4) floordiv 32) * 32)>652 653// CHECK-LABEL: func.func @simple_fill(654func.func @simple_fill(%arg0: memref<128xf32>) -> memref<128xf32> {655 %c0 = arith.constant 0 : index656 %cst = arith.constant dense<0.000000e+00> : vector<32xf32>657// CHECK: %[[C1:.*]] = arith.constant 1 : index658// CHECK: %[[C4:.*]] = arith.constant 4 : index659// CHECK: %[[C8:.*]] = arith.constant 8 : index660// CHECK: gpu.launch661 scf.forall (%arg1) in (1) {662// CHECK: %[[BIDX:.*]] = gpu.block_id x663// CHECK: %[[BLX:.*]] = affine.apply #[[$MAPB]]()[%[[BIDX]]]664 %0 = affine.apply #map(%arg1)665 %subview = memref.subview %arg0[%0] [128] [1] : memref<128xf32> to memref<128xf32, strided<[1], offset: ?>>666 scf.forall (%arg2) in (4) {667// CHECK: %[[TIDX:.*]] = gpu.thread_id x668// CHECK: %[[TIDY:.*]] = gpu.thread_id y669// CHECK: %[[TIDZ:.*]] = gpu.thread_id z670// CHECK: %[[THX:.*]] = affine.apply #[[$MAPW]]()[%[[TIDX]], %[[TIDY]], %[[TIDZ]]]671// CHECK-NOT: scf.if672// CHECK: memref.subview %{{.*}}[%[[THX]]]673 %1 = affine.apply #map1(%arg2)674 %subview_0 = memref.subview %subview[%1] [32] [1] : memref<128xf32, strided<[1], offset: ?>> to memref<32xf32, strided<[1], offset: ?>>675 vector.transfer_write %cst, %subview_0[%c0] {in_bounds = [true]} : vector<32xf32>, memref<32xf32, strided<[1], offset: ?>>676 memref.copy %subview_0, %subview_0 : memref<32xf32, strided<[1], offset: ?>> to memref<32xf32, strided<[1], offset: ?>>677 } {mapping = [#gpu.warp<linear_dim_0>]}678 memref.copy %subview, %subview : memref<128xf32, strided<[1], offset: ?>> to memref<128xf32, strided<[1], offset: ?>>679 } {mapping = [#gpu.block<x>]}680 return %arg0 : memref<128xf32>681}682 683module attributes {transform.with_named_sequence} {684 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {685 %func = transform.structured.match ops{["func.func"]} in %module_op686 : (!transform.any_op) -> !transform.any_op687 %gpu_launch = transform.gpu.map_forall_to_blocks %func generate_gpu_launch688 : (!transform.any_op) -> !transform.any_op689 transform.gpu.map_nested_forall_to_threads %gpu_launch block_dims = [4, 8, 4]690 : (!transform.any_op) -> !transform.any_op691 transform.yield692 }693}694 695// -----696 697#map = affine_map<(d0) -> (d0 * 128)>698#map1 = affine_map<(d0) -> (d0 * 32)>699 700// CHECK-DAG: #[[$MAPB:.*]] = affine_map<()[s0] -> (s0 * 128)>701// CHECK-DAG: #[[$MAPLANE:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 73) mod 32)>702// CHECK-DAG: #[[$MAPI:.*]] = affine_map<()[s0, s1] -> (s0 * 32 + s1 * 2336 - ((s0 + s1 * 73) floordiv 2) * 64)>703// CHECK-DAG: #[[$MAPJ:.*]] = affine_map<()[s0, s1] -> ((((s0 + s1 * 73) mod 32) floordiv 2) * 32)>704 705// CHECK-LABEL: func.func @simple_fill(706func.func @simple_fill(%arg0: memref<128x256xf32>) -> memref<128x256xf32> {707 %c0 = arith.constant 0 : index708 %cst = arith.constant dense<0.000000e+00> : vector<16x32xf32>709 // CHECK: %[[C6:.*]] = arith.constant 6 : index710 // CHECK: gpu.launch711 scf.forall (%arg1) in (1) {712 // CHECK: %[[BIDX:.*]] = gpu.block_id x713 // CHECK: %[[BLX:.*]] = affine.apply #[[$MAPB]]()[%[[BIDX]]]714 %0 = affine.apply #map(%arg1)715 %subview = memref.subview %arg0[%0, 0] [128, 256] [1, 1]716 : memref<128x256xf32> to memref<128x256xf32, strided<[256, 1], offset: ?>>717 718 // %arg2 and %arg3 map to lanes [0, 6) and are turned into epxressions719 // involving threadIdx.x/y by the map_nested_forall_to_threads720 // transformation. This results in a if (linear_thread_id < 6) conditional.721 scf.forall (%arg2, %arg3) in (2, 3) {722 // CHECK: %[[TIDX:.*]] = gpu.thread_id x723 // CHECK: %[[TIDY:.*]] = gpu.thread_id y724 // CHECK: %[[LID:.*]] = affine.apply #[[$MAPLANE]]()[%[[TIDX]], %[[TIDY]]]725 // CHECK: %[[COND:.*]] = arith.cmpi ult, %[[LID]], %[[C6]]726 // CHECK: scf.if %[[COND]]727 // CHECK: %[[I:.*]] = affine.apply #[[$MAPI]]()[%[[TIDX]], %[[TIDY]]]728 // CHECK: %[[J:.*]] = affine.apply #[[$MAPJ]]()[%[[TIDX]], %[[TIDY]]]729 // CHECK: memref.subview %{{.*}}[%[[I]], %[[J]]]730 %1 = affine.apply #map1(%arg2)731 %2 = affine.apply #map1(%arg3)732 %subview_0 = memref.subview %subview[%1, %2] [16, 32] [1, 1] 733 : memref<128x256xf32, strided<[256, 1], offset: ?>> to memref<16x32xf32, strided<[256, 1], offset: ?>>734 vector.transfer_write %cst, %subview_0[%c0, %c0] {in_bounds = [true, true]} 735 : vector<16x32xf32>, memref<16x32xf32, strided<[256, 1], offset: ?>>736 737 // This could be obtained e.g. if a previous transformation mapped this loop738 // to lanes. This can aslo be written by hand as valid IR.739 } {mapping = [#gpu.lane<linear_dim_0>, #gpu.lane<linear_dim_1>]}740 } {mapping = [#gpu.block<x>]}741 return %arg0 : memref<128x256xf32>742}743 744module attributes {transform.with_named_sequence} {745 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {746 %func = transform.structured.match ops{["func.func"]} in %module_op747 : (!transform.any_op) -> !transform.any_op748 %gpu_launch = transform.gpu.map_forall_to_blocks %func generate_gpu_launch749 : (!transform.any_op) -> !transform.any_op750 751 // This transformation maps scf.forall ivs to a particular mapping of thread752 // ids (laneid, threadid, warpid or warpgroupid).753 transform.gpu.map_nested_forall_to_threads %gpu_launch block_dims = [73, 5, 1]754 : (!transform.any_op) -> !transform.any_op755 transform.yield756 }757}758 759// -----760 761#map = affine_map<(d0) -> (d0 * 128)>762#map1 = affine_map<(d0) -> (d0 * 32)>763 764// CHECK-DAG: #[[$MAPB:.*]] = affine_map<()[s0] -> (s0 * 128)>765// CHECK-DAG: #[[$MAP_LIN_W:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 73) floordiv 32)>766// CHECK-DAG: #[[$MAP_W0:.*]] = affine_map<()[s0] -> (s0 * 32 - (s0 floordiv 2) * 64)>767// CHECK-DAG: #[[$MAP_W1:.*]] = affine_map<()[s0] -> ((s0 floordiv 2) * 32)>768 769// CHECK-LABEL: func.func @simple_fill(770func.func @simple_fill(%arg0: memref<128xf32>) -> memref<128xf32> {771 %c0 = arith.constant 0 : index772 %cst = arith.constant dense<0.000000e+00> : vector<32xf32>773// CHECK-DAG: %[[C0_i64:.*]] = arith.constant 0 : i64774// CHECK-DAG: %[[C1_i64:.*]] = arith.constant 1 : i64775/// 0x2f1 is 753776// CHECK-DAG: %[[C753_i64:.*]] = arith.constant 753 : i64777 778// CHECK: gpu.launch779 scf.forall (%arg1) in (1) {780// CHECK: %[[BIDX:.*]] = gpu.block_id x781// CHECK: %[[BLX:.*]] = affine.apply #[[$MAPB]]()[%[[BIDX]]]782 %0 = affine.apply #map(%arg1)783 %subview = memref.subview %arg0[%0] [128] [1] : memref<128xf32> to memref<128xf32, strided<[1], offset: ?>>784 785 // %arg2 and %arg3 map to lanes [0, 6) and are turned into epxressions786 // involving threadIdx.x/y by the map_nested_forall_to_threads787 // transformation. This results in a if (linear_thread_id < 6) conditional.788 scf.forall (%arg2, %arg3) in (2, 3) {789 // CHECK: %[[TIDX:.*]] = gpu.thread_id x790 // CHECK: %[[TIDY:.*]] = gpu.thread_id y791 792 // CHECK: %[[LIN_W:.*]] = affine.apply #[[$MAP_LIN_W]]()[%[[TIDX]], %[[TIDY]]]793 //794 // Compute the active warps below using the mask + popcnt795 // CHECK: %[[LIN_W_i64:.*]] = arith.index_castui %[[LIN_W]] : index to i64796 // CHECK: %[[TWO_POW_W:.*]] = arith.shli %[[C1_i64]], %[[LIN_W_i64]] : i64797 // CHECK: %[[FILTER_TILL_W:.*]] = arith.subi %[[TWO_POW_W]], %[[C1_i64]] : i64798 // CHECK: %[[ACTIVE_TILL_W:.*]] = arith.andi %[[FILTER_TILL_W]], %[[C753_i64]] : i64799 // CHECK: %[[LOGICAL_ID_W_i64:.*]] = math.ctpop %[[ACTIVE_TILL_W]] : i64800 // CHECK: %[[LOGICAL_ID_W:.*]] = arith.index_castui %[[LOGICAL_ID_W_i64]] : i64 to index801 //802 // Dynamically compute whether this warp is active below using the mask + popcnt803 // CHECK: %[[IS_ACTIVE_W_MASK:.*]] = arith.andi %[[TWO_POW_W]], %[[C753_i64]] : i64804 // CHECK: %[[IS_ACTIVE_W:.*]] = arith.cmpi ne, %[[IS_ACTIVE_W_MASK]], %[[C0_i64]] : i64805 // CHECK: scf.if %[[IS_ACTIVE_W]] {806 807 // CHECK: %[[W0:.*]] = affine.apply #[[$MAP_W0]]()[%[[LOGICAL_ID_W]]]808 // CHECK: %[[W1:.*]] = affine.apply #[[$MAP_W1]]()[%[[LOGICAL_ID_W]]]809 // CHECK: memref.subview %{{.*}}[%[[W0]]] [%[[W1]]]810 %1 = affine.apply #map1(%arg2)811 %2 = affine.apply #map1(%arg3)812 %subview_0 = memref.subview %subview[%1] [%2] [1] : memref<128xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>813 vector.transfer_write %cst, %subview_0[%c0] {in_bounds = [true]} : vector<32xf32>, memref<?xf32, strided<[1], offset: ?>>814 815 // This could be obtained e.g. if a previous transformation mapped this loop816 // to lanes. This can aslo be written by hand as valid IR.817 // This additionally uses the hex mask: 0x 10 1111 0001818 } {mapping = [#gpu.warp<linear_dim_0>, #gpu.warp<linear_dim_1>, #gpu.mask<0x2f1>]}819 820 memref.copy %subview, %subview : memref<128xf32, strided<[1], offset: ?>> to memref<128xf32, strided<[1], offset: ?>>821 } {mapping = [#gpu.block<x>]}822 return %arg0 : memref<128xf32>823}824 825module attributes {transform.with_named_sequence} {826 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {827 %func = transform.structured.match ops{["func.func"]} in %module_op828 : (!transform.any_op) -> !transform.any_op829 %gpu_launch = transform.gpu.map_forall_to_blocks %func generate_gpu_launch830 : (!transform.any_op) -> !transform.any_op831 832 // This transformation maps scf.forall ivs to a particular mapping of thread833 // ids (laneid, threadid, warpid or warpgroupid).834 transform.gpu.map_nested_forall_to_threads %gpu_launch block_dims = [73, 5, 1]835 : (!transform.any_op) -> !transform.any_op836 transform.yield837 }838}839