|
| 1 | +// RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 \ |
| 2 | +// RUN: -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm | \ |
| 3 | +// RUN: mlir-cpu-runner -e main -entry-point-result=void \ |
| 4 | +// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext | \ |
| 5 | +// RUN: FileCheck %s |
| 6 | + |
| 7 | +func @print_memref_f32(memref<*xf32>) |
| 8 | + |
| 9 | +func @alloc_1d_filled_inc_f32(%arg0: index, %arg1: f32) -> memref<?xf32> { |
| 10 | + %c0 = constant 0 : index |
| 11 | + %c1 = constant 1 : index |
| 12 | + %0 = alloc(%arg0) : memref<?xf32> |
| 13 | + scf.for %arg2 = %c0 to %arg0 step %c1 { |
| 14 | + %tmp = index_cast %arg2 : index to i32 |
| 15 | + %tmp1 = sitofp %tmp : i32 to f32 |
| 16 | + %tmp2 = addf %tmp1, %arg1 : f32 |
| 17 | + store %tmp2, %0[%arg2] : memref<?xf32> |
| 18 | + } |
| 19 | + return %0 : memref<?xf32> |
| 20 | +} |
| 21 | + |
| 22 | +func @vector_add_cycle(%id : index, %A: memref<?xf32>, %B: memref<?xf32>, %C: memref<?xf32>) { |
| 23 | + %c0 = constant 0 : index |
| 24 | + %cf0 = constant 0.0 : f32 |
| 25 | + %a = vector.transfer_read %A[%c0], %cf0: memref<?xf32>, vector<64xf32> |
| 26 | + %b = vector.transfer_read %B[%c0], %cf0: memref<?xf32>, vector<64xf32> |
| 27 | + %acc = addf %a, %b: vector<64xf32> |
| 28 | + vector.transfer_write %acc, %C[%c0]: vector<64xf32>, memref<?xf32> |
| 29 | + return |
| 30 | +} |
| 31 | + |
| 32 | +// Loop over a function containinng a large add vector and distribute it so that |
| 33 | +// each iteration of the loop process part of the vector operation. |
| 34 | +func @main() { |
| 35 | + %cf1 = constant 1.0 : f32 |
| 36 | + %cf2 = constant 2.0 : f32 |
| 37 | + %c0 = constant 0 : index |
| 38 | + %c1 = constant 1 : index |
| 39 | + %c64 = constant 64 : index |
| 40 | + %out = alloc(%c64) : memref<?xf32> |
| 41 | + %in1 = call @alloc_1d_filled_inc_f32(%c64, %cf1) : (index, f32) -> memref<?xf32> |
| 42 | + %in2 = call @alloc_1d_filled_inc_f32(%c64, %cf2) : (index, f32) -> memref<?xf32> |
| 43 | + scf.for %arg5 = %c0 to %c64 step %c1 { |
| 44 | + call @vector_add_cycle(%arg5, %in1, %in2, %out) : (index, memref<?xf32>, memref<?xf32>, memref<?xf32>) -> () |
| 45 | + } |
| 46 | + %converted = memref_cast %out : memref<?xf32> to memref<*xf32> |
| 47 | + call @print_memref_f32(%converted): (memref<*xf32>) -> () |
| 48 | + // CHECK: Unranked{{.*}}data = |
| 49 | + // CHECK: [ |
| 50 | + // CHECK-SAME: 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, |
| 51 | + // CHECK-SAME: 29, 31, 33, 35, 37, 39, 41, 43, 45, 47, 49, 51, |
| 52 | + // CHECK-SAME: 53, 55, 57, 59, 61, 63, 65, 67, 69, 71, 73, 75, |
| 53 | + // CHECK-SAME: 77, 79, 81, 83, 85, 87, 89, 91, 93, 95, 97, 99, |
| 54 | + // CHECK-SAME: 101, 103, 105, 107, 109, 111, 113, 115, 117, 119, |
| 55 | + // CHECK-SAME: 121, 123, 125, 127, 129] |
| 56 | + dealloc %out : memref<?xf32> |
| 57 | + dealloc %in1 : memref<?xf32> |
| 58 | + dealloc %in2 : memref<?xf32> |
| 59 | + return |
| 60 | +} |
0 commit comments