xref: /llvm-project/mlir/test/Examples/Toy/Ch1/ast.toy (revision 444bb1f1bbbd7450fd30acc4ac1091fe916a740f)
138b71d6bSMehdi Amini# RUN: toyc-ch1 %s -emit=ast 2>&1 | FileCheck %s
238b71d6bSMehdi Amini
3e88dbc8cSJacques Pienaar# User defined generic function that operates on unknown shaped arguments.
438b71d6bSMehdi Aminidef multiply_transpose(a, b) {
5941a1c43SRiver Riddle  return transpose(a) * transpose(b);
638b71d6bSMehdi Amini}
738b71d6bSMehdi Amini
838b71d6bSMehdi Aminidef main() {
9213dda68SMehdi Amini  # Define a variable `a` with shape <2, 3>, initialized with the literal value.
10213dda68SMehdi Amini  # The shape is inferred from the supplied literal.
11213dda68SMehdi Amini  var a = [[1, 2, 3], [4, 5, 6]];
12fde21c6fSRiver Riddle  # b is identical to a, the literal array is implicitly reshaped: defining new
13e88dbc8cSJacques Pienaar  # variables is the way to reshape arrays (element count in literal must match
14e88dbc8cSJacques Pienaar  # the size of specified shape).
1538b71d6bSMehdi Amini  var b<2, 3> = [1, 2, 3, 4, 5, 6];
16e88dbc8cSJacques Pienaar
1738b71d6bSMehdi Amini  # This call will specialize `multiply_transpose` with <2, 3> for both
18bfedf169SThomas Preud'homme  # arguments and deduce a return type of <3, 2> in initialization of `c`.
1938b71d6bSMehdi Amini  var c = multiply_transpose(a, b);
2038b71d6bSMehdi Amini  # A second call to `multiply_transpose` with <2, 3> for both arguments will
21bfedf169SThomas Preud'homme  # reuse the previously specialized and inferred version and return `<3, 2>`
2238b71d6bSMehdi Amini  var d = multiply_transpose(b, a);
23bfedf169SThomas Preud'homme  # A new call with `<3, 2>` for both dimension will trigger another
2438b71d6bSMehdi Amini  # specialization of `multiply_transpose`.
25bfedf169SThomas Preud'homme  var e = multiply_transpose(c, d);
26*444bb1f1SAndrey Portnoy  # Finally, calling into `multiply_transpose` with incompatible shapes
27*444bb1f1SAndrey Portnoy  # (<2, 3> and <3, 2>) will trigger a shape inference error.
28*444bb1f1SAndrey Portnoy  var f = multiply_transpose(a, c);
2938b71d6bSMehdi Amini}
3038b71d6bSMehdi Amini
3138b71d6bSMehdi Amini
3238b71d6bSMehdi Amini# CHECK: Module:
3338b71d6bSMehdi Amini# CHECK-NEXT:     Function
345633813bSRahul Joshi# CHECK-NEXT:       Proto 'multiply_transpose' @{{.*}}ast.toy:4:1
3538b71d6bSMehdi Amini# CHECK-NEXT:       Params: [a, b]
3638b71d6bSMehdi 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
4438b71d6bSMehdi Amini# CHECK-NEXT:             ]
4538b71d6bSMehdi Amini# CHECK-NEXT:       } // Block
4638b71d6bSMehdi Amini# CHECK-NEXT:     Function
475633813bSRahul Joshi# CHECK-NEXT:       Proto 'main' @{{.*}}ast.toy:8:1
4838b71d6bSMehdi Amini# CHECK-NEXT:       Params: []
4938b71d6bSMehdi Amini# CHECK-NEXT:       Block {
50e88dbc8cSJacques Pienaar# CHECK-NEXT:         VarDecl a<> @{{.*}}ast.toy:11:3
51e88dbc8cSJacques Pienaar# 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
54e88dbc8cSJacques Pienaar# CHECK-NEXT:         VarDecl c<> @{{.*}}ast.toy:19:3
55e88dbc8cSJacques Pienaar# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:19:11
56e88dbc8cSJacques Pienaar# CHECK-NEXT:             var: a @{{.*}}ast.toy:19:30
57e88dbc8cSJacques Pienaar# CHECK-NEXT:             var: b @{{.*}}ast.toy:19:33
5838b71d6bSMehdi Amini# CHECK-NEXT:           ]
59e88dbc8cSJacques Pienaar# CHECK-NEXT:         VarDecl d<> @{{.*}}ast.toy:22:3
60e88dbc8cSJacques Pienaar# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:22:11
61e88dbc8cSJacques Pienaar# CHECK-NEXT:             var: b @{{.*}}ast.toy:22:30
62e88dbc8cSJacques Pienaar# CHECK-NEXT:             var: a @{{.*}}ast.toy:22:33
6338b71d6bSMehdi Amini# CHECK-NEXT:           ]
64e88dbc8cSJacques Pienaar# CHECK-NEXT:         VarDecl e<> @{{.*}}ast.toy:25:3
65e88dbc8cSJacques Pienaar# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:25:11
66bfedf169SThomas Preud'homme# CHECK-NEXT:             var: c @{{.*}}ast.toy:25:30
67bfedf169SThomas Preud'homme# CHECK-NEXT:             var: d @{{.*}}ast.toy:25:33
6838b71d6bSMehdi Amini# CHECK-NEXT:           ]
699f89c34fSAart Bik# CHECK-NEXT:         VarDecl f<> @{{.*}}ast.toy:28:3
70e88dbc8cSJacques Pienaar# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}ast.toy:28:11
71*444bb1f1SAndrey Portnoy# CHECK-NEXT:             var: a @{{.*}}ast.toy:28:30
72*444bb1f1SAndrey Portnoy# CHECK-NEXT:             var: c @{{.*}}ast.toy:28:33
7338b71d6bSMehdi Amini# CHECK-NEXT:           ]
74*444bb1f1SAndrey Portnoy# CHECK-NEXT:       } // Block
75