mlir/test/Transforms/parametric-mapping.mlir

   1 // RUN: mlir-opt -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(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)
   8 func.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   return
  19 }
  20
  21 //      CHECK: func @map2d
  22 // CHECK-SAME: (%[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
  23 func.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.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 }