blob: 37981e37bbdcfda09c70dabce57dcf4b20721069 [file] [log] [blame]
// RUN: iree-run-mlir --iree-input-type=mhlo --iree-hal-target-backends=vmvx %s | FileCheck %s
// RUN: iree-run-mlir --iree-input-type=mhlo --iree-hal-target-backends=llvm-cpu %s | FileCheck %s
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo --iree-hal-target-backends=vulkan-spirv %s | FileCheck %s)
// CHECK-LABEL: EXEC @entry
func.func @entry() -> tensor<5x5xf32> {
%arg0 = util.unfoldable_constant dense<0.0> : tensor<f32>
%arg1 = util.unfoldable_constant dense<[[1.0],[-2.0],[-3.0],[4.0],[-5.0]]> : tensor<5x1xf32>
%arg2 = util.unfoldable_constant dense<1.0> : tensor<f32>
%arg3 = util.unfoldable_constant dense<[[3.46499,-7.64389,-5.72249,5.98053,17.6892],[2.9707,-6.20734,-4.25962,4.76055,13.8784],[2.47641,-4.77079,-2.79675,3.54056,10.0675],[1.98212,-3.33424,-1.33388,2.32058,6.25666],[1.48783,-1.8977,0.12899,1.1006,2.4458]]> : tensor<5x5xf32>
%arg4 = util.unfoldable_constant dense<0.0> : tensor<5xf32>
%ret0 = call @_entry(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<f32>, tensor<5x1xf32>, tensor<f32>, tensor<5x5xf32>, tensor<5xf32>) -> tensor<5x5xf32>
return %ret0 : tensor<5x5xf32>
}
func.func private @_entry(
%0: tensor<f32>,
%1: tensor<5x1xf32>,
%2: tensor<f32>,
%3: tensor<5x5xf32>,
%4: tensor<5xf32>) -> tensor<5x5xf32> {
%5 = "mhlo.broadcast_in_dim"(%1) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>, name = "broadcast.44"} : (tensor<5x1xf32>) -> tensor<5x1x5xf32>
%6 = "mhlo.broadcast_in_dim"(%2) {broadcast_dimensions = dense<[]> : tensor<0xi64>, name = "broadcast.9"} : (tensor<f32>) -> tensor<5x1x5xf32>
%7 = mhlo.multiply %5, %6 : tensor<5x1x5xf32>
%8 = "mhlo.broadcast_in_dim"(%0) {broadcast_dimensions = dense<[]> : tensor<0xi64>, name = "broadcast.47"} : (tensor<f32>) -> tensor<5x1x5xf32>
%9 = "mhlo.compare"(%7, %8) {comparison_direction = #mhlo<comparison_direction GT>} : (tensor<5x1x5xf32>, tensor<5x1x5xf32>) -> tensor<5x1x5xi1>
%10 = "mhlo.broadcast_in_dim"(%0) {broadcast_dimensions = dense<[]> : tensor<0xi64>, name = "broadcast.11"} : (tensor<f32>) -> tensor<5x1x5xf32>
%11 = "mhlo.broadcast_in_dim"(%0) {broadcast_dimensions = dense<[]> : tensor<0xi64>, name = "broadcast.67"} : (tensor<f32>) -> tensor<5x5xf32>
%12 = "mhlo.broadcast_in_dim"(%4) {broadcast_dimensions = dense<1> : tensor<1xi64>, name = "broadcast.64"} : (tensor<5xf32>) -> tensor<5x5xf32>
%13 = mhlo.add %3, %12 : tensor<5x5xf32>
%14 = mhlo.maximum %11, %13 {name = "maximum.68"} : tensor<5x5xf32>
%15 = "mhlo.reshape"(%14) {name = "reshape.70"} : (tensor<5x5xf32>) -> tensor<5x1x5xf32>
%16 = "mhlo.select"(%9, %10, %15) {name = "select.71"} : (tensor<5x1x5xi1>, tensor<5x1x5xf32>, tensor<5x1x5xf32>) -> tensor<5x1x5xf32>
%17 = "mhlo.copy"(%16) {name = "copy.4"} : (tensor<5x1x5xf32>) -> tensor<5x1x5xf32>
%18 = "mhlo.reshape"(%17) {name = "reshape.72"} : (tensor<5x1x5xf32>) -> tensor<5x5xf32>
return %18 : tensor<5x5xf32>
}
// On separate lines to avoid "[[" which FileCheck interprets as substitutions
// CHECK: 5x5xf32=
// CHECK-SAME: [0 0 0 0 0]
// CHECK-SAME: [2.97{{[0-9]+}} 0 0 4.76{{[0-9]+}} 13.87{{[0-9]+}}]
// CHECK-SAME: [2.47{{[0-9]+}} 0 0 3.54{{[0-9]+}} 10.06{{[0-9]+}}]
// CHECK-SAME: [0 0 0 0 0]
// CHECK-SAME: [1.48{{[0-9]+}} 0 0.12{{[0-9]+}} 1.10{{[0-9]+}} 2.44{{[0-9]+}}]