208 lines · plain
1// DEFINE: %{entry_point} = entry2// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm3// DEFINE: %{run} = %mcr_aarch64_cmd \4// DEFINE: -march=aarch64 -mattr=+sve,+sme \5// DEFINE: -e %{entry_point} -entry-point-result=void \6// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib7 8// RUN: %{compile} | %{run} | FileCheck %s9 10// 2-D vector load (SME tile).11func.func @transfer_read_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {12 %c4 = arith.constant 4 : index13 %pad = arith.constant 0.0 : f3214 %0 = vector.transfer_read %A[%base1, %base2], %pad {in_bounds=[true, true]} :15 memref<?x?xf32>, vector<[4]x[4]xf32>16 17 vector.print str "TILE BEGIN:\n"18 vector.print %0: vector<[4]x[4]xf32>19 20 return21}22 23// 2-D vector load (SME tile) + transpose.24func.func @transfer_read_2d_transposed(%A : memref<?x?xf32>, %base1: index, %base2: index) {25 %pad = arith.constant 0.0 : f3226 %0 = vector.transfer_read %A[%base1, %base2], %pad27 {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds=[true, true]}28 : memref<?x?xf32>, vector<[4]x[4]xf32>29 30 vector.print str "TILE BEGIN:\n"31 vector.print %0 : vector<[4]x[4]xf32>32 33 return34}35 36// 2-D vector load (SME tile) with mask and pad of zero.37func.func @transfer_read_2d_mask(%A : memref<?x?xf32>, %base1: index, %base2: index) {38 %c2 = arith.constant 2 : index39 %c3 = arith.constant 3 : index40 %pad = arith.constant 0.0 : f3241 %mask = vector.create_mask %c2, %c3 : vector<[4]x[4]xi1>42 %0 = vector.transfer_read %A[%base1, %base2], %pad, %mask43 {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>44 45 vector.print str "TILE BEGIN:\n"46 vector.print %0: vector<[4]x[4]xf32>47 48 return49}50 51// 2-D vector load (SME tile) with mask and pad of zero + transpose.52func.func @transfer_read_2d_mask_transposed(%A : memref<?x?xf32>, %base1: index, %base2: index) {53 %c2 = arith.constant 2 : index54 %c3 = arith.constant 3 : index55 %pad = arith.constant 0.0 : f3256 %mask = vector.create_mask %c2, %c3 : vector<[4]x[4]xi1>57 %0 = vector.transfer_read %A[%base1, %base2], %pad, %mask58 {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds=[true, true]}59 : memref<?x?xf32>, vector<[4]x[4]xf32>60 61 vector.print str "TILE BEGIN:\n"62 vector.print %0: vector<[4]x[4]xf32>63 64 return65}66 67// 2-D vector load (SME tile) with mask and non-zero pad.68func.func @transfer_read_2d_mask_non_zero_pad(%A : memref<?x?xf32>, %base1: index, %base2: index) {69 %c2 = arith.constant 2 : index70 %c3 = arith.constant 3 : index71 %pad = arith.constant -42.0 : f3272 %mask = vector.create_mask %c2, %c3 : vector<[4]x[4]xi1>73 %0 = vector.transfer_read %A[%base1, %base2], %pad, %mask74 {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>75 76 vector.print str "TILE BEGIN:\n"77 vector.print %0: vector<[4]x[4]xf32>78 79 return80}81 82// 2-D vector load (SME tile) with mask and non-zero pad + transpose.83func.func @transfer_read_2d_mask_non_zero_pad_transposed(%A : memref<?x?xf32>, %base1: index, %base2: index) {84 %c2 = arith.constant 2 : index85 %c3 = arith.constant 3 : index86 %pad = arith.constant -42.0 : f3287 %mask = vector.create_mask %c2, %c3 : vector<[4]x[4]xi1>88 %0 = vector.transfer_read %A[%base1, %base2], %pad, %mask89 {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds=[true, true]}90 : memref<?x?xf32>, vector<[4]x[4]xf32>91 92 vector.print str "TILE BEGIN:\n"93 vector.print %0: vector<[4]x[4]xf32>94 95 return96}97 98// Allocate heap memory of size 'd0' x 'd1' and initialize.99//100// Example:101//102// initialize_memory(%c4, %c5)103//104// 0, 1, 2, 3, 4105// 10, 11, 12, 13, 14106// 20, 21, 22, 23, 24107// 30, 31, 32, 33, 34108//109// Returns dynamic memref. It's the callers responsiblity to free the returned110// memref.111func.func @initialize_memory(%d0 : index, %d1 : index) -> memref<?x?xf32> {112 %c0 = arith.constant 0 : index113 %c1 = arith.constant 1 : index114 %c1_f32 = arith.constant 1.0 : f32115 %c10_f32 = arith.constant 10.0 : f32116 117 %A = memref.alloc(%d0, %d1) : memref<?x?xf32>118 119 %init = arith.constant 0.0 : f32120 scf.for %i = %c0 to %d0 step %c1 iter_args(%val = %init) -> f32 {121 scf.for %j = %c0 to %d1 step %c1 iter_args(%inner_val = %val) -> f32 {122 memref.store %inner_val, %A[%i, %j] : memref<?x?xf32>123 %inner_val_next = arith.addf %inner_val, %c1_f32 : f32124 scf.yield %inner_val_next : f32125 }126 %val_next = arith.addf %val, %c10_f32 : f32127 scf.yield %val_next : f32128 }129 130 return %A : memref<?x?xf32>131}132 133func.func @entry() {134 %c0 = arith.constant 0 : index135 %c1 = arith.constant 1 : index136 %c2 = arith.constant 2 : index137 %c4 = arith.constant 4 : index138 139 // Allocate enough memory to load a 32-bit tile plus a tiny bit more to test140 // non-zero offsets while remaining inbounds.141 %svl_s = arm_sme.streaming_vl <word>142 %svl_s_plus_two = arith.addi %svl_s, %c2 : index143 144 %A = call @initialize_memory(%svl_s_plus_two, %svl_s_plus_two) : (index, index) -> memref<?x?xf32>145 146 // 1.a. Read 2D vector from 2D memref.147 //148 // CHECK-LABEL: TILE BEGIN:149 // CHECK-NEXT: ( 0, 1, 2, 3150 // CHECK-NEXT: ( 10, 11, 12, 13151 // CHECK-NEXT: ( 20, 21, 22, 23152 // CHECK-NEXT: ( 30, 31, 32, 33153 call @transfer_read_2d(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()154 155 // 1.b. Same as 1.a., but with non-zero offsets.156 //157 // CHECK-LABEL: TILE BEGIN:158 // CHECK-NEXT: ( 12, 13, 14, 15159 // CHECK-NEXT: ( 22, 23, 24, 25160 // CHECK-NEXT: ( 32, 33, 34, 35161 // CHECK-NEXT: ( 42, 43, 44, 45162 call @transfer_read_2d(%A, %c1, %c2) : (memref<?x?xf32>, index, index) -> ()163 164 // 2. Same as 1.a., but with mask and a pad of constant zero.165 // CHECK-LABEL: TILE BEGIN:166 // CHECK-NEXT: ( 0, 1, 2, 0167 // CHECK-NEXT: ( 10, 11, 12, 0168 // CHECK-NEXT: ( 0, 0, 0, 0169 // CHECK-NEXT: ( 0, 0, 0, 0170 call @transfer_read_2d_mask(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()171 172 // 3. Same as 1.a., but with mask and non-zero pad.173 // CHECK-LABEL: TILE BEGIN:174 // CHECK-NEXT: ( 0, 1, 2, -42175 // CHECK-NEXT: ( 10, 11, 12, -42176 // CHECK-NEXT: ( -42, -42, -42, -42177 // CHECK-NEXT: ( -42, -42, -42, -42178 call @transfer_read_2d_mask_non_zero_pad(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()179 180 // 4. Same as 1.a., but transpose the result.181 // CHECK-LABEL: TILE BEGIN:182 // CHECK-NEXT: ( 0, 10, 20, 30183 // CHECK-NEXT: ( 1, 11, 21, 31184 // CHECK-NEXT: ( 2, 12, 22, 32185 // CHECK-NEXT: ( 3, 13, 23, 33186 call @transfer_read_2d_transposed(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()187 188 // 5. Same as 2., but transpose the result.189 // CHECK-LABEL: TILE BEGIN:190 // CHECK-NEXT: ( 0, 10, 0, 0191 // CHECK-NEXT: ( 1, 11, 0, 0192 // CHECK-NEXT: ( 2, 12, 0, 0193 // CHECK-NEXT: ( 0, 0, 0, 0194 call @transfer_read_2d_mask_transposed(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()195 196 // 5. Same as 3, but transpose the result.197 // CHECK-LABEL: TILE BEGIN:198 // CHECK-NEXT: ( 0, 10, -42, -42199 // CHECK-NEXT: ( 1, 11, -42, -42200 // CHECK-NEXT: ( 2, 12, -42, -42201 // CHECK-NEXT: ( -42, -42, -42, -42202 call @transfer_read_2d_mask_non_zero_pad_transposed(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()203 204 memref.dealloc %A : memref<?x?xf32>205 206 return207}208