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