mlir/test/Dialect/Affine/SuperVectorize/vectorize_3d.mlir

   1 // RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=32,64,256 test-fastest-varying=2,1,0" | FileCheck %s
   2
   3 func.func @vec3d(%A : memref<?x?x?xf32>) {
   4   %c0 = arith.constant 0 : index
   5   %c1 = arith.constant 1 : index
   6   %c2 = arith.constant 2 : index
   7   %0 = memref.dim %A, %c0 : memref<?x?x?xf32>
   8   %1 = memref.dim %A, %c1 : memref<?x?x?xf32>
   9   %2 = memref.dim %A, %c2 : memref<?x?x?xf32>
  10   // CHECK: affine.for %{{.*}} = 0 to %{{.*}} {
  11   // CHECK:   affine.for %{{.*}} = 0 to %{{.*}} {
  12   // CHECK:     affine.for %{{.*}} = 0 to %{{.*}} step 32 {
  13   // CHECK:       affine.for %{{.*}} = 0 to %{{.*}} step 64 {
  14   // CHECK:         affine.for %{{.*}} = 0 to %{{.*}} step 256 {
  15   // CHECK:           %{{.*}} = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}], %{{.*}} : memref<?x?x?xf32>, vector<32x64x256xf32>
  16   affine.for %t0 = 0 to %0 {
  17     affine.for %t1 = 0 to %0 {
  18       affine.for %i0 = 0 to %0 {
  19         affine.for %i1 = 0 to %1 {
  20           affine.for %i2 = 0 to %2 {
  21             %a2 = affine.load %A[%i0, %i1, %i2] : memref<?x?x?xf32>
  22           }
  23         }
  24       }
  25     }
  26   }
  27   return
  28 }