72 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 18 // Allocate memory.19 %mem1 = memref.alloca(%svl_s, %svl_s) : memref<?x?xi32>20 21 // Fill each "row" of "mem1" with row number.22 //23 // For example, assuming an SVL of 128-bits:24 //25 // 0, 0, 0, 026 // 1, 1, 1, 127 // 2, 2, 2, 228 // 3, 3, 3, 329 //30 %init_0 = arith.constant 0 : i3231 scf.for %i = %c0 to %svl_s step %c1 iter_args(%val = %init_0) -> (i32) {32 %splat_val = vector.broadcast %val : i32 to vector<[4]xi32>33 vector.store %splat_val, %mem1[%i, %c0] : memref<?x?xi32>, vector<[4]xi32>34 %val_next = arith.addi %val, %c1_i32 : i3235 scf.yield %val_next : i3236 }37 38 // Load tile from "mem1".39 %tile = vector.load %mem1[%c0, %c0] : memref<?x?xi32>, vector<[4]x[4]xi32>40 41 // Transpose tile.42 %transposed_tile = vector.transpose %tile, [1, 0] : vector<[4]x[4]xi32> to vector<[4]x[4]xi32>43 44 // Dump the original tile. The smallest SVL is 128-bits so the tile will be at45 // least 4x4xi32.46 //47 // CHECK: TILE BEGIN48 // CHECK-NEXT: ( 0, 0, 0, 049 // CHECK-NEXT: ( 1, 1, 1, 150 // CHECK-NEXT: ( 2, 2, 2, 251 // CHECK-NEXT: ( 3, 3, 3, 352 // CHECK: TILE END53 vector.print str "TILE BEGIN\n"54 vector.print %tile : vector<[4]x[4]xi32>55 vector.print str "TILE END\n"56 57 // Dump the transposed tile. The smallest SVL is 128-bits so the tile will be58 // at least 4x4xi32.59 //60 // CHECK: TILE BEGIN61 // CHECK-NEXT: ( 0, 1, 2, 362 // CHECK-NEXT: ( 0, 1, 2, 363 // CHECK-NEXT: ( 0, 1, 2, 364 // CHECK-NEXT: ( 0, 1, 2, 365 // CHECK: TILE END66 vector.print str "TILE BEGIN\n"67 vector.print %transposed_tile : vector<[4]x[4]xi32>68 vector.print str "TILE END\n"69 70 return71}72