mlir/test/Dialect/Affine/affine-expand-index-ops.mlir

   1 // RUN: mlir-opt %s -affine-expand-index-ops -split-input-file | FileCheck %s
   2
   3 //   CHECK-DAG:   #[[$map0:.+]] = affine_map<()[s0] -> (s0 floordiv 50176)>
   4 //   CHECK-DAG:   #[[$map1:.+]] = affine_map<()[s0] -> ((s0 mod 50176) floordiv 224)>
   5 //   CHECK-DAG:   #[[$map2:.+]] = affine_map<()[s0] -> (s0 mod 224)>
   6
   7 // CHECK-LABEL: @static_basis
   8 //  CHECK-SAME:    (%[[IDX:.+]]: index)
   9 //       CHECK:   %[[N:.+]] = affine.apply #[[$map0]]()[%[[IDX]]]
  10 //       CHECK:   %[[P:.+]] = affine.apply #[[$map1]]()[%[[IDX]]]
  11 //       CHECK:   %[[Q:.+]] = affine.apply #[[$map2]]()[%[[IDX]]]
  12 //       CHECK:   return %[[N]], %[[P]], %[[Q]]
  13 func.func @static_basis(%linear_index: index) -> (index, index, index) {
  14   %b0 = arith.constant 16 : index
  15   %b1 = arith.constant 224 : index
  16   %b2 = arith.constant 224 : index
  17   %1:3 = affine.delinearize_index %linear_index into (%b0, %b1, %b2) : index, index, index
  18   return %1#0, %1#1, %1#2 : index, index, index
  19 }
  20
  21 // -----
  22
  23 //   CHECK-DAG:   #[[$map0:.+]] = affine_map<()[s0, s1, s2] -> (s2 floordiv (s0 * s1))>
  24 //   CHECK-DAG:   #[[$map1:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) floordiv s1)>
  25 //   CHECK-DAG:   #[[$map2:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) mod s1)>
  26
  27 // CHECK-LABEL: @dynamic_basis
  28 //  CHECK-SAME:    (%[[IDX:.+]]: index, %[[MEMREF:.+]]: memref
  29 //   CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
  30 //   CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
  31 //        CHECK:  %[[DIM1:.+]] = memref.dim %[[MEMREF]], %[[C1]] :
  32 //        CHECK:  %[[DIM2:.+]] = memref.dim %[[MEMREF]], %[[C2]] :
  33 //       CHECK:   %[[N:.+]] = affine.apply #[[$map0]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
  34 //       CHECK:   %[[P:.+]] = affine.apply #[[$map1]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
  35 //       CHECK:   %[[Q:.+]] = affine.apply #[[$map2]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
  36 //       CHECK:   return %[[N]], %[[P]], %[[Q]]
  37 func.func @dynamic_basis(%linear_index: index, %src: memref<?x?x?xf32>) -> (index, index, index) {
  38   %c0 = arith.constant 0 : index
  39   %c1 = arith.constant 1 : index
  40   %c2 = arith.constant 2 : index
  41   %b0 = memref.dim %src, %c0 : memref<?x?x?xf32>
  42   %b1 = memref.dim %src, %c1 : memref<?x?x?xf32>
  43   %b2 = memref.dim %src, %c2 : memref<?x?x?xf32>
  44   %1:3 = affine.delinearize_index %linear_index into (%b0, %b1, %b2) : index, index, index
  45   return %1#0, %1#1, %1#2 : index, index, index
  46 }