75 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 10func.func @entry() {11 %c0 = arith.constant 0 : index12 %c1 = arith.constant 1 : index13 %c1_i32 = arith.constant 1 : i3214 15 // Calculate the size of a 32-bit tile, e.g. ZA{n}.s.16 %svl_s = arm_sme.streaming_vl <word>17 %za_s_size = arith.muli %svl_s, %svl_s : index18 19 // Allocate memory.20 %mem1 = memref.alloca(%svl_s, %svl_s) : memref<?x?xi32>21 22 // Fill each "row" of "mem1" with row number.23 //24 // For example, assuming an SVL of 128-bits:25 //26 // 0, 0, 0, 027 // 1, 1, 1, 128 // 2, 2, 2, 229 // 3, 3, 3, 330 //31 %init_0 = arith.constant 0 : i3232 scf.for %i = %c0 to %svl_s step %c1 iter_args(%val = %init_0) -> (i32) {33 %splat_val = vector.broadcast %val : i32 to vector<[4]xi32>34 vector.store %splat_val, %mem1[%i, %c0] : memref<?x?xi32>, vector<[4]xi32>35 %val_next = arith.addi %val, %c1_i32 : i3236 scf.yield %val_next : i3237 }38 39 // Load tile from "mem1" vertically.40 %0 = arm_sme.tile_load %mem1[%c0, %c0] layout<vertical> : memref<?x?xi32>, vector<[4]x[4]xi32>41 42 // 1. ORIGINAL HORIZONTAL LAYOUT43 // Dump "mem1". The smallest SVL is 128-bits so the tile will be at least44 // 4x4xi32.45 //46 // CHECK: TILE BEGIN47 // CHECK-NEXT: ( 0, 0, 0, 048 // CHECK-NEXT: ( 1, 1, 1, 149 // CHECK-NEXT: ( 2, 2, 2, 250 // CHECK-NEXT: ( 3, 3, 3, 351 // CHECK: TILE END52 vector.print str "TILE BEGIN\n"53 scf.for %i = %c0 to %svl_s step %c1 {54 %tileslice = vector.load %mem1[%i, %c0] : memref<?x?xi32>, vector<[4]xi32>55 vector.print %tileslice : vector<[4]xi32>56 }57 vector.print str "TILE END\n"58 59 // 2. VERTICAL LAYOUT60 // Dump "mem2". The smallest SVL is 128-bits so the tile will be at least61 // 4x4xi32.62 //63 // CHECK: TILE BEGIN64 // CHECK-NEXT: ( 0, 1, 2, 365 // CHECK-NEXT: ( 0, 1, 2, 366 // CHECK-NEXT: ( 0, 1, 2, 367 // CHECK-NEXT: ( 0, 1, 2, 368 // CHECK: TILE END69 vector.print str "TILE BEGIN\n"70 vector.print %0 : vector<[4]x[4]xi32>71 vector.print str "TILE END\n"72 73 return74}75