xref: /llvm-project/mlir/test/Examples/Toy/Ch5/ast.toy (revision 5633813bf376ef12056cc8ce34c03c445d0dbce5)
1c39592b0SMehdi Amini# RUN: toyc-ch5 %s -emit=ast 2>&1 | FileCheck %s
2c39592b0SMehdi Amini
31ba9bb05SRiver Riddle# User defined generic function that operates on unknown shaped arguments.
4c39592b0SMehdi Aminidef multiply_transpose(a, b) {
5941a1c43SRiver Riddle  return transpose(a) * transpose(b);
6c39592b0SMehdi Amini}
7c39592b0SMehdi Amini
8c39592b0SMehdi Aminidef main() {
9c39592b0SMehdi Amini  # Define a variable `a` with shape <2, 3>, initialized with the literal value.
10c39592b0SMehdi Amini  # The shape is inferred from the supplied literal.
11c39592b0SMehdi Amini  var a = [[1, 2, 3], [4, 5, 6]];
121ba9bb05SRiver Riddle  # b is identical to a, the literal array is implicitly reshaped: defining new
131ba9bb05SRiver Riddle  # variables is the way to reshape arrays (element count in literal must match
141ba9bb05SRiver Riddle  # the size of specified shape).
15c39592b0SMehdi Amini  var b<2, 3> = [1, 2, 3, 4, 5, 6];
161ba9bb05SRiver Riddle
17c39592b0SMehdi Amini  # This call will specialize `multiply_transpose` with <2, 3> for both
18c39592b0SMehdi Amini  # arguments and deduce a return type of <2, 2> in initialization of `c`.
19c39592b0SMehdi Amini  var c = multiply_transpose(a, b);
20c39592b0SMehdi Amini  # A second call to `multiply_transpose` with <2, 3> for both arguments will
21c39592b0SMehdi Amini  # reuse the previously specialized and inferred version and return `<2, 2>`
22c39592b0SMehdi Amini  var d = multiply_transpose(b, a);
23c39592b0SMehdi Amini  # A new call with `<2, 2>` for both dimension will trigger another
24c39592b0SMehdi Amini  # specialization of `multiply_transpose`.
25c39592b0SMehdi Amini  var e = multiply_transpose(b, c);
26c39592b0SMehdi Amini  # Finally, calling into `multiply_transpose` with incompatible shape will
27c39592b0SMehdi Amini  # trigger a shape inference error.
289f89c34fSAart Bik  var f = multiply_transpose(transpose(a), c);
29c39592b0SMehdi Amini}
30c39592b0SMehdi Amini
31c39592b0SMehdi Amini
32c39592b0SMehdi Amini# CHECK: Module:
33c39592b0SMehdi Amini# CHECK-NEXT:     Function
34*5633813bSRahul Joshi# CHECK-NEXT:       Proto 'multiply_transpose' @{{.*}}ast.toy:4:1
35c39592b0SMehdi Amini# CHECK-NEXT:       Params: [a, b]
36c39592b0SMehdi Amini# CHECK-NEXT:       Block {
371bdfc9e7SRiver Riddle# CHECK-NEXT:         Return
38941a1c43SRiver Riddle# CHECK-NEXT:           BinOp: * @{{.*}}ast.toy:5:25
39941a1c43SRiver Riddle# CHECK-NEXT:             Call 'transpose' [ @{{.*}}ast.toy:5:10
40941a1c43SRiver Riddle# CHECK-NEXT:               var: a @{{.*}}ast.toy:5:20
41941a1c43SRiver Riddle# CHECK-NEXT:             ]
42941a1c43SRiver Riddle# CHECK-NEXT:             Call 'transpose' [ @{{.*}}ast.toy:5:25
43941a1c43SRiver Riddle# CHECK-NEXT:               var: b @{{.*}}ast.toy:5:35
44c39592b0SMehdi Amini# CHECK-NEXT:             ]
45c39592b0SMehdi Amini# CHECK-NEXT:       } // Block
46c39592b0SMehdi Amini# CHECK-NEXT:     Function
47*5633813bSRahul Joshi# CHECK-NEXT:       Proto 'main' @{{.*}}ast.toy:8:1
48c39592b0SMehdi Amini# CHECK-NEXT:       Params: []
49c39592b0SMehdi Amini# CHECK-NEXT:       Block {
501ba9bb05SRiver Riddle# CHECK-NEXT:         VarDecl a<> @{{.*}}ast.toy:11:3
511ba9bb05SRiver Riddle# CHECK-NEXT:           Literal: <2, 3>[ <3>[ 1.000000e+00, 2.000000e+00, 3.000000e+00], <3>[ 4.000000e+00, 5.000000e+00, 6.000000e+00]] @{{.*}}ast.toy:11:11
5206b0db53SRiver Riddle# CHECK-NEXT:         VarDecl b<2, 3> @{{.*}}ast.toy:15:3
5306b0db53SRiver Riddle# CHECK-NEXT:           Literal: <6>[ 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00] @{{.*}}ast.toy:15:17
541ba9bb05SRiver Riddle# CHECK-NEXT:         VarDecl c<> @{{.*}}ast.toy:19:3
551ba9bb05SRiver Riddle# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:19:11
561ba9bb05SRiver Riddle# CHECK-NEXT:             var: a @{{.*}}ast.toy:19:30
571ba9bb05SRiver Riddle# CHECK-NEXT:             var: b @{{.*}}ast.toy:19:33
58c39592b0SMehdi Amini# CHECK-NEXT:           ]
591ba9bb05SRiver Riddle# CHECK-NEXT:         VarDecl d<> @{{.*}}ast.toy:22:3
601ba9bb05SRiver Riddle# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:22:11
611ba9bb05SRiver Riddle# CHECK-NEXT:             var: b @{{.*}}ast.toy:22:30
621ba9bb05SRiver Riddle# CHECK-NEXT:             var: a @{{.*}}ast.toy:22:33
63c39592b0SMehdi Amini# CHECK-NEXT:           ]
641ba9bb05SRiver Riddle# CHECK-NEXT:         VarDecl e<> @{{.*}}ast.toy:25:3
651ba9bb05SRiver Riddle# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:25:11
661ba9bb05SRiver Riddle# CHECK-NEXT:             var: b @{{.*}}ast.toy:25:30
671ba9bb05SRiver Riddle# CHECK-NEXT:             var: c @{{.*}}ast.toy:25:33
68c39592b0SMehdi Amini# CHECK-NEXT:           ]
699f89c34fSAart Bik# CHECK-NEXT:         VarDecl f<> @{{.*}}ast.toy:28:3
701ba9bb05SRiver Riddle# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:28:11
711ba9bb05SRiver Riddle# CHECK-NEXT:             Call 'transpose' [ @{{.*}}ast.toy:28:30
721ba9bb05SRiver Riddle# CHECK-NEXT:               var: a @{{.*}}ast.toy:28:40
73c39592b0SMehdi Amini# CHECK-NEXT:             ]
741ba9bb05SRiver Riddle# CHECK-NEXT:             var: c @{{.*}}ast.toy:28:44
75c39592b0SMehdi Amini# CHECK-NEXT:           ]
76c39592b0SMehdi Amini
77