164 lines · plain
1// RUN: mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm | \2// RUN: %mcr_aarch64_cmd \3// RUN: -e=main -entry-point-result=void \4// RUN: -march=aarch64 -mattr="+sve,+sme" \5// RUN: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib,%native_mlir_arm_runner_utils | \6// RUN: FileCheck %s7 8#transpose = affine_map<(d0, d1) -> (d1, d0)>9 10func.func @fill2DMemrefRows(%memref: memref<?x?xf32>) {11 %c0 = arith.constant 0 : index12 %c1 = arith.constant 1 : index13 %rows = memref.dim %memref, %c0 : memref<?x?xf32>14 %cols = memref.dim %memref, %c1 : memref<?x?xf32>15 scf.for %i = %c0 to %rows step %c1 {16 scf.for %j = %c0 to %cols step %c1 {17 %n = arith.addi %i, %c1 : index18 %val = arith.index_cast %n : index to i3219 %val_f32 = arith.sitofp %val : i32 to f3220 memref.store %val_f32, %memref[%i, %j] : memref<?x?xf32>21 }22 }23 return24}25 26func.func @testTransposedReadWithMask(%maskRows: index, %maskCols: index) {27 %in = memref.alloca() : memref<4x16xf32>28 %out = memref.alloca() : memref<16x4xf32>29 30 %inDyn = memref.cast %in : memref<4x16xf32> to memref<?x?xf32>31 %outDyn = memref.cast %out : memref<16x4xf32> to memref<?x?xf32>32 33 func.call @fill2DMemrefRows(%inDyn) : (memref<?x?xf32>) -> ()34 35 /// A mask so we only read the first maskRows x maskCols portion of %in.36 %mask = vector.create_mask %maskRows, %maskCols : vector<[4]x[16]xi1>37 %pad = arith.constant 0.0 : f3238 %c0 = arith.constant 0 : index39 40 /// A vector.transfer_read with a transpose permutation map. So the input data41 /// (and mask) have a [4]x[16] shape, but the output is [16]x[4].42 %readTransposed = vector.transfer_read %inDyn[%c0, %c0], %pad, %mask43 {permutation_map = #transpose, in_bounds = [true, true]} : memref<?x?xf32>, vector<[16]x[4]xf32>44 45 /// Write the vector back to a memref (that also has a transposed shape).46 vector.transfer_write %readTransposed, %outDyn[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[4]xf32>, memref<?x?xf32>47 48 /// Print the input memref.49 vector.print str "Input memref:\n"50 %inUnranked = memref.cast %inDyn : memref<?x?xf32> to memref<*xf32>51 call @printMemrefF32(%inUnranked) : (memref<*xf32>) -> ()52 53 /// Print the result memref.54 vector.print str "Masked transposed result:\n"55 %outUnranked = memref.cast %outDyn : memref<?x?xf32> to memref<*xf32>56 call @printMemrefF32(%outUnranked) : (memref<*xf32>) -> ()57 58 return59}60 61func.func @testTransposedWriteWithMask(%maskRows: index, %maskCols: index) {62 %in = memref.alloca() : memref<16x4xf32>63 %out = memref.alloca() : memref<4x16xf32>64 65 %c0_f32 = arith.constant 0.0 : f3266 linalg.fill ins(%c0_f32 : f32) outs(%out : memref<4x16xf32>)67 68 %inDyn = memref.cast %in : memref<16x4xf32> to memref<?x?xf32>69 %outDyn = memref.cast %out : memref<4x16xf32> to memref<?x?xf32>70 71 func.call @fill2DMemrefRows(%inDyn) : (memref<?x?xf32>) -> ()72 73 /// A regular read.74 %c0 = arith.constant 0 : index75 %read = vector.transfer_read %inDyn[%c0, %c0], %c0_f32 {in_bounds = [true, true]}76 : memref<?x?xf32>, vector<[16]x[4]xf32>77 78 /// A mask so we only write the first maskRows x maskCols portion of transpose(%in).79 %mask = vector.create_mask %maskRows, %maskCols : vector<[4]x[16]xi1>80 81 /// Write out the data with a transpose. Here (like the read test) the mask82 /// matches the shape of the memory, not the vector.83 vector.transfer_write %read, %outDyn[%c0, %c0], %mask {permutation_map = #transpose, in_bounds = [true, true]}84 : vector<[16]x[4]xf32>, memref<?x?xf32>85 86 /// Print the input memref.87 vector.print str "Input memref:\n"88 %inUnranked = memref.cast %inDyn : memref<?x?xf32> to memref<*xf32>89 call @printMemrefF32(%inUnranked) : (memref<*xf32>) -> ()90 91 /// Print the result memref.92 vector.print str "Masked transposed result:\n"93 %outUnranked = memref.cast %outDyn : memref<?x?xf32> to memref<*xf32>94 call @printMemrefF32(%outUnranked) : (memref<*xf32>) -> ()95 96 return97}98 99func.func @main() {100 /// Set the SVL to 128-bits (i.e. vscale = 1).101 /// This test is for the use of multiple tiles rather than scalability.102 %c128 = arith.constant 128 : i32103 func.call @setArmSVLBits(%c128) : (i32) -> ()104 105 // CHECK: Input memref:106 // CHECK: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]107 // CHECK-NEXT: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]108 // CHECK-NEXT: [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]109 // CHECK-NEXT: [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4]110 //111 // CHECK: Masked transposed result:112 // CHECK: [1, 2, 0, 0]113 // CHECK-NEXT: [1, 2, 0, 0]114 // CHECK-NEXT: [1, 2, 0, 0]115 // CHECK-NEXT: [1, 2, 0, 0]116 // CHECK-NEXT: [1, 2, 0, 0]117 // CHECK-NEXT: [1, 2, 0, 0]118 // CHECK-NEXT: [1, 2, 0, 0]119 // CHECK-NEXT: [1, 2, 0, 0]120 // CHECK-NEXT: [1, 2, 0, 0]121 // CHECK-NEXT: [1, 2, 0, 0]122 // CHECK-NEXT: [1, 2, 0, 0]123 // CHECK-NEXT: [1, 2, 0, 0]124 // CHECK-NEXT: [1, 2, 0, 0]125 // CHECK-NEXT: [1, 2, 0, 0]126 // CHECK-NEXT: [1, 2, 0, 0]127 // CHECK-NEXT: [0, 0, 0, 0]128 %readMaskRows = arith.constant 2 : index129 %readMaskCols = arith.constant 15 : index130 func.call @testTransposedReadWithMask(%readMaskRows, %readMaskCols) : (index, index) -> ()131 132 // CHECK: Input memref:133 // CHECK: [1, 1, 1, 1]134 // CHECK-NEXT: [2, 2, 2, 2]135 // CHECK-NEXT: [3, 3, 3, 3]136 // CHECK-NEXT: [4, 4, 4, 4]137 // CHECK-NEXT: [5, 5, 5, 5]138 // CHECK-NEXT: [6, 6, 6, 6]139 // CHECK-NEXT: [7, 7, 7, 7]140 // CHECK-NEXT: [8, 8, 8, 8]141 // CHECK-NEXT: [9, 9, 9, 9]142 // CHECK-NEXT: [10, 10, 10, 10]143 // CHECK-NEXT: [11, 11, 11, 11]144 // CHECK-NEXT: [12, 12, 12, 12]145 // CHECK-NEXT: [13, 13, 13, 13]146 // CHECK-NEXT: [14, 14, 14, 14]147 // CHECK-NEXT: [15, 15, 15, 15]148 // CHECK-NEXT: [16, 16, 16, 16]149 //150 // CHECK: Masked transposed result:151 // CHECK: [1, 2, 3, 4, 5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0]152 // CHECK-NEXT: [1, 2, 3, 4, 5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0]153 // CHECK-NEXT: [1, 2, 3, 4, 5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0]154 // CHECK-NEXT: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]155 %writeMaskRows = arith.constant 3 : index156 %writeMaskCols = arith.constant 8 : index157 func.call @testTransposedWriteWithMask(%writeMaskRows, %writeMaskCols) : (index, index) -> ()158 159 return160}161 162func.func private @printMemrefF32(%ptr : memref<*xf32>)163func.func private @setArmSVLBits(%bits : i32)164