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