77 lines · plain
1// RUN: mlir-opt %s -convert-linalg-to-loops -convert-scf-to-cf -expand-strided-metadata -lower-affine -convert-arith-to-llvm -finalize-memref-to-llvm -convert-func-to-llvm -convert-cf-to-llvm -reconcile-unrealized-casts | \2// RUN: mlir-runner -O3 -e main -entry-point-result=void \3// RUN: -shared-libs=%mlir_runner_utils \4// RUN: | FileCheck %s5 6func.func private @printMemrefF32(memref<*xf32>)7 8func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {9 %c0 = arith.constant 0 : index10 %c1 = arith.constant 1 : index11 %f0 = arith.constant 0.0 : f3212 %x = memref.dim %A, %c0 : memref<?x?xf32>13 %y = memref.dim %B, %c1 : memref<?x?xf32>14 %C = memref.alloc(%x, %y) : memref<?x?xf32>15 linalg.fill ins(%f0 : f32) outs(%C : memref<?x?xf32>)16 linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)17 outs(%C: memref<?x?xf32>)18 return %C : memref<?x?xf32>19}20 21func.func @matvec(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {22 %c0 = arith.constant 0 : index23 %c1 = arith.constant 1 : index24 %f0 = arith.constant 0.0 : f3225 %m = memref.dim %A, %c0 : memref<?x?xf32>26 %x = memref.dim %A, %c1 : memref<?x?xf32>27 %n = memref.dim %B, %c1 : memref<?x?xf32>28 %C = memref.alloc(%m, %n) : memref<?x?xf32>29 linalg.fill ins(%f0 : f32) outs(%C : memref<?x?xf32>)30 scf.for %i = %c0 to %n step %c1 {31 %b = memref.subview %B[0, %i][%x, 1][1, 1] : memref<?x?xf32> to memref<?xf32, strided<[?], offset: ?>>32 %c = memref.subview %C[0, %i][%m, 1][1, 1] : memref<?x?xf32> to memref<?xf32, strided<[?], offset: ?>>33 linalg.matvec ins(%A, %b: memref<?x?xf32>, memref<?xf32, strided<[?], offset: ?>>)34 outs(%c: memref<?xf32, strided<[?], offset: ?>>)35 }36 return %C : memref<?x?xf32>37}38 39func.func @main() {40 %c0 = arith.constant 0 : index41 %c1 = arith.constant 1 : index42 %m = arith.constant 5 : index43 %x = arith.constant 3 : index44 %n = arith.constant 2 : index45 %val1 = arith.constant 13.0 : f3246 %val2 = arith.constant 17.0 : f3247 %A = memref.alloc(%m, %x) : memref<?x?xf32>48 %B = memref.alloc(%x, %n) : memref<?x?xf32>49 linalg.fill ins(%val1 : f32) outs(%A : memref<?x?xf32>)50 linalg.fill ins(%val2 : f32) outs(%B : memref<?x?xf32>)51 memref.store %val1, %B[%c0, %c0] : memref<?x?xf32>52 %C1 = call @matmul(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>53 %C2 = call @matvec(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>54 scf.for %i = %c0 to %m step %c1 {55 scf.for %j = %c0 to %n step %c1 {56 %e1 = memref.load %C1[%i, %j] : memref<?x?xf32>57 %e2 = memref.load %C2[%i, %j] : memref<?x?xf32>58 %c = arith.cmpf oeq, %e1, %e2 : f3259 cf.assert %c, "Matmul does not produce same output as matvec"60 }61 }62 %C2_ = memref.cast %C2 : memref<?x?xf32> to memref<*xf32>63 call @printMemrefF32(%C2_) : (memref<*xf32>) -> ()64 memref.dealloc %C1 : memref<?x?xf32>65 memref.dealloc %C2 : memref<?x?xf32>66 return67}68 69// CHECK: Unranked Memref base@ = {{.*}} rank = 2 offset = 0 sizes = [5, 2] strides = [2, 1] data =70// CHECK-NEXT: [71// CHECK-SAME: [611, 663],72// CHECK-NEXT: [611, 663],73// CHECK-NEXT: [611, 663],74// CHECK-NEXT: [611, 663],75// CHECK-NEXT: [611, 663]76// CHECK-SAME: ]77