67 lines · plain
1// DEFINE: %{entry_point} = entry2// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm3// DEFINE: %{run} = %mcr_aarch64_cmd \4// DEFINE: -march=aarch64 -mattr=+sve,+sme \5// DEFINE: -e %{entry_point} -entry-point-result=i32 \6// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils7 8// NOTE: To run this test, your CPU must support SME.9 10// RUN: %{compile} | %{run} | FileCheck %s11 12// VLA memcopy in streaming mode.13func.func @streaming_kernel_copy(%src : memref<?xi64>, %dst : memref<?xi64>, %size : index) attributes {arm_streaming} {14 %c0 = arith.constant 0 : index15 %c2 = arith.constant 2 : index16 %vscale = vector.vscale17 %step = arith.muli %c2, %vscale : index18 scf.for %i = %c0 to %size step %step {19 %0 = vector.load %src[%i] : memref<?xi64>, vector<[2]xi64>20 vector.store %0, %dst[%i] : memref<?xi64>, vector<[2]xi64>21 }22 return23}24 25func.func @entry() -> i32 {26 %i0 = arith.constant 0: i6427 %r0 = arith.constant 0: i3228 %c0 = arith.constant 0: index29 %c4 = arith.constant 4: index30 %c32 = arith.constant 32: index31 32 // Set up memory.33 %a = memref.alloc() : memref<32xi64>34 %a_copy = memref.alloc() : memref<32xi64>35 %a_data = arith.constant dense<[1 , 2, 3 , 4 , 5, 6, 7, 8,36 9, 10, 11, 12, 13, 14, 15, 16,37 17, 18, 19, 20, 21, 22, 23, 24,38 25, 26, 27, 28, 29, 30, 31, 32]> : vector<32xi64>39 vector.transfer_write %a_data, %a[%c0] : vector<32xi64>, memref<32xi64>40 41 // Call kernel.42 %0 = memref.cast %a : memref<32xi64> to memref<?xi64>43 %1 = memref.cast %a_copy : memref<32xi64> to memref<?xi64>44 call @streaming_kernel_copy(%0, %1, %c32) : (memref<?xi64>, memref<?xi64>, index) -> ()45 46 // Print and verify.47 //48 // CHECK: ( 1, 2, 3, 4 )49 // CHECK-NEXT: ( 5, 6, 7, 8 )50 // CHECK-NEXT: ( 9, 10, 11, 12 )51 // CHECK-NEXT: ( 13, 14, 15, 16 )52 // CHECK-NEXT: ( 17, 18, 19, 20 )53 // CHECK-NEXT: ( 21, 22, 23, 24 )54 // CHECK-NEXT: ( 25, 26, 27, 28 )55 // CHECK-NEXT: ( 29, 30, 31, 32 )56 scf.for %i = %c0 to %c32 step %c4 {57 %cv = vector.transfer_read %a_copy[%i], %i0 : memref<32xi64>, vector<4xi64>58 vector.print %cv : vector<4xi64>59 }60 61 // Release resources.62 memref.dealloc %a : memref<32xi64>63 memref.dealloc %a_copy : memref<32xi64>64 65 return %r0 : i3266}67