The code and documentation for this chapter of the tutorial have been updated to follow the new flow. The toy 'array' type has been replaced by usages of the MLIR tensor type. The code has also been cleaned up and modernized. Closes tensorflow/mlir#101 PiperOrigin-RevId: 265744086
32 lines
1.8 KiB
Plaintext
32 lines
1.8 KiB
Plaintext
# RUN: toyc-ch2 %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(
|
|
# CHECK-SAME: [[VAL_0:%.*]]: tensor<*xf64>, [[VAL_1:%.*]]: tensor<*xf64>)
|
|
# CHECK-NEXT: attributes {toy.generic} {
|
|
# CHECK-NEXT: [[VAL_2:%.*]] = "toy.transpose"([[VAL_1]]) : (tensor<*xf64>) -> tensor<*xf64>
|
|
# CHECK-NEXT: [[VAL_3:%.*]] = "toy.mul"([[VAL_0]], [[VAL_2]]) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64>
|
|
# CHECK-NEXT: "toy.return"([[VAL_3]]) : (tensor<*xf64>) -> ()
|
|
|
|
# CHECK-LABEL: func @main() {
|
|
# CHECK-NEXT: [[VAL_4:%.*]] = "toy.constant"() {value = dense<{{\[\[}}1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>
|
|
# CHECK-NEXT: [[VAL_5:%.*]] = "toy.reshape"([[VAL_4]]) : (tensor<2x3xf64>) -> tensor<2x3xf64>
|
|
# CHECK-NEXT: [[VAL_6:%.*]] = "toy.constant"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>} : () -> tensor<6xf64>
|
|
# CHECK-NEXT: [[VAL_7:%.*]] = "toy.reshape"([[VAL_6]]) : (tensor<6xf64>) -> tensor<2x3xf64>
|
|
# CHECK-NEXT: [[VAL_8:%.*]] = "toy.generic_call"([[VAL_5]], [[VAL_7]]) {callee = @multiply_transpose} : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
|
|
# CHECK-NEXT: [[VAL_9:%.*]] = "toy.generic_call"([[VAL_7]], [[VAL_5]]) {callee = @multiply_transpose} : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>
|
|
# CHECK-NEXT: "toy.print"([[VAL_9]]) : (tensor<*xf64>) -> ()
|
|
# CHECK-NEXT: "toy.return"() : () -> ()
|