1// RUN: mlir-opt --gpu-kernel-outlining --convert-gpu-to-nvvm %s | FileCheck %s
2
3func @main() {
4  %data = alloc() : memref<2x6xf32>
5  %sum = alloc() : memref<2xf32>
6  %mul = alloc() : memref<2xf32>
7  %c1 = constant 1 : index
8
9  // ADD + MUL
10  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
11             threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1, %block_z = %c1) {
12    %val = load %data[%bx, %tx] : memref<2x6xf32>
13    %reduced0 = "gpu.all_reduce"(%val) ({}) { op = "add" } : (f32) -> (f32)
14    store %reduced0, %sum[%bx] : memref<2xf32>
15    %reduced1 = "gpu.all_reduce"(%val) ({}) { op = "mul" } : (f32) -> (f32)
16    store %reduced1, %mul[%bx] : memref<2xf32>
17    gpu.terminator
18  }
19
20// CHECK:      gpu.module @main_kernel {
21// CHECK-NEXT:   llvm.mlir.global internal @{{.*}}() {addr_space = 3 : i32} : !llvm.array<32 x float>
22// CHECK-NEXT:   llvm.mlir.global internal @{{.*}}() {addr_space = 3 : i32} : !llvm.array<32 x float>
23
24  return
25}
26