53 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(test-mapping-to-processing-elements))" %s | FileCheck %s2 3// CHECK: #[[mul_map:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>4// CHECK: #[[add_map:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>5 6// CHECK: func @map1d7// CHECK-SAME: (%[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)8func.func @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 return19}20 21// CHECK: func @map2d22// CHECK-SAME: (%[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)23func.func @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.x31 // CHECK: %[[bidxXbdimx:.+]] = affine.apply #[[mul_map]]()[%[[blocks]]#0, %[[threads]]#1]32 //33 // threadIdx.x + blockIdx.x * blockDim.x34 // 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_offset40 // CHECK: %[[new_lb:.+]] = affine.apply #[[add_map]]()[%[[thread_offset]], %[[lb]]]41 //42 // stepXgdimx = step * gridDim.x43 // CHECK: %[[stepXgdimx:.+]] = affine.apply #[[mul_map]]()[%[[blocks]]#1, %[[step]]]44 //45 // new_step = step * gridDim.x * blockDim.x46 // 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 return52}53