mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir

   1 // RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=8 test-fastest-varying=0" -split-input-file | FileCheck %s
   2
   3 // CHECK-DAG: #[[$MAP_ID0:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 12)>
   4 // CHECK-DAG: #[[$MAP_ID1:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16)>
   5
   6 // CHECK-LABEL: vec_affine_apply
   7 // CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {
   8 func.func @vec_affine_apply(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
   9 // CHECK:       affine.for %[[ARG2:.*]] = 0 to 8 {
  10 // CHECK-NEXT:    affine.for %[[ARG3:.*]] = 0 to 24 {
  11 // CHECK-NEXT:      affine.for %[[ARG4:.*]] = 0 to 48 step 8 {
  12 // CHECK-NEXT:        %[[S0:.*]] = affine.apply #[[$MAP_ID0]](%[[ARG3]])
  13 // CHECK-NEXT:        %[[S1:.*]] = affine.apply #[[$MAP_ID1]](%[[ARG4]])
  14 // CHECK-NEXT:        %[[CST:.*]] = arith.constant 0.000000e+00 : f32
  15 // CHECK-NEXT:        %[[S2:.*]] = vector.transfer_read %[[ARG0]][%[[ARG2]], %[[S0]], %[[S1]]], %[[CST]] : memref<8x12x16xf32>, vector<8xf32>
  16 // CHECK-NEXT:        vector.transfer_write %[[S2]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : vector<8xf32>, memref<8x24x48xf32>
  17 // CHECK-NEXT:      }
  18 // CHECK-NEXT:    }
  19 // CHECK-NEXT:  }
  20 // CHECK-NEXT:  return
  21   affine.for %arg2 = 0 to 8 {
  22     affine.for %arg3 = 0 to 24 {
  23       affine.for %arg4 = 0 to 48 {
  24         %0 = affine.apply affine_map<(d0) -> (d0 mod 12)>(%arg3)
  25         %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
  26         %2 = affine.load %arg0[%arg2, %0, %1] : memref<8x12x16xf32>
  27         affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
  28       }
  29     }
  30   }
  31   return
  32 }
  33
  34 // -----
  35
  36 // CHECK-DAG: #[[$MAP_ID2:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16 + 1)>
  37
  38 // CHECK-LABEL: vec_affine_apply_2
  39 // CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {
  40 func.func @vec_affine_apply_2(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
  41 // CHECK:      affine.for %[[ARG2:.*]] = 0 to 8 {
  42 // CHECK-NEXT:   affine.for %[[ARG3:.*]] = 0 to 12 {
  43 // CHECK-NEXT:     affine.for %[[ARG4:.*]] = 0 to 48 step 8 {
  44 // CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID2]](%[[ARG4]])
  45 // CHECK-NEXT:       %[[CST:.*]] = arith.constant 0.000000e+00 : f32
  46 // CHECK-NEXT:       %[[S1:.*]] = vector.transfer_read %[[ARG0]][%[[ARG2]], %[[ARG3]], %[[S0]]], %[[CST]] : memref<8x12x16xf32>, vector<8xf32>
  47 // CHECK-NEXT:       vector.transfer_write %[[S1]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : vector<8xf32>, memref<8x24x48xf32>
  48 // CHECK-NEXT:     }
  49 // CHECK-NEXT:   }
  50 // CHECK-NEXT: }
  51   affine.for %arg2 = 0 to 8 {
  52     affine.for %arg3 = 0 to 12 {
  53       affine.for %arg4 = 0 to 48 {
  54         %1 = affine.apply affine_map<(d0) -> (d0 mod 16 + 1)>(%arg4)
  55         %2 = affine.load %arg0[%arg2, %arg3, %1] : memref<8x12x16xf32>
  56         affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
  57       }
  58     }
  59   }
  60   return
  61 }
  62
  63 // -----
  64
  65 // CHECK-LABEL: no_vec_affine_apply
  66 // CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xi32>, %[[ARG1:.*]]: memref<8x24x48xi32>) {
  67 func.func @no_vec_affine_apply(%arg0: memref<8x12x16xi32>, %arg1: memref<8x24x48xi32>) {
  68 // CHECK:      affine.for %[[ARG2:.*]] = 0 to 8 {
  69 // CHECK-NEXT:   affine.for %[[ARG3:.*]] = 0 to 24 {
  70 // CHECK-NEXT:     affine.for %[[ARG4:.*]] = 0 to 48 {
  71 // CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID0]](%[[ARG3]])
  72 // CHECK-NEXT:       %[[S1:.*]] = affine.apply #[[$MAP_ID1]](%[[ARG4]])
  73 // CHECK-NEXT:       %[[S2:.*]] = affine.load %[[ARG0]][%[[ARG2]], %[[S0]], %[[S1]]] : memref<8x12x16xi32>
  74 // CHECK-NEXT:       %[[S3:.*]] = arith.index_cast %[[S2]] : i32 to index
  75 // CHECK-NEXT:       %[[S4:.*]] = affine.apply #[[$MAP_ID1]](%[[S3]])
  76 // CHECK-NEXT:       %[[S5:.*]] = arith.index_cast %[[S4]] : index to i32
  77 // CHECK-NEXT:       affine.store %[[S5]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : memref<8x24x48xi32>
  78 // CHECK-NEXT:     }
  79 // CHECK-NEXT:   }
  80 // CHECK-NEXT: }
  81 // CHECK-NEXT: return
  82   affine.for %arg2 = 0 to 8 {
  83     affine.for %arg3 = 0 to 24 {
  84       affine.for %arg4 = 0 to 48 {
  85         %0 = affine.apply affine_map<(d0) -> (d0 mod 12)>(%arg3)
  86         %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
  87         %2 = affine.load %arg0[%arg2, %0, %1] : memref<8x12x16xi32>
  88         %3 = arith.index_cast %2 : i32 to index
  89         %4 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%3)
  90         %5 = arith.index_cast %4 : index to i32
  91         affine.store %5, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xi32>
  92       }
  93     }
  94   }
  95   return
  96 }
  97
  98 // -----
  99
 100 // CHECK-DAG: #[[$MAP_ID1:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16)>
 101
 102 // CHECK-LABEL: affine_map_with_expr
 103 // CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {
 104 func.func @affine_map_with_expr(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
 105 // CHECK:      affine.for %[[ARG2:.*]] = 0 to 8 {
 106 // CHECK-NEXT:   affine.for %[[ARG3:.*]] = 0 to 12 {
 107 // CHECK-NEXT:     affine.for %[[ARG4:.*]] = 0 to 48 {
 108 // CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID1]](%[[ARG4]])
 109 // CHECK-NEXT:       %[[S1:.*]] = affine.load %[[ARG0]][%[[ARG2]], %[[ARG3]], %[[S0]] + 1] : memref<8x12x16xf32>
 110 // CHECK-NEXT:       affine.store %[[S1]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : memref<8x24x48xf32>
 111 // CHECK-NEXT:     }
 112 // CHECK-NEXT:   }
 113 // CHECK-NEXT: }
 114 // CHECK-NEXT: return
 115   affine.for %arg2 = 0 to 8 {
 116     affine.for %arg3 = 0 to 12 {
 117       affine.for %arg4 = 0 to 48 {
 118         %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
 119         %2 = affine.load %arg0[%arg2, %arg3, %1 + 1] : memref<8x12x16xf32>
 120         affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
 121       }
 122     }
 123   }
 124   return
 125 }
 126
 127 // -----
 128
 129 // CHECK-DAG: #[[$MAP_ID3:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (d0)>
 130 // CHECK-DAG: #[[$MAP_ID4:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (d1)>
 131 // CHECK-DAG: #[[$MAP_ID5:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (d2 + 1)>
 132 // CHECK-DAG: #[[$MAP_ID6:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (0)>
 133
 134 // CHECK-LABEL: affine_map_with_expr_2
 135 // CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>, %[[I0:.*]]: index) {
 136 func.func @affine_map_with_expr_2(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>, %i: index) {
 137 // CHECK:      affine.for %[[ARG3:.*]] = 0 to 8 {
 138 // CHECK-NEXT:   affine.for %[[ARG4:.*]] = 0 to 12 {
 139 // CHECK-NEXT:     affine.for %[[ARG5:.*]] = 0 to 48 step 8 {
 140 // CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID3]](%[[ARG3]], %[[ARG4]], %[[I0]])
 141 // CHECK-NEXT:       %[[S1:.*]] = affine.apply #[[$MAP_ID4]](%[[ARG3]], %[[ARG4]], %[[I0]])
 142 // CHECK-NEXT:       %[[S2:.*]] = affine.apply #[[$MAP_ID5]](%[[ARG3]], %[[ARG4]], %[[I0]])
 143 // CHECK-NEXT:       %[[CST:.*]] = arith.constant 0.000000e+00 : f32
 144 // CHECK-NEXT:       %[[S3:.*]] = vector.transfer_read %[[ARG0]][%[[S0]], %[[S1]], %[[S2]]], %[[CST]] {permutation_map = #[[$MAP_ID6]]} : memref<8x12x16xf32>, vector<8xf32>
 145 // CHECK-NEXT:       vector.transfer_write %[[S3]], %[[ARG1]][%[[ARG3]], %[[ARG4]], %[[ARG5]]] : vector<8xf32>, memref<8x24x48xf32>
 146 // CHECK-NEXT:     }
 147 // CHECK-NEXT:   }
 148 // CHECK-NEXT: }
 149 // CHECK-NEXT: return
 150   affine.for %arg2 = 0 to 8 {
 151     affine.for %arg3 = 0 to 12 {
 152       affine.for %arg4 = 0 to 48 {
 153         %2 = affine.load %arg0[%arg2, %arg3, %i + 1] : memref<8x12x16xf32>
 154         affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
 155       }
 156     }
 157   }
 158   return
 159 }