348 lines · plain
1// RUN: mlir-opt -int-range-optimizations -split-input-file %s | FileCheck %s2 3// CHECK-LABEL: func @launch_func4func.func @launch_func(%arg0 : index) {5 %0 = test.with_bounds {6 umin = 3 : index, umax = 5 : index,7 smin = 3 : index, smax = 5 : index8 } : index9 %1 = test.with_bounds {10 umin = 7 : index, umax = 11 : index,11 smin = 7 : index, smax = 11 : index12 } : index13 gpu.launch blocks(%block_id_x, %block_id_y, %block_id_z) in (%grid_dim_x = %0, %grid_dim_y = %1, %grid_dim_z = %arg0)14 threads(%thread_id_x, %thread_id_y, %thread_id_z) in (%block_dim_x = %arg0, %block_dim_y = %0, %block_dim_z = %1) {15 16 // CHECK: test.reflect_bounds {smax = 5 : index, smin = 3 : index, umax = 5 : index, umin = 3 : index}17 // CHECK: test.reflect_bounds {smax = 11 : index, smin = 7 : index, umax = 11 : index, umin = 7 : index}18 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}19 %grid_dim_x0 = test.reflect_bounds %grid_dim_x : index20 %grid_dim_y0 = test.reflect_bounds %grid_dim_y : index21 %grid_dim_z0 = test.reflect_bounds %grid_dim_z : index22 23 // CHECK: test.reflect_bounds {smax = 4 : index, smin = 0 : index, umax = 4 : index, umin = 0 : index}24 // CHECK: test.reflect_bounds {smax = 10 : index, smin = 0 : index, umax = 10 : index, umin = 0 : index}25 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}26 %block_id_x0 = test.reflect_bounds %block_id_x : index27 %block_id_y0 = test.reflect_bounds %block_id_y : index28 %block_id_z0 = test.reflect_bounds %block_id_z : index29 30 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}31 // CHECK: test.reflect_bounds {smax = 5 : index, smin = 3 : index, umax = 5 : index, umin = 3 : index}32 // CHECK: test.reflect_bounds {smax = 11 : index, smin = 7 : index, umax = 11 : index, umin = 7 : index}33 %block_dim_x0 = test.reflect_bounds %block_dim_x : index34 %block_dim_y0 = test.reflect_bounds %block_dim_y : index35 %block_dim_z0 = test.reflect_bounds %block_dim_z : index36 37 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}38 // CHECK: test.reflect_bounds {smax = 4 : index, smin = 0 : index, umax = 4 : index, umin = 0 : index}39 // CHECK: test.reflect_bounds {smax = 10 : index, smin = 0 : index, umax = 10 : index, umin = 0 : index}40 %thread_id_x0 = test.reflect_bounds %thread_id_x : index41 %thread_id_y0 = test.reflect_bounds %thread_id_y : index42 %thread_id_z0 = test.reflect_bounds %thread_id_z : index43 44 // The launch bounds are not constant, and so this can't infer anything45 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}46 %thread_id_op = gpu.thread_id y47 %thread_id_op0 = test.reflect_bounds %thread_id_op : index48 gpu.terminator49 }50 51 func.return52}53 54// -----55 56// CHECK-LABEL: func @kernel57module attributes {gpu.container_module} {58 gpu.module @gpu_module {59 llvm.func @kernel() attributes {gpu.kernel} {60 61 %grid_dim_x = gpu.grid_dim x62 %grid_dim_y = gpu.grid_dim y63 %grid_dim_z = gpu.grid_dim z64 65 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}66 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}67 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}68 %grid_dim_x0 = test.reflect_bounds %grid_dim_x : index69 %grid_dim_y0 = test.reflect_bounds %grid_dim_y : index70 %grid_dim_z0 = test.reflect_bounds %grid_dim_z : index71 72 %block_id_x = gpu.block_id x73 %block_id_y = gpu.block_id y74 %block_id_z = gpu.block_id z75 76 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}77 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}78 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}79 %block_id_x0 = test.reflect_bounds %block_id_x : index80 %block_id_y0 = test.reflect_bounds %block_id_y : index81 %block_id_z0 = test.reflect_bounds %block_id_z : index82 83 %block_dim_x = gpu.block_dim x84 %block_dim_y = gpu.block_dim y85 %block_dim_z = gpu.block_dim z86 87 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}88 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}89 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}90 %block_dim_x0 = test.reflect_bounds %block_dim_x : index91 %block_dim_y0 = test.reflect_bounds %block_dim_y : index92 %block_dim_z0 = test.reflect_bounds %block_dim_z : index93 94 %thread_id_x = gpu.thread_id x95 %thread_id_y = gpu.thread_id y96 %thread_id_z = gpu.thread_id z97 98 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}99 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}100 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}101 %thread_id_x0 = test.reflect_bounds %thread_id_x : index102 %thread_id_y0 = test.reflect_bounds %thread_id_y : index103 %thread_id_z0 = test.reflect_bounds %thread_id_z : index104 105 %global_id_x = gpu.global_id x106 %global_id_y = gpu.global_id y107 %global_id_z = gpu.global_id z108 109 // CHECK: test.reflect_bounds {smax = 9223372036854775807 : index, smin = -9223372036854775808 : index, umax = -8589934592 : index, umin = 0 : index}110 // CHECK: test.reflect_bounds {smax = 9223372036854775807 : index, smin = -9223372036854775808 : index, umax = -8589934592 : index, umin = 0 : index}111 // CHECK: test.reflect_bounds {smax = 9223372036854775807 : index, smin = -9223372036854775808 : index, umax = -8589934592 : index, umin = 0 : index}112 %global_id_x0 = test.reflect_bounds %global_id_x : index113 %global_id_y0 = test.reflect_bounds %global_id_y : index114 %global_id_z0 = test.reflect_bounds %global_id_z : index115 116 %subgroup_size = gpu.subgroup_size : index117 %lane_id = gpu.lane_id118 %num_subgroups = gpu.num_subgroups : index119 %subgroup_id = gpu.subgroup_id : index120 121 // CHECK: test.reflect_bounds {smax = 128 : index, smin = 1 : index, umax = 128 : index, umin = 1 : index}122 // CHECK: test.reflect_bounds {smax = 127 : index, smin = 0 : index, umax = 127 : index, umin = 0 : index}123 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}124 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}125 %subgroup_size0 = test.reflect_bounds %subgroup_size : index126 %lane_id0 = test.reflect_bounds %lane_id : index127 %num_subgroups0 = test.reflect_bounds %num_subgroups : index128 %subgroup_id0 = test.reflect_bounds %subgroup_id : index129 130 llvm.return131 }132 }133}134 135// -----136 137// CHECK-LABEL: func @annotated_kernel138module attributes {gpu.container_module} {139 gpu.module @gpu_module {140 gpu.func @annotated_kernel() kernel141 attributes {known_block_size = array<i32: 8, 12, 16>,142 known_grid_size = array<i32: 20, 24, 28>} {143 144 %grid_dim_x = gpu.grid_dim x145 %grid_dim_y = gpu.grid_dim y146 %grid_dim_z = gpu.grid_dim z147 148 // CHECK: test.reflect_bounds {smax = 20 : index, smin = 20 : index, umax = 20 : index, umin = 20 : index}149 // CHECK: test.reflect_bounds {smax = 24 : index, smin = 24 : index, umax = 24 : index, umin = 24 : index}150 // CHECK: test.reflect_bounds {smax = 28 : index, smin = 28 : index, umax = 28 : index, umin = 28 : index}151 %grid_dim_x0 = test.reflect_bounds %grid_dim_x : index152 %grid_dim_y0 = test.reflect_bounds %grid_dim_y : index153 %grid_dim_z0 = test.reflect_bounds %grid_dim_z : index154 155 %block_id_x = gpu.block_id x156 %block_id_y = gpu.block_id y157 %block_id_z = gpu.block_id z158 159 // CHECK: test.reflect_bounds {smax = 19 : index, smin = 0 : index, umax = 19 : index, umin = 0 : index}160 // CHECK: test.reflect_bounds {smax = 23 : index, smin = 0 : index, umax = 23 : index, umin = 0 : index}161 // CHECK: test.reflect_bounds {smax = 27 : index, smin = 0 : index, umax = 27 : index, umin = 0 : index}162 %block_id_x0 = test.reflect_bounds %block_id_x : index163 %block_id_y0 = test.reflect_bounds %block_id_y : index164 %block_id_z0 = test.reflect_bounds %block_id_z : index165 166 %block_dim_x = gpu.block_dim x167 %block_dim_y = gpu.block_dim y168 %block_dim_z = gpu.block_dim z169 170 // CHECK: test.reflect_bounds {smax = 8 : index, smin = 8 : index, umax = 8 : index, umin = 8 : index}171 // CHECK: test.reflect_bounds {smax = 12 : index, smin = 12 : index, umax = 12 : index, umin = 12 : index}172 // CHECK: test.reflect_bounds {smax = 16 : index, smin = 16 : index, umax = 16 : index, umin = 16 : index}173 %block_dim_x0 = test.reflect_bounds %block_dim_x : index174 %block_dim_y0 = test.reflect_bounds %block_dim_y : index175 %block_dim_z0 = test.reflect_bounds %block_dim_z : index176 177 %thread_id_x = gpu.thread_id x178 %thread_id_y = gpu.thread_id y179 %thread_id_z = gpu.thread_id z180 181 // CHECK: test.reflect_bounds {smax = 7 : index, smin = 0 : index, umax = 7 : index, umin = 0 : index}182 // CHECK: test.reflect_bounds {smax = 11 : index, smin = 0 : index, umax = 11 : index, umin = 0 : index}183 // CHECK: test.reflect_bounds {smax = 15 : index, smin = 0 : index, umax = 15 : index, umin = 0 : index}184 %thread_id_x0 = test.reflect_bounds %thread_id_x : index185 %thread_id_y0 = test.reflect_bounds %thread_id_y : index186 %thread_id_z0 = test.reflect_bounds %thread_id_z : index187 188 %global_id_x = gpu.global_id x189 %global_id_y = gpu.global_id y190 %global_id_z = gpu.global_id z191 192 // CHECK: test.reflect_bounds {smax = 159 : index, smin = 0 : index, umax = 159 : index, umin = 0 : index}193 // CHECK: test.reflect_bounds {smax = 287 : index, smin = 0 : index, umax = 287 : index, umin = 0 : index}194 // CHECK: test.reflect_bounds {smax = 447 : index, smin = 0 : index, umax = 447 : index, umin = 0 : index}195 %global_id_x0 = test.reflect_bounds %global_id_x : index196 %global_id_y0 = test.reflect_bounds %global_id_y : index197 %global_id_z0 = test.reflect_bounds %global_id_z : index198 199 %subgroup_size = gpu.subgroup_size : index200 %lane_id = gpu.lane_id201 %num_subgroups = gpu.num_subgroups : index202 %subgroup_id = gpu.subgroup_id : index203 204 // CHECK: test.reflect_bounds {smax = 128 : index, smin = 1 : index, umax = 128 : index, umin = 1 : index}205 // CHECK: test.reflect_bounds {smax = 127 : index, smin = 0 : index, umax = 127 : index, umin = 0 : index}206 // CHECK: test.reflect_bounds {smax = 4294967295 : index, smin = 1 : index, umax = 4294967295 : index, umin = 1 : index}207 // CHECK: test.reflect_bounds {smax = 4294967294 : index, smin = 0 : index, umax = 4294967294 : index, umin = 0 : index}208 %subgroup_size0 = test.reflect_bounds %subgroup_size : index209 %lane_id0 = test.reflect_bounds %lane_id : index210 %num_subgroups0 = test.reflect_bounds %num_subgroups : index211 %subgroup_id0 = test.reflect_bounds %subgroup_id : index212 213 gpu.return214 }215 }216}217 218// -----219 220// CHECK-LABEL: func @annotated_kernel221module {222 func.func @annotated_kernel()223 attributes {gpu.known_block_size = array<i32: 8, 12, 16>,224 gpu.known_grid_size = array<i32: 20, 24, 28>} {225 226 %block_id_x = gpu.block_id x227 %block_id_y = gpu.block_id y228 %block_id_z = gpu.block_id z229 230 // CHECK: test.reflect_bounds {smax = 19 : index, smin = 0 : index, umax = 19 : index, umin = 0 : index}231 // CHECK: test.reflect_bounds {smax = 23 : index, smin = 0 : index, umax = 23 : index, umin = 0 : index}232 // CHECK: test.reflect_bounds {smax = 27 : index, smin = 0 : index, umax = 27 : index, umin = 0 : index}233 %block_id_x0 = test.reflect_bounds %block_id_x : index234 %block_id_y0 = test.reflect_bounds %block_id_y : index235 %block_id_z0 = test.reflect_bounds %block_id_z : index236 237 %thread_id_x = gpu.thread_id x238 %thread_id_y = gpu.thread_id y239 %thread_id_z = gpu.thread_id z240 241 // CHECK: test.reflect_bounds {smax = 7 : index, smin = 0 : index, umax = 7 : index, umin = 0 : index}242 // CHECK: test.reflect_bounds {smax = 11 : index, smin = 0 : index, umax = 11 : index, umin = 0 : index}243 // CHECK: test.reflect_bounds {smax = 15 : index, smin = 0 : index, umax = 15 : index, umin = 0 : index}244 %thread_id_x0 = test.reflect_bounds %thread_id_x : index245 %thread_id_y0 = test.reflect_bounds %thread_id_y : index246 %thread_id_z0 = test.reflect_bounds %thread_id_z : index247 248 return249 }250}251 252// -----253 254// CHECK-LABEL: func @local_bounds_kernel255module attributes {gpu.container_module} {256 gpu.module @gpu_module {257 gpu.func @local_bounds_kernel() kernel {258 259 %grid_dim_x = gpu.grid_dim x upper_bound 20260 %grid_dim_y = gpu.grid_dim y upper_bound 24261 %grid_dim_z = gpu.grid_dim z upper_bound 28262 263 // CHECK: test.reflect_bounds {smax = 20 : index, smin = 1 : index, umax = 20 : index, umin = 1 : index}264 // CHECK: test.reflect_bounds {smax = 24 : index, smin = 1 : index, umax = 24 : index, umin = 1 : index}265 // CHECK: test.reflect_bounds {smax = 28 : index, smin = 1 : index, umax = 28 : index, umin = 1 : index}266 %grid_dim_x0 = test.reflect_bounds %grid_dim_x : index267 %grid_dim_y0 = test.reflect_bounds %grid_dim_y : index268 %grid_dim_z0 = test.reflect_bounds %grid_dim_z : index269 270 %block_id_x = gpu.block_id x upper_bound 20271 %block_id_y = gpu.block_id y upper_bound 24272 %block_id_z = gpu.block_id z upper_bound 28273 274 // CHECK: test.reflect_bounds {smax = 19 : index, smin = 0 : index, umax = 19 : index, umin = 0 : index}275 // CHECK: test.reflect_bounds {smax = 23 : index, smin = 0 : index, umax = 23 : index, umin = 0 : index}276 // CHECK: test.reflect_bounds {smax = 27 : index, smin = 0 : index, umax = 27 : index, umin = 0 : index}277 %block_id_x0 = test.reflect_bounds %block_id_x : index278 %block_id_y0 = test.reflect_bounds %block_id_y : index279 %block_id_z0 = test.reflect_bounds %block_id_z : index280 281 %block_dim_x = gpu.block_dim x upper_bound 8282 %block_dim_y = gpu.block_dim y upper_bound 12283 %block_dim_z = gpu.block_dim z upper_bound 16284 285 // CHECK: test.reflect_bounds {smax = 8 : index, smin = 1 : index, umax = 8 : index, umin = 1 : index}286 // CHECK: test.reflect_bounds {smax = 12 : index, smin = 1 : index, umax = 12 : index, umin = 1 : index}287 // CHECK: test.reflect_bounds {smax = 16 : index, smin = 1 : index, umax = 16 : index, umin = 1 : index}288 %block_dim_x0 = test.reflect_bounds %block_dim_x : index289 %block_dim_y0 = test.reflect_bounds %block_dim_y : index290 %block_dim_z0 = test.reflect_bounds %block_dim_z : index291 292 %thread_id_x = gpu.thread_id x upper_bound 8293 %thread_id_y = gpu.thread_id y upper_bound 12294 %thread_id_z = gpu.thread_id z upper_bound 16295 296 // CHECK: test.reflect_bounds {smax = 7 : index, smin = 0 : index, umax = 7 : index, umin = 0 : index}297 // CHECK: test.reflect_bounds {smax = 11 : index, smin = 0 : index, umax = 11 : index, umin = 0 : index}298 // CHECK: test.reflect_bounds {smax = 15 : index, smin = 0 : index, umax = 15 : index, umin = 0 : index}299 %thread_id_x0 = test.reflect_bounds %thread_id_x : index300 %thread_id_y0 = test.reflect_bounds %thread_id_y : index301 %thread_id_z0 = test.reflect_bounds %thread_id_z : index302 303 %global_id_x = gpu.global_id x upper_bound 160304 %global_id_y = gpu.global_id y upper_bound 288305 %global_id_z = gpu.global_id z upper_bound 448306 307 // CHECK: test.reflect_bounds {smax = 159 : index, smin = 0 : index, umax = 159 : index, umin = 0 : index}308 // CHECK: test.reflect_bounds {smax = 287 : index, smin = 0 : index, umax = 287 : index, umin = 0 : index}309 // CHECK: test.reflect_bounds {smax = 447 : index, smin = 0 : index, umax = 447 : index, umin = 0 : index}310 %global_id_x0 = test.reflect_bounds %global_id_x : index311 %global_id_y0 = test.reflect_bounds %global_id_y : index312 %global_id_z0 = test.reflect_bounds %global_id_z : index313 314 %subgroup_size = gpu.subgroup_size upper_bound 32 : index315 %subgroup_id = gpu.subgroup_id upper_bound 32 : index316 %num_subgroups = gpu.num_subgroups upper_bound 8 : index317 %lane_id = gpu.lane_id upper_bound 64318 319 // CHECK: test.reflect_bounds {smax = 32 : index, smin = 1 : index, umax = 32 : index, umin = 1 : index}320 // CHECK: test.reflect_bounds {smax = 31 : index, smin = 0 : index, umax = 31 : index, umin = 0 : index}321 // CHECK: test.reflect_bounds {smax = 8 : index, smin = 1 : index, umax = 8 : index, umin = 1 : index}322 // CHECK: test.reflect_bounds {smax = 63 : index, smin = 0 : index, umax = 63 : index, umin = 0 : index}323 %subgroup_size0 = test.reflect_bounds %subgroup_size : index324 %subgroup_id0 = test.reflect_bounds %subgroup_id : index325 %num_subgroups0 = test.reflect_bounds %num_subgroups : index326 %lane_id0 = test.reflect_bounds %lane_id : index327 328 gpu.return329 }330 }331}332 333// -----334 335// CHECK-LABEL: func @broadcast336func.func @broadcast(%idx: i32) {337 %0 = test.with_bounds { umin = 0 : index, umax = 10 : index, smin = 0 : index, smax = 10 : index } : index338 %1 = gpu.subgroup_broadcast %0, first_active_lane : index339 %2 = gpu.subgroup_broadcast %0, specific_lane %idx : index340 341 // CHECK: test.reflect_bounds {smax = 10 : index, smin = 0 : index, umax = 10 : index, umin = 0 : index}342 // CHECK: test.reflect_bounds {smax = 10 : index, smin = 0 : index, umax = 10 : index, umin = 0 : index}343 344 %4 = test.reflect_bounds %1 : index345 %5 = test.reflect_bounds %2 : index346 return347}348