160 lines · plain
1// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=8 test-fastest-varying=0" -split-input-file | FileCheck %s2 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_apply7// CHECK-SAME: (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {8func.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:.*]] = ub.poison : f3215// 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: return21 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 return32}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_239// CHECK-SAME: (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {40func.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:.*]] = ub.poison : f3246// 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 return61}62 63// -----64 65// CHECK-LABEL: no_vec_affine_apply66// CHECK-SAME: (%[[ARG0:.*]]: memref<8x12x16xi32>, %[[ARG1:.*]]: memref<8x24x48xi32>) {67func.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 index75// CHECK-NEXT: %[[S4:.*]] = affine.apply #[[$MAP_ID1]](%[[S3]])76// CHECK-NEXT: %[[S5:.*]] = arith.index_cast %[[S4]] : index to i3277// CHECK-NEXT: affine.store %[[S5]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : memref<8x24x48xi32>78// CHECK-NEXT: }79// CHECK-NEXT: }80// CHECK-NEXT: }81// CHECK-NEXT: return82 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 index89 %4 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%3)90 %5 = arith.index_cast %4 : index to i3291 affine.store %5, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xi32>92 }93 }94 }95 return96}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_expr103// CHECK-SAME: (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {104func.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: return115 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 return125}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_2135// CHECK-SAME: (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>, %[[I0:.*]]: index) {136func.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:.*]] = ub.poison : f32144// 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: return150 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 return159}160