mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir

   1 // RUN: mlir-opt %s -affine-super-vectorizer-test="vector-shape-ratio=4,8" 2>&1 | FileCheck %s
   2 // RUN: mlir-opt %s -affine-super-vectorizer-test="vector-shape-ratio=2,5,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
   5 func.func @vector_add_2d(%arg0: index, %arg1: index) -> f32 {
   6   // Nothing should be matched in this first block.
   7   // CHECK-NOT:matched: {{.*}} = memref.alloc{{.*}}
   8   // CHECK-NOT:matched: {{.*}} = arith.constant 0{{.*}}
   9   // CHECK-NOT:matched: {{.*}} = arith.constant 1{{.*}}
  10   %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
  11   %1 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
  12   %2 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
  13   %c0 = arith.constant 0 : index
  14   %cst = arith.constant 1.000000e+00 : f32
  15
  16   // CHECK:matched: {{.*}} arith.constant dense{{.*}} with shape ratio: 2, 32
  17   %cst_1 = arith.constant dense<1.000000e+00> : vector<8x256xf32>
  18   // CHECK:matched: {{.*}} arith.constant dense{{.*}} with shape ratio: 1, 3, 7, 2, 1
  19   %cst_a = arith.constant dense<1.000000e+00> : vector<1x3x7x8x8xf32>
  20   // CHECK-NOT:matched: {{.*}} arith.constant dense{{.*}} with shape ratio: 1, 3, 7, 1{{.*}}
  21   %cst_b = arith.constant dense<1.000000e+00> : vector<1x3x7x4x4xf32>
  22   // TEST-3x4x5x8:matched: {{.*}} arith.constant dense{{.*}} with shape ratio: 3, 2, 1, 4
  23   %cst_c = arith.constant dense<1.000000e+00> : vector<3x4x5x8xf32>
  24   // TEST-3x4x5x8-NOT:matched: {{.*}} arith.constant dense{{.*}} with shape ratio{{.*}}
  25   %cst_d = arith.constant dense<1.000000e+00> : vector<3x4x4x8xf32>
  26   // TEST-3x4x5x8:matched: {{.*}} arith.constant dense{{.*}} with shape ratio: 1, 1, 2, 16
  27   %cst_e = arith.constant dense<1.000000e+00> : vector<1x2x10x32xf32>
  28
  29   // Nothing should be matched in this last block.
  30   // CHECK-NOT:matched: {{.*}} = arith.constant 7{{.*}}
  31   // CHECK-NOT:matched: {{.*}} = arith.constant 42{{.*}}
  32   // CHECK-NOT:matched: {{.*}} = memref.load{{.*}}
  33   // CHECK-NOT:matched: return {{.*}}
  34   %c7 = arith.constant 7 : index
  35   %c42 = arith.constant 42 : index
  36   %9 = memref.load %2[%c7, %c42] : memref<?x?xf32>
  37   return %9 : f32
  38 }
  39
  40 // VECNEST-LABEL: func @double_loop_nest
  41 func.func @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:       }