[mlir][gpu] Fix cuda integration tests

[Hardcode84]

Summary

https://reviews.llvm.org/D138758 has added uniform flag to gpu reduce ops, update integration tests.

[akuegel]

Thanks :)

mlir/test/Integration/GPU/CUDA/all-reduce-and.mlir

func.func @main() {
   memref.store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   memref.store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   memref.store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // AND
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xi32>
-    %reduced = gpu.all_reduce and %val {} : (i32) -> (i32)
+    %reduced = gpu.all_reduce and %val uniform {} : (i32) -> (i32)
     memref.store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
   call @printMemrefI32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
 func.func private @printMemrefI32(memref<*xi32>)
 

mlir/test/Integration/GPU/CUDA/all-reduce-max.mlir

func.func @main() {
   memref.store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   memref.store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   memref.store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // MAX
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xi32>
-    %reduced = gpu.all_reduce max %val {} : (i32) -> (i32)
+    %reduced = gpu.all_reduce max %val uniform {} : (i32) -> (i32)
     memref.store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
   call @printMemrefI32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [16, 11]
 
   return
 }
 
 func.func private @printMemrefI32(memref<*xi32>)
 

mlir/test/Integration/GPU/CUDA/all-reduce-min.mlir

func.func @main() {
   memref.store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   memref.store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   memref.store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // MIN
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xi32>
-    %reduced = gpu.all_reduce min %val {} : (i32) -> (i32)
+    %reduced = gpu.all_reduce min %val uniform {} : (i32) -> (i32)
     memref.store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
   call @printMemrefI32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
 func.func private @printMemrefI32(memref<*xi32>)
 

mlir/test/Integration/GPU/CUDA/all-reduce-op.mlir

func.func @main() {
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %sy, %block_z = %sz) {
     %t0 = arith.muli %tz, %block_y : index
     %t1 = arith.addi %ty, %t0 : index
     %t2 = arith.muli %t1, %block_x : index
     %idx = arith.addi %tx, %t2 : index
     %t3 = arith.index_cast %idx : index to i32
     %val = arith.sitofp %t3 : i32 to f32
-    %sum = gpu.all_reduce add %val {} : (f32) -> (f32)
+    %sum = gpu.all_reduce add %val uniform {} : (f32) -> (f32)
     memref.store %sum, %dst[%tz, %ty, %tx] : memref<?x?x?xf32>
     gpu.terminator
   }
   call @printMemrefF32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
 func.func private @printMemrefF32(%ptr : memref<*xf32>)

mlir/test/Integration/GPU/CUDA/all-reduce-or.mlir

func.func @main() {
   memref.store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   memref.store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   memref.store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // OR
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xi32>
-    %reduced = gpu.all_reduce or %val {} : (i32) -> (i32)
+    %reduced = gpu.all_reduce or %val uniform {} : (i32) -> (i32)
     memref.store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
   call @printMemrefI32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 15]
 
   return
 }
 
 func.func private @printMemrefI32(memref<*xi32>)
 

mlir/test/Integration/GPU/CUDA/all-reduce-region.mlir

func.func @main() {
   %one = arith.constant 1 : index
   %c0 = arith.constant 0 : index
   %sx = memref.dim %dst, %c0 : memref<?xf32>
   %cast_dst = memref.cast %dst : memref<?xf32> to memref<*xf32>
   gpu.host_register %cast_dst : memref<*xf32>
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %val = arith.index_cast %tx : index to i32
-    %xor = gpu.all_reduce %val {
+    %xor = gpu.all_reduce %val uniform {
     ^bb(%lhs : i32, %rhs : i32):
       %xor = arith.xori %lhs, %rhs : i32
       "gpu.yield"(%xor) : (i32) -> ()
     } : (i32) -> (i32)
     %res = arith.sitofp %xor : i32 to f32
     memref.store %res, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
   call @printMemrefF32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
 func.func private @printMemrefF32(memref<*xf32>)

mlir/test/Integration/GPU/CUDA/all-reduce-xor.mlir

func.func @main() {
   memref.store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   memref.store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   memref.store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // XOR
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xi32>
-    %reduced = gpu.all_reduce xor %val {} : (i32) -> (i32)
+    %reduced = gpu.all_reduce xor %val uniform {} : (i32) -> (i32)
     memref.store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
   call @printMemrefI32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 1]
 
   return
 }
 
 func.func private @printMemrefI32(memref<*xi32>)
 

mlir/test/Integration/GPU/CUDA/multiple-all-reduce.mlir

func.func @main() {
   memref.store %cst7, %data[%c1, %c3] : memref<2x6xf32>
   memref.store %cst10, %data[%c1, %c4] : memref<2x6xf32>
   memref.store %cst11, %data[%c1, %c5] : memref<2x6xf32>
 
   // ADD + MUL
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xf32>
-    %reduced0 = gpu.all_reduce add %val {} : (f32) -> (f32)
+    %reduced0 = gpu.all_reduce add %val uniform {} : (f32) -> (f32)
     memref.store %reduced0, %sum[%bx] : memref<2xf32>
-    %reduced1 = gpu.all_reduce mul %val {} : (f32) -> (f32)
+    %reduced1 = gpu.all_reduce mul %val uniform {} : (f32) -> (f32)
     memref.store %reduced1, %mul[%bx] : memref<2xf32>
     gpu.terminator
   }
 
   call @printMemrefF32(%cast_sum) : (memref<*xf32>) -> ()
   // CHECK: [31, 39]
 
   call @printMemrefF32(%cast_mul) : (memref<*xf32>) -> ()
   // CHECK: [0, 27720]
 
   return
 }
 
 func.func private @printMemrefF32(memref<*xf32>)