mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir

   1 // RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" -split-input-file | FileCheck %s
   2
   3 // Specific tests to check vectorization of uniform/divergent values.
   4
   5 // CHECK-LABEL: @uniform_arg
   6 // CHECK-SAME:  %[[in:.*]]: memref<512xf32>,
   7 // CHECK-SAME:  %[[uniform:.*]]: f32
   8 func.func @uniform_arg(%in : memref<512xf32>, %uniform : f32) {
   9   affine.for %i = 0 to 512 {
  10     %ld = affine.load %in[%i] : memref<512xf32>
  11     %add = arith.addf %ld, %uniform : f32
  12   }
  13   return
  14 }
  15
  16 // CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform]] : f32 to vector<128xf32>
  17 // CHECK-NEXT: affine.for
  18 // CHECK:        arith.addf %{{.*}}, %[[bcast]] : vector<128xf32>
  19
  20 // -----
  21
  22 // CHECK-LABEL: @multi_use_uniform_arg
  23 // CHECK-SAME:  %[[in:.*]]: memref<512xf32>
  24 // CHECK-SAME:  %[[uniform:.*]]: f32
  25 func.func @multi_use_uniform_arg(%in : memref<512xf32>, %uniform : f32) {
  26   affine.for %i = 0 to 512 {
  27     %ld = affine.load %in[%i] : memref<512xf32>
  28     %user0 = arith.addf %ld, %uniform : f32
  29     %user1 = arith.addf %ld, %uniform : f32
  30   }
  31   return
  32 }
  33
  34 // CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform]] : f32 to vector<128xf32>
  35 // CHECK-NOT:  vector.broadcast
  36 // CHECK-NEXT: affine.for
  37 // CHECK:        arith.addf %{{.*}}, %[[bcast]] : vector<128xf32>
  38 // CHECK:        arith.addf %{{.*}}, %[[bcast]] : vector<128xf32>
  39
  40 // -----
  41
  42 // CHECK-LABEL: @uniform_load
  43 func.func @uniform_load(%A : memref<?x?xf32>, %C : memref<?x?xf32>) {
  44   %c0 = arith.constant 0 : index
  45   %N = memref.dim %A, %c0 : memref<?x?xf32>
  46   affine.for %i = 0 to %N {
  47     %uniform_ld = affine.load %A[%i, %i] : memref<?x?xf32>
  48     affine.for %j = 0 to %N {
  49       %b = affine.load %A[%i, %j] : memref<?x?xf32>
  50       %c = arith.addf %uniform_ld, %b : f32
  51     }
  52   }
  53   return
  54 }
  55
  56 // CHECK:      affine.for
  57 // CHECK-NEXT:   %[[uniform_ld:.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
  58 // CHECK-NEXT:   %[[bcast:.*]] = vector.broadcast %[[uniform_ld]] : f32 to vector<128xf32>
  59 // CHECK-NEXT:   affine.for
  60 // CHECK:          arith.addf %[[bcast]], %{{.*}} : vector<128xf32>