Add Chapter 4 for the Toy tutorial: shape inference, function specialization, and basic combines

--

PiperOrigin-RevId: 242050514

6 files changed, 173 insertions, 0 deletions

diff --git a/mlir/test/Examples/Toy/Ch4/ast.toy b/mlir/test/Examples/Toy/Ch4/ast.toy
new file mode 100644
index 00000000000..dbc9d927abe
--- /dev/null
+++ b/mlir/test/Examples/Toy/Ch4/ast.toy
@@ -0,0 +1,73 @@
+# RUN: toyc-ch4 %s -emit=ast 2>&1 | FileCheck %s
+
+
+# User defined generic function that operates solely on 
+def multiply_transpose(a, b) {
+  return a * transpose(b);
+}
+
+def main() {
+  # Define a variable `a` with shape <2, 3>, initialized with the literal value.
+  # The shape is inferred from the supplied literal.
+  var a = [[1, 2, 3], [4, 5, 6]];
+  # b is identical to a, the literal array is implicitely reshaped: defining new
+  # variables is the way to reshape arrays (element count must match).
+  var b<2, 3> = [1, 2, 3, 4, 5, 6];
+  # This call will specialize `multiply_transpose` with <2, 3> for both
+  # arguments and deduce a return type of <2, 2> in initialization of `c`.
+  var c = multiply_transpose(a, b);
+  # A second call to `multiply_transpose` with <2, 3> for both arguments will
+  # reuse the previously specialized and inferred version and return `<2, 2>`
+  var d = multiply_transpose(b, a);
+  # A new call with `<2, 2>` for both dimension will trigger another
+  # specialization of `multiply_transpose`.
+  var e = multiply_transpose(b, c);
+  # Finally, calling into `multiply_transpose` with incompatible shape will
+  # trigger a shape inference error.
+  var e = multiply_transpose(transpose(a), c);
+}
+
+
+# CHECK: Module:
+# CHECK-NEXT:     Function
+# CHECK-NEXT:       Proto 'multiply_transpose' @{{.*}}Toy/Ch4/ast.toy:5:1'
+# CHECK-NEXT:       Params: [a, b]
+# CHECK-NEXT:       Block {
+# CHECK-NEXT:         Retur
+# CHECK-NEXT:           BinOp: * @{{.*}}Toy/Ch4/ast.toy:6:14
+# CHECK-NEXT:             var: a @{{.*}}Toy/Ch4/ast.toy:6:10
+# CHECK-NEXT:             Call 'transpose' [ @{{.*}}Toy/Ch4/ast.toy:6:14
+# CHECK-NEXT:               var: b @{{.*}}Toy/Ch4/ast.toy:6:24
+# CHECK-NEXT:             ]
+# CHECK-NEXT:       } // Block
+# CHECK-NEXT:     Function
+# CHECK-NEXT:       Proto 'main' @{{.*}}Toy/Ch4/ast.toy:9:1'
+# CHECK-NEXT:       Params: []
+# CHECK-NEXT:       Block {
+# CHECK-NEXT:         VarDecl a<> @{{.*}}Toy/Ch4/ast.toy:12:3
+# 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]] @{{.*}}Toy/Ch4/ast.toy:12:11
+# CHECK-NEXT:         VarDecl b<2, 3> @{{.*}}Toy/Ch4/ast.toy:15:3
+# CHECK-NEXT:           Literal: <6>[ 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00] @{{.*}}Toy/Ch4/ast.toy:15:17
+# CHECK-NEXT:         VarDecl c<> @{{.*}}Toy/Ch4/ast.toy:18:3
+# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}Toy/Ch4/ast.toy:18:11
+# CHECK-NEXT:             var: a @{{.*}}Toy/Ch4/ast.toy:18:30
+# CHECK-NEXT:             var: b @{{.*}}Toy/Ch4/ast.toy:18:33
+# CHECK-NEXT:           ]
+# CHECK-NEXT:         VarDecl d<> @{{.*}}Toy/Ch4/ast.toy:21:3
+# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}Toy/Ch4/ast.toy:21:11
+# CHECK-NEXT:             var: b @{{.*}}Toy/Ch4/ast.toy:21:30
+# CHECK-NEXT:             var: a @{{.*}}Toy/Ch4/ast.toy:21:33
+# CHECK-NEXT:           ]
+# CHECK-NEXT:         VarDecl e<> @{{.*}}Toy/Ch4/ast.toy:24:3
+# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}Toy/Ch4/ast.toy:24:11
+# CHECK-NEXT:             var: b @{{.*}}Toy/Ch4/ast.toy:24:30
+# CHECK-NEXT:             var: c @{{.*}}Toy/Ch4/ast.toy:24:33
+# CHECK-NEXT:           ]
+# CHECK-NEXT:         VarDecl e<> @{{.*}}Toy/Ch4/ast.toy:27:3
+# CHECK-NEXT:           Call 'multiply_transpose' [ @{{.*}}Toy/Ch4/ast.toy:27:11
+# CHECK-NEXT:             Call 'transpose' [ @{{.*}}Toy/Ch4/ast.toy:27:30
+# CHECK-NEXT:               var: a @{{.*}}Toy/Ch4/ast.toy:27:40
+# CHECK-NEXT:             ]
+# CHECK-NEXT:             var: c @{{.*}}Toy/Ch4/ast.toy:27:44
+# CHECK-NEXT:           ]
+
diff --git a/mlir/test/Examples/Toy/Ch4/codegen.toy b/mlir/test/Examples/Toy/Ch4/codegen.toy
new file mode 100644
index 00000000000..c0497a1851a
--- /dev/null
+++ b/mlir/test/Examples/Toy/Ch4/codegen.toy
@@ -0,0 +1,32 @@
+# RUN: toyc-ch4 %s -emit=mlir 2>&1 | FileCheck %s
+
+# User defined generic function that operates on unknown shaped arguments
+def multiply_transpose(a, b) {
+  return a * transpose(b);
+}
+
+def main() {
+  var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
+  var b<2, 3> = [1, 2, 3, 4, 5, 6];
+  var c = multiply_transpose(a, b);
+  var d = multiply_transpose(b, a);
+  print(d);
+}
+
+# CHECK-LABEL: func @multiply_transpose(%arg0: !toy<"array">, %arg1: !toy<"array">)
+# CHECK-NEXT:   attributes  {toy.generic: true} {
+# CHECK-NEXT:   %0 = "toy.transpose"(%arg1) : (!toy<"array">) -> !toy<"array">
+# CHECK-NEXT:   %1 = "toy.mul"(%arg0, %0) : (!toy<"array">, !toy<"array">) -> !toy<"array">
+# CHECK-NEXT:   "toy.return"(%1) : (!toy<"array">) -> ()
+# CHECK-NEXT: }
+
+# CHECK-LABEL: func @main() {
+# CHECK-NEXT:   %0 = "toy.constant"() {value: dense<tensor<2x3xf64>, {{\[\[}}1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]>} : () -> !toy<"array<2, 3>">
+# CHECK-NEXT:   %1 = "toy.reshape"(%0) : (!toy<"array<2, 3>">) -> !toy<"array<2, 3>">
+# CHECK-NEXT:   %2 = "toy.constant"() {value: dense<tensor<6xf64>, [1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]>} : () -> !toy<"array<6>">
+# CHECK-NEXT:   %3 = "toy.reshape"(%2) : (!toy<"array<6>">) -> !toy<"array<2, 3>">
+# CHECK-NEXT:   %4 = "toy.generic_call"(%1, %3) {callee: "multiply_transpose"} : (!toy<"array<2, 3>">, !toy<"array<2, 3>">) -> !toy<"array">
+# CHECK-NEXT:   %5 = "toy.generic_call"(%3, %1) {callee: "multiply_transpose"} : (!toy<"array<2, 3>">, !toy<"array<2, 3>">) -> !toy<"array">
+# CHECK-NEXT:   "toy.print"(%5) : (!toy<"array">) -> ()
+# CHECK-NEXT:   "toy.return"() : () -> ()
+
diff --git a/mlir/test/Examples/Toy/Ch4/invalid.mlir b/mlir/test/Examples/Toy/Ch4/invalid.mlir
new file mode 100644
index 00000000000..0e33397730e
--- /dev/null
+++ b/mlir/test/Examples/Toy/Ch4/invalid.mlir
@@ -0,0 +1,11 @@
+// RUN: not toyc-ch4 %s -emit=mlir 2>&1
+
+
+// This IR is not "valid":
+// - toy.print should not return a value.
+// - toy.print should take an argument.
+// - There should be a block terminator.
+// This all round-trip since this is opaque for MLIR.
+func @main() {
+  %0 = "toy.print"()  : () -> !toy<"array<2, 3>">
+}
diff --git a/mlir/test/Examples/Toy/Ch4/scalar.toy b/mlir/test/Examples/Toy/Ch4/scalar.toy
new file mode 100644
index 00000000000..51fbcaf8b7a
--- /dev/null
+++ b/mlir/test/Examples/Toy/Ch4/scalar.toy
@@ -0,0 +1,14 @@
+# RUN: toyc-ch4 %s -emit=mlir 2>&1 | FileCheck %s
+
+def main() {
+  var a<2, 2> = 5.5;
+  print(a);
+}
+
+# CHECK-LABEL: func @main() {
+# CHECK-NEXT:    %0 = "toy.constant"() {value: dense<tensor<1xf64>, [5.500000e+00]>} : () -> !toy<"array<1>">
+# CHECK-NEXT:    %1 = "toy.reshape"(%0) : (!toy<"array<1>">) -> !toy<"array<2, 2>">
+# CHECK-NEXT:    "toy.print"(%1) : (!toy<"array<2, 2>">) -> ()
+# CHECK-NEXT:    "toy.return"() : () -> ()
+# CHECK-NEXT:  }
+
diff --git a/mlir/test/Examples/Toy/Ch4/transpose_transpose.toy b/mlir/test/Examples/Toy/Ch4/transpose_transpose.toy
new file mode 100644
index 00000000000..31399eee53f
--- /dev/null
+++ b/mlir/test/Examples/Toy/Ch4/transpose_transpose.toy
@@ -0,0 +1,19 @@
+# RUN: toyc-ch4 %s -emit=mlir 2>&1 | FileCheck %s
+# RUN: toyc-ch4 %s -emit=mlir -opt 2>&1 | FileCheck %s --check-prefix=OPT
+
+def transpose_transpose(x) {
+  return transpose(transpose(x));
+}
+
+def main() {
+  print(transpose_transpose([[1, 2], [3, 4]]));
+}
+
+#CHECK-LABEL: func @transpose_transpose
+#CHECK: transpose
+#CHECK-LABEL: main
+
+
+#OPT-LABEL: func @transpose_transpose
+#OPT-NOT: transpose
+
diff --git a/mlir/test/Examples/Toy/Ch4/trivialReshape.toy b/mlir/test/Examples/Toy/Ch4/trivialReshape.toy
new file mode 100644
index 00000000000..c7a805d89ef
--- /dev/null
+++ b/mlir/test/Examples/Toy/Ch4/trivialReshape.toy
@@ -0,0 +1,24 @@
+# RUN: toyc-ch4 %s -emit=mlir 2>&1 | FileCheck %s
+# RUN: toyc-ch4 %s -emit=mlir -opt 2>&1 | FileCheck %s --check-prefix=OPT
+
+# We expect no reshape in this function with optimizations enabled
+def foo(a) {
+  var b<2,1> = a;
+  var c<2,1> = b;
+  print(c);
+}
+
+def main() {
+  var a<2, 1> = [1, 2];
+  foo(a);
+}
+
+# without optimizations, match the reshape
+#CHECK-LABEL: func @foo
+#CHECK: reshape
+#CHECK-LABEL: main
+
+# with optimizations, ensure no reshape
+#OPT-LABEL: main
+#OPT-LABEL: func @foo_2x1
+#OPT-NOT: reshape