157 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// Vector store.11func.func @transfer_write_2d(%A : memref<?x?xf32>, %base1: index, %base2: index) {12 %c0 = arith.constant 0.0 : f3213 %zero = vector.broadcast %c0 : f32 to vector<[4]x[4]xf32>14 vector.transfer_write %zero, %A[%base1, %base2] {in_bounds=[true, true]} :15 vector<[4]x[4]xf32>, memref<?x?xf32>16 return17}18 19// Masked vector store.20func.func @transfer_write_2d_mask(%A : memref<?x?xf32>, %base1: index, %base2: index) {21 %c0 = arith.constant 0.0 : f3222 %c2 = arith.constant 2 : index23 %c3 = arith.constant 3 : index24 %mask = vector.create_mask %c2, %c3 : vector<[4]x[4]xi1>25 %zero = vector.broadcast %c0 : f32 to vector<[4]x[4]xf32>26 vector.transfer_write %zero, %A[%base1, %base2], %mask {in_bounds=[true, true]} :27 vector<[4]x[4]xf32>, memref<?x?xf32>28 return29}30 31// Vector transpose + store.32func.func @transfer_write_2d_transposed(%A : memref<?x?xf32>, %base1: index, %base2: index) {33 %0 = vector.load %A[%base1, %base2] : memref<?x?xf32>, vector<[4]x[4]xf32>34 vector.transfer_write %0, %A[%base1, %base2] {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds=[true, true]} :35 vector<[4]x[4]xf32>, memref<?x?xf32>36 return37}38 39// Vector transpose + masked store.40func.func @transfer_write_2d_mask_transposed(%A : memref<?x?xf32>, %base1: index, %base2: index) {41 %c2 = arith.constant 2 : index42 %c4 = arith.constant 4 : index43 %mask = vector.create_mask %c4, %c2 : vector<[4]x[4]xi1>44 %0 = vector.load %A[%base1, %base2] : memref<?x?xf32>, vector<[4]x[4]xf32>45 vector.transfer_write %0, %A[%base1, %base2], %mask {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds=[true, true]} :46 vector<[4]x[4]xf32>, memref<?x?xf32>47 return48}49 50// Vector load + print.51func.func @load_and_print(%A : memref<?x?xf32>, %base1: index, %base2: index) {52 %0 = vector.load %A[%base1, %base2] : memref<?x?xf32>, vector<[4]x[4]xf32>53 54 vector.print str "TILE BEGIN:\n"55 vector.print %0: vector<[4]x[4]xf32>56 57 return58}59 60// Allocate heap memory of size 'd0' x 'd1' and initialize.61//62// Example:63//64// initialize_memory(%c4, %c5)65//66// 0, 1, 2, 3, 467// 10, 11, 12, 13, 1468// 20, 21, 22, 23, 2469// 30, 31, 32, 33, 3470//71// Returns dynamic memref. It's the callers responsiblity to free the returned72// memref.73func.func @initialize_memory(%d0 : index, %d1 : index) -> memref<?x?xf32> {74 %c0 = arith.constant 0 : index75 %c1 = arith.constant 1 : index76 %c1_f32 = arith.constant 1.0 : f3277 %c10_f32 = arith.constant 10.0 : f3278 79 %A = memref.alloc(%d0, %d1) : memref<?x?xf32>80 81 %init = arith.constant 0.0 : f3282 scf.for %i = %c0 to %d0 step %c1 iter_args(%val = %init) -> f32 {83 scf.for %j = %c0 to %d1 step %c1 iter_args(%inner_val = %val) -> f32 {84 memref.store %inner_val, %A[%i, %j] : memref<?x?xf32>85 %inner_val_next = arith.addf %inner_val, %c1_f32 : f3286 scf.yield %inner_val_next : f3287 }88 %val_next = arith.addf %val, %c10_f32 : f3289 scf.yield %val_next : f3290 }91 92 return %A : memref<?x?xf32>93}94 95func.func @entry() {96 %c0 = arith.constant 0 : index97 %c2 = arith.constant 2 : index98 %c4 = arith.constant 4 : index99 100 // 1. Initialize memory101 //102 // Allocate enough memory to load a 32-bit tile plus a tiny bit more to test103 // non-zero offsets while remaining inbounds.104 %svl_s = arm_sme.streaming_vl <word>105 %svl_s_plus_two = arith.addi %svl_s, %c2 : index106 %A = call @initialize_memory(%svl_s_plus_two, %svl_s_plus_two) : (index, index) -> memref<?x?xf32>107 108 // CHECK-LABEL: TILE BEGIN:109 // CHECK-NEXT: ( 0, 1, 2, 3110 // CHECK-NEXT: ( 10, 11, 12, 13111 // CHECK-NEXT: ( 20, 21, 22, 23112 // CHECK-NEXT: ( 30, 31, 32, 33113 call @load_and_print(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()114 115 // 2. Write 2-D vector of zeroes to 1. at offset [2, 2].116 // CHECK-LABEL: TILE BEGIN:117 // CHECK-NEXT: ( 0, 1, 2, 3118 // CHECK-NEXT: ( 10, 11, 12, 13119 // CHECK-NEXT: ( 20, 21, 0, 0120 // CHECK-NEXT: ( 30, 31, 0, 0121 call @transfer_write_2d(%A, %c2, %c2) : (memref<?x?xf32>, index, index) -> ()122 call @load_and_print(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()123 124 // 3. Write 2-D vector of zeroes to 2. but with mask (nrows=2, ncols=3).125 // CHECK-LABEL: TILE BEGIN:126 // CHECK-NEXT: ( 0, 0, 0, 3127 // CHECK-NEXT: ( 0, 0, 0, 13128 // CHECK-NEXT: ( 20, 21, 0, 0129 // CHECK-NEXT: ( 30, 31, 0, 0130 call @transfer_write_2d_mask(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()131 call @load_and_print(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()132 133 // 4. Reload 3. + transpose + store.134 // CHECK-LABEL: TILE BEGIN:135 // CHECK-NEXT: ( 0, 0, 20, 30136 // CHECK-NEXT: ( 0, 0, 21, 31137 // CHECK-NEXT: ( 0, 0, 0, 0138 // CHECK-NEXT: ( 3, 13, 0, 0139 call @transfer_write_2d_transposed(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()140 call @load_and_print(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()141 142 // 5. Reload 4. + transpose + masked store (nrows=4, ncols=2).143 // The mask applies after permutation. Columns 2 and 3 (from 4.) are144 // preserved.145 // CHECK-LABEL: TILE BEGIN:146 // CHECK-NEXT: ( 0, 0, 20, 30147 // CHECK-NEXT: ( 0, 0, 21, 31148 // CHECK-NEXT: ( 20, 21, 0, 0149 // CHECK-NEXT: ( 30, 31, 0, 0150 call @transfer_write_2d_mask_transposed(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()151 call @load_and_print(%A, %c0, %c0) : (memref<?x?xf32>, index, index) -> ()152 153 memref.dealloc %A : memref<?x?xf32>154 155 return156}157