xref: /llvm-project/flang/test/HLFIR/matmul.fir (revision 09472ba315043748db01b092fdaca54daf19a513) 
1// Test hlfir.matmul operation parse, verify (no errors), and unparse
2
3// RUN: fir-opt %s | fir-opt | FileCheck %s
4
5// arguments are expressions of known shape
6func.func @matmul0(%arg0: !hlfir.expr<2x2xi32>, %arg1: !hlfir.expr<2x2xi32>) {
7  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<2x2xi32>, !hlfir.expr<2x2xi32>) -> !hlfir.expr<2x2xi32>
8  return
9}
10// CHECK-LABEL: func.func @matmul0
11// CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x2xi32>,
12// CHECK:           %[[ARG1:.*]]: !hlfir.expr<2x2xi32>) {
13// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<2x2xi32>, !hlfir.expr<2x2xi32>) -> !hlfir.expr<2x2xi32>
14// CHECK-NEXT:    return
15// CHECK-NEXT:  }
16
17// arguments are expressions of assumed shape
18func.func @matmul1(%arg0: !hlfir.expr<?x?xi32>, %arg1: !hlfir.expr<?x?xi32>) {
19  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<?x?xi32>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?x?xi32>
20  return
21}
22// CHECK-LABEL: func.func @matmul1
23// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?xi32>,
24// CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x?xi32>) {
25// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x?xi32>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?x?xi32>
26// CHECK-NEXT:    return
27// CHECK-NEXT:  }
28
29// arguments are expressions where only some dimensions are known #1
30func.func @matmul2(%arg0: !hlfir.expr<2x?xi32>, %arg1: !hlfir.expr<?x2xi32>) {
31  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<2x?xi32>, !hlfir.expr<?x2xi32>) -> !hlfir.expr<2x2xi32>
32  return
33}
34// CHECK-LABEL: func.func @matmul2
35// CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x?xi32>,
36// CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x2xi32>) {
37// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<2x?xi32>, !hlfir.expr<?x2xi32>) -> !hlfir.expr<2x2xi32>
38// CHECK-NEXT:    return
39// CHECK-NEXT:  }
40
41// arguments are expressions where only some dimensions are known #2
42func.func @matmul3(%arg0: !hlfir.expr<?x2xi32>, %arg1: !hlfir.expr<2x?xi32>) {
43  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<?x2xi32>, !hlfir.expr<2x?xi32>) -> !hlfir.expr<?x?xi32>
44  return
45}
46// CHECK-LABEL: func.func @matmul3
47// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x2xi32>,
48// CHECK:           %[[ARG1:.*]]: !hlfir.expr<2x?xi32>) {
49// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x2xi32>, !hlfir.expr<2x?xi32>) -> !hlfir.expr<?x?xi32>
50// CHECK-NEXT:    return
51// CHECK-NEXT:  }
52
53// arguments are logicals
54func.func @matmul4(%arg0: !hlfir.expr<?x?x!fir.logical<4>>, %arg1: !hlfir.expr<?x?x!fir.logical<4>>) {
55  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<?x?x!fir.logical<4>>, !hlfir.expr<?x?x!fir.logical<4>>) -> !hlfir.expr<?x?x!fir.logical<4>>
56  return
57}
58// CHECK-LABEL: func.func @matmul4
59// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?x!fir.logical<4>>,
60// CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x?x!fir.logical<4>>) {
61// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x?x!fir.logical<4>>, !hlfir.expr<?x?x!fir.logical<4>>) -> !hlfir.expr<?x?x!fir.logical<4>>
62// CHECK-NEXT:    return
63// CHECK-NEXT:  }
64
65// lhs is rank 1
66func.func @matmul5(%arg0: !hlfir.expr<?xi32>, %arg1: !hlfir.expr<?x?xi32>) {
67  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<?xi32>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?xi32>
68  return
69}
70// CHECK-LABEL: func.func @matmul5
71// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?xi32>,
72// CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x?xi32>) {
73// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?xi32>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?xi32>
74// CHECK-NEXT:    return
75// CHECK-NEXT:  }
76
77// rhs is rank 1
78func.func @matmul6(%arg0: !hlfir.expr<?x?xi32>, %arg1: !hlfir.expr<?xi32>) {
79  %res = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<?x?xi32>, !hlfir.expr<?xi32>) -> !hlfir.expr<?xi32>
80  return
81}
82// CHECK-LABEL: func.func @matmul6
83// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?xi32>,
84// CHECK:           %[[ARG1:.*]]: !hlfir.expr<?xi32>) {
85// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x?xi32>, !hlfir.expr<?xi32>) -> !hlfir.expr<?xi32>
86// CHECK-NEXT:    return
87// CHECK-NEXT:  }
88
89// arguments are boxed arrays
90func.func @matmul7(%arg0: !fir.box<!fir.array<2x2xf32>>, %arg1: !fir.box<!fir.array<2x2xf32>>) {
91  %res = hlfir.matmul %arg0 %arg1 : (!fir.box<!fir.array<2x2xf32>>, !fir.box<!fir.array<2x2xf32>>) -> !hlfir.expr<2x2xf32>
92  return
93}
94// CHECK-LABEL: func.func @matmul7
95// CHECK:           %[[ARG0:.*]]: !fir.box<!fir.array<2x2xf32>>,
96// CHECK:           %[[ARG1:.*]]: !fir.box<!fir.array<2x2xf32>>) {
97// CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul %[[ARG0]] %[[ARG1]] : (!fir.box<!fir.array<2x2xf32>>, !fir.box<!fir.array<2x2xf32>>) -> !hlfir.expr<2x2xf32>
98// CHECK-NEXT:    return
99// CHECK-NEXT:  }
100