mlir/test/Conversion/VectorToSCF/unrolled-vector-to-loops.mlir

   1 // RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file -allow-unregistered-dialect | FileCheck %s
   2
   3 // CHECK-LABEL: func @transfer_read_inbounds
   4 func @transfer_read_inbounds(%A : memref<?x?x?xf32>) -> (vector<2x3x4xf32>) {
   5   %f0 = arith.constant 0.0: f32
   6   %c0 = arith.constant 0: index
   7
   8   // CHECK:      vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
   9   // CHECK-NEXT: vector.insert {{.*}} [0, 0] : vector<4xf32> into vector<2x3x4xf32>
  10   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  11   // CHECK-NEXT: vector.insert {{.*}} [0, 1] : vector<4xf32> into vector<2x3x4xf32>
  12   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  13   // CHECK-NEXT: vector.insert {{.*}} [0, 2] : vector<4xf32> into vector<2x3x4xf32>
  14   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  15   // CHECK-NEXT: vector.insert {{.*}} [1, 0] : vector<4xf32> into vector<2x3x4xf32>
  16   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  17   // CHECK-NEXT: vector.insert {{.*}} [1, 1] : vector<4xf32> into vector<2x3x4xf32>
  18   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  19   // CHECK-NEXT: vector.insert {{.*}} [1, 2] : vector<4xf32> into vector<2x3x4xf32>
  20   // CHECK-NOT: scf.if
  21   // CHECK-NOT: scf.for
  22   %vec = vector.transfer_read %A[%c0, %c0, %c0], %f0 {in_bounds = [true, true, true]} : memref<?x?x?xf32>, vector<2x3x4xf32>
  23   return %vec : vector<2x3x4xf32>
  24 }
  25
  26 // -----
  27
  28 // CHECK-LABEL: func @transfer_read_out_of_bounds
  29 func @transfer_read_out_of_bounds(%A : memref<?x?x?xf32>) -> (vector<2x3x4xf32>) {
  30   %f0 = arith.constant 0.0: f32
  31   %c0 = arith.constant 0: index
  32
  33   // CHECK: scf.if
  34   // CHECK: scf.if
  35   // CHECK: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  36   // CHECK: vector.insert {{.*}} [0, 0] : vector<4xf32> into vector<2x3x4xf32>
  37   // CHECK: scf.if
  38   // CHECK: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  39   // CHECK: vector.insert {{.*}} [0, 1] : vector<4xf32> into vector<2x3x4xf32>
  40   // CHECK: scf.if
  41   // CHECK: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  42   // CHECK: vector.insert {{.*}} [0, 2] : vector<4xf32> into vector<2x3x4xf32>
  43   // CHECK: scf.if
  44   // CHECK: scf.if
  45   // CHECK: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  46   // CHECK: vector.insert {{.*}} [1, 0] : vector<4xf32> into vector<2x3x4xf32>
  47   // CHECK: scf.if
  48   // CHECK: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  49   // CHECK: vector.insert {{.*}} [1, 1] : vector<4xf32> into vector<2x3x4xf32>
  50   // CHECK: scf.if
  51   // CHECK: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  52   // CHECK: vector.insert {{.*}} [1, 2] : vector<4xf32> into vector<2x3x4xf32>
  53   // CHECK-NOT: scf.for
  54   %vec = vector.transfer_read %A[%c0, %c0, %c0], %f0 : memref<?x?x?xf32>, vector<2x3x4xf32>
  55   return %vec : vector<2x3x4xf32>
  56 }
  57
  58 // -----
  59
  60 func @transfer_read_mask(%A : memref<?x?x?xf32>, %mask : vector<2x3x4xi1>) -> (vector<2x3x4xf32>) {
  61   %f0 = arith.constant 0.0: f32
  62   %c0 = arith.constant 0: index
  63
  64   // CHECK:      vector.extract %{{.*}}[0, 0] : vector<2x3x4xi1>
  65   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  66   // CHECK-NEXT: vector.insert {{.*}} [0, 0] : vector<4xf32> into vector<2x3x4xf32>
  67   // CHECK-NEXT: vector.extract %{{.*}}[0, 1] : vector<2x3x4xi1>
  68   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  69   // CHECK-NEXT: vector.insert {{.*}} [0, 1] : vector<4xf32> into vector<2x3x4xf32>
  70   // CHECK-NEXT: vector.extract %{{.*}}[0, 2] : vector<2x3x4xi1>
  71   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  72   // CHECK-NEXT: vector.insert {{.*}} [0, 2] : vector<4xf32> into vector<2x3x4xf32>
  73   // CHECK-NEXT: vector.extract %{{.*}}[1, 0] : vector<2x3x4xi1>
  74   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  75   // CHECK-NEXT: vector.insert {{.*}} [1, 0] : vector<4xf32> into vector<2x3x4xf32>
  76   // CHECK-NEXT: vector.extract %{{.*}}[1, 1] : vector<2x3x4xi1>
  77   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  78   // CHECK-NEXT: vector.insert {{.*}} [1, 1] : vector<4xf32> into vector<2x3x4xf32>
  79   // CHECK-NEXT: vector.extract %{{.*}}[1, 2] : vector<2x3x4xi1>
  80   // CHECK-NEXT: vector.transfer_read {{.*}} : memref<?x?x?xf32>, vector<4xf32>
  81   // CHECK-NEXT: vector.insert {{.*}} [1, 2] : vector<4xf32> into vector<2x3x4xf32>
  82   // CHECK-NOT: scf.if
  83   // CHECK-NOT: scf.for
  84   %vec = vector.transfer_read %A[%c0, %c0, %c0], %f0, %mask {in_bounds = [true, true, true]}: memref<?x?x?xf32>, vector<2x3x4xf32>
  85   return %vec : vector<2x3x4xf32>
  86 }