54 lines · plain
1// Tests multiple kernels running concurrently. Runs two kernels, which2// increment a global atomic counter and wait for the counter to reach 2.3//4// RUN: mlir-opt %s \5// RUN: | mlir-opt -gpu-lower-to-nvvm-pipeline="cubin-format=%gpu_compilation_format allow-pattern-rollback=0" \6// RUN: | env CUDA_MODULE_LOADING=EAGER mlir-runner \7// RUN: --shared-libs=%mlir_cuda_runtime \8// RUN: --shared-libs=%mlir_runner_utils \9// RUN: --entry-point-result=void10 11// CUDA_MODULE_LOADING=EAGER avoids an implicit context synchronization on first12// use of each kernel. It is technically not needed for this test, because13// there is only one kernel.14 15module attributes {gpu.container_module} {16 17gpu.module @kernels {18 gpu.func @kernel(%memref: memref<i32>) kernel {19 %c0 = arith.constant 0 : i3220 %c1 = arith.constant 1 : i3221 %c2 = arith.constant 2 : i3222 %block = memref.atomic_rmw addi %c1, %memref[] : (i32, memref<i32>) -> i3223 scf.while: () -> () {24 %value = memref.atomic_rmw addi %c0, %memref[] : (i32, memref<i32>) -> i3225 %cond = arith.cmpi slt, %value, %c2 : i3226 scf.condition(%cond)27 } do {28 scf.yield29 }30 gpu.return31 }32}33 34func.func @main() {35 %c0 = arith.constant 0 : i3236 %c1 = arith.constant 1 : index37 %memref = gpu.alloc host_shared () : memref<i32>38 memref.store %c0, %memref[] : memref<i32>39 %0 = gpu.wait async40 %1 = gpu.wait async41 %2 = gpu.launch_func async [%0] @kernels::@kernel42 blocks in (%c1, %c1, %c1)43 threads in (%c1, %c1, %c1)44 args(%memref: memref<i32>)45 %3 = gpu.launch_func async [%1] @kernels::@kernel46 blocks in (%c1, %c1, %c1)47 threads in (%c1, %c1, %c1)48 args(%memref: memref<i32>)49 gpu.wait [%2, %3]50 return51}52 53}54