1// RUN: mlir-opt %s -affine-super-vectorizer-test -vector-shape-ratio 4 -vector-shape-ratio 8 2>&1 | FileCheck %s
2// RUN: mlir-opt %s -affine-super-vectorizer-test -vector-shape-ratio 2 -vector-shape-ratio 5 -vector-shape-ratio 2 2>&1 | FileCheck %s -check-prefix=TEST-3x4x5x8
3// RUN: mlir-opt %s -affine-super-vectorizer-test -vectorize-affine-loop-nest 2>&1 | FileCheck %s -check-prefix=VECNEST
4
5func @vector_add_2d(%arg0: index, %arg1: index) -> f32 {
6  // Nothing should be matched in this first block.
7  // CHECK-NOT:matched: {{.*}} = alloc{{.*}}
8  // CHECK-NOT:matched: {{.*}} = constant 0{{.*}}
9  // CHECK-NOT:matched: {{.*}} = constant 1{{.*}}
10  %0 = alloc(%arg0, %arg1) : memref<?x?xf32>
11  %1 = alloc(%arg0, %arg1) : memref<?x?xf32>
12  %2 = alloc(%arg0, %arg1) : memref<?x?xf32>
13  %c0 = constant 0 : index
14  %cst = constant 1.000000e+00 : f32
15
16  // CHECK:matched: {{.*}} constant dense{{.*}} with shape ratio: 2, 32
17  %cst_1 = constant dense<1.000000e+00> : vector<8x256xf32>
18  // CHECK:matched: {{.*}} constant dense{{.*}} with shape ratio: 1, 3, 7, 2, 1
19  %cst_a = constant dense<1.000000e+00> : vector<1x3x7x8x8xf32>
20  // CHECK-NOT:matched: {{.*}} constant dense{{.*}} with shape ratio: 1, 3, 7, 1{{.*}}
21  %cst_b = constant dense<1.000000e+00> : vector<1x3x7x4x4xf32>
22  // TEST-3x4x5x8:matched: {{.*}} constant dense{{.*}} with shape ratio: 3, 2, 1, 4
23  %cst_c = constant dense<1.000000e+00> : vector<3x4x5x8xf32>
24  // TEST-3x4x4x8-NOT:matched: {{.*}} constant dense{{.*}} with shape ratio{{.*}}
25  %cst_d = constant dense<1.000000e+00> : vector<3x4x4x8xf32>
26  // TEST-3x4x4x8:matched: {{.*}} constant dense{{.*}} with shape ratio: 1, 1, 2, 16
27  %cst_e = constant dense<1.000000e+00> : vector<1x2x10x32xf32>
28
29  // Nothing should be matched in this last block.
30  // CHECK-NOT:matched: {{.*}} = constant 7{{.*}}
31  // CHECK-NOT:matched: {{.*}} = constant 42{{.*}}
32  // CHECK-NOT:matched: {{.*}} = load{{.*}}
33  // CHECK-NOT:matched: return {{.*}}
34  %c7 = constant 7 : index
35  %c42 = constant 42 : index
36  %9 = load %2[%c7, %c42] : memref<?x?xf32>
37  return %9 : f32
38}
39
40// VECNEST-LABEL: func @double_loop_nest
41func @double_loop_nest(%a: memref<20x30xf32>, %b: memref<20xf32>) {
42
43  affine.for %i = 0 to 20 {
44    %b_ld = affine.load %b[%i] : memref<20xf32>
45    affine.for %j = 0 to 30 {
46      %a_ld = affine.load %a[%i, %j] : memref<20x30xf32>
47      affine.store %a_ld, %a[%i, %j] : memref<20x30xf32>
48    }
49    affine.store %b_ld, %b[%i] : memref<20xf32>
50  }
51
52  return
53}
54
55// VECNEST:       affine.for %{{.*}} = 0 to 20 step 4 {
56// VECNEST:         vector.transfer_read
57// VECNEST-NEXT:    affine.for %{{.*}} = 0 to 30 {
58// VECNEST:           vector.transfer_read
59// VECNEST-NEXT:      vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] {permutation_map = #{{.*}}}
60// VECNEST-NEXT:    }
61// VECNEST-NEXT:    vector.transfer_write
62// VECNEST:       }
63