205 lines · plain
1// Test that dist_schedule gets correctly translated with the correct schedule type and chunk size where appropriate while using workshare loops.2 3// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s4 5llvm.func @distribute_wsloop_dist_schedule_chunked_schedule_chunked(%n: i32, %teams: i32, %threads: i32, %dcs: i32) {6 %0 = llvm.mlir.constant(0 : i32) : i327 %1 = llvm.mlir.constant(1 : i32) : i328 %scs = llvm.mlir.constant(64 : i32) : i329 10 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {11 omp.parallel {12 omp.distribute dist_schedule_static dist_schedule_chunk_size(%dcs : i32) {13 omp.wsloop schedule(static = %scs : i32) {14 omp.loop_nest (%i) : i32 = (%0) to (%n) step (%1) {15 omp.yield16 }17 } {omp.composite}18 } {omp.composite}19 omp.terminator20 } {omp.composite}21 omp.terminator22 }23 llvm.return24}25// CHECK: define internal void @distribute_wsloop_dist_schedule_chunked_schedule_chunked..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {26// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 33, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i32 1, i32 64)27// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 91, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i32 1, i32 %3)28 29llvm.func @distribute_wsloop_dist_schedule_chunked_schedule_chunked_i64(%n: i32, %teams: i32, %threads: i32) {30 %0 = llvm.mlir.constant(0 : i64) : i6431 %1 = llvm.mlir.constant(1 : i64) : i6432 %dcs = llvm.mlir.constant(1024 : i64) : i6433 %scs = llvm.mlir.constant(64 : i64) : i6434 %n64 = llvm.zext %n : i32 to i6435 36 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {37 omp.parallel {38 omp.distribute dist_schedule_static dist_schedule_chunk_size(%dcs : i64) {39 omp.wsloop schedule(static = %scs : i64) {40 omp.loop_nest (%i) : i64 = (%0) to (%n64) step (%1) {41 omp.yield42 }43 } {omp.composite}44 } {omp.composite}45 omp.terminator46 } {omp.composite}47 omp.terminator48 }49 llvm.return50}51// CHECK: define internal void @distribute_wsloop_dist_schedule_chunked_schedule_chunked_i64..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {52// CHECK: call void @__kmpc_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 33, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i64 1, i64 64)53// call void @__kmpc_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 91, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i64 1, i64 1024)54 55// -----56 57llvm.func @distribute_wsloop_dist_schedule_chunked(%n: i32, %teams: i32, %threads: i32) {58 %0 = llvm.mlir.constant(0 : i32) : i3259 %1 = llvm.mlir.constant(1 : i32) : i3260 %dcs = llvm.mlir.constant(1024 : i32) : i3261 62 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {63 omp.parallel {64 omp.distribute dist_schedule_static dist_schedule_chunk_size(%dcs : i32) {65 omp.wsloop schedule(static) {66 omp.loop_nest (%i) : i32 = (%0) to (%n) step (%1) {67 omp.yield68 }69 } {omp.composite}70 } {omp.composite}71 omp.terminator72 } {omp.composite}73 omp.terminator74 }75 llvm.return76}77// CHECK: define internal void @distribute_wsloop_dist_schedule_chunked..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {78// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 34, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i32 1, i32 0)79// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 91, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i32 1, i32 1024)80 81llvm.func @distribute_wsloop_dist_schedule_chunked_i64(%n: i32, %teams: i32, %threads: i32) {82 %0 = llvm.mlir.constant(0 : i64) : i6483 %1 = llvm.mlir.constant(1 : i64) : i6484 %dcs = llvm.mlir.constant(1024 : i64) : i6485 %n64 = llvm.zext %n : i32 to i6486 87 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {88 omp.parallel {89 omp.distribute dist_schedule_static dist_schedule_chunk_size(%dcs : i64) {90 omp.wsloop schedule(static) {91 omp.loop_nest (%i) : i64 = (%0) to (%n64) step (%1) {92 omp.yield93 }94 } {omp.composite}95 } {omp.composite}96 omp.terminator97 } {omp.composite}98 omp.terminator99 }100 llvm.return101}102// CHECK: define internal void @distribute_wsloop_dist_schedule_chunked_i64..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {103// CHECK: call void @__kmpc_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 34, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i64 1, i64 0)104// CHECK: call void @__kmpc_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 91, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i64 1, i64 1024)105 106// -----107 108llvm.func @distribute_wsloop_schedule_chunked(%n: i32, %teams: i32, %threads: i32) {109 %0 = llvm.mlir.constant(0 : i32) : i32110 %1 = llvm.mlir.constant(1 : i32) : i32111 %scs = llvm.mlir.constant(64 : i32) : i32112 113 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {114 omp.parallel {115 omp.distribute dist_schedule_static {116 omp.wsloop schedule(static = %scs : i32) {117 omp.loop_nest (%i) : i32 = (%0) to (%n) step (%1) {118 omp.yield119 }120 } {omp.composite}121 } {omp.composite}122 omp.terminator123 } {omp.composite}124 omp.terminator125 }126 llvm.return127}128// CHECK: define internal void @distribute_wsloop_schedule_chunked..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {129// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 33, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i32 1, i32 64)130// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 92, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i32 1, i32 0)131 132llvm.func @distribute_wsloop_schedule_chunked_i64(%n: i32, %teams: i32, %threads: i32) {133 %0 = llvm.mlir.constant(0 : i64) : i64134 %1 = llvm.mlir.constant(1 : i64) : i64135 %scs = llvm.mlir.constant(64 : i64) : i64136 %n64 = llvm.zext %n : i32 to i64137 138 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {139 omp.parallel {140 omp.distribute dist_schedule_static {141 omp.wsloop schedule(static = %scs : i64) {142 omp.loop_nest (%i) : i64 = (%0) to (%n64) step (%1) {143 omp.yield144 }145 } {omp.composite}146 } {omp.composite}147 omp.terminator148 } {omp.composite}149 omp.terminator150 }151 llvm.return152}153 154// CHECK: define internal void @distribute_wsloop_schedule_chunked_i64..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {155// CHECK: call void @__kmpc_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 33, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i64 1, i64 64)156// CHECK: call void @__kmpc_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 92, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.stride, i64 1, i64 0)157 158// -----159 160llvm.func @distribute_wsloop_no_chunks(%n: i32, %teams: i32, %threads: i32) {161 %0 = llvm.mlir.constant(0 : i32) : i32162 %1 = llvm.mlir.constant(1 : i32) : i32163 164 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {165 omp.parallel {166 omp.distribute dist_schedule_static {167 omp.wsloop schedule(static) {168 omp.loop_nest (%i) : i32 = (%0) to (%n) step (%1) {169 omp.yield170 }171 } {omp.composite}172 } {omp.composite}173 omp.terminator174 } {omp.composite}175 omp.terminator176 }177 llvm.return178}179// CHECK: define internal void @distribute_wsloop_no_chunks..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {180// CHECK: call void @__kmpc_dist_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 34, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.distupperbound, ptr %p.stride, i32 1, i32 0)181// CHECK: call void @__kmpc_dist_for_static_init_4u(ptr @1, i32 %omp_global_thread_num9, i32 92, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.distupperbound10, ptr %p.stride, i32 1, i32 0)182 183llvm.func @distribute_wsloop_no_chunks_i64(%n: i32, %teams: i32, %threads: i32) {184 %0 = llvm.mlir.constant(0 : i64) : i64185 %1 = llvm.mlir.constant(1 : i64) : i64186 %n64 = llvm.zext %n : i32 to i64187 188 omp.teams num_teams(to %teams : i32) thread_limit(%threads : i32) {189 omp.parallel {190 omp.distribute dist_schedule_static {191 omp.wsloop schedule(static) {192 omp.loop_nest (%i) : i64 = (%0) to (%n64) step (%1) {193 omp.yield194 }195 } {omp.composite}196 } {omp.composite}197 omp.terminator198 } {omp.composite}199 omp.terminator200 }201 llvm.return202}203// CHECK: define internal void @distribute_wsloop_no_chunks_i64..omp_par(ptr noalias %tid.addr, ptr noalias %zero.addr, ptr %0) #0 {204// CHECK: call void @__kmpc_dist_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 34, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.distupperbound, ptr %p.stride, i64 1, i64 0)205// CHECK: call void @__kmpc_dist_for_static_init_8u(ptr @1, i32 %omp_global_thread_num9, i32 92, ptr %p.lastiter, ptr %p.lowerbound, ptr %p.upperbound, ptr %p.distupperbound10, ptr %p.stride, i64 1, i64 0)