| // RUN: mlir-opt --transform-interpreter --split-input-file -canonicalize -cse %s | FileCheck %s |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-LABEL: func.func @blocks_3d( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @blocks_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| %one = arith.constant 1 : index |
| // CHECK: gpu.launch |
| // CHECK: %[[BLKX:.*]] = gpu.block_id x |
| // CHECK: %[[BLKY:.*]] = gpu.block_id y |
| // CHECK: memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]]] |
| // CHECK: memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]]] |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c7, %c9) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.block<x>, #gpu.block<y>]} |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_forall_to_blocks %funcop grid_dims = [12, 9, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0) -> (d0 floordiv 128)> |
| |
| // CHECK-LABEL: func.func @warpgroup_3d( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @warpgroup_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c1 = arith.constant 1 : index |
| %c3 = arith.constant 3 : index |
| %one = arith.constant 1 : index |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index |
| // CHECK-DAG: %[[C384:.*]] = arith.constant 384 : index |
| // CHECK-DAG: %[[C512:.*]] = arith.constant 512 : index |
| |
| // CHECK: gpu.launch |
| // CHECK: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK-DAG: %[[WG:.*]] = affine.apply #[[$MAP]](%[[TIDX]]) |
| // CHECK-DAG: %[[CMPX:.*]] = arith.cmpi ult, %[[TIDX]], %[[C384]] : index |
| // CHECK-DAG: %[[CMPY:.*]] = arith.cmpi ult, %[[TIDY]], %[[C1]] : index |
| // CHECK: %[[COND:.*]] = arith.andi %[[CMPX]], %[[CMPY]] : i1 |
| // CHECK: scf.if %[[COND]] |
| // CHECK: memref.load %[[ARGX]][%[[WG]], %[[TIDY]]] |
| // CHECK: memref.load %[[ARGY]][%[[WG]], %[[TIDY]]] |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c3, %c1) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.warpgroup<x>, #gpu.warpgroup<y>]} |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [512, 2, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0) -> (d0 floordiv 16)> |
| |
| // CHECK-LABEL: func.func @warp_3d( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @warp_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c2 = arith.constant 2 : index |
| %c3 = arith.constant 3 : index |
| %one = arith.constant 1 : index |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index |
| // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index |
| // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index |
| // CHECK-DAG: %[[c64:.*]] = arith.constant 64 : index |
| |
| // CHECK: gpu.launch |
| // CHECK: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK-DAG: %[[W:.*]] = affine.apply #[[$MAP]](%[[TIDX]]) |
| // CHECK-DAG: %[[CMPX:.*]] = arith.cmpi ult, %[[TIDX]], %[[C32]] : index |
| // CHECK-DAG: %[[CMPY:.*]] = arith.cmpi ult, %[[TIDY]], %[[C3]] : index |
| // CHECK: %[[COND:.*]] = arith.andi %[[CMPX]], %[[CMPY]] : i1 |
| // CHECK: scf.if %[[COND]] |
| // CHECK: memref.load %[[ARGX]][%[[W]], %[[TIDY]]] |
| // CHECK: memref.load %[[ARGY]][%[[W]], %[[TIDY]]] |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j, %k) in (%c2, %c3, %c3) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.warp<x>, #gpu.warp<y>, #gpu.warp<z>]} |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [64, 4, 3] warp_size = 16: (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-LABEL: func.func @threads_3d( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @threads_3d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %one = arith.constant 1 : index |
| %c12 = arith.constant 12 : index |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| // CHECK: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK: %[[C12:.*]] = arith.constant 12 : index |
| // CHECK: %[[C9:.*]] = arith.constant 9 : index |
| // CHECK: %[[C7:.*]] = arith.constant 7 : index |
| // CHECK: gpu.launch async [%{{.*}}] blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C1]], %{{.*}} = %[[C1]], %{{.*}} = %[[C1]]) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C12]], %{{.*}} = %[[C9]], %{{.*}} = %[[C1]]) |
| // CHECK: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK: arith.cmpi ult, %[[TIDX]], %[[C9]] : index |
| // CHECK: arith.cmpi ult, %[[TIDY]], %[[C7]] : index |
| // CHECK: memref.load %[[ARGX]][%[[TIDY]], %[[TIDX]]] |
| // CHECK: memref.load %[[ARGY]][%[[TIDY]], %[[TIDX]]] |
| // CHECK: gpu.barrier |
| // CHECK: arith.cmpi ult, %[[TIDY]], %[[C1]] : index |
| // CHECK: memref.load %[[ARGT]][%[[TIDX]]] |
| // CHECK: gpu.barrier |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c7, %c9) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.thread<y>, #gpu.thread<x>]} |
| scf.forall (%i) in (%c12) { |
| %7 = memref.load %t[%i] : !type1d |
| %8 = arith.addf %alpha, %7 : f32 |
| memref.store %8, %t[%i] : !type1d |
| } {mapping = [#gpu.thread<x>] } |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type4d = memref<32x64x4x32xf32> |
| |
| // CHECK-LABEL: func.func @saxpy4d( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<32x64x4x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<32x64x4x32xf32> |
| func.func @saxpy4d(%x: !type4d, %y: !type4d, %alpha : f32) -> !type4d { |
| %c32 = arith.constant 32 : index |
| %c64 = arith.constant 64 : index |
| %c4 = arith.constant 4 : index |
| // CHECK: %[[C32:.*]] = arith.constant 32 : index |
| // CHECK: %[[C64:.*]] = arith.constant 64 : index |
| // CHECK: %[[C4:.*]] = arith.constant 4 : index |
| // CHECK: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C64]], %{{.*}} = %[[C1]]) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C4]], %{{.*}} = %[[C1]]) |
| // CHECK: %[[BLKX:.*]] = gpu.block_id x |
| // CHECK: %[[BLKY:.*]] = gpu.block_id y |
| // CHECK: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK: memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]] |
| // CHECK: memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]] |
| scf.forall (%i, %j) in (%c32, %c64) { |
| scf.forall (%k, %l) in (%c4, %c32) { |
| %4 = memref.load %x[%i, %j, %k, %l] : !type4d |
| %5 = memref.load %y[%i, %j, %k, %l] : !type4d |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j, %k, %l] : !type4d |
| } { mapping = [#gpu.thread<y>, #gpu.thread<x>] } |
| } { mapping = [#gpu.block<x>, #gpu.block<y>] } |
| return %y : !type4d |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| %gpuLaunch = transform.gpu.map_forall_to_blocks %funcop { generate_gpu_launch } : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %gpuLaunch block_dims = [32, 4, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-LABEL: func.func @saxpy2d_no_barrier( |
| func.func @saxpy2d_no_barrier(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %one = arith.constant 1 : index |
| %c12 = arith.constant 12 : index |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| // CHECK-NOT: gpu.barrier |
| // CHECK: return |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c7, %c9) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.thread<y>, #gpu.thread<x>] } |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] sync_after_distribute = false : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<32x32xf32> |
| // CHECK-LABEL: func.func @saxpy2d_singleloop( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<32x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<32x32xf32> |
| func.func @saxpy2d_singleloop(%x: !type, %y: !type, %stream : !gpu.async.token) -> !type { |
| %c32 = arith.constant 32 : index |
| %one = arith.constant 1 : index |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| // CHECK: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK: memref.load %[[ARGX]][%[[TIDX]], %[[TIDX]]] |
| // CHECK: memref.load %[[ARGY]][%[[TIDX]], %[[TIDX]]] |
| scf.forall (%i) in (%c32) { |
| %4 = memref.load %x[%i, %i] : !type |
| %5 = memref.load %y[%i, %i] : !type |
| %6 = arith.mulf %4, %5 : f32 |
| memref.store %6, %y[%i, %i] : !type |
| } { mapping = [#gpu.thread<x>] } |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [32, 1, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<3 x 2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-LABEL: func.func @saxpy3d_fold_id_z( |
| func.func @saxpy3d_fold_id_z(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %one = arith.constant 1 : index |
| %c12 = arith.constant 12 : index |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| // CHECK: %[[C0:.+]] = arith.constant 0 : index |
| // CHECK-NOT: gpu.thread_id z |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j, %k) in (%one, %c7, %c9) { |
| // CHECK: memref.load %{{.*}}[%[[C0]], |
| // CHECK: memref.load %{{.*}}[%[[C0]], |
| %4 = memref.load %x[%i, %j, %k] : !type |
| %5 = memref.load %y[%i, %j, %k] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| // CHECK: memref.store %{{.*}}, %{{.*}}[%[[C0]] |
| memref.store %6, %y[%i, %j, %k] : !type |
| } { mapping = [#gpu.thread<z>, #gpu.thread<y>, #gpu.thread<x>] } |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] sync_after_distribute = false : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAPWGLIN:.*]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 32 + d2 * 256)> |
| // CHECK-DAG: #[[$MAPWGX:.*]] = affine_map<(d0, d1) -> (((d0 + d1 * 32) floordiv 128) mod 2)> |
| // CHECK-DAG: #[[$MAPWGY:.*]] = affine_map<(d0, d1, d2) -> (d2 + ((d0 + d1 * 32) floordiv 128) floordiv 2)> |
| |
| // CHECK-LABEL: func.func @warpgroup_linear( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @warpgroup_linear(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c2 = arith.constant 2 : index |
| %c3 = arith.constant 3 : index |
| %one = arith.constant 1 : index |
| |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C768:.*]] = arith.constant 768 : index |
| // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index |
| // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index |
| // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index |
| |
| // CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK-DAG: %[[TIDZ:.*]] = gpu.thread_id z |
| // CHECK-DAG: %[[WIDLIN:.*]] = affine.apply #[[$MAPWGLIN]](%[[TIDX]], %[[TIDY]], %[[TIDZ]]) |
| // CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWGX]](%[[TIDX]], %[[TIDY]]) |
| // CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWGY]](%[[TIDX]], %[[TIDY]], %[[TIDZ]]) |
| // CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[WIDLIN]], %[[C768]] : index |
| // CHECK: scf.if %[[CMPLIN]] |
| // CHECK: memref.load %[[ARGX]][%[[WIDX]], %[[WIDY]]] |
| // CHECK: memref.load %[[ARGY]][%[[WIDX]], %[[WIDY]]] |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c2, %c3) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.warpgroup<linear_dim_0>, #gpu.warpgroup<linear_dim_1>]} |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [32, 8, 4] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAPWLIN:.*]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 32 + d2 * 256)> |
| // CHECK-DAG: #[[$MAPWX:.*]] = affine_map<(d0, d1, d2) -> ((d1 + d2 * 8 + d0 floordiv 32) mod 2)> |
| // CHECK-DAG: #[[$MAPWY:.*]] = affine_map<(d0, d1, d2) -> ((d1 + d2 * 8 + d0 floordiv 32) floordiv 2)> |
| |
| // CHECK-LABEL: func.func @warp_linear( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @warp_linear(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c2 = arith.constant 2 : index |
| %c3 = arith.constant 3 : index |
| %one = arith.constant 1 : index |
| |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index |
| // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index |
| // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index |
| // CHECK-DAG: %[[C192:.*]] = arith.constant 192 : index |
| |
| // CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK-DAG: %[[TIDZ:.*]] = gpu.thread_id z |
| // CHECK-DAG: %[[WIDLIN:.*]] = affine.apply #[[$MAPWLIN]](%[[TIDX]], %[[TIDY]], %[[TIDZ]]) |
| // CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWX]](%[[TIDX]], %[[TIDY]], %[[TIDZ]]) |
| // CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWY]](%[[TIDX]], %[[TIDY]], %[[TIDZ]]) |
| // CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[WIDLIN]], %[[C192]] : index |
| // CHECK: scf.if %[[CMPLIN]] |
| // CHECK: memref.load %[[ARGX]][%[[WIDX]], %[[WIDY]]] |
| // CHECK: memref.load %[[ARGY]][%[[WIDX]], %[[WIDY]]] |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c2, %c3) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.warp<linear_dim_0>, #gpu.warp<linear_dim_1>]} |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop block_dims = [32, 8, 4] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAPWX:.*]] = affine_map<(d0, d1) -> (((d0 + d1 * 18) floordiv 32) mod 3)> |
| // CHECK-DAG: #[[$MAPWY:.*]] = affine_map<(d0, d1) -> ((((d0 + d1 * 18) floordiv 32) mod 6) floordiv 3)> |
| |
| // CHECK-DAG: #[[$MAPLIN:.*]] = affine_map<(d0, d1) -> (d0 + d1 * 18)> |
| // CHECK-DAG: #[[$MAPLX:.*]] = affine_map<(d0, d1) -> ((d0 + d1 * 18) mod 10)> |
| // CHECK-DAG: #[[$MAPLY:.*]] = affine_map<(d0, d1) -> ((d0 + d1 * 18) floordiv 10)> |
| |
| // CHECK-LABEL: func.func @map_multi_level_linear( |
| func.func @map_multi_level_linear(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %one = arith.constant 1 : index |
| %c10 = arith.constant 10 : index |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c3 = arith.constant 3 : index |
| |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C11:.*]] = arith.constant 11 : index |
| // CHECK-DAG: %[[C18:.*]] = arith.constant 18 : index |
| // CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index |
| // CHECK-DAG: %[[C192:.*]] = arith.constant 192 : index |
| |
| // check that both the thread level and the warp level got distributed. |
| // CHECK-NOT: #gpu.thread |
| // CHECK-NOT: #gpu.warp |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| // CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id y |
| scf.forall (%i, %j) in (%c7, %c9) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.thread<y>, #gpu.thread<x>]} |
| |
| // CHECK-DAG: %[[LIN:.*]] = affine.apply #[[$MAPLIN]](%[[TIDX]], %[[TIDY]]) |
| // CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWX]](%[[TIDX]], %[[TIDY]]) |
| // CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWY]](%[[TIDX]], %[[TIDY]]) |
| // CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[LIN]], %[[C192]] : index |
| // CHECK: scf.if %[[CMPLIN]] |
| scf.forall (%i, %j, %k) in (%c3, %c2, %c1) { |
| %7 = memref.load %x[%i, %j] : !type |
| %8 = arith.addf %alpha, %7 : f32 |
| memref.store %8, %y[%i, %j] : !type |
| } {mapping = [#gpu.warp<linear_dim_0>, #gpu.warp<linear_dim_1>, #gpu.warp<linear_dim_2>] } |
| |
| // CHECK-DAG: %[[LIDX:.*]] = affine.apply #[[$MAPLX]](%[[TIDX]], %[[TIDY]]) |
| // CHECK-DAG: %[[LIDY:.*]] = affine.apply #[[$MAPLY]](%[[TIDX]], %[[TIDY]]) |
| // CHECK-DAG: %[[COND:.*]] = arith.cmpi ult, %[[LIN]], %[[C20]] : index |
| // CHECK: scf.if %[[COND]] |
| // CHECK: memref.load %{{.*}}[%[[LIDX]]] : memref<32xf32> |
| // CHECK: memref.store %{{.*}}[%[[LIDY]]] : memref<32xf32> |
| scf.forall (%i, %j) in (%c10, %c2) { |
| %7 = memref.load %t[%i] : !type1d |
| %8 = arith.addf %alpha, %7 : f32 |
| memref.store %8, %t[%j] : !type1d |
| } {mapping = [#gpu.thread<linear_dim_0>, #gpu.thread<linear_dim_1>] } |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %funcop |
| block_dims = [18, 11, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAPBLIN:.*]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 12 + d2 * 108)> |
| // CHECK-DAG: #[[$MAPBX:.*]] = affine_map<(d0, d1, d2) -> ((d0 + d1 * 12 + d2 * 108) mod 7)> |
| // CHECK-DAG: #[[$MAPBY:.*]] = affine_map<(d0, d1, d2) -> ((d0 + d1 * 12 + d2 * 108) floordiv 7)> |
| |
| // CHECK-LABEL: func.func @block_linear_existing_launch( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @block_linear_existing_launch( |
| %x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| %one = arith.constant 1 : index |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C9:.*]] = arith.constant 9 : index |
| // CHECK-DAG: %[[C12:.*]] = arith.constant 12 : index |
| // CHECK-DAG: %[[C63:.*]] = arith.constant 63 : index |
| // CHECK: gpu.launch async [{{.*}}] blocks({{.*}}) in (%{{.*}} = %[[C12]], %{{.*}} = %[[C9]], %{{.*}} = %[[C1]]) threads |
| // CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x |
| // CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y |
| // CHECK-DAG: %[[BIDZ:.*]] = gpu.block_id z |
| // CHECK-DAG: %[[BIDLIN:.*]] = affine.apply #[[$MAPBLIN]](%[[BIDX]], %[[BIDY]], %[[BIDZ]]) |
| // CHECK-DAG: %[[BLX:.*]] = affine.apply #[[$MAPBX]](%[[BIDX]], %[[BIDY]], %[[BIDZ]]) |
| // CHECK-DAG: %[[BLY:.*]] = affine.apply #[[$MAPBY]](%[[BIDX]], %[[BIDY]], %[[BIDZ]]) |
| // CHECK-DAG: %[[CMPLIN:.*]] = arith.cmpi ult, %[[BIDLIN]], %[[C63]] : index |
| // CHECK: scf.if %[[CMPLIN]] |
| // CHECK: memref.load %[[ARGX]][%[[BLX]], %[[BLY]]] |
| // CHECK: memref.load %[[ARGY]][%[[BLX]], %[[BLY]]] |
| %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) |
| threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) |
| { |
| scf.forall (%i, %j) in (%c7, %c9) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.block<linear_dim_0>, #gpu.block<linear_dim_1>]} |
| gpu.terminator |
| } |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_forall_to_blocks %funcop grid_dims = [12, 9, 1] : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| !type = memref<2 x 32 x f32> |
| !type1d = memref<32 x f32> |
| |
| // CHECK-DAG: #[[$MAPBX:.*]] = affine_map<(d0) -> (d0 mod 7)> |
| // CHECK-DAG: #[[$MAPBY:.*]] = affine_map<(d0, d1, d2) -> (d1 + d2 * 9 + d0 floordiv 7)> |
| |
| // CHECK-LABEL: func.func @block_linear_generate_launch( |
| // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> |
| // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> |
| func.func @block_linear_generate_launch( |
| %x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { |
| %c9 = arith.constant 9 : index |
| %c7 = arith.constant 7 : index |
| %one = arith.constant 1 : index |
| |
| // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index |
| // CHECK-DAG: %[[C9:.*]] = arith.constant 9 : index |
| // CHECK: gpu.launch blocks({{.*}}) in (%{{.*}} = %[[C7]], %{{.*}} = %[[C9]], %{{.*}} = %[[C1]]) threads |
| // CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x |
| // CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y |
| // CHECK-DAG: %[[BIDZ:.*]] = gpu.block_id z |
| // CHECK-DAG: %[[BLX:.*]] = affine.apply #[[$MAPBX]](%[[BIDX]]) |
| // CHECK-DAG: %[[BLY:.*]] = affine.apply #[[$MAPBY]](%[[BIDX]], %[[BIDY]], %[[BIDZ]]) |
| // CHECK: memref.load %[[ARGX]][%[[BLX]], %[[BLY]]] |
| // CHECK: memref.load %[[ARGY]][%[[BLX]], %[[BLY]]] |
| scf.forall (%i, %j) in (%c7, %c9) { |
| %4 = memref.load %x[%i, %j] : !type |
| %5 = memref.load %y[%i, %j] : !type |
| %6 = math.fma %alpha, %4, %5 : f32 |
| memref.store %6, %y[%i, %j] : !type |
| } { mapping = [#gpu.block<linear_dim_0>, #gpu.block<linear_dim_1>]} |
| |
| return %y : !type |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { |
| %funcop = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_forall_to_blocks %funcop generate_gpu_launch : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| // ----- |
| |
| #map = affine_map<(d0) -> (d0 * 128)> |
| #map1 = affine_map<(d0) -> (d0 * 32)> |
| |
| // CHECK-DAG: #[[$MAPB:.*]] = affine_map<(d0) -> (d0 * 128)> |
| // CHECK-DAG: #[[$MAPW:.*]] = affine_map<(d0, d1, d2) -> (d2 * 32 + ((d0 + d1 * 4) floordiv 32) * 32)> |
| |
| // CHECK-LABEL: func.func @simple_fill( |
| func.func @simple_fill(%arg0: memref<128xf32>) -> memref<128xf32> { |
| %c0 = arith.constant 0 : index |
| %cst = arith.constant dense<0.000000e+00> : vector<32xf32> |
| // CHECK: %[[C1:.*]] = arith.constant 1 : index |
| // CHECK: %[[C4:.*]] = arith.constant 4 : index |
| // CHECK: %[[C8:.*]] = arith.constant 8 : index |
| // CHECK: gpu.launch |
| scf.forall (%arg1) in (1) { |
| // CHECK: %[[BIDX:.*]] = gpu.block_id x |
| // CHECK: %[[BLX:.*]] = affine.apply #[[$MAPB]](%[[BIDX]]) |
| %0 = affine.apply #map(%arg1) |
| %subview = memref.subview %arg0[%0] [128] [1] : memref<128xf32> to memref<128xf32, strided<[1], offset: ?>> |
| scf.forall (%arg2) in (4) { |
| // CHECK: %[[TIDX:.*]] = gpu.thread_id x |
| // CHECK: %[[TIDY:.*]] = gpu.thread_id y |
| // CHECK: %[[TIDZ:.*]] = gpu.thread_id z |
| // CHECK: %[[THX:.*]] = affine.apply #[[$MAPW]](%[[TIDX]], %[[TIDY]], %[[TIDZ]]) |
| // CHECK-NOT: scf.if |
| // CHECK: memref.subview %{{.*}}[%[[THX]]] |
| %1 = affine.apply #map1(%arg2) |
| %subview_0 = memref.subview %subview[%1] [32] [1] : memref<128xf32, strided<[1], offset: ?>> to memref<32xf32, strided<[1], offset: ?>> |
| vector.transfer_write %cst, %subview_0[%c0] {in_bounds = [true]} : vector<32xf32>, memref<32xf32, strided<[1], offset: ?>> |
| memref.copy %subview_0, %subview_0 : memref<32xf32, strided<[1], offset: ?>> to memref<32xf32, strided<[1], offset: ?>> |
| } {mapping = [#gpu.warp<linear_dim_0>]} |
| memref.copy %subview, %subview : memref<128xf32, strided<[1], offset: ?>> to memref<128xf32, strided<[1], offset: ?>> |
| } {mapping = [#gpu.block<x>]} |
| return %arg0 : memref<128xf32> |
| } |
| |
| module attributes {transform.with_named_sequence} { |
| transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) { |
| %func = transform.structured.match ops{["func.func"]} in %module_op |
| : (!transform.any_op) -> !transform.any_op |
| %gpu_launch = transform.gpu.map_forall_to_blocks %func generate_gpu_launch |
| : (!transform.any_op) -> !transform.any_op |
| transform.gpu.map_nested_forall_to_threads %gpu_launch block_dims = [4, 8, 4] |
| : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
