1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
|
// RUN: mlir-opt %s -convert-linalg-to-loops -convert-scf-to-std -convert-linalg-to-llvm -convert-memref-to-llvm -convert-std-to-llvm | \
// RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s
func private @print_memref_f32(memref<*xf32>)
func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
%c0 = constant 0 : index
%c1 = constant 1 : index
%f0 = constant 0.0 : f32
%x = memref.dim %A, %c0 : memref<?x?xf32>
%y = memref.dim %B, %c1 : memref<?x?xf32>
%C = memref.alloc(%x, %y) : memref<?x?xf32>
linalg.fill(%f0, %C) : f32, memref<?x?xf32>
linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
outs(%C: memref<?x?xf32>)
return %C : memref<?x?xf32>
}
func @matvec(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
%c0 = constant 0 : index
%c1 = constant 1 : index
%f0 = constant 0.0 : f32
%m = memref.dim %A, %c0 : memref<?x?xf32>
%x = memref.dim %A, %c1 : memref<?x?xf32>
%n = memref.dim %B, %c1 : memref<?x?xf32>
%C = memref.alloc(%m, %n) : memref<?x?xf32>
linalg.fill(%f0, %C) : f32, memref<?x?xf32>
scf.for %i = %c0 to %n step %c1 {
%b = memref.subview %B[0, %i][%x, 1][1, 1] : memref<?x?xf32> to memref<?xf32, offset: ?, strides: [?]>
%c = memref.subview %C[0, %i][%m, 1][1, 1] : memref<?x?xf32> to memref<?xf32, offset: ?, strides: [?]>
linalg.matvec ins(%A, %b: memref<?x?xf32>, memref<?xf32, offset: ?, strides: [?]>)
outs(%c: memref<?xf32, offset: ?, strides: [?]>)
}
return %C : memref<?x?xf32>
}
func @main() {
%c0 = constant 0 : index
%c1 = constant 1 : index
%m = constant 5 : index
%x = constant 3 : index
%n = constant 2 : index
%val1 = constant 13.0 : f32
%val2 = constant 17.0 : f32
%A = memref.alloc(%m, %x) : memref<?x?xf32>
%B = memref.alloc(%x, %n) : memref<?x?xf32>
linalg.fill(%val1, %A) : f32, memref<?x?xf32>
linalg.fill(%val2, %B) : f32, memref<?x?xf32>
memref.store %val1, %B[%c0, %c0] : memref<?x?xf32>
%C1 = call @matmul(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>
%C2 = call @matvec(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>
scf.for %i = %c0 to %m step %c1 {
scf.for %j = %c0 to %n step %c1 {
%e1 = memref.load %C1[%i, %j] : memref<?x?xf32>
%e2 = memref.load %C2[%i, %j] : memref<?x?xf32>
%c = cmpf oeq, %e1, %e2 : f32
assert %c, "Matmul does not produce same output as matvec"
}
}
%C2_ = memref.cast %C2 : memref<?x?xf32> to memref<*xf32>
call @print_memref_f32(%C2_) : (memref<*xf32>) -> ()
return
}
// CHECK: Unranked Memref base@ = {{.*}} rank = 2 offset = 0 sizes = [5, 2] strides = [2, 1] data =
// CHECK-NEXT: [
// CHECK-SAME: [611, 663],
// CHECK-NEXT: [611, 663],
// CHECK-NEXT: [611, 663],
// CHECK-NEXT: [611, 663],
// CHECK-NEXT: [611, 663]
// CHECK-SAME: ]
|
