| // RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" | FileCheck %s |
| |
| // A per-thread array accumulator (memref.alloca) for a block+thread reduction is |
| // reduced element-by-element across the parallel dimensions with gpu.all_reduce - |
| // the array analog of the scalar acc.reduction_accumulate. |
| |
| // CHECK-LABEL: func.func @array_reduction |
| // CHECK: gpu.launch |
| // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32> |
| // CHECK-NOT: acc.reduction_accumulate_array |
| // CHECK-NOT: acc.bounds |
| // CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { |
| // CHECK: %[[ELT:.*]] = memref.load %[[ALLOCA]][%[[IV]]] : memref<2xi32> |
| // CHECK: %[[RED:.*]] = gpu.all_reduce add %[[ELT]] |
| // Per-thread alloca: the all_reduce result is stored unpredicated. |
| // CHECK: memref.store %[[RED]], %[[ALLOCA]][%[[IV]]] : memref<2xi32> |
| // CHECK: } |
| |
| func.func @array_reduction(%arg0: memref<2xi32>) { |
| %0 = acc.copyin varPtr(%arg0 : memref<2xi32>) dataClause(acc_reduction) implicit(true) name("r") -> memref<2xi32> |
| acc.kernel_environment dataOperands(%0 : memref<2xi32>) { |
| %c1_pw = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1_pw par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg2 = %0) : (memref<2xi32>) { |
| %c2 = arith.constant 2 : index |
| %c0_i32 = arith.constant 0 : i32 |
| %c1_i32 = arith.constant 1 : i32 |
| %c1 = arith.constant 1 : index |
| %c0 = arith.constant 0 : index |
| %2 = acc.reduction_init %arg2 <add> : memref<2xi32> { |
| %alloca = memref.alloca() : memref<2xi32> |
| scf.parallel (%i) = (%c0) to (%c2) step (%c1) { |
| memref.store %c0_i32, %alloca[%i] : memref<2xi32> |
| scf.reduce |
| } {acc.par_dims = #acc<par_dims[thread_x]>} |
| acc.yield %alloca : memref<2xi32> |
| } |
| scf.parallel (%bx_iv) = (%c0) to (%kbx) step (%c1) { |
| scf.parallel (%tx_iv) = (%c0) to (%ktx) step (%c1) { |
| %3 = memref.load %2[%c0] : memref<2xi32> |
| %4 = arith.addi %3, %c1_i32 : i32 |
| memref.store %4, %2[%c0] : memref<2xi32> |
| scf.reduce |
| } {acc.par_dims = #acc<par_dims[thread_x]>} |
| scf.reduce |
| } {acc.par_dims = #acc<par_dims[block_x]>} |
| %b = acc.bounds extent(%c2 : index) |
| acc.reduction_accumulate_array %2 bounds(%b) <add> par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2xi32> |
| acc.reduction_combine_region %2 into %arg2 : memref<2xi32> { |
| scf.for %i = %c0 to %c2 step %c1 { |
| %3 = memref.load %2[%i] : memref<2xi32> |
| %4 = memref.load %arg2[%i] : memref<2xi32> |
| %5 = arith.addi %3, %4 : i32 |
| memref.store %5, %arg2[%i] : memref<2xi32> |
| } |
| } |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| } |
| acc.copyout accPtr(%0 : memref<2xi32>) to varPtr(%arg0 : memref<2xi32>) dataClause(acc_reduction) implicit(true) name("r") |
| return |
| } |
| |
| // CHECK-LABEL: func.func @array_reduction_small_shared |
| // CHECK: memref.alloc() : memref<2xi32> |
| // CHECK-NOT: gpu.all_reduce |
| // CHECK-NOT: acc.reduction_accumulate_array |
| func.func @array_reduction_small_shared() { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) { |
| %c2 = arith.constant 2 : index |
| %shared = memref.alloc() : memref<2xi32> |
| %bounds = acc.bounds extent(%c2 : index) |
| acc.reduction_accumulate_array %shared bounds(%bounds) <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2xi32> |
| memref.dealloc %shared : memref<2xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // CHECK-LABEL: func.func @array_reduction_strided_extent |
| // CHECK: gpu.launch |
| // CHECK: %[[LB:.*]] = arith.constant 1 : index |
| // CHECK: %[[STEP:.*]] = arith.constant 2 : index |
| // CHECK: %[[EXTENT:.*]] = arith.constant 3 : index |
| // CHECK: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[STEP]] : index |
| // CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SPAN]] : index |
| // CHECK: scf.for %{{.*}} = %[[LB]] to %[[UB]] step %[[STEP]] |
| func.func @array_reduction_strided_extent() { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) { |
| %c1_b = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c3 = arith.constant 3 : index |
| %local = memref.alloca() : memref<8xi32> |
| %bounds = acc.bounds lowerbound(%c1_b : index) extent(%c3 : index) |
| stride(%c2 : index) |
| acc.reduction_accumulate_array %local bounds(%bounds) <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : memref<8xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // A dynamically-shaped accumulator (a strided view whose type conveys no size) |
| // is classified per-thread from par_dims: a thread dimension means per-thread |
| // storage, so lowering emits the per-element gpu.all_reduce. |
| // |
| // CHECK-LABEL: func.func @array_reduction_dynamic_par_dims |
| // CHECK: scf.for |
| // CHECK: gpu.all_reduce add |
| func.func @array_reduction_dynamic_par_dims(%buf: memref<?xi32>, %n: index) { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg0 = %buf, %ext = %n) : (memref<?xi32>, index) { |
| %view = memref.reinterpret_cast %arg0 to offset: [0], sizes: [%ext], strides: [1] |
| : memref<?xi32> to memref<?xi32, strided<[1]>> |
| %bounds = acc.bounds extent(%ext : index) |
| acc.reduction_accumulate_array %view bounds(%bounds) <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : memref<?xi32, strided<[1]>> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // CHECK-LABEL: func.func @rank_two_array_reduction |
| // CHECK: %[[R2_ALLOCA:.*]] = memref.alloca() : memref<2x3xi32> |
| // CHECK: scf.for %[[R2_I:.*]] = |
| // CHECK: scf.for %[[R2_J:.*]] = |
| // CHECK: memref.store %{{.*}}, %[[R2_ALLOCA]][%[[R2_I]], %[[R2_J]]] : memref<2x3xi32> |
| // CHECK: scf.for %[[R2_LINEAR:.*]] = |
| // CHECK: %[[R2_J_IDX:.*]] = arith.remui %[[R2_LINEAR]], %{{.*}} : index |
| // CHECK: %[[R2_ROW:.*]] = arith.divui %[[R2_LINEAR]], %{{.*}} : index |
| // CHECK: %[[R2_I_IDX:.*]] = arith.remui %[[R2_ROW]], %{{.*}} : index |
| // CHECK: memref.load %[[R2_ALLOCA]][%[[R2_I_IDX]], %[[R2_J_IDX]]] : memref<2x3xi32> |
| // CHECK: %[[R2_RESULT:.*]] = gpu.all_reduce add |
| // CHECK: memref.store %[[R2_RESULT]], %[[R2_ALLOCA]][%[[R2_I_IDX]], %[[R2_J_IDX]]] |
| |
| func.func @rank_two_array_reduction() { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| %private = acc.privatize par_dims(#acc<par_dims[block_x, thread_x]>) : () -> !acc.private_type<memref<2x3xi32>> |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg0 = %private) : (!acc.private_type<memref<2x3xi32>>) { |
| %c6 = arith.constant 6 : index |
| %local = acc.private_local %arg0 {acc.par_dims = #acc<par_dims[block_x, thread_x]>} : (!acc.private_type<memref<2x3xi32>>) -> memref<2x3xi32> |
| %bounds = acc.bounds extent(%c6 : index) |
| acc.reduction_accumulate_array %local bounds(%bounds) <add> par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2x3xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // CHECK-LABEL: func.func @rank_three_array_reduction |
| // CHECK: %[[R3_ALLOCA:.*]] = memref.alloca() : memref<2x2x2xi32> |
| // CHECK: scf.for %[[R3_I:.*]] = |
| // CHECK: scf.for %[[R3_J:.*]] = |
| // CHECK: scf.for %[[R3_K:.*]] = |
| // CHECK: memref.store %{{.*}}, %[[R3_ALLOCA]][%[[R3_I]], %[[R3_J]], %[[R3_K]]] : memref<2x2x2xi32> |
| // CHECK: scf.for %[[R3_LINEAR:.*]] = |
| // CHECK: %[[R3_K_IDX:.*]] = arith.remui %[[R3_LINEAR]], %{{.*}} : index |
| // CHECK: %[[R3_PLANE:.*]] = arith.divui %[[R3_LINEAR]], %{{.*}} : index |
| // CHECK: %[[R3_J_IDX:.*]] = arith.remui %[[R3_PLANE]], %{{.*}} : index |
| // CHECK: %[[R3_ROW:.*]] = arith.divui %[[R3_PLANE]], %{{.*}} : index |
| // CHECK: %[[R3_I_IDX:.*]] = arith.remui %[[R3_ROW]], %{{.*}} : index |
| // CHECK: memref.load %[[R3_ALLOCA]][%[[R3_I_IDX]], %[[R3_J_IDX]], %[[R3_K_IDX]]] : memref<2x2x2xi32> |
| // CHECK: %[[R3_RESULT:.*]] = gpu.all_reduce add |
| // CHECK: memref.store %[[R3_RESULT]], %[[R3_ALLOCA]][%[[R3_I_IDX]], %[[R3_J_IDX]], %[[R3_K_IDX]]] |
| |
| func.func @rank_three_array_reduction() { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| %private = acc.privatize par_dims(#acc<par_dims[block_x, thread_x]>) : () -> !acc.private_type<memref<2x2x2xi32>> |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg0 = %private) : (!acc.private_type<memref<2x2x2xi32>>) { |
| %c8 = arith.constant 8 : index |
| %local = acc.private_local %arg0 {acc.par_dims = #acc<par_dims[block_x, thread_x]>} : (!acc.private_type<memref<2x2x2xi32>>) -> memref<2x2x2xi32> |
| %bounds = acc.bounds extent(%c8 : index) |
| acc.reduction_accumulate_array %local bounds(%bounds) <add> par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2x2x2xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // CHECK-LABEL: func.func @dynamic_rank_two_array_reduction |
| // CHECK: scf.for %[[LINEAR:.*]] = |
| // CHECK: %[[DIM1:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<2x?xi32> |
| // CHECK: %[[IDX1:.*]] = arith.remui %[[LINEAR]], %[[DIM1]] : index |
| // CHECK: %[[ROW:.*]] = arith.divui %[[LINEAR]], %[[DIM1]] : index |
| // CHECK: %[[C2:.*]] = arith.constant 2 : index |
| // CHECK: %[[IDX0:.*]] = arith.remui %[[ROW]], %[[C2]] : index |
| // CHECK: %[[ELEM:.*]] = memref.load %{{.*}}[%[[IDX0]], %[[IDX1]]] : memref<2x?xi32> |
| // CHECK: %[[RESULT:.*]] = gpu.all_reduce add %[[ELEM]] |
| // CHECK: memref.store %[[RESULT]], %{{.*}}[%[[IDX0]], %[[IDX1]]] : memref<2x?xi32> |
| func.func @dynamic_rank_two_array_reduction( |
| %local: memref<2x?xi32>, %n: index) { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) |
| ins(%arg0 = %local, %arg1 = %n) |
| : (memref<2x?xi32>, index) { |
| %c2 = arith.constant 2 : index |
| %extent = arith.muli %c2, %arg1 : index |
| %bounds = acc.bounds extent(%extent : index) |
| acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2x?xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // CHECK-LABEL: func.func @rank_two_partial_bounds_strided_layout |
| // CHECK: %[[LB:.*]] = arith.constant 5 : index |
| // CHECK: %[[STEP:.*]] = arith.constant 2 : index |
| // CHECK: %[[EXTENT:.*]] = arith.constant 3 : index |
| // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<3x4xi32, strided<[8, 2]>> |
| // CHECK: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[STEP]] : index |
| // CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SPAN]] : index |
| // CHECK: scf.for %[[LINEAR:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] |
| // CHECK: %[[COL:.*]] = arith.remui %[[LINEAR]], %{{.*}} : index |
| // CHECK: %[[ROW_LINEAR:.*]] = arith.divui %[[LINEAR]], %{{.*}} : index |
| // CHECK: %[[ROW:.*]] = arith.remui %[[ROW_LINEAR]], %{{.*}} : index |
| // CHECK: memref.load %[[ALLOCA]][%[[ROW]], %[[COL]]] : memref<3x4xi32, strided<[8, 2]>> |
| func.func @rank_two_partial_bounds_strided_layout() { |
| %c1 = arith.constant 1 : index |
| %c128 = arith.constant 128 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) { |
| %lb = arith.constant 5 : index |
| %step = arith.constant 2 : index |
| %extent = arith.constant 3 : index |
| %local = memref.alloca() : memref<3x4xi32, strided<[8, 2]>> |
| %bounds = acc.bounds lowerbound(%lb : index) extent(%extent : index) |
| stride(%step : index) |
| acc.reduction_accumulate_array %local bounds(%bounds) <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : memref<3x4xi32, strided<[8, 2]>> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // CHECK-LABEL: func.func @partial_thread_x_reduction |
| // CHECK: %[[C16:.*]] = arith.constant 16 : index |
| // CHECK-NOT: arith.constant 31 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C16]], |
| // CHECK: gpu.all_reduce add |
| func.func @partial_thread_x_reduction() { |
| %c1 = arith.constant 1 : index |
| %c16 = arith.constant 16 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c16 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx) { |
| %c2 = arith.constant 2 : index |
| %local = memref.alloca() : memref<2xi32> |
| %bounds = acc.bounds extent(%c2 : index) |
| acc.reduction_accumulate_array %local bounds(%bounds) <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // A worker-only launch keeps its (1, N, 1) shape: every worker owns one thread, |
| // so the workers are already the subgroup lanes. Padding ThreadX would fold the |
| // workers into lanes and serialize them. |
| // |
| // CHECK-LABEL: func.func @thread_y_reduction_single_thread_rows |
| // CHECK: %[[C16_ROWS:.*]] = arith.constant 16 : index |
| // CHECK-NOT: arith.constant 32 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %[[C16_ROWS]], |
| func.func @thread_y_reduction_single_thread_rows() { |
| %c1 = arith.constant 1 : index |
| %c16 = arith.constant 16 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %ty = acc.par_width %c16 par_dim(#acc.par_dim<thread_y>) |
| acc.compute_region launch(%kbx = %bx, %kty = %ty) { |
| %c0_i32 = arith.constant 0 : i32 |
| %local = memref.alloca() : memref<i32> |
| acc.reduction_accumulate %c0_i32 to %local <add> |
| par_dims(#acc<par_dims[block_x, thread_y]>) : i32 -> memref<i32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // A worker-only launch with several threads per worker keeps its shape too: |
| // the partials are combined in the lowest ThreadY threads of the block. |
| // |
| // CHECK-LABEL: func.func @thread_y_reduction_narrow_rows |
| // CHECK: %[[C8_NARROW:.*]] = arith.constant 8 : index |
| // CHECK: %[[C16_NARROW:.*]] = arith.constant 16 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C8_NARROW]], %{{.*}} = %[[C16_NARROW]], |
| func.func @thread_y_reduction_narrow_rows() { |
| %c1 = arith.constant 1 : index |
| %c8 = arith.constant 8 : index |
| %c16 = arith.constant 16 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c8 par_dim(#acc.par_dim<thread_x>) |
| %ty = acc.par_width %c16 par_dim(#acc.par_dim<thread_y>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx, %kty = %ty) { |
| %c0_i32 = arith.constant 0 : i32 |
| %local = memref.alloca() : memref<i32> |
| acc.reduction_accumulate %c0_i32 to %local <add> |
| par_dims(#acc<par_dims[block_x, thread_y]>) : i32 -> memref<i32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // A ThreadX reduction in the same region shuffles within a row, so the rows |
| // are aligned again and blockDim.y is divided by the same factor. |
| // |
| // CHECK-LABEL: func.func @thread_y_reduction_with_thread_x_reduction |
| // CHECK: %[[C32_BOTH:.*]] = arith.constant 32 : index |
| // CHECK: %[[C4_BOTH:.*]] = arith.constant 4 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32_BOTH]], %{{.*}} = %[[C4_BOTH]], |
| func.func @thread_y_reduction_with_thread_x_reduction() { |
| %c1 = arith.constant 1 : index |
| %c8 = arith.constant 8 : index |
| %c16 = arith.constant 16 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c8 par_dim(#acc.par_dim<thread_x>) |
| %ty = acc.par_width %c16 par_dim(#acc.par_dim<thread_y>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx, %kty = %ty) { |
| %c0_i32 = arith.constant 0 : i32 |
| %worker = memref.alloca() : memref<i32> |
| %vector = memref.alloca() : memref<i32> |
| acc.reduction_accumulate %c0_i32 to %worker <add> |
| par_dims(#acc<par_dims[block_x, thread_y]>) : i32 -> memref<i32> |
| acc.reduction_accumulate %c0_i32 to %vector <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : i32 -> memref<i32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // More workers than a subgroup still get aligned: the worker-indexed shared |
| // reduction buffer only holds subgroupSize entries. |
| // |
| // CHECK-LABEL: func.func @thread_y_reduction_more_workers_than_subgroup |
| // CHECK: %[[C32_WIDE:.*]] = arith.constant 32 : index |
| // CHECK: %[[C2_WIDE:.*]] = arith.constant 2 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32_WIDE]], %{{.*}} = %[[C2_WIDE]], |
| func.func @thread_y_reduction_more_workers_than_subgroup() { |
| %c1 = arith.constant 1 : index |
| %c64 = arith.constant 64 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %ty = acc.par_width %c64 par_dim(#acc.par_dim<thread_y>) |
| acc.compute_region launch(%kbx = %bx, %kty = %ty) { |
| %c0_i32 = arith.constant 0 : i32 |
| %local = memref.alloca() : memref<i32> |
| acc.reduction_accumulate %c0_i32 to %local <add> |
| par_dims(#acc<par_dims[block_x, thread_y]>) : i32 -> memref<i32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // A ThreadX-only reduction still needs aligned rows when ThreadY is greater |
| // than one, because a physical subgroup must not contain multiple logical rows. |
| // |
| // CHECK-LABEL: func.func @thread_x_reduction_with_thread_y_width |
| // CHECK: %[[C32_ROWS:.*]] = arith.constant 32 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32_ROWS]], |
| func.func @thread_x_reduction_with_thread_y_width() { |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c16 = arith.constant 16 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c16 par_dim(#acc.par_dim<thread_x>) |
| %ty = acc.par_width %c2 par_dim(#acc.par_dim<thread_y>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx, %kty = %ty) { |
| %c0_i32 = arith.constant 0 : i32 |
| %local = memref.alloca() : memref<i32> |
| acc.reduction_accumulate %c0_i32 to %local <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : i32 -> memref<i32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // ThreadZ rows have the same subgroup-packing constraint as ThreadY rows. |
| // |
| // CHECK-LABEL: func.func @thread_x_reduction_with_thread_z_width |
| // CHECK: %[[C352_Z_ROWS:.*]] = arith.constant 352 : index |
| // CHECK: %[[C2_Z_ROWS:.*]] = arith.constant 2 : index |
| // CHECK: gpu.launch |
| // CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C352_Z_ROWS]], %{{.*}} = %{{.*}}, %{{.*}} = %[[C2_Z_ROWS]]) |
| func.func @thread_x_reduction_with_thread_z_width() { |
| %c1 = arith.constant 1 : index |
| %c3 = arith.constant 3 : index |
| %c341 = arith.constant 341 : index |
| %bx = acc.par_width %c1 par_dim(#acc.par_dim<block_x>) |
| %tx = acc.par_width %c341 par_dim(#acc.par_dim<thread_x>) |
| %tz = acc.par_width %c3 par_dim(#acc.par_dim<thread_z>) |
| acc.compute_region launch(%kbx = %bx, %ktx = %tx, %ktz = %tz) { |
| %c0_i32 = arith.constant 0 : i32 |
| %local = memref.alloca() : memref<i32> |
| acc.reduction_accumulate %c0_i32 to %local <add> |
| par_dims(#acc<par_dims[block_x, thread_x]>) : i32 -> memref<i32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
| |
| // Thread-only array accumulate is well-defined when the region launches no |
| // block dim (single block): a within-block all_reduce is a complete reduction. |
| // CHECK-LABEL: func.func @thread_only_array_reduction_single_block |
| // CHECK: gpu.launch |
| // CHECK-NOT: acc.reduction_accumulate_array |
| // CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { |
| // CHECK: %[[ELT:.*]] = memref.load %[[ALLOCA:.*]][%[[IV]]] : memref<8xi32> |
| // CHECK: %[[RED:.*]] = gpu.all_reduce add %[[ELT]] |
| // CHECK: memref.store %[[RED]], %[[ALLOCA]][%[[IV]]] : memref<8xi32> |
| // CHECK: } |
| func.func @thread_only_array_reduction_single_block() { |
| %c128 = arith.constant 128 : index |
| %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) |
| acc.compute_region launch(%ktx = %tx) { |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c8 = arith.constant 8 : index |
| %c0_i32 = arith.constant 0 : i32 |
| %local = memref.alloca() : memref<8xi32> |
| scf.for %i = %c0 to %c8 step %c1 { |
| memref.store %c0_i32, %local[%i] : memref<8xi32> |
| } |
| %bounds = acc.bounds extent(%c8 : index) |
| acc.reduction_accumulate_array %local bounds(%bounds) <add> |
| par_dims(#acc<par_dims[thread_x]>) : memref<8xi32> |
| acc.yield |
| } <{origin = "acc.parallel"}> |
| return |
| } |
