From b09bf2a4dcae38663a4bec58cb61347bba71d662 Mon Sep 17 00:00:00 2001 From: Stephan Herhut Date: Thu, 30 Jul 2020 01:04:37 -0700 Subject: [PATCH] Move hlo end to end tests to the hlo directory tree. PiperOrigin-RevId: 323955773 --- tests/end2end/broadcast.mlir | 213 +++++++++++++++++++++++++++++++++++ tests/end2end/reshape.mlir | 190 +++++++++++++++++++++++++++++++ 2 files changed, 403 insertions(+) create mode 100644 tests/end2end/broadcast.mlir create mode 100644 tests/end2end/reshape.mlir diff --git a/tests/end2end/broadcast.mlir b/tests/end2end/broadcast.mlir new file mode 100644 index 0000000..43fb7c9 --- /dev/null +++ b/tests/end2end/broadcast.mlir @@ -0,0 +1,213 @@ +// RUN: mlir-hlo-opt %s -mhlo-test-chlo-legalize-to-hlo -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -lhlo-copy-removal -canonicalize -cse -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse -convert-linalg-to-llvm -test-lhlo-legalize-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | FileCheck %s + +func @main() -> () { + call @trivial_broadcast_wrapper() : () -> () + call @broadcast_in_X_dim_wrapper() : () -> () + call @broadcast_in_Y_dim_wrapper() : () -> () + call @broadcast_in_X_dim_transpose_wrapper() : () -> () + call @broadcast_in_Y_dim_transpose_wrapper() : () -> () + call @broadcast_scalar_1d_wrapper() : () -> () + call @broadcast_scalar_2d_wrapper() : () -> () + return +} + +func @print_memref_i8(memref<*xi8>) attributes { llvm.emit_c_interface } +func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface } + +func @trivial_broadcast_wrapper() { + %input = alloc() : memref<3xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0] : memref<3xf32> + %c2f32 = constant 2.0 : f32 + %c1 = constant 1 : index + store %c2f32, %input[%c1] : memref<3xf32> + %c3f32 = constant 3.0 : f32 + %c2 = constant 2 : index + store %c3f32, %input[%c2] : memref<3xf32> + %input_tensor = tensor_load %input : memref<3xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<0> : tensor<1xi64> + } : (tensor<3xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [2, 2, 2, 2] +// CHECK: [3, 3, 3, 3] + +func @broadcast_in_X_dim_wrapper() { + %input = alloc() : memref<1x4xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0, %c0] : memref<1x4xf32> + %c2f32 = constant 2.0 : f32 + %c1 = constant 1 : index + store %c2f32, %input[%c0, %c1] : memref<1x4xf32> + %c3f32 = constant 3.0 : f32 + %c2 = constant 2 : index + store %c3f32, %input[%c0, %c2] : memref<1x4xf32> + %c4f32 = constant 4.0 : f32 + %c3 = constant 3 : index + store %c4f32, %input[%c0, %c3] : memref<1x4xf32> + %input_tensor = tensor_load %input : memref<1x4xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<[0, 1]> : tensor<2xi64> + } : (tensor<1x4xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 2, 3, 4] +// CHECK: [1, 2, 3, 4] +// CHECK: [1, 2, 3, 4] + +func @broadcast_in_Y_dim_wrapper() { + %input = alloc() : memref<3x1xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0, %c0] : memref<3x1xf32> + %c2f32 = constant 2.0 : f32 + %c1 = constant 1 : index + store %c2f32, %input[%c1, %c0] : memref<3x1xf32> + %c3f32 = constant 3.0 : f32 + %c2 = constant 2 : index + store %c3f32, %input[%c2, %c0] : memref<3x1xf32> + %input_tensor = tensor_load %input : memref<3x1xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<[0, 1]> : tensor<2xi64> + } : (tensor<3x1xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [2, 2, 2, 2] +// CHECK: [3, 3, 3, 3] + +func @broadcast_in_X_dim_transpose_wrapper() { + %input = alloc() : memref<4x1xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0, %c0] : memref<4x1xf32> + %c2f32 = constant 2.0 : f32 + %c1 = constant 1 : index + store %c2f32, %input[%c1, %c0] : memref<4x1xf32> + %c3f32 = constant 3.0 : f32 + %c2 = constant 2 : index + store %c3f32, %input[%c2, %c0] : memref<4x1xf32> + %c4f32 = constant 4.0 : f32 + %c3 = constant 3 : index + store %c4f32, %input[%c3, %c0] : memref<4x1xf32> + %input_tensor = tensor_load %input : memref<4x1xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<[1, 0]> : tensor<2xi64> + } : (tensor<4x1xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 2, 3, 4] +// CHECK: [1, 2, 3, 4] +// CHECK: [1, 2, 3, 4] + +func @broadcast_in_Y_dim_transpose_wrapper() { + %input = alloc() : memref<1x3xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0, %c0] : memref<1x3xf32> + %c2f32 = constant 2.0 : f32 + %c1 = constant 1 : index + store %c2f32, %input[%c0, %c1] : memref<1x3xf32> + %c3f32 = constant 3.0 : f32 + %c2 = constant 2 : index + store %c3f32, %input[%c0, %c2] : memref<1x3xf32> + %input_tensor = tensor_load %input : memref<1x3xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<[1, 0]> : tensor<2xi64> + } : (tensor<1x3xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [2, 2, 2, 2] +// CHECK: [3, 3, 3, 3] + +func @broadcast_scalar_1d_wrapper() { + %input = alloc() : memref<1xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0] : memref<1xf32> + %input_tensor = tensor_load %input : memref<1xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<0> : tensor<1xi64> + } : (tensor<1xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [1, 1, 1, 1] + +func @broadcast_scalar_2d_wrapper() { + %input = alloc() : memref<1x1xf32> + %c1f32 = constant 1.0 : f32 + %c0 = constant 0 : index + store %c1f32, %input[%c0, %c0] : memref<1x1xf32> + %input_tensor = tensor_load %input : memref<1x1xf32> + + %output_tensor = "mhlo.broadcast_in_dim"(%input_tensor) { + broadcast_dimensions = dense<[0, 1]> : tensor<2xi64> + } : (tensor<1x1xf32>) -> tensor<3x4xf32> + + %output = alloc() : memref<3x4xf32> + tensor_store %output_tensor, %output : memref<3x4xf32> + + %cast_for_print = memref_cast %output : memref<3x4xf32> to memref<*xf32> + call @print_memref_f32(%cast_for_print) : (memref<*xf32>) -> () + return +} +// CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [1, 1, 1, 1] +// CHECK: [1, 1, 1, 1] + diff --git a/tests/end2end/reshape.mlir b/tests/end2end/reshape.mlir new file mode 100644 index 0000000..3d98ce5 --- /dev/null +++ b/tests/end2end/reshape.mlir @@ -0,0 +1,190 @@ +// RUN: mlir-hlo-opt %s -mhlo-test-chlo-legalize-to-hlo -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -lhlo-copy-removal -canonicalize -cse -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops -convert-scf-to-std -canonicalize -cse -test-lhlo-legalize-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext | FileCheck %s + +func @main() -> () { + call @reshape_with_static_shape_size_matrix_to_1D() : () -> () + call @reshape_with_static_shape_size_matrix_to_3D() : () -> () + call @reshape_with_dynamic_shape_size_matrix_to_1D() : () -> () + call @reshape_with_dynamic_shape_size_matrix_to_3D() : () -> () + return +} + +func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface } + +func @reshape_with_static_shape_size_matrix_to_1D() { + %c0 = constant 0 : index + %c1 = constant 1 : index + + // Initialize input. + %input = alloc() : memref<2x3xf32> + %dim_x = dim %input, %c0 : memref<2x3xf32> + %dim_y = dim %input, %c1 : memref<2x3xf32> + scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) { + %i_i64 = index_cast %i : index to i64 + %i_f32 = sitofp %i_i64 : i64 to f32 + store %i_f32, %input[%i, %j] : memref<2x3xf32> + } + %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> + call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> () + // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] + // CHECK: [0, 0, 0] + // CHECK: [1, 1, 1] + + // Initialize shape. + %shape = alloc() : memref<1xi64> + %num_elements = muli %dim_x, %dim_y : index + %num_elements_i64 = index_cast %num_elements : index to i64 + store %num_elements_i64, %shape[%c0] : memref<1xi64> + + // 1. Ranked input, ranked output. + %output_1 = lmhlo.reshape_memref_cast %input(%shape) + : (memref<2x3xf32>, memref<1xi64>) -> memref<6xf32> + %unranked_output_1 = memref_cast %output_1 : memref<6xf32> to memref<*xf32> + call @print_memref_f32(%unranked_output_1) : (memref<*xf32>) -> () + // CHECK: rank = 1 offset = 0 sizes = [6] strides = [1] + // CHECK: [0, 0, 0, 1, 1, 1] + + // 2. Ranked input, unranked output. + %output_2 = lmhlo.reshape_memref_cast %input(%shape) + : (memref<2x3xf32>, memref<1xi64>) -> memref<*xf32> + call @print_memref_f32(%output_2) : (memref<*xf32>) -> () + // CHECK: rank = 1 offset = 0 sizes = [6] strides = [1] + // CHECK: [0, 0, 0, 1, 1, 1] + + // 3. Unranked input, ranked output. + %output_3 = lmhlo.reshape_memref_cast %unranked_input(%shape) + : (memref<*xf32>, memref<1xi64>) -> memref + %unranked_output_3 = memref_cast %output_3 : memref to memref<*xf32> + call @print_memref_f32(%unranked_output_3) : (memref<*xf32>) -> () + // CHECK: rank = 1 offset = 0 sizes = [6] strides = [1] + // CHECK: [0, 0, 0, 1, 1, 1] + + // 4. Unranked input, unranked output. + %output_4 = lmhlo.reshape_memref_cast %unranked_input(%shape) + : (memref<*xf32>, memref<1xi64>) -> memref<*xf32> + call @print_memref_f32(%output_4) : (memref<*xf32>) -> () + // CHECK: rank = 1 offset = 0 sizes = [6] strides = [1] + // CHECK: [0, 0, 0, 1, 1, 1] + return +} + +func @reshape_with_static_shape_size_matrix_to_3D() { + %c0 = constant 0 : index + %c1 = constant 1 : index + %c2 = constant 2 : index + + // Initialize input. + %input = alloc() : memref<2x3xf32> + %dim_x = dim %input, %c0 : memref<2x3xf32> + %dim_y = dim %input, %c1 : memref<2x3xf32> + scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) { + %i_i64 = index_cast %i : index to i64 + %i_f32 = sitofp %i_i64 : i64 to f32 + store %i_f32, %input[%i, %j] : memref<2x3xf32> + } + %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> + call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> () + // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] + // CHECK: [0, 0, 0] + // CHECK: [1, 1, 1] + + // Initialize shape. + %shape = alloc() : memref<3xi64> + %c1_i64 = constant 1 : i64 + %c2_i64 = constant 2 : i64 + %c3_i64 = constant 3 : i64 + store %c3_i64, %shape[%c0] : memref<3xi64> + store %c1_i64, %shape[%c1] : memref<3xi64> + store %c2_i64, %shape[%c2] : memref<3xi64> + + // Static shape input and shape, dynamic output. + %unranked_output = lmhlo.reshape_memref_cast %input(%shape) + : (memref<2x3xf32>, memref<3xi64>) -> memref<*xf32> + call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> () + // CHECK: rank = 3 offset = 0 sizes = [3, 1, 2] strides = [2, 2, 1] + // CHECK: {{\[}}{{\[}}[0, 0]], + // CHECK: {{\[}}[0, 1]], + // CHECK: {{\[}}[1, 1]]] + return +} + +func @reshape_with_dynamic_shape_size_matrix_to_1D() { + %c0 = constant 0 : index + %c1 = constant 1 : index + + // Initialize input. + %input = alloc() : memref<2x3xf32> + %dim_x = dim %input, %c0 : memref<2x3xf32> + %dim_y = dim %input, %c1 : memref<2x3xf32> + scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) { + %i_i64 = index_cast %i : index to i64 + %i_f32 = sitofp %i_i64 : i64 to f32 + store %i_f32, %input[%i, %j] : memref<2x3xf32> + } + %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> + call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> () + // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] + // CHECK: [0, 0, 0] + // CHECK: [1, 1, 1] + + // Initialize shape. + %shape = alloc(%c1) : memref + %num_elements = muli %dim_x, %dim_y : index + %num_elements_i64 = index_cast %num_elements : index to i64 + store %num_elements_i64, %shape[%c0] : memref + + // 1. Ranked input, unranked output. + %output_2 = lmhlo.reshape_memref_cast %input(%shape) + : (memref<2x3xf32>, memref) -> memref<*xf32> + call @print_memref_f32(%output_2) : (memref<*xf32>) -> () + // CHECK: rank = 1 offset = 0 sizes = [6] strides = [1] + // CHECK: [0, 0, 0, 1, 1, 1] + + // 2. Unranked input, unranked output. + %output_4 = lmhlo.reshape_memref_cast %unranked_input(%shape) + : (memref<*xf32>, memref) -> memref<*xf32> + call @print_memref_f32(%output_4) : (memref<*xf32>) -> () + // CHECK: rank = 1 offset = 0 sizes = [6] strides = [1] + // CHECK: [0, 0, 0, 1, 1, 1] + return +} + +func @reshape_with_dynamic_shape_size_matrix_to_3D() { + %c0 = constant 0 : index + %c1 = constant 1 : index + %c2 = constant 2 : index + %c3 = constant 3 : index + + // Initialize input. + %input = alloc() : memref<2x3xf32> + %dim_x = dim %input, %c0 : memref<2x3xf32> + %dim_y = dim %input, %c1 : memref<2x3xf32> + scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) { + %i_i64 = index_cast %i : index to i64 + %i_f32 = sitofp %i_i64 : i64 to f32 + store %i_f32, %input[%i, %j] : memref<2x3xf32> + } + %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> + call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> () + // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] + // CHECK: [0, 0, 0] + // CHECK: [1, 1, 1] + + // Initialize shape. + %shape = alloc(%c3) : memref + %c1_i64 = constant 1 : i64 + %c2_i64 = constant 2 : i64 + %c3_i64 = constant 3 : i64 + store %c3_i64, %shape[%c0] : memref + store %c1_i64, %shape[%c1] : memref + store %c2_i64, %shape[%c2] : memref + + // Static shape input, dynamic output and shape. + %unranked_output = lmhlo.reshape_memref_cast %input(%shape) + : (memref<2x3xf32>, memref) -> memref<*xf32> + call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> () + // CHECK: rank = 3 offset = 0 sizes = [3, 1, 2] strides = [2, 2, 1] + // CHECK: {{\[}}{{\[}}[0, 0]], + // CHECK: {{\[}}[0, 1]], + // CHECK: {{\[}}[1, 1]]] + return +}