1// RUN: mlir-opt -allow-unregistered-dialect -test-mapping-to-processing-elements %s | FileCheck %s
2
3// CHECK: #[[mul_map:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
4// CHECK: #[[add_map:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
5
6//      CHECK: func @map1d
7// CHECK-SAME: (%[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
8func @map1d(%lb: index, %ub: index, %step: index) {
9  // CHECK: %[[threads:.*]]:2 = "new_processor_id_and_range"() : () -> (index, index)
10  %0:2 = "new_processor_id_and_range"() : () -> (index, index)
11
12  // CHECK: %[[thread_offset:.+]] = affine.apply #[[mul_map]]()[%[[threads]]#0, %[[step]]]
13  // CHECK: %[[new_lb:.+]] = affine.apply #[[add_map]]()[%[[thread_offset]], %[[lb]]]
14  // CHECK: %[[new_step:.+]] = affine.apply #[[mul_map]]()[%[[threads]]#1, %[[step]]]
15
16  // CHECK: scf.for %{{.*}} = %[[new_lb]] to %[[ub]] step %[[new_step]] {
17  scf.for %i = %lb to %ub step %step {}
18  return
19}
20
21//      CHECK: func @map2d
22// CHECK-SAME: (%[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
23func @map2d(%lb : index, %ub : index, %step : index) {
24  // CHECK: %[[blocks:.*]]:2 = "new_processor_id_and_range"() : () -> (index, index)
25  %0:2 = "new_processor_id_and_range"() : () -> (index, index)
26
27  // CHECK: %[[threads:.*]]:2 = "new_processor_id_and_range"() : () -> (index, index)
28  %1:2 = "new_processor_id_and_range"() : () -> (index, index)
29
30  // blockIdx.x * blockDim.x
31  // CHECK: %[[bidxXbdimx:.+]] = affine.apply #[[mul_map]]()[%[[blocks]]#0, %[[threads]]#1]
32  //
33  // threadIdx.x + blockIdx.x * blockDim.x
34  // CHECK: %[[tidxpbidxXbdimx:.+]] = affine.apply #[[add_map]]()[%[[bidxXbdimx]], %[[threads]]#0]
35  //
36  // thread_offset = step * (threadIdx.x + blockIdx.x * blockDim.x)
37  // CHECK: %[[thread_offset:.+]] = affine.apply #[[mul_map]]()[%[[tidxpbidxXbdimx]], %[[step]]]
38  //
39  // new_lb = lb + thread_offset
40  // CHECK: %[[new_lb:.+]] = affine.apply #[[add_map]]()[%[[thread_offset]], %[[lb]]]
41  //
42  // stepXgdimx = step * gridDim.x
43  // CHECK: %[[stepXgdimx:.+]] = affine.apply #[[mul_map]]()[%[[blocks]]#1, %[[step]]]
44  //
45  // new_step = step * gridDim.x * blockDim.x
46  // CHECK: %[[new_step:.+]] = affine.apply #[[mul_map]]()[%[[threads]]#1, %[[stepXgdimx]]]
47  //
48  // CHECK: scf.for %{{.*}} = %[[new_lb]] to %[[ub]] step %[[new_step]] {
49
50  scf.for %i = %lb to %ub step %step {}
51  return
52}
53