118 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-scf -arm-sve-legalize-vector-storage -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm | \2// RUN: %mcr_aarch64_cmd -e=entry -entry-point-result=void --march=aarch64 --mattr="+sve" -shared-libs=%native_mlir_c_runner_utils | \3// RUN: FileCheck %s4 5/// This tests basic functionality of arrays of scalable vectors, which in MLIR6/// are vectors with a single trailing scalable dimension. This test requires7/// the -arm-sve-legalize-vector-storage pass to ensure the loads/stores done8/// here are be legal for the LLVM backend.9 10func.func @read_and_print_2d_vector(%memref: memref<3x?xf32>) {11 %cst = arith.constant 0.000000e+00 : f3212 %c0 = arith.constant 0 : index13 %c1 = arith.constant 1 : index14 %c2 = arith.constant 2 : index15 %dim = memref.dim %memref, %c1 : memref<3x?xf32>16 %mask = vector.create_mask %c2, %dim : vector<3x[8]xi1>17 %vector = vector.transfer_read %memref[%c0,%c0], %cst, %mask {in_bounds = [true, true]} : memref<3x?xf32>, vector<3x[8]xf32>18 19 /// TODO: Support vector.print for arrays of scalable vectors.20 %row0 = vector.extract %vector[0] : vector<[8]xf32> from vector<3x[8]xf32>21 %row1 = vector.extract %vector[1] : vector<[8]xf32> from vector<3x[8]xf32>22 %row2 = vector.extract %vector[2] : vector<[8]xf32> from vector<3x[8]xf32>23 24 /// Print each of the vectors.25 /// vscale is >= 1, so at least 8 elements will be printed.26 27 vector.print str "read_and_print_2d_vector()\n"28 // CHECK-LABEL: read_and_print_2d_vector()29 // CHECK: ( 8, 8, 8, 8, 8, 8, 8, 830 vector.print %row0 : vector<[8]xf32>31 // CHECK: ( 8, 8, 8, 8, 8, 8, 8, 832 vector.print %row1 : vector<[8]xf32>33 /// This last row is all zero due to our mask.34 // CHECK: ( 0, 0, 0, 0, 0, 0, 0, 035 vector.print %row2 : vector<[8]xf32>36 37 return38}39 40func.func @print_1x2xVSCALExf32(%vector: vector<1x2x[4]xf32>) {41 /// TODO: Support vector.print for arrays of scalable vectors.42 %slice0 = vector.extract %vector[0, 1] : vector<[4]xf32> from vector<1x2x[4]xf32>43 %slice1 = vector.extract %vector[0, 1] : vector<[4]xf32> from vector<1x2x[4]xf32>44 vector.print %slice0 : vector<[4]xf32>45 vector.print %slice1 : vector<[4]xf32>46 return47}48 49func.func @add_arrays_of_scalable_vectors(%a: memref<1x2x?xf32>, %b: memref<1x2x?xf32>) {50 %c0 = arith.constant 0 : index51 %c2 = arith.constant 2 : index52 %c3 = arith.constant 2 : index53 %cst = arith.constant 0.000000e+00 : f3254 %dim_a = memref.dim %a, %c2 : memref<1x2x?xf32>55 %dim_b = memref.dim %b, %c2 : memref<1x2x?xf32>56 %mask_a = vector.create_mask %c2, %c3, %dim_a : vector<1x2x[4]xi1>57 %mask_b = vector.create_mask %c2, %c3, %dim_b : vector<1x2x[4]xi1>58 59 /// Print each of the vectors.60 /// vscale is >= 1, so at least 4 elements will be printed.61 62 // CHECK-LABEL: Vector A63 // CHECK-NEXT: ( 5, 5, 5, 564 // CHECK-NEXT: ( 5, 5, 5, 565 vector.print str "\nVector A\n"66 %vector_a = vector.transfer_read %a[%c0, %c0, %c0], %cst, %mask_a {in_bounds = [true, true, true]} : memref<1x2x?xf32>, vector<1x2x[4]xf32>67 func.call @print_1x2xVSCALExf32(%vector_a) : (vector<1x2x[4]xf32>) -> ()68 69 // CHECK-LABEL: Vector B70 // CHECK-NEXT: ( 4, 4, 4, 471 // CHECK-NEXT: ( 4, 4, 4, 472 vector.print str "\nVector B\n"73 %vector_b = vector.transfer_read %b[%c0, %c0, %c0], %cst, %mask_b {in_bounds = [true, true, true]} : memref<1x2x?xf32>, vector<1x2x[4]xf32>74 func.call @print_1x2xVSCALExf32(%vector_b) : (vector<1x2x[4]xf32>) -> ()75 76 // CHECK-LABEL: Sum77 // CHECK-NEXT: ( 9, 9, 9, 978 // CHECK-NEXT: ( 9, 9, 9, 979 vector.print str "\nSum\n"80 %sum = arith.addf %vector_a, %vector_b : vector<1x2x[4]xf32>81 func.call @print_1x2xVSCALExf32(%sum) : (vector<1x2x[4]xf32>) -> ()82 83 return84}85 86func.func @entry() {87 %vscale = vector.vscale88 89 %c4 = arith.constant 4 : index90 %c8 = arith.constant 8 : index91 %f32_8 = arith.constant 8.0 : f3292 %f32_5 = arith.constant 5.0 : f3293 %f32_4 = arith.constant 4.0 : f3294 95 %test_1_memref_size = arith.muli %vscale, %c8 : index96 %test_1_memref = memref.alloca(%test_1_memref_size) : memref<3x?xf32>97 98 linalg.fill ins(%f32_8 : f32) outs(%test_1_memref :memref<3x?xf32>)99 100 vector.print str "=> Print and read 2D arrays of scalable vectors:\n"101 func.call @read_and_print_2d_vector(%test_1_memref) : (memref<3x?xf32>) -> ()102 103 vector.print str "\n====================\n"104 105 %test_2_memref_size = arith.muli %vscale, %c4 : index106 %test_2_memref_a = memref.alloca(%test_2_memref_size) : memref<1x2x?xf32>107 %test_2_memref_b = memref.alloca(%test_2_memref_size) : memref<1x2x?xf32>108 109 linalg.fill ins(%f32_5 : f32) outs(%test_2_memref_a :memref<1x2x?xf32>)110 linalg.fill ins(%f32_4 : f32) outs(%test_2_memref_b :memref<1x2x?xf32>)111 112 vector.print str "=> Reading and adding two 3D arrays of scalable vectors:\n"113 func.call @add_arrays_of_scalable_vectors(114 %test_2_memref_a, %test_2_memref_b) : (memref<1x2x?xf32>, memref<1x2x?xf32>) -> ()115 116 return117}118