wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/mlir/lite/stablehlo/tests/BUILD
+37
View File
@@ -0,0 +1,37 @@
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//visibility:private",
],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
driver = "//tensorflow/compiler/mlir/lite/stablehlo:run_lit.sh",
size_override = {
"legalize-skip-quantization-ops.mlir": "medium",
},
test_file_exts = [
"mlir",
"cc",
],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir/lite:flatbuffer_translate",
"//tensorflow/compiler/mlir/lite:tf_tfl_translate",
"//tensorflow/compiler/mlir/lite/stablehlo:odml-to-stablehlo-opt",
"//tensorflow/compiler/mlir/lite/stablehlo:odml_to_stablehlo",
"@llvm-project//llvm:FileCheck",
"@llvm-project//mlir:run_lit.sh",
],
)
tensorflow/compiler/mlir/lite/stablehlo/tests/call_xla_module_to_stablehlo.mlir
+42
View File
@@ -0,0 +1,42 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1660 : i32}} {
func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "args_tf_0", outputs = "Identity"}} {
%0 = tf_executor.graph {
%outputs, %control = tf_executor.island wraps "tf.Identity"(%arg0) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32>
%outputs_0, %control_1 = tf_executor.island wraps "tf.XlaSharding"(%outputs) {_XlaSharding = "", device = "", sharding = "", unspecified_dims = []} : (tensor<2x3xi32>) -> tensor<2x3xi32>
%outputs_2, %control_3 = tf_executor.island wraps "tf.XlaCallModule"(%outputs_0) {Sout = [#tf_type.shape<2x3>], device = "", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "ML\EFR\01StableHLO_v0.9.0\00\01\17\05\01\03\01\03\05\03\07\07\09\0B\03]?\0B\01)\07\0F\0B+\0B\0F\0B\0B\0B3\0B\0B\0B\0B\0F\0B\0F\0B\13\0B\03\17\0F\13\0B\0B\0B\0F\13\0B\0B\0B\0B\01\05\0B\0F\03\07\17\17\07\02\D7\1F\11\03\05\05\0D\03\09\09\0B\0D\03\0F\03\05\11\05\0F\11\01\00\05\11\05\13\05\15\03\0B\15)\171\193\05;\1B=\05\17\05\19\05\1B\05\1D\1D\1F\01\05\1F\1D#%\05!\17'\A9\01\05#\03\03+\0D\03-/\1D%\1D'#\07\03\035\0D\0379\1D)\1D+\1D-\1D/\01\09\01\02\02)\05\09\0D\09\11\03\05\03\05\1B\04C\05\01\11\01\07\07\03\01\05\03\11\01\13\07\03\05\0B\03\05\1D\05\06!\03\05\05\01\01\07\04\01\03\03\06\03\01\05\01\00f\051\0F\0B\03!\1B\1D[;\05\1F\15\1D\15\1D%)9\13\15\19\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00multiply_v1\00return_v1\00sym_name\00jax.uses_shape_polymorphism\00mhlo.num_partitions\00mhlo.num_replicas\00jit_jax_model\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00x\00jit(jax_model)/jit(main)/mul\00experimental/users/ypang/lite/convert_ulm.py\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00", platforms = ["CPU"], version = 8 : i64} : (tensor<2x3xi32>) -> tensor<2x3xi32>
%control_4 = tf_executor.island(%control_3) wraps "tf.NoOp"() {device = ""} : () -> ()
%outputs_5, %control_6 = tf_executor.island wraps "tf.PreventGradient"(%outputs_2) {device = "", message = "The jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients"} : (tensor<2x3xi32>) -> tensor<2x3xi32>
%outputs_7, %control_8 = tf_executor.island wraps "tf.Identity"(%outputs_5) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32>
%outputs_9, %control_10 = tf_executor.island(%control_4) wraps "tf.Identity"(%outputs_7) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32>
tf_executor.fetch %outputs_9 : tensor<2x3xi32>
}
return %0 : tensor<2x3xi32>
}
}
// CHECK: module attributes
// CHECK-SAME: tfl.metadata = {{{.*}}keep_stablehlo_constant = "true"{{.*}}}
// CHECK-NEXT: func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> attributes {tf.entry_function = {inputs = "args_tf_0", outputs = "Identity"}} {
// CHECK-NEXT: %0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32>
// CHECK-NEXT: %1 = stablehlo.multiply %0, %0 : tensor<2x3xi32>
// CHECK-NEXT: return %1 : tensor<2x3xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/compose-uniform-quantized-type.mlir
+556
View File
@@ -0,0 +1,556 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt --compose-uniform-quantized-type \
// RUN: --split-input-file --verify-diagnostics %s | FileCheck %s
module {
// CHECK-LABEL: quantized_conv_op
// CHECK-SAME: %[[ARG:.*]]: tensor<1x3x3x4xf32>
func.func @quantized_conv_op(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1x1x1x1xf32> // Input inverse scale.
%2 = stablehlo.constant dense<-128> : tensor<1x1x1x1xi8> // Input zero point.
%3 = stablehlo.constant dense<1> : tensor<3x3x4x4xi8> // Quantized filter tensor.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x1x4xf32>
%5 = stablehlo.constant dense<4.000000e+03> : tensor<1x1x1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<127> : tensor<1x1x1x1xi8> // Output zero point.
%7 = call @uniform_quantize_0(%arg0, %1, %2) : (tensor<1x3x3x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8>
%8 = stablehlo.convert %7 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
%9 = stablehlo.convert %3 : (tensor<3x3x4x4xi8>) -> tensor<3x3x4x4xf32>
%10 = stablehlo.convolution(%8, %9) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
%11 = stablehlo.reshape %2 : (tensor<1x1x1x1xi8>) -> tensor<1xi8>
%12 = stablehlo.broadcast_in_dim %11, dims = [0] : (tensor<1xi8>) -> tensor<1x3x3x4xi8>
%13 = stablehlo.convert %12 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
%14 = stablehlo.convert %3 : (tensor<3x3x4x4xi8>) -> tensor<3x3x4x4xf32>
%15 = stablehlo.convolution(%13, %14) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
%16 = stablehlo.subtract %10, %15 : tensor<1x3x3x4xf32>
%17 = stablehlo.broadcast_in_dim %4, dims = [0, 1, 2, 3] : (tensor<1x1x1x4xf32>) -> tensor<1x3x3x4xf32>
%18 = stablehlo.multiply %16, %17 : tensor<1x3x3x4xf32>
%19 = call @uniform_quantize_1(%18, %5, %6) : (tensor<1x3x3x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8>
%20 = call @uniform_dequantize_0(%19, %5, %6) : (tensor<1x3x3x4xi8>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32>
return %20 : tensor<1x3x3x4xf32>
}
// CHECK: %[[FILTER:.*]] = stablehlo.constant() <{value = dense<1> : tensor<3x3x4x4xi8>}> : () -> tensor<3x3x4x4x!quant.uniform<i8:f32:3, {{{.*}}}>>
// CHECK: %[[QUANT_ARG:.*]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%[[QUANT_ARG]], %[[FILTER]]) {{.*}} : (tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>, tensor<3x3x4x4x!quant.uniform<i8:f32:3, {{.*}}>>) -> tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]] : (tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3x3x4xf32>
// CHECK: return %[[DEQUANT]] : tensor<1x3x3x4xf32>
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xi8>
return %0 : tensor<1x3x3x4xi8>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xi8>
return %0 : tensor<1x3x3x4xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x3x3x4xi8>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
return %0 : tensor<1x3x3x4xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// Tests a variant where there is no stablehlo.convert op in between the
// filter constant and the convolution op.
//
// `filter (f32) -> convolution`
//
// instead of:
//
// `filter (i8) -> convert (i8 -> f32) -> convolution`
module {
// CHECK-LABEL: quantized_conv_op_with_no_filter_convert
// CHECK-SAME: %[[ARG:.*]]: tensor<1x3x3x4xf32>
func.func @quantized_conv_op_with_no_filter_convert(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1x1x1x1xf32> // Input inverse scale.
%2 = stablehlo.constant dense<-128> : tensor<1x1x1x1xi8> // Input zero point.
%3 = stablehlo.constant dense<2.000000e+01> : tensor<3x3x4x4xf32> // Quantized filter tensor.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x1x4xf32>
%5 = stablehlo.constant dense<4.000000e+03> : tensor<1x1x1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<127> : tensor<1x1x1x1xi8> // Output zero point.
%7 = call @uniform_quantize_0(%arg0, %1, %2) : (tensor<1x3x3x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8>
%8 = stablehlo.convert %7 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
%9 = stablehlo.convolution(%8, %3) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
%10 = stablehlo.reshape %2 : (tensor<1x1x1x1xi8>) -> tensor<1xi8>
%11 = stablehlo.broadcast_in_dim %10, dims = [0] : (tensor<1xi8>) -> tensor<1x3x3x4xi8>
%12 = stablehlo.convert %11 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
%13 = stablehlo.convolution(%12, %3) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
%14 = stablehlo.subtract %9, %13 : tensor<1x3x3x4xf32>
%15 = stablehlo.broadcast_in_dim %4, dims = [0, 1, 2, 3] : (tensor<1x1x1x4xf32>) -> tensor<1x3x3x4xf32>
%16 = stablehlo.multiply %14, %15 : tensor<1x3x3x4xf32>
%17 = call @uniform_quantize_1(%16, %5, %6) : (tensor<1x3x3x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8>
%18 = call @uniform_dequantize_0(%17, %5, %6) : (tensor<1x3x3x4xi8>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32>
return %18 : tensor<1x3x3x4xf32>
}
// CHECK: %[[FILTER:.*]] = stablehlo.constant() <{value = dense<20> : tensor<3x3x4x4xi8>}> : () -> tensor<3x3x4x4x!quant.uniform<i8:f32:3, {{{.*}}}>>
// CHECK: %[[QUANT_ARG:.*]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%[[QUANT_ARG]], %[[FILTER]]) {{.*}} : (tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>, tensor<3x3x4x4x!quant.uniform<i8:f32:3, {{.*}}>>) -> tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]] : (tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3x3x4xf32>
// CHECK: return %[[DEQUANT]] : tensor<1x3x3x4xf32>
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xi8>
return %0 : tensor<1x3x3x4xi8>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xi8>
return %0 : tensor<1x3x3x4xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x3x3x4xi8>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
return %0 : tensor<1x3x3x4xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// The pattern should not match when there are no `uniform_quantize` call
// for the input.
module {
// CHECK-LABEL: conv_no_input_uniform_quantize_call
func.func @conv_no_input_uniform_quantize_call(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1x1x1x1xf32> // Input inverse scale.
%2 = stablehlo.constant dense<-128> : tensor<1x1x1x1xi8> // Input zero point.
%3 = stablehlo.constant dense<2.000000e+01> : tensor<3x3x4x4xf32> // Quantized filter tensor.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x1x4xf32>
%5 = stablehlo.constant dense<4.000000e+03> : tensor<1x1x1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<127> : tensor<1x1x1x1xi8> // Output zero point.
%10 = stablehlo.convolution(%arg0, %3) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
%11 = stablehlo.reshape %2 : (tensor<1x1x1x1xi8>) -> tensor<1xi8>
%12 = stablehlo.broadcast_in_dim %11, dims = [0] : (tensor<1xi8>) -> tensor<1x3x3x4xi8>
%13 = stablehlo.convert %12 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
%15 = stablehlo.convolution(%13, %3) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
%16 = stablehlo.subtract %10, %15 : tensor<1x3x3x4xf32>
%17 = stablehlo.broadcast_in_dim %4, dims = [0, 1, 2, 3] : (tensor<1x1x1x4xf32>) -> tensor<1x3x3x4xf32>
%18 = stablehlo.multiply %16, %17 : tensor<1x3x3x4xf32>
%19 = call @uniform_quantize_1(%18, %5, %6) : (tensor<1x3x3x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8>
%20 = call @uniform_dequantize_0(%19, %5, %6) : (tensor<1x3x3x4xi8>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32>
return %20 : tensor<1x3x3x4xf32>
}
// CHECK-NOT: stablehlo.uniform_quantize
// CHECK-NOT: stablehlo.uniform_dequantize
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_1(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xi8>
return %0 : tensor<1x3x3x4xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x3x3x4xi8>, %arg1: tensor<1x1x1x1xf32>, %arg2: tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
return %0 : tensor<1x3x3x4xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
module {
// CHECK-LABEL: quantized_dot_general
// CHECK-SAME: %[[ARG:.*]]: tensor<1x4x2xf32>
func.func @quantized_dot_general(%arg0: tensor<1x4x2xf32>) -> tensor<1x4x3xf32> {
%0 = stablehlo.constant dense<3.000000e+00> : tensor<1x1x1xf32> // Input inverse scale.
%1 = stablehlo.constant dense<1> : tensor<1x1x1xi8> // Input zero point.
%2 = stablehlo.constant dense<5> : tensor<2x3xi8> // Quantized filter.
%3 = stablehlo.constant dense<4> : tensor<1x1x3xi32> // Precalculated q2 * z1.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x3xf32> // Merged scale: s1 * s2.
%5 = stablehlo.constant dense<2.000000e+02> : tensor<1x1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<2> : tensor<1x1x1xi8> // Output zero point.
%7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x4x2xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x2xi8>
%8 = stablehlo.convert %7 : (tensor<1x4x2xi8>) -> tensor<1x4x2xf32>
%9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
%10 = stablehlo.dot_general %8, %9, contracting_dims = [2] x [0] : (tensor<1x4x2xf32>, tensor<2x3xf32>) -> tensor<1x4x3xf32>
%11 = stablehlo.convert %3 : (tensor<1x1x3xi32>) -> tensor<1x1x3xf32>
%12 = stablehlo.broadcast_in_dim %11, dims = [0, 1, 2] : (tensor<1x1x3xf32>) -> tensor<1x4x3xf32> // Optional
%13 = stablehlo.subtract %10, %12 : tensor<1x4x3xf32> // Precalculated zp_neg.
%14 = stablehlo.broadcast_in_dim %4, dims = [0, 1, 2] : (tensor<1x1x3xf32>) -> tensor<1x4x3xf32> // Optional
%15 = stablehlo.multiply %13, %14 : tensor<1x4x3xf32> // s1 * s2
%16 = call @uniform_quantize_1(%15, %5, %6) : (tensor<1x4x3xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x3xi8>
%17 = call @uniform_dequantize_0(%16, %5, %6) : (tensor<1x4x3xi8>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x3xf32>
return %17 : tensor<1x4x3xf32>
}
// Quantization dimension == 1 because it is the output feature dimension.
// CHECK: %[[FILTER:.*]] = stablehlo.constant() <{value = dense<5> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8:f32:1, {{{.*}}}>>
// CHECK: %[[QUANT_ARG:.*]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x4x2xf32>) -> tensor<1x4x2x!quant.uniform<i8:f32, {{.*}}:1>>
// CHECK: %[[CONV:.*]] = stablehlo.dot_general %[[QUANT_ARG]], %[[FILTER]], contracting_dims = [2] x [0] : (tensor<1x4x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8:f32:1, {{.*}}>>) -> tensor<1x4x3x!quant.uniform<i8:f32, {{.*}}:2>>
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]] : (tensor<1x4x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x4x3xf32>
// CHECK: return %[[DEQUANT]] : tensor<1x4x3xf32>
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x4x2xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<1x4x2xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x4x2xf32>) -> tensor<1x4x2xi8>
return %0 : tensor<1x4x2xi8>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x4x3xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<1x4x3xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x4x3xf32>) -> tensor<1x4x3xi8>
return %0 : tensor<1x4x3xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x4x3xi8>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<1x4x3xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x4x3xi8>) -> tensor<1x4x3xf32>
return %0 : tensor<1x4x3xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// Tests that when dot_general's filter comes from an f32 constant
// it is cast to i8 after the conversion.
module {
// CHECK-LABEL: quantized_dot_general_float_filter
// CHECK-SAME: %[[ARG:.*]]: tensor<1x4x2xf32>
func.func @quantized_dot_general_float_filter(%arg0: tensor<1x4x2xf32>) -> tensor<1x4x3xf32> {
%0 = stablehlo.constant dense<3.000000e+00> : tensor<1x1x1xf32> // Input inverse scale.
%1 = stablehlo.constant dense<1> : tensor<1x1x1xi8> // Input zero point.
// Filter, disguised as f32 but the values are actually i8.
%2 = stablehlo.constant dense<5.000000e+00> : tensor<2x3xf32>
%3 = stablehlo.constant dense<4> : tensor<1x1x3xi32> // Precalculated q2 * z1.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x3xf32> // Merged scale: s1 * s2.
%5 = stablehlo.constant dense<2.000000e+02> : tensor<1x1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<2> : tensor<1x1x1xi8> // Output zero point.
%7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x4x2xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x2xi8>
%8 = stablehlo.convert %7 : (tensor<1x4x2xi8>) -> tensor<1x4x2xf32>
%9 = stablehlo.dot_general %8, %2, contracting_dims = [2] x [0] : (tensor<1x4x2xf32>, tensor<2x3xf32>) -> tensor<1x4x3xf32>
%10 = stablehlo.convert %3 : (tensor<1x1x3xi32>) -> tensor<1x1x3xf32>
%11 = stablehlo.broadcast_in_dim %10, dims = [0, 1, 2] : (tensor<1x1x3xf32>) -> tensor<1x4x3xf32> // Optional
%12 = stablehlo.subtract %9, %11 : tensor<1x4x3xf32> // Precalculated zp_neg.
%13 = stablehlo.broadcast_in_dim %4, dims = [0, 1, 2] : (tensor<1x1x3xf32>) -> tensor<1x4x3xf32> // Optional
%14 = stablehlo.multiply %12, %13 : tensor<1x4x3xf32> // s1 * s2
%15 = call @uniform_quantize_1(%14, %5, %6) : (tensor<1x4x3xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x3xi8>
%16 = call @uniform_dequantize_0(%15, %5, %6) : (tensor<1x4x3xi8>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<1x4x3xf32>
return %16 : tensor<1x4x3xf32>
}
// Quantization dimension == 1 because it is the output feature dimension.
// Quantized filter values (from f32 constant) are cast to i8.
// CHECK: %[[FILTER:.*]] = stablehlo.constant() <{value = dense<5> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8:f32:1, {{{.*}}}>>
// CHECK: %[[QUANT_ARG:.*]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x4x2xf32>) -> tensor<1x4x2x!quant.uniform<i8:f32, {{.*}}:1>>
// CHECK: %[[CONV:.*]] = stablehlo.dot_general %[[QUANT_ARG]], %[[FILTER]], contracting_dims = [2] x [0] : (tensor<1x4x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8:f32:1, {{.*}}>>) -> tensor<1x4x3x!quant.uniform<i8:f32, {{.*}}:2>>
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]] : (tensor<1x4x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x4x3xf32>
// CHECK: return %[[DEQUANT]] : tensor<1x4x3xf32>
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x4x2xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<1x4x2xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x4x2xf32>) -> tensor<1x4x2xi8>
return %0 : tensor<1x4x2xi8>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x4x3xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<1x4x3xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x4x3xf32>) -> tensor<1x4x3xi8>
return %0 : tensor<1x4x3xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x4x3xi8>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<1x4x3xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x4x3xi8>) -> tensor<1x4x3xf32>
return %0 : tensor<1x4x3xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// Tests that the conversion is successful even when there are no
// broadcast_in_dim ops for the second arguments of the subtract op and
// multiply op.
module {
// CHECK-LABEL: quantized_dot_general_no_broadcast
// CHECK-SAME: %[[ARG:.*]]: tensor<1x2xf32>
func.func @quantized_dot_general_no_broadcast(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<3.000000e+00> : tensor<1x1xf32> // Input inverse scale.
%1 = stablehlo.constant dense<1> : tensor<1x1xi8> // Input zero point.
%2 = stablehlo.constant dense<5> : tensor<2x3xi8> // Quantized filter.
%3 = stablehlo.constant dense<4> : tensor<1x3xi32> // Precalculated z1 * q2.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x3xf32> // Merged scale: s1 * s2.
%5 = stablehlo.constant dense<2.000000e+02> : tensor<1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<2> : tensor<1x1xi8> // Output zero point.
%7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x2xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x2xi8>
%8 = stablehlo.convert %7 : (tensor<1x2xi8>) -> tensor<1x2xf32>
%9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
%10 = stablehlo.dot_general %8, %9, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%11 = stablehlo.convert %3 : (tensor<1x3xi32>) -> tensor<1x3xf32>
%12 = stablehlo.subtract %10, %11 : tensor<1x3xf32> // q1 * q2 - z1 * q2
%13 = stablehlo.multiply %12, %4 : tensor<1x3xf32> // s1 * s2
%14 = call @uniform_quantize_1(%13, %5, %6) : (tensor<1x3xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xi8>
%15 = call @uniform_dequantize_0(%14, %5, %6) : (tensor<1x3xi8>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xf32>
return %15 : tensor<1x3xf32>
}
// Quantization dimension == 1 because it is the output feature dimension.
// CHECK: %[[FILTER:.*]] = stablehlo.constant() <{value = dense<5> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8:f32:1, {{{.*}}}>>
// CHECK: %[[QUANT_ARG:.*]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}:1>>
// CHECK: %[[CONV:.*]] = stablehlo.dot_general %[[QUANT_ARG]], %[[FILTER]], contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}:2>>
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3xf32>
// CHECK: return %[[DEQUANT]] : tensor<1x3xf32>
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi8>
return %0 : tensor<1x2xi8>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3xf32>) -> tensor<1x3xi8>
return %0 : tensor<1x3xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x3xi8>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x3xi8>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// Tests that the conversion doesn't happen when the uniform_quantize function
// outputs a i32 storage type.
module {
// CHECK-LABEL: quantized_dot_general_uniform_quantize_to_i32
// CHECK-SAME: %[[ARG:.*]]: tensor<1x2xf32>
func.func @quantized_dot_general_uniform_quantize_to_i32(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<3.000000e+00> : tensor<1x1xf32> // Input inverse scale.
%1 = stablehlo.constant dense<1> : tensor<1x1xi8> // Input zero point.
%2 = stablehlo.constant dense<5> : tensor<2x3xi8> // Quantized filter.
%3 = stablehlo.constant dense<4> : tensor<1x3xi32> // Precalculated z1 * q2.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x3xf32> // Merged scale: s1 * s2.
%5 = stablehlo.constant dense<2.000000e+02> : tensor<1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<2> : tensor<1x1xi8> // Output zero point.
// This uniform_quantize function is expected to output i8 instead of i32.
%7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x2xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x2xi32>
%8 = stablehlo.convert %7 : (tensor<1x2xi32>) -> tensor<1x2xf32>
%9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
%10 = stablehlo.dot_general %8, %9, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%11 = stablehlo.convert %3 : (tensor<1x3xi32>) -> tensor<1x3xf32>
%12 = stablehlo.subtract %10, %11 : tensor<1x3xf32> // q1 * q2 - z1 * q2
%13 = stablehlo.multiply %12, %4 : tensor<1x3xf32> // s1 * s2
%14 = call @uniform_quantize_1(%13, %5, %6) : (tensor<1x3xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xi8>
%15 = call @uniform_dequantize_0(%14, %5, %6) : (tensor<1x3xi8>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xf32>
return %15 : tensor<1x3xf32>
}
// CHECK-NOT: stablehlo.uniform_quantize
// CHECK-NOT: !quant.uniform
// CHECK: stablehlo.dot_general
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi32> {
%0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi32>
return %0 : tensor<1x2xi32>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3xf32>) -> tensor<1x3xi8>
return %0 : tensor<1x3xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x3xi8>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x3xi8>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// Tests that the conversion doesn't happen when the filter tensor is i32.
module {
// CHECK-LABEL: quantized_dot_general_filter_i32
// CHECK-SAME: %[[ARG:.*]]: tensor<1x2xf32>
func.func @quantized_dot_general_filter_i32(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<3.000000e+00> : tensor<1x1xf32> // Input inverse scale.
%1 = stablehlo.constant dense<1> : tensor<1x1xi8> // Input zero point.
%2 = stablehlo.constant dense<5> : tensor<2x3xi32> // Quantized filter - the pattern expects i8 but i32 is given.
%3 = stablehlo.constant dense<4> : tensor<1x3xi32> // Precalculated z1 * q2.
%4 = stablehlo.constant dense<3.000000e+03> : tensor<1x3xf32> // Merged scale: s1 * s2.
%5 = stablehlo.constant dense<2.000000e+02> : tensor<1x1xf32> // Output inverse scale.
%6 = stablehlo.constant dense<2> : tensor<1x1xi8> // Output zero point.
%7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x2xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x2xi8>
%8 = stablehlo.convert %7 : (tensor<1x2xi8>) -> tensor<1x2xf32>
%9 = stablehlo.convert %2 : (tensor<2x3xi32>) -> tensor<2x3xf32>
%10 = stablehlo.dot_general %8, %9, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%11 = stablehlo.convert %3 : (tensor<1x3xi32>) -> tensor<1x3xf32>
%12 = stablehlo.subtract %10, %11 : tensor<1x3xf32> // q1 * q2 - z1 * q2
%13 = stablehlo.multiply %12, %4 : tensor<1x3xf32> // s1 * s2
%14 = call @uniform_quantize_1(%13, %5, %6) : (tensor<1x3xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xi8>
%15 = call @uniform_dequantize_0(%14, %5, %6) : (tensor<1x3xi8>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x3xf32>
return %15 : tensor<1x3xf32>
}
// CHECK-NOT: stablehlo.uniform_quantize
// CHECK-NOT: !quant.uniform
// CHECK: stablehlo.dot_general
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi8>
return %0 : tensor<1x2xi8>
}
// CHECK: @uniform_quantize_0
func.func private @uniform_quantize_1(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xi8> {
%0 = stablehlo.convert %arg0 : (tensor<1x3xf32>) -> tensor<1x3xi8>
return %0 : tensor<1x3xi8>
}
// CHECK: @uniform_quantize_1
func.func private @uniform_dequantize_0(%arg0: tensor<1x3xi8>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xf32> {
%0 = stablehlo.convert %arg0 : (tensor<1x3xi8>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: @uniform_dequantize_0
}
// -----
// Tests that a quantized dot_general op is composed when both operands are
// actiavations.
// CHECK-LABEL: dot_general_with_two_activations
// CHECK-SAME: %[[ARG_0:.*]]: tensor<8x16x16xf32>
// CHECK-SAME: %[[ARG_1:.*]]: tensor<8x16x4xf32>
module {
func.func @dot_general_with_two_activations(%arg0: tensor<8x16x16xf32>, %arg1: tensor<8x16x4xf32>) -> tensor<8x16x4xf32> {
%1 = stablehlo.constant dense<2.000000e-01> : tensor<1x1x1xf32> // Input 1 inverse scale (1 / s1).
%2 = stablehlo.constant dense<-128> : tensor<1x1x1xi8> // Input 1 zero point (z1).
%3 = stablehlo.constant dense<-128> : tensor<1x1x1xi32> // Input 1 zero point (z1) (upcast & folded into i32).
%4 = stablehlo.constant dense<4.000000e-01> : tensor<1x1x1xf32> // Input 2 inverse scale (1 / s2).
%5 = stablehlo.constant dense<0> : tensor<1x1x1xi8> // Input 2 zero point (z2).
%6 = stablehlo.constant dense<0> : tensor<1x1x1xi32> // Input 2 zero point (z2) (upcast & folded into i32).
%7 = stablehlo.constant dense<5.000000e-01> : tensor<1x1x1xf32> // Output inverse scale (1 / s3).
%8 = stablehlo.constant dense<-5> : tensor<1x1x1xi8> // Output zero point (z3).
%9 = stablehlo.constant dense<1.250000e+01> : tensor<1x1x1xf32> // Merged scale (s1 * s2).
%10 = call @uniform_quantize(%arg0, %1, %2) : (tensor<8x16x16xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x16xi8> // q1
%11 = call @uniform_quantize_0(%arg1, %4, %5) : (tensor<8x16x4xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xi8> // q2
%12 = stablehlo.convert %10 : (tensor<8x16x16xi8>) -> tensor<8x16x16xi32>
%13 = stablehlo.broadcast_in_dim %3, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x16xi32>
%14 = stablehlo.subtract %12, %13 : tensor<8x16x16xi32> // q1 - z1
%15 = stablehlo.convert %11 : (tensor<8x16x4xi8>) -> tensor<8x16x4xi32>
%16 = stablehlo.broadcast_in_dim %6, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x4xi32>
%17 = stablehlo.subtract %15, %16 : tensor<8x16x4xi32> // q2 - z2
// Corresponds to einsum expression: b i j, b j d -> b i d
%18 = stablehlo.dot_general %14, %17, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16xi32>, tensor<8x16x4xi32>) -> tensor<8x16x4xi32>
%19 = stablehlo.convert %18 : (tensor<8x16x4xi32>) -> tensor<8x16x4xf32>
%20 = stablehlo.broadcast_in_dim %9, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<8x16x4xf32>
%21 = stablehlo.multiply %19, %20 : tensor<8x16x4xf32> // * s1 s2
%22 = call @uniform_quantize_1(%21, %7, %8) : (tensor<8x16x4xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xi8>
%23 = call @uniform_dequantize(%22, %7, %8) : (tensor<8x16x4xi8>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xf32>
return %23 : tensor<8x16x4xf32>
}
// CHECK: %[[UQ_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<8x16x16xf32>) -> tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>
// CHECK: %[[UQ_1:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<8x16x4xf32>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[UQ_0]], %[[UQ_1]], batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>
// CHECK: %[[DQ_0:.*]] = stablehlo.uniform_dequantize %[[DOT_GENERAL]] : (tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>) -> tensor<8x16x4xf32>
// CHECK: return %[[DQ_0]]
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize(%arg0: tensor<8x16x16xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x16xi8> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x16xf32>) -> tensor<8x16x16xi8>
return %0 : tensor<8x16x16xi8>
}
func.func private @uniform_quantize_0(%arg0: tensor<8x16x4xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x4xf32>) -> tensor<8x16x4xi8>
return %0 : tensor<8x16x4xi8>
}
func.func private @uniform_quantize_1(%arg0: tensor<8x16x4xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x4xf32>) -> tensor<8x16x4xi8>
return %0 : tensor<8x16x4xi8>
}
func.func private @uniform_dequantize(%arg0: tensor<8x16x4xi8>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x4xf32> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x4xi8>) -> tensor<8x16x4xf32>
return %0 : tensor<8x16x4xf32>
}
}
// -----
// Tests that a quantized dot_general op is composed when both operands are
// activations, where input zero points are not folded into i32 constants.
// CHECK-LABEL: dot_general_with_two_activations
// CHECK-SAME: %[[ARG_0:.*]]: tensor<8x16x16xf32>
// CHECK-SAME: %[[ARG_1:.*]]: tensor<8x16x4xf32>
module {
func.func @dot_general_with_two_activations(%arg0: tensor<8x16x16xf32>, %arg1: tensor<8x16x4xf32>) -> tensor<8x16x4xf32> {
%1 = stablehlo.constant dense<2.000000e-01> : tensor<1x1x1xf32> // Input 1 inverse scale (1 / s1).
%2 = stablehlo.constant dense<-128> : tensor<1x1x1xi8> // Input 1 zero point (z1).
%3 = stablehlo.constant dense<4.000000e-01> : tensor<1x1x1xf32> // Input 2 inverse scale (1 / s2).
%4 = stablehlo.constant dense<0> : tensor<1x1x1xi8> // Input 2 zero point (z2).
%5 = stablehlo.constant dense<5.000000e-01> : tensor<1x1x1xf32> // Output inverse scale (1 / s3).
%6 = stablehlo.constant dense<-5> : tensor<1x1x1xi8> // Output zero point (z3).
%7 = stablehlo.constant dense<1.250000e+01> : tensor<1x1x1xf32> // Merged scale (s1 * s2).
%8 = call @uniform_quantize(%arg0, %1, %2) : (tensor<8x16x16xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x16xi8> // q1
%9 = call @uniform_quantize_0(%arg1, %3, %4) : (tensor<8x16x4xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xi8> // q2
%10 = stablehlo.convert %8 : (tensor<8x16x16xi8>) -> tensor<8x16x16xi32>
%11 = stablehlo.convert %2 : (tensor<1x1x1xi8>) -> tensor<1x1x1xi32>
%12 = stablehlo.broadcast_in_dim %11, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x16xi32>
%13 = stablehlo.subtract %10, %12 : tensor<8x16x16xi32> // q1 - z1
%14 = stablehlo.convert %9 : (tensor<8x16x4xi8>) -> tensor<8x16x4xi32>
%15 = stablehlo.convert %4 : (tensor<1x1x1xi8>) -> tensor<1x1x1xi32>
%16 = stablehlo.broadcast_in_dim %15, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x4xi32>
%17 = stablehlo.subtract %14, %16 : tensor<8x16x4xi32> // q2 - z2
// Corresponds to einsum expression: b i j, b j d -> b i d
%18 = stablehlo.dot_general %13, %17, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16xi32>, tensor<8x16x4xi32>) -> tensor<8x16x4xi32>
%19 = stablehlo.convert %18 : (tensor<8x16x4xi32>) -> tensor<8x16x4xf32>
%20 = stablehlo.broadcast_in_dim %7, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<8x16x4xf32>
%21 = stablehlo.multiply %19, %20 : tensor<8x16x4xf32> // * s1 s2
%22 = call @uniform_quantize_1(%21, %5, %6) : (tensor<8x16x4xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xi8>
%23 = call @uniform_dequantize(%22, %5, %6) : (tensor<8x16x4xi8>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xf32>
return %23 : tensor<8x16x4xf32>
}
// CHECK: %[[UQ_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<8x16x16xf32>) -> tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>
// CHECK: %[[UQ_1:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<8x16x4xf32>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[UQ_0]], %[[UQ_1]], batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>
// CHECK: %[[DQ_0:.*]] = stablehlo.uniform_dequantize %[[DOT_GENERAL]] : (tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>) -> tensor<8x16x4xf32>
// CHECK: return %[[DQ_0]]
// The following uniform_quantize & uniform_dequantize functions do NOT have
// the correct body. Only the type signatures matter for testing.
func.func private @uniform_quantize(%arg0: tensor<8x16x16xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x16xi8> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x16xf32>) -> tensor<8x16x16xi8>
return %0 : tensor<8x16x16xi8>
}
func.func private @uniform_quantize_0(%arg0: tensor<8x16x4xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x4xf32>) -> tensor<8x16x4xi8>
return %0 : tensor<8x16x4xi8>
}
func.func private @uniform_quantize_1(%arg0: tensor<8x16x4xf32>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x4xi8> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x4xf32>) -> tensor<8x16x4xi8>
return %0 : tensor<8x16x4xi8>
}
func.func private @uniform_dequantize(%arg0: tensor<8x16x4xi8>, %arg1: tensor<1x1x1xf32>, %arg2: tensor<1x1x1xi8>) -> tensor<8x16x4xf32> {
%0 = stablehlo.convert %arg0 : (tensor<8x16x4xi8>) -> tensor<8x16x4xf32>
return %0 : tensor<8x16x4xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/composite-lowering.mlir
+493
View File
@@ -0,0 +1,493 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt -composite-lowering -verify-diagnostics %s | FileCheck %s
func.func @hardswish(%arg0: tensor<2xf32>) -> (tensor<*xf32>) {
%0 = mhlo.composite "aten.hardswish.default" %arg0 {decomposition = @XlaCallModule_aten.hardswish.default.impl_0} : (tensor<2xf32>) -> tensor<2xf32>
%1 = "tf.Identity"(%0) {device = ""} : (tensor<2xf32>) -> tensor<*xf32>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
return %2 : tensor<*xf32>
}
func.func private @XlaCallModule_aten.hardswish.default.impl_0(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = mhlo.constant dense<6.000000e+00> : tensor<f32>
%1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<2xf32>
%2 = mhlo.constant dense<3.40282347E+38> : tensor<f32>
%3 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<2xf32>
%4 = mhlo.constant dense<3.000000e+00> : tensor<f32>
%5 = "mhlo.broadcast_in_dim"(%4) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<2xf32>
%6 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%7 = "mhlo.broadcast_in_dim"(%6) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<2xf32>
%8 = mhlo.constant dense<-3.40282347E+38> : tensor<f32>
%9 = "mhlo.broadcast_in_dim"(%8) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<2xf32>
%10 = mhlo.add %arg0, %5 : tensor<2xf32>
%11 = mhlo.clamp %7, %10, %3 : tensor<2xf32>
%12 = mhlo.clamp %9, %11, %1 : tensor<2xf32>
%13 = mhlo.multiply %arg0, %12 : tensor<2xf32>
%14 = mhlo.divide %13, %1 : tensor<2xf32>
return %14 : tensor<2xf32>
}
// CHECK-LABEL: func.func @hardswish(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<2xf32>) -> tensor<*xf32> {
// CHECK: %[[VAL_1:.*]] = "tfl.hard_swish"(%[[VAL_0]]) : (tensor<2xf32>) -> tensor<2xf32>
// CHECK: %[[VAL_2:.*]] = "tf.Identity"(%[[VAL_1]]) {device = ""} : (tensor<2xf32>) -> tensor<*xf32>
// CHECK: %[[VAL_3:.*]] = "tf.Identity"(%[[VAL_2]]) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
// CHECK: return %[[VAL_3]] : tensor<*xf32>
// CHECK: }
func.func @avg_pool2d_1(%arg0: tensor<1x3x6x6xf32>) -> (tensor<*xf32>) {
%0 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = false, count_include_pad = true, divisor_override = "py_None", kernel_size = dense<3> : tensor<2xi64>, padding = dense<0> : tensor<2xi64>, stride = dense<1> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_0} : (tensor<1x3x6x6xf32>) -> tensor<1x3x4x4xf32>
%1 = "tf.Identity"(%0) {device = ""} : (tensor<1x3x4x4xf32>) -> tensor<*xf32>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
return %2 : tensor<*xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_0(%arg0: tensor<1x3x6x6xf32>) -> tensor<1x3x4x4xf32> {
%0 = mhlo.constant dense<1.000000e+00> : tensor<f32>
%1 = "mhlo.broadcast_in_dim"(%0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<6x6xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.reduce_window"(%arg0, %2) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%7 = mhlo.add %arg1, %arg2 : tensor<f32>
mhlo.return %7 : tensor<f32>
}) {window_dimensions = dense<[1, 1, 3, 3]> : tensor<4xi64>} : (tensor<1x3x6x6xf32>, tensor<f32>) -> tensor<1x3x4x4xf32>
%4 = "mhlo.reduce_window"(%1, %2) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%7 = mhlo.add %arg1, %arg2 : tensor<f32>
mhlo.return %7 : tensor<f32>
}) {window_dimensions = dense<3> : tensor<2xi64>} : (tensor<6x6xf32>, tensor<f32>) -> tensor<4x4xf32>
%5 = "mhlo.broadcast_in_dim"(%4) {broadcast_dimensions = dense<[2, 3]> : tensor<2xi64>} : (tensor<4x4xf32>) -> tensor<1x3x4x4xf32>
%6 = mhlo.divide %3, %5 : tensor<1x3x4x4xf32>
return %6 : tensor<1x3x4x4xf32>
}
// CHECK-LABEL: func.func @avg_pool2d_1(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x3x6x6xf32>) -> tensor<*xf32> {
// CHECK: %[[VAL_1:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %[[VAL_2:.*]] = "tfl.transpose"(%[[VAL_0]], %[[VAL_1]]) : (tensor<1x3x6x6xf32>, tensor<4xi32>) -> tensor<1x6x6x3xf32>
// CHECK: %[[VAL_3:.*]] = arith.constant dense<0> : tensor<4x2xi32>
// CHECK: %[[VAL_4:.*]] = "tfl.pad"(%[[VAL_2]], %[[VAL_3]]) : (tensor<1x6x6x3xf32>, tensor<4x2xi32>) -> tensor<1x6x6x3xf32>
// CHECK: %[[VAL_5:.*]] = "tfl.average_pool_2d"(%[[VAL_4]]) <{filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x6x6x3xf32>) -> tensor<1x4x4x3xf32>
// CHECK: %[[VAL_6:.*]] = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %[[VAL_7:.*]] = "tfl.transpose"(%[[VAL_5]], %[[VAL_6]]) : (tensor<1x4x4x3xf32>, tensor<4xi32>) -> tensor<1x3x4x4xf32>
// CHECK: %[[VAL_8:.*]] = "tf.Identity"(%[[VAL_7]]) {device = ""} : (tensor<1x3x4x4xf32>) -> tensor<*xf32>
// CHECK: %[[VAL_9:.*]] = "tf.Identity"(%[[VAL_8]]) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
// CHECK: return %[[VAL_9]] : tensor<*xf32>
func.func @avg_pool2d_2(%arg0: tensor<1x3x6x6xf32>) -> (tensor<*xf32>) {
%0 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = false, count_include_pad = false, divisor_override = "py_None", kernel_size = dense<3> : tensor<2xi64>, padding = dense<1> : tensor<2xi64>, stride = dense<1> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_1} : (tensor<1x3x6x6xf32>) -> tensor<1x3x6x6xf32>
%1 = "tf.Identity"(%0) {device = ""} : (tensor<1x3x6x6xf32>) -> tensor<*xf32>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
return %2 : tensor<*xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_1(%arg0: tensor<1x3x6x6xf32>) -> tensor<1x3x6x6xf32> {
%0 = mhlo.constant dense<[[0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00], [0.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00], [0.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00], [0.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00], [0.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00], [0.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00], [0.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 0.000000e+00], [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]]> : tensor<8x8xf32>
%1 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%2 = "mhlo.pad"(%arg0, %1) {edge_padding_high = dense<[0, 0, 1, 1]> : tensor<4xi64>, edge_padding_low = dense<[0, 0, 1, 1]> : tensor<4xi64>, interior_padding = dense<0> : tensor<4xi64>} : (tensor<1x3x6x6xf32>, tensor<f32>) -> tensor<1x3x8x8xf32>
%3 = "mhlo.reduce_window"(%2, %1) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%7 = mhlo.add %arg1, %arg2 : tensor<f32>
mhlo.return %7 : tensor<f32>
}) {window_dimensions = dense<[1, 1, 3, 3]> : tensor<4xi64>} : (tensor<1x3x8x8xf32>, tensor<f32>) -> tensor<1x3x6x6xf32>
%4 = "mhlo.reduce_window"(%0, %1) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%7 = mhlo.add %arg1, %arg2 : tensor<f32>
mhlo.return %7 : tensor<f32>
}) {window_dimensions = dense<3> : tensor<2xi64>} : (tensor<8x8xf32>, tensor<f32>) -> tensor<6x6xf32>
%5 = "mhlo.broadcast_in_dim"(%4) {broadcast_dimensions = dense<[2, 3]> : tensor<2xi64>} : (tensor<6x6xf32>) -> tensor<1x3x6x6xf32>
%6 = mhlo.divide %3, %5 : tensor<1x3x6x6xf32>
return %6 : tensor<1x3x6x6xf32>
}
// CHECK-LABEL: func.func @avg_pool2d_2(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x3x6x6xf32>) -> tensor<*xf32> {
// CHECK: %[[VAL_1:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %[[VAL_2:.*]] = "tfl.transpose"(%[[VAL_0]], %[[VAL_1]]) : (tensor<1x3x6x6xf32>, tensor<4xi32>) -> tensor<1x6x6x3xf32>
// CHECK: %[[VAL_3:.*]] = "tfl.average_pool_2d"(%[[VAL_2]]) <{filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x6x6x3xf32>) -> tensor<1x6x6x3xf32>
// CHECK: %[[VAL_4:.*]] = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %[[VAL_5:.*]] = "tfl.transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor<1x6x6x3xf32>, tensor<4xi32>) -> tensor<1x3x6x6xf32>
// CHECK: %[[VAL_6:.*]] = "tf.Identity"(%[[VAL_5]]) {device = ""} : (tensor<1x3x6x6xf32>) -> tensor<*xf32>
// CHECK: %[[VAL_7:.*]] = "tf.Identity"(%[[VAL_6]]) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
// CHECK: return %[[VAL_7]] : tensor<*xf32>
func.func @avg_pool2d_3(%arg0: tensor<1x1x1x8xf32>) -> (tensor<1x1x1x4xf32>) {
%2 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = false, count_include_pad = true, divisor_override = "py_None", kernel_size = dense<[1, 3]> : tensor<2xi64>, padding = dense<[0, 1]> : tensor<2xi64>, stride = dense<[1, 2]> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_2} : (tensor<1x1x1x8xf32>) -> tensor<1x1x1x4xf32>
return %2 : tensor<1x1x1x4xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_2(%arg0: tensor<1x1x1x8xf32>) -> tensor<1x1x1x4xf32>
// CHECK-LABEL: avg_pool2d_3
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x8xf32>, tensor<4xi32>) -> tensor<1x1x8x1xf32>
// CHECK{LITERAL}: %cst_0 = arith.constant dense<[[0, 0], [0, 0], [1, 1], [0, 0]]> : tensor<4x2xi32>
// CHECK: %1 = "tfl.pad"(%0, %cst_0) : (tensor<1x1x8x1xf32>, tensor<4x2xi32>) -> tensor<1x1x10x1xf32>
// CHECK: %2 = "tfl.average_pool_2d"(%1) <{filter_height = 1 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 2 : i32}> : (tensor<1x1x10x1xf32>) -> tensor<1x1x4x1xf32>
// CHECK: %cst_1 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %3 = "tfl.transpose"(%2, %cst_1) : (tensor<1x1x4x1xf32>, tensor<4xi32>) -> tensor<1x1x1x4xf32>
// CHECK: return %3 : tensor<1x1x1x4xf32>
func.func @avg_pool2d_4(%arg0: tensor<1x1x1x9xf32>) -> (tensor<1x1x1x4xf32>) {
%2 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = false, count_include_pad = false, divisor_override = "py_None", kernel_size = dense<[1, 3]> : tensor<2xi64>, padding = dense<[0, 0]> : tensor<2xi64>, stride = dense<[1, 2]> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_3} : (tensor<1x1x1x9xf32>) -> tensor<1x1x1x4xf32>
return %2 : tensor<1x1x1x4xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_3(%arg0: tensor<1x1x1x9xf32>) -> tensor<1x1x1x4xf32>
// CHECK-LABEL: avg_pool2d_4
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x9xf32>, tensor<4xi32>) -> tensor<1x1x9x1xf32>
// CHECK: %1 = "tfl.average_pool_2d"(%0) <{filter_height = 1 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 2 : i32}> : (tensor<1x1x9x1xf32>) -> tensor<1x1x4x1xf32>
// CHECK: %cst_0 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %2 = "tfl.transpose"(%1, %cst_0) : (tensor<1x1x4x1xf32>, tensor<4xi32>) -> tensor<1x1x1x4xf32>
// CHECK: return %2 : tensor<1x1x1x4xf32>
func.func @avg_pool2d_5(%arg0: tensor<1x1x3x3xf32>) -> (tensor<1x1x2x2xf32>) {
%0 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = true, count_include_pad = true, divisor_override = "py_None", kernel_size = dense<[2, 2]> : tensor<2xi64>, padding = dense<[0, 0]> : tensor<2xi64>, stride = dense<[2, 2]> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_4} : (tensor<1x1x3x3xf32>) -> tensor<1x1x2x2xf32>
return %0 : tensor<1x1x2x2xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_4(%arg0: tensor<1x1x3x3xf32>) -> tensor<1x1x2x2xf32>
// CHECK-LABEL: avg_pool2d_5
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x3x3xf32>, tensor<4xi32>) -> tensor<1x3x3x1xf32>
// CHECK{LITERAL}: %cst_0 = arith.constant dense<[[0, 0], [0, 1], [0, 1], [0, 0]]> : tensor<4x2xi32>
// CHECK: %1 = "tfl.pad"(%0, %cst_0) : (tensor<1x3x3x1xf32>, tensor<4x2xi32>) -> tensor<1x4x4x1xf32>
// CHECK: %2 = "tfl.average_pool_2d"(%1) <{filter_height = 2 : i32, filter_width = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32>
// CHECK{LITERAL}: %cst_1 = arith.constant dense<[[[[1.000000e+00], [2.000000e+00]], [[2.000000e+00], [4.000000e+00]]]]> : tensor<1x2x2x1xf32>
// CHECK: %3 = tfl.mul %2, %cst_1 {fused_activation_function = "NONE"} : tensor<1x2x2x1xf32>
// CHECK: %cst_2 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %4 = "tfl.transpose"(%3, %cst_2) : (tensor<1x2x2x1xf32>, tensor<4xi32>) -> tensor<1x1x2x2xf32>
// CHECK: return %4 : tensor<1x1x2x2xf32>
func.func @avg_pool2d_6(%arg0: tensor<1x1x1x7xf32>) -> (tensor<1x1x1x2xf32>) {
%0 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = true, count_include_pad = true, divisor_override = "py_None", kernel_size = dense<[1, 5]> : tensor<2xi64>, padding = dense<[0, 0]> : tensor<2xi64>, stride = dense<[1, 3]> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_5} : (tensor<1x1x1x7xf32>) -> tensor<1x1x1x2xf32>
return %0 : tensor<1x1x1x2xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_5(%arg0: tensor<1x1x1x7xf32>) -> tensor<1x1x1x2xf32>
// CHECK-LABEL: avg_pool2d_6
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x7xf32>, tensor<4xi32>) -> tensor<1x1x7x1xf32>
// CHECK{LITERAL}: %cst_0 = arith.constant dense<[[0, 0], [0, 0], [0, 1], [0, 0]]> : tensor<4x2xi32>
// CHECK: %1 = "tfl.pad"(%0, %cst_0) : (tensor<1x1x7x1xf32>, tensor<4x2xi32>) -> tensor<1x1x8x1xf32>
// CHECK: %2 = "tfl.average_pool_2d"(%1) <{filter_height = 1 : i32, filter_width = 5 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 3 : i32}> : (tensor<1x1x8x1xf32>) -> tensor<1x1x2x1xf32>
// CHECK{LITERAL}: %cst_1 = arith.constant dense<[[[[1.000000e+00], [1.250000e+00]]]]> : tensor<1x1x2x1xf32>
// CHECK: %3 = tfl.mul %2, %cst_1 {fused_activation_function = "NONE"} : tensor<1x1x2x1xf32>
// CHECK: %cst_2 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %4 = "tfl.transpose"(%3, %cst_2) : (tensor<1x1x2x1xf32>, tensor<4xi32>) -> tensor<1x1x1x2xf32>
func.func @avg_pool2d_7(%arg0: tensor<1x1x1x8xf32>) -> (tensor<1x1x1x5xf32>) {
%0 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = true, count_include_pad = true, divisor_override = "py_None", kernel_size = dense<[1, 3]> : tensor<2xi64>, padding = dense<[0, 1]> : tensor<2xi64>, stride = dense<[1, 2]> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_6} : (tensor<1x1x1x8xf32>) -> tensor<1x1x1x5xf32>
return %0 : tensor<1x1x1x5xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_6(%arg0: tensor<1x1x1x8xf32>) -> tensor<1x1x1x5xf32>
// CHECK-LABEL: avg_pool2d_7
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK{LITERAL}: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x8xf32>, tensor<4xi32>) -> tensor<1x1x8x1xf32>
// CHECK{LITERAL}: %cst_0 = arith.constant dense<[[0, 0], [0, 0], [1, 2], [0, 0]]> : tensor<4x2xi32>
// CHECK: %1 = "tfl.pad"(%0, %cst_0) : (tensor<1x1x8x1xf32>, tensor<4x2xi32>) -> tensor<1x1x11x1xf32>
// CHECK: %2 = "tfl.average_pool_2d"(%1) <{filter_height = 1 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 2 : i32}> : (tensor<1x1x11x1xf32>) -> tensor<1x1x5x1xf32>
// CHECK{LITERAL}: %cst_1 = arith.constant dense<[[[[1.000000e+00], [1.000000e+00], [1.000000e+00], [1.000000e+00], [1.500000e+00]]]]> : tensor<1x1x5x1xf32>
// CHECK: %3 = tfl.mul %2, %cst_1 {fused_activation_function = "NONE"} : tensor<1x1x5x1xf32>
// CHECK: %cst_2 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %4 = "tfl.transpose"(%3, %cst_2) : (tensor<1x1x5x1xf32>, tensor<4xi32>) -> tensor<1x1x1x5xf32>
func.func @upsample_bilinear2d(%arg0: tensor<1x64x16x16xf32>) -> (tensor<1x64x32x32xf32>) {
%0 = mhlo.composite "odml.upsample_bilinear2d" %arg0 {composite_attributes = {is_nchw_op = true, align_corners = false, size = dense<32> : tensor<2xi64>}, decomposition = @XlaCallModule_odml.upsample_bilinear2d.impl_21_0} : (tensor<1x64x16x16xf32>) -> tensor<1x64x32x32xf32>
return %0 : tensor<1x64x32x32xf32>
}
func.func private @XlaCallModule_odml.upsample_bilinear2d.impl_21_0(%arg0: tensor<1x64x16x16xf32>) -> tensor<1x64x32x32xf32> {
%0 = mhlo.constant dense<[[0.000000e+00], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01], [7.500000e-01], [2.500000e-01]]> : tensor<32x1xf32>
%1 = mhlo.constant dense<[0.000000e+00, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01, 7.500000e-01, 2.500000e-01]> : tensor<32xf32>
%2 = mhlo.constant dense<[1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13, 14, 14, 15, 15, 15]> : tensor<32xi64>
%3 = mhlo.constant dense<16> : tensor<i64>
%4 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i64>) -> tensor<32x32xi64>
%5 = mhlo.constant dense<0> : tensor<i64>
%6 = "mhlo.broadcast_in_dim"(%5) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i64>) -> tensor<32x32xi64>
%7 = mhlo.constant dense<[0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13, 14, 14, 15]> : tensor<32xi64>
%8 = "mhlo.broadcast_in_dim"(%7) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<32xi64>) -> tensor<32x32xi64>
%9 = mhlo.compare LT, %8, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
%10 = mhlo.add %8, %4 : tensor<32x32xi64>
%11 = mhlo.select %9, %10, %8 : tensor<32x32xi1>, tensor<32x32xi64>
%12 = mhlo.reshape %11 : (tensor<32x32xi64>) -> tensor<32x32x1xi64>
%13 = "mhlo.broadcast_in_dim"(%7) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<32xi64>) -> tensor<32x32xi64>
%14 = mhlo.compare LT, %13, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
%15 = mhlo.add %13, %4 : tensor<32x32xi64>
%16 = mhlo.select %14, %15, %13 : tensor<32x32xi1>, tensor<32x32xi64>
%17 = mhlo.reshape %16 : (tensor<32x32xi64>) -> tensor<32x32x1xi64>
%18 = "mhlo.concatenate"(%12, %17) <{dimension = 2 : i64}> : (tensor<32x32x1xi64>, tensor<32x32x1xi64>) -> tensor<32x32x2xi64>
%19 = "mhlo.gather"(%arg0, %18) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1], collapsed_slice_dims = [2, 3], start_index_map = [2, 3], index_vector_dim = 2>, slice_sizes = dense<[1, 64, 1, 1]> : tensor<4xi64>}> : (tensor<1x64x16x16xf32>, tensor<32x32x2xi64>) -> tensor<1x64x32x32xf32>
%20 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<32xi64>) -> tensor<32x32xi64>
%21 = mhlo.compare LT, %20, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
%22 = mhlo.add %20, %4 : tensor<32x32xi64>
%23 = mhlo.select %21, %22, %20 : tensor<32x32xi1>, tensor<32x32xi64>
%24 = mhlo.reshape %23 : (tensor<32x32xi64>) -> tensor<32x32x1xi64>
%25 = "mhlo.concatenate"(%12, %24) <{dimension = 2 : i64}> : (tensor<32x32x1xi64>, tensor<32x32x1xi64>) -> tensor<32x32x2xi64>
%26 = "mhlo.gather"(%arg0, %25) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1], collapsed_slice_dims = [2, 3], start_index_map = [2, 3], index_vector_dim = 2>, slice_sizes = dense<[1, 64, 1, 1]> : tensor<4xi64>}> : (tensor<1x64x16x16xf32>, tensor<32x32x2xi64>) -> tensor<1x64x32x32xf32>
%27 = mhlo.subtract %26, %19 : tensor<1x64x32x32xf32>
%28 = "mhlo.broadcast_in_dim"(%1) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<32xf32>) -> tensor<1x64x32x32xf32>
%29 = mhlo.multiply %27, %28 : tensor<1x64x32x32xf32>
%30 = mhlo.add %19, %29 : tensor<1x64x32x32xf32>
%31 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<32xi64>) -> tensor<32x32xi64>
%32 = mhlo.compare LT, %31, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
%33 = mhlo.add %31, %4 : tensor<32x32xi64>
%34 = mhlo.select %32, %33, %31 : tensor<32x32xi1>, tensor<32x32xi64>
%35 = mhlo.reshape %34 : (tensor<32x32xi64>) -> tensor<32x32x1xi64>
%36 = "mhlo.concatenate"(%35, %17) <{dimension = 2 : i64}> : (tensor<32x32x1xi64>, tensor<32x32x1xi64>) -> tensor<32x32x2xi64>
%37 = "mhlo.gather"(%arg0, %36) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1], collapsed_slice_dims = [2, 3], start_index_map = [2, 3], index_vector_dim = 2>, slice_sizes = dense<[1, 64, 1, 1]> : tensor<4xi64>}> : (tensor<1x64x16x16xf32>, tensor<32x32x2xi64>) -> tensor<1x64x32x32xf32>
%38 = "mhlo.concatenate"(%35, %24) <{dimension = 2 : i64}> : (tensor<32x32x1xi64>, tensor<32x32x1xi64>) -> tensor<32x32x2xi64>
%39 = "mhlo.gather"(%arg0, %38) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1], collapsed_slice_dims = [2, 3], start_index_map = [2, 3], index_vector_dim = 2>, slice_sizes = dense<[1, 64, 1, 1]> : tensor<4xi64>}> : (tensor<1x64x16x16xf32>, tensor<32x32x2xi64>) -> tensor<1x64x32x32xf32>
%40 = mhlo.subtract %39, %37 : tensor<1x64x32x32xf32>
%41 = mhlo.multiply %40, %28 : tensor<1x64x32x32xf32>
%42 = mhlo.add %37, %41 : tensor<1x64x32x32xf32>
%43 = mhlo.subtract %42, %30 : tensor<1x64x32x32xf32>
%44 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<[2, 3]> : tensor<2xi64>}> : (tensor<32x1xf32>) -> tensor<1x64x32x1xf32>
%45 = mhlo.reshape %44 : (tensor<1x64x32x1xf32>) -> tensor<1x64x32xf32>
%46 = "mhlo.broadcast_in_dim"(%45) <{broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}> : (tensor<1x64x32xf32>) -> tensor<1x64x32x32xf32>
%47 = mhlo.multiply %43, %46 : tensor<1x64x32x32xf32>
%48 = mhlo.add %30, %47 : tensor<1x64x32x32xf32>
return %48 : tensor<1x64x32x32xf32>
}
// CHECK-LABEL: func.func @upsample_bilinear2d(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x64x16x16xf32>) -> tensor<1x64x32x32xf32> {
// CHECK-DAG: %[[VAL_1:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_2:.*]] = "tfl.transpose"(%[[VAL_0]], %[[VAL_1]]) : (tensor<1x64x16x16xf32>, tensor<4xi32>) -> tensor<1x16x16x64xf32>
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<32> : tensor<2xi32>
// CHECK-DAG: %[[VAL_4:.*]] = "tfl.resize_bilinear"(%[[VAL_2]], %[[VAL_3]]) <{align_corners = false, half_pixel_centers = true}> : (tensor<1x16x16x64xf32>, tensor<2xi32>) -> tensor<1x32x32x64xf32>
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK-DAG: %[[VAL_6:.*]] = "tfl.transpose"(%[[VAL_4]], %[[VAL_5]]) : (tensor<1x32x32x64xf32>, tensor<4xi32>) -> tensor<1x64x32x32xf32>
// CHECK: return %[[VAL_6]] : tensor<1x64x32x32xf32>
// CHECK: }
func.func private @XlaCallModule_tfl.gelu.impl_0(%arg0: tensor<1x4x4x1xf32>) -> (tensor<1x4x4x1xf32>)
func.func @jax_gelu_approx(%arg0: tensor<1x4x4x1xf32>) -> (tensor<1x4x4x1xf32>) {
%2 = mhlo.composite "tfl.gelu" %arg0 {composite_attributes = {approximate = true}, decomposition = @XlaCallModule_tfl.gelu.impl_0} : (tensor<1x4x4x1xf32>) -> tensor<1x4x4x1xf32>
return %2 : tensor<1x4x4x1xf32>
}
// CHECK-LABEL: jax_gelu_approx
// CHECK: %0 = "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<1x4x4x1xf32>) -> tensor<1x4x4x1xf32>
func.func private @XlaCallModule_tfl.gelu.impl_1(%arg0: tensor<1x4x4x1xf32>) -> (tensor<1x4x4x1xf32>)
func.func @jax_gelu(%arg0: tensor<1x4x4x1xf32>) -> (tensor<1x4x4x1xf32>) {
%2 = mhlo.composite "tfl.gelu" %arg0 {composite_attributes = {approximate = false}, decomposition = @XlaCallModule_tfl.gelu.impl_1} : (tensor<1x4x4x1xf32>) -> tensor<1x4x4x1xf32>
return %2 : tensor<1x4x4x1xf32>
}
// CHECK-LABEL: jax_gelu
// CHECK: %0 = "tfl.gelu"(%arg0) <{approximate = false}> : (tensor<1x4x4x1xf32>) -> tensor<1x4x4x1xf32>
func.func private @gelu_decomp_1(%arg0: tensor<5x10xf32>) -> tensor<5x10xf32>
func.func @gelu_aten(%arg0: tensor<5x10xf32>) -> (tensor<*xf32>) {
%0 = mhlo.composite "aten.gelu.default" %arg0 {composite_attributes = {approximate = "none"}, decomposition = @gelu_decomp_1} : (tensor<5x10xf32>) -> tensor<5x10xf32>
%1 = "tf.Identity"(%0) {device = ""} : (tensor<5x10xf32>) -> tensor<*xf32>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
return %2 : tensor<*xf32>
}
// CHECK-LABEL: gelu_aten
// CHECK: %0 = "tfl.gelu"(%arg0) <{approximate = false}> : (tensor<5x10xf32>) -> tensor<5x10xf32>
func.func private @gelu_decomp_2(%arg0: tensor<5x10xf32>) -> tensor<5x10xf32>
func.func @gelu_aten_approximate(%arg0: tensor<5x10xf32>) -> (tensor<*xf32>) {
%0 = mhlo.composite "aten.gelu.default" %arg0 {composite_attributes = {approximate = "tanh"}, decomposition = @gelu_decomp_2} : (tensor<5x10xf32>) -> tensor<5x10xf32>
%1 = "tf.Identity"(%0) {device = ""} : (tensor<5x10xf32>) -> tensor<*xf32>
%2 = "tf.Identity"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
return %2 : tensor<*xf32>
}
// CHECK-LABEL: gelu_aten_approximate
// CHECK: %0 = "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<5x10xf32>) -> tensor<5x10xf32>
// CHECK-LABEL func.func @jax_image_resize_nearest
func.func @jax_image_resize_nearest(%arg0: tensor<1x2x2x10xf32>) -> (tensor<1x4x4x10xf32>) {
%1 = mhlo.composite "tfl.resize_nearest_neighbor" %arg0 {composite_attributes = {is_nchw_op = false, size = dense<4> : tensor<2xi64>}, decomposition = @XlaCallModule_tfl.resize_nearest_neighbor.impl_0} : (tensor<1x2x2x10xf32>) -> tensor<1x4x4x10xf32>
return %1 : tensor<1x4x4x10xf32>
}
func.func private @XlaCallModule_tfl.resize_nearest_neighbor.impl_0(%arg0: tensor<1x2x2x10xf32>) -> tensor<1x4x4x10xf32> {
%0 = call @XlaCallModule__resize_0(%arg0) : (tensor<1x2x2x10xf32>) -> tensor<1x4x4x10xf32>
return %0 : tensor<1x4x4x10xf32>
}
func.func private @XlaCallModule__resize_0(%arg0: tensor<1x2x2x10xf32>) -> (tensor<1x4x4x10xf32>) {
%0 = mhlo.constant dense<2> : tensor<i32>
%1 = mhlo.constant dense<0> : tensor<i32>
%2 = mhlo.constant dense<4.000000e+00> : tensor<f32>
%3 = mhlo.constant dense<2.000000e+00> : tensor<f32>
%4 = mhlo.constant dense<5.000000e-01> : tensor<f32>
%5 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<4xf32>
%6 = "mhlo.broadcast_in_dim"(%4) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
%7 = mhlo.add %5, %6 : tensor<4xf32>
%8 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
%9 = mhlo.multiply %7, %8 : tensor<4xf32>
%10 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
%11 = mhlo.divide %9, %10 : tensor<4xf32>
%12 = mhlo.floor %11 : tensor<4xf32>
%13 = mhlo.convert %12 : (tensor<4xf32>) -> tensor<4xi32>
%14 = "mhlo.broadcast_in_dim"(%1) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
%15 = mhlo.compare LT, %13, %14, SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
%16 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
%17 = mhlo.add %13, %16 : tensor<4xi32>
%18 = mhlo.select %15, %17, %13 : tensor<4xi1>, tensor<4xi32>
%19 = "mhlo.broadcast_in_dim"(%18) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<4x1xi32>
%20 = "mhlo.gather"(%arg0, %19) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 2, 3], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>, slice_sizes = dense<[1, 1, 2, 10]> : tensor<4xi64>}> : (tensor<1x2x2x10xf32>, tensor<4x1xi32>) -> tensor<1x4x2x10xf32>
%21 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<4xf32>
%22 = "mhlo.broadcast_in_dim"(%4) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
%23 = mhlo.add %21, %22 : tensor<4xf32>
%24 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
%25 = mhlo.multiply %23, %24 : tensor<4xf32>
%26 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
%27 = mhlo.divide %25, %26 : tensor<4xf32>
%28 = mhlo.floor %27 : tensor<4xf32>
%29 = mhlo.convert %28 : (tensor<4xf32>) -> tensor<4xi32>
%30 = "mhlo.broadcast_in_dim"(%1) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
%31 = mhlo.compare LT, %29, %30, SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
%32 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
%33 = mhlo.add %29, %32 : tensor<4xi32>
%34 = mhlo.select %31, %33, %29 : tensor<4xi1>, tensor<4xi32>
%35 = "mhlo.broadcast_in_dim"(%34) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<4x1xi32>
%36 = "mhlo.gather"(%20, %35) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1, 3], collapsed_slice_dims = [2], start_index_map = [2], index_vector_dim = 1>, slice_sizes = dense<[1, 4, 1, 10]> : tensor<4xi64>}> : (tensor<1x4x2x10xf32>, tensor<4x1xi32>) -> tensor<1x4x4x10xf32>
return %36 : tensor<1x4x4x10xf32>
}
// CHECK: %cst = arith.constant dense<4> : tensor<2xi32>
// CHECK: %0 = "tfl.resize_nearest_neighbor"(%arg0, %cst) <{align_corners = false, half_pixel_centers = true}> : (tensor<1x2x2x10xf32>, tensor<2xi32>) -> tensor<1x4x4x10xf32>
// CHECK: return %0 : tensor<1x4x4x10xf32>
// CHECK-LABEL func.func @jax_image_resize_bilinear
func.func @jax_image_resize_bilinear(%arg0: tensor<1x44x44x128xf32>) -> (tensor<1x88x88x128xf32> {jax.result_info = ""}) {
%1 = mhlo.composite "odml.upsample_bilinear2d" %arg0 {composite_attributes = {is_nchw_op = false, align_corners = false, size = dense<88> : tensor<2xi64>}, decomposition = @XlaCallModule_tfl.resize_bilinear.impl_0} : (tensor<1x44x44x128xf32>) -> tensor<1x88x88x128xf32>
return %1 : tensor<1x88x88x128xf32>
}
func.func private @XlaCallModule_tfl.resize_bilinear.impl_0(%arg0: tensor<1x44x44x128xf32>) -> tensor<1x88x88x128xf32> {
%0 = call @XlaCallModule__resize_bilinear_0(%arg0) : (tensor<1x44x44x128xf32>) -> tensor<1x88x88x128xf32>
return %0 : tensor<1x88x88x128xf32>
}
func.func private @XlaCallModule__resize_bilinear_0(%arg0: tensor<1x44x44x128xf32>) -> tensor<1x88x88x128xf32> {
%cst = arith.constant dense<88> : tensor<2xi32>
// Because the decomposition desn't matter here, we just use the
// resize_bilinear op to create a correct graph.
%1 = "tfl.resize_bilinear"(%arg0, %cst) <{align_corners = false, half_pixel_centers = true}> : (tensor<1x44x44x128xf32>, tensor<2xi32>) -> tensor<1x88x88x128xf32>
return %1 : tensor<1x88x88x128xf32>
}
// CHECK: %cst = arith.constant dense<88> : tensor<2xi32>
// CHECK: %0 = "tfl.resize_bilinear"(%arg0, %cst) <{align_corners = false, half_pixel_centers = true}> : (tensor<1x44x44x128xf32>, tensor<2xi32>) -> tensor<1x88x88x128xf32>
// CHECK: return %0 : tensor<1x88x88x128xf32>
// CHECK-LABEL func.func @jax_image_resize_nearest_nchw
func.func @jax_image_resize_nearest_nchw(%arg0: tensor<4x8x32x32xf32>) -> (tensor<4x8x64x64xf32>) {
%0 = call @XlaCallModule_tfl.resize_nearest_neighbor.impl_1(%arg0) : (tensor<4x8x32x32xf32>) -> tensor<4x8x64x64xf32>
%1 = mhlo.composite "tfl.resize_nearest_neighbor" %arg0 {composite_attributes = {is_nchw_op = true, size = dense<64> : tensor<2xi64>}, decomposition = @XlaCallModule_tfl.resize_nearest_neighbor.impl_1} : (tensor<4x8x32x32xf32>) -> tensor<4x8x64x64xf32>
return %1 : tensor<4x8x64x64xf32>
}
func.func private @XlaCallModule_tfl.resize_nearest_neighbor.impl_1(%arg0: tensor<4x8x32x32xf32>) -> tensor<4x8x64x64xf32> {
%0 = call @XlaCallModule__resize_1(%arg0) : (tensor<4x8x32x32xf32>) -> tensor<4x8x64x64xf32>
return %0 : tensor<4x8x64x64xf32>
}
func.func private @XlaCallModule__resize_1(%arg0: tensor<4x8x32x32xf32>) -> (tensor<4x8x64x64xf32>) {
%0 = mhlo.constant dense<32> : tensor<64xi32>
%1 = mhlo.constant dense<0> : tensor<64xi32>
%2 = mhlo.constant dense<6.400000e+01> : tensor<64xf32>
%3 = mhlo.constant dense<3.200000e+01> : tensor<64xf32>
%4 = mhlo.constant dense<5.000000e-01> : tensor<64xf32>
%5 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<64xf32>
%6 = mhlo.add %5, %4 : tensor<64xf32>
%7 = mhlo.multiply %6, %3 : tensor<64xf32>
%8 = mhlo.divide %7, %2 : tensor<64xf32>
%9 = mhlo.floor %8 : tensor<64xf32>
%10 = mhlo.convert %9 : (tensor<64xf32>) -> tensor<64xi32>
%11 = mhlo.compare LT, %10, %1, SIGNED : (tensor<64xi32>, tensor<64xi32>) -> tensor<64xi1>
%12 = mhlo.add %10, %0 : tensor<64xi32>
%13 = mhlo.select %11, %12, %10 : tensor<64xi1>, tensor<64xi32>
%14 = mhlo.reshape %13 : (tensor<64xi32>) -> tensor<64x1xi32>
%15 = "mhlo.gather"(%arg0, %14) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1, 3], collapsed_slice_dims = [2], start_index_map = [2], index_vector_dim = 1>, slice_sizes = dense<[4, 8, 1, 32]> : tensor<4xi64>}> : (tensor<4x8x32x32xf32>, tensor<64x1xi32>) -> tensor<4x8x64x32xf32>
%16 = "mhlo.gather"(%15, %14) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1, 2], collapsed_slice_dims = [3], start_index_map = [3], index_vector_dim = 1>, slice_sizes = dense<[4, 8, 64, 1]> : tensor<4xi64>}> : (tensor<4x8x64x32xf32>, tensor<64x1xi32>) -> tensor<4x8x64x64xf32>
return %16 : tensor<4x8x64x64xf32>
}
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %1 = "tfl.transpose"(%arg0, %cst) : (tensor<4x8x32x32xf32>, tensor<4xi32>) -> tensor<4x32x32x8xf32>
// CHECK: %cst_0 = arith.constant dense<64> : tensor<2xi32>
// CHECK: %2 = "tfl.resize_nearest_neighbor"(%1, %cst_0) <{align_corners = false, half_pixel_centers = true}> : (tensor<4x32x32x8xf32>, tensor<2xi32>) -> tensor<4x64x64x8xf32>
// CHECK: %cst_1 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %3 = "tfl.transpose"(%2, %cst_1) : (tensor<4x64x64x8xf32>, tensor<4xi32>) -> tensor<4x8x64x64xf32>
// CHECK: return %3 : tensor<4x8x64x64xf32>
func.func @embedding_lookup(%arg0: tensor<1xi32>, %arg1: tensor<32000x2048xf32>) -> tensor<1x2048xf32> {
%0 = mhlo.composite "odml.embedding_lookup" %arg0, %arg1 {decomposition = @XlaCallModule_odml.embedding_lookup.impl_0} : (tensor<1xi32>, tensor<32000x2048xf32>) -> tensor<1x2048xf32>
return %0 : tensor<1x2048xf32>
}
func.func private @XlaCallModule_odml.embedding_lookup.impl_0(%arg0: tensor<1xi32>, %arg1: tensor<32000x2048xf32>) -> tensor<1x2048xf32> {
%0 = "mhlo.gather"(%arg1, %arg0) <{dimension_numbers = #mhlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, slice_sizes = dense<[1, 2048]> : tensor<2xi64>}> : (tensor<32000x2048xf32>, tensor<1xi32>) -> tensor<1x2048xf32>
return %0 : tensor<1x2048xf32>
}
// CHECK-LABEL: func.func @embedding_lookup(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1xi32>, %[[ARG_1:.*]]: tensor<32000x2048xf32>) -> tensor<1x2048xf32> {
// CHECK: %[[VAL_1:.*]] = "tfl.embedding_lookup"(%[[ARG_0]], %[[ARG_1]]) : (tensor<1xi32>, tensor<32000x2048xf32>) -> tensor<1x2048xf32>
// CHECK: return %[[VAL_1]] : tensor<1x2048xf32>
// CHECK: }
func.func @embedding_lookup_dynamic(%arg0: tensor<1xi32>, %arg1: tensor<32000x2048xf32>, %arg2: tensor<i32>) -> tensor<1x2048xf32> {
%0 = mhlo.composite "odml.embedding_lookup" %arg2, %arg0, %arg1 {decomposition = @XlaCallModule_odml.embedding_lookup.impl_1} : (tensor<i32>, tensor<1xi32>, tensor<32000x2048xf32>) -> tensor<1x2048xf32>
return %0 : tensor<1x2048xf32>
}
func.func private @XlaCallModule_odml.embedding_lookup.impl_1(%arg2: tensor<i32>, %arg0: tensor<1xi32>, %arg1: tensor<32000x2048xf32>) -> tensor<1x2048xf32> {
%0 = "mhlo.gather"(%arg1, %arg0) <{dimension_numbers = #mhlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, slice_sizes = dense<[1, 2048]> : tensor<2xi64>}> : (tensor<32000x2048xf32>, tensor<1xi32>) -> tensor<1x2048xf32>
return %0 : tensor<1x2048xf32>
}
// CHECK-LABEL: func.func @embedding_lookup_dynamic(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1xi32>, %[[ARG_1:.*]]: tensor<32000x2048xf32>, %[[ARG_2:.*]]: tensor<i32>) -> tensor<1x2048xf32> {
// CHECK: %[[VAL_1:.*]] = "tfl.embedding_lookup"(%[[ARG_0]], %[[ARG_1]]) : (tensor<1xi32>, tensor<32000x2048xf32>) -> tensor<1x2048xf32>
// CHECK: return %[[VAL_1]] : tensor<1x2048xf32>
// CHECK: }
func.func @random_uniform(%arg0: tensor<3xi32>) -> tensor<1x2x3xf32> {
%0 = mhlo.composite "odml.random_uniform" %arg0 {composite_attributes = {seed = 0 : i64, seed2 = 1: i64}, decomposition = @XlaCallModule_odml.random_uniform.impl_0} : (tensor<3xi32>) -> tensor<1x2x3xf32>
return %0 : tensor<1x2x3xf32>
}
func.func private @XlaCallModule_odml.random_uniform.impl_0(%arg0: tensor<3xi32>) -> tensor<1x2x3xf32> {
%0 = mhlo.constant dense<1.000000e+00> : tensor<1x2x3xf32>
return %0 : tensor<1x2x3xf32>
}
// CHECK-LABEL func.func @random_uniform
// CHECK: %0 = "tfl.random_uniform"(%arg0) <{seed = 0 : i64, seed2 = 1 : i64}> : (tensor<3xi32>) -> tensor<1x2x3xf32>
// CHECK: return %0 : tensor<1x2x3xf32>
func.func @random_standard_normal(%arg0: tensor<3xi32>) -> tensor<1x2x3xf32> {
%0 = mhlo.composite "odml.random_standard_normal" %arg0 {composite_attributes = {seed = 0 : i64, seed2 = 1: i64}, decomposition = @XlaCallModule_odml.random_standard_normal.impl_0} : (tensor<3xi32>) -> tensor<1x2x3xf32>
return %0 : tensor<1x2x3xf32>
}
func.func private @XlaCallModule_odml.random_standard_normal.impl_0(%arg0: tensor<3xi32>) -> tensor<1x2x3xf32> {
%0 = mhlo.constant dense<1.000000e+00> : tensor<1x2x3xf32>
return %0 : tensor<1x2x3xf32>
}
// CHECK-LABEL func.func @random_standard_normal
// CHECK: %0 = "tfl.random_standard_normal"(%arg0) <{seed = 0 : i64, seed2 = 1 : i64}> : (tensor<3xi32>) -> tensor<1x2x3xf32>
// CHECK: return %0 : tensor<1x2x3xf32>
func.func private @XlaCallModule_tfl.unpack.impl_0(%arg0: tensor<1x3x4x1xf32>) -> (tensor<1x4x1xf32>, tensor<1x4x1xf32>, tensor<1x4x1xf32>)
func.func @jax_unstack(%arg0: tensor<1x3x4x1xf32>) -> (tensor<1x4x1xf32>, tensor<1x4x1xf32>, tensor<1x4x1xf32>) {
%0:3 = mhlo.composite "tfl.unpack" %arg0 {composite_attributes = {num = 3 : i32, axis = 1 : i32}, decomposition = @XlaCallModule_tfl.unpack.impl_0} : (tensor<1x3x4x1xf32>) -> (tensor<1x4x1xf32>, tensor<1x4x1xf32>, tensor<1x4x1xf32>)
return %0#0, %0#1, %0#2 : tensor<1x4x1xf32>, tensor<1x4x1xf32>, tensor<1x4x1xf32>
}
// CHECK-LABEL: jax_unstack
// CHECK: %0:3 = "tfl.unpack"(%arg0) <{axis = 1 : i32, num = 3 : i32}> : (tensor<1x3x4x1xf32>) -> (tensor<1x4x1xf32>, tensor<1x4x1xf32>, tensor<1x4x1xf32>)
tensorflow/compiler/mlir/lite/stablehlo/tests/fold_broadcast.mlir
+71
View File
@@ -0,0 +1,71 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -constant-fold-broadcast-pass -cse -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @foldBroadcastInDimBeforeMulOp_bcast_dim_1D_float
func.func @foldBroadcastInDimBeforeMulOp_bcast_dim_1D_float() -> (tensor<1x1x2x4xf32>) {
// CHECK-DAG: %[[RES:.*]] = mhlo.constant dense<{{\[\[\[\[}}1.000000e+00, 4.000000e+00, 9.000000e+00, 1.600000e+01], [5.000000e+00, 1.200000e+01, 2.100000e+01, 3.200000e+01]]]]> : tensor<1x1x2x4xf32>
%cst0 = mhlo.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
%cst1 = mhlo.constant dense<[[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]]> : tensor<1x1x2x4xf32>
%0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xf32>) -> tensor<1x1x2x4xf32>
%1 = mhlo.multiply %0, %cst1 : tensor<1x1x2x4xf32>
// CHECK: return %[[RES]] : tensor<1x1x2x4xf32>
func.return %1 : tensor<1x1x2x4xf32>
}
// CHECK-LABEL: @foldBroadcastInDimBeforeMulOp_bcast_dim_2D_float
func.func @foldBroadcastInDimBeforeMulOp_bcast_dim_2D_float() -> (tensor<1x2x2x3xf32>) {
// CHECK-DAG: %[[RES:.*]] = mhlo.constant dense<{{\[\[\[\[}}1.000000e+00, 4.000000e+00, 9.000000e+00], [4.000000e+00, 1.000000e+01, 1.800000e+01]], {{\[\[}}2.800000e+01, 4.000000e+01, 5.400000e+01], [4.000000e+01, 5.500000e+01, 7.200000e+01]]]]> : tensor<1x2x2x3xf32>
%cst0 = mhlo.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf32>
%cst1 = mhlo.constant dense<[[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]]]> : tensor<1x2x2x3xf32>
%0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<[1, 3]> : tensor<2xi64>}> : (tensor<2x3xf32>) -> tensor<1x2x2x3xf32>
%1 = mhlo.multiply %0, %cst1 : tensor<1x2x2x3xf32>
// CHECK: return %[[RES]] : tensor<1x2x2x3xf32>
func.return %1 : tensor<1x2x2x3xf32>
}
// CHECK-LABEL: @foldBroadcastInDimBeforeMulOp_bcast_dim_1D_int
func.func @foldBroadcastInDimBeforeMulOp_bcast_dim_1D_int() -> (tensor<1x1x2x4xi32>) {
// CHECK-DAG: %[[RES:.*]] = mhlo.constant dense<{{\[\[\[\[}}1, 4, 9, 16], [5, 12, 21, 32]]]]> : tensor<1x1x2x4xi32>
%cst0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
%cst1 = mhlo.constant dense<[[[[1, 2, 3, 4], [5, 6, 7, 8]]]]> : tensor<1x1x2x4xi32>
%0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
%1 = mhlo.multiply %0, %cst1 : tensor<1x1x2x4xi32>
// CHECK: return %[[RES]] : tensor<1x1x2x4xi32>
func.return %1 : tensor<1x1x2x4xi32>
}
// CHECK-LABEL: @foldBroadcastInDimBeforeMulOp_bcast_dim_4D_int
func.func @foldBroadcastInDimBeforeMulOp_bcast_dim_4D_int() -> tensor<1x2x1x4xi32> {
// CHECK-DAG: %[[RES:.*]] = mhlo.constant dense<{{\[\[\[\[}}0, 1, 4, 9]], {{\[\[}}0, 1, 4, 9]]]]> : tensor<1x2x1x4xi32>
%0 = mhlo.constant dense<[[[[0, 1, 2, 3]]]]> : tensor<1x1x1x4xi32>
%1 = mhlo.constant dense<[[[[0, 1, 2, 3]], [[0, 1, 2, 3]]]]> : tensor<1x2x1x4xi32>
%2 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x1x1x4xi32>) -> tensor<1x2x1x4xi32>
%3 = mhlo.multiply %1, %2 : tensor<1x2x1x4xi32>
// CHECK: return %[[RES]] : tensor<1x2x1x4xi32>
return %3 : tensor<1x2x1x4xi32>
}
// CHECK: @notFoldBroadcastInDimBeforeMulOpWhenArgIsNonConst_bcast_dim_1D_int(%[[ARG:.*]]: tensor<1x1x2x4xi32>) -> tensor<1x1x2x4xi32>
func.func @notFoldBroadcastInDimBeforeMulOpWhenArgIsNonConst_bcast_dim_1D_int(%arg0: tensor<1x1x2x4xi32>) -> (tensor<1x1x2x4xi32>) {
// CHECK-DAG: %[[CONST:.*]] = mhlo.constant dense<{{\[}}1, 2, 3, 4]> : tensor<4xi32>
%cst0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
// CHECK: %[[BROADCAST:.*]] = "mhlo.broadcast_in_dim"(%[[CONST]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
%0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
// CHECK: %[[MUL:.*]] = mhlo.multiply %[[BROADCAST]], %[[ARG]] : tensor<1x1x2x4xi32>
%1 = mhlo.multiply %0, %arg0 : tensor<1x1x2x4xi32>
// CHECK: return %[[MUL]] : tensor<1x1x2x4xi32>
func.return %1 : tensor<1x1x2x4xi32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/fuse_mhlo_convolution.mlir
+59
View File
@@ -0,0 +1,59 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -fuse-mhlo-convolution-pass -cse | FileCheck %s
// CHECK-LABEL: @fuseMulAndConv2D
// CHECK-SAME: %[[INPUT:[^:[:space:]]+]]
func.func @fuseMulAndConv2D(%input: tensor<1x256x256x3xf32>) -> (tensor<1x256x256x2xf32>) {
// CHECK-DAG: %[[FILTER:.+]] = mhlo.constant dense<{{\[\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]]]> : tensor<1x1x3x2xf32>
// CHECK-DAG: %[[CST:.+]] = mhlo.constant dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>
// CHECK-DAG: %[[CST_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[CST]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>) -> tensor<1x1x3x2xf32>
// CHECK-DAG: %[[NEW_FILTER:.+]] = mhlo.multiply %[[CST_BCAST]], %[[FILTER]] : tensor<1x1x3x2xf32>
// CHECK-DAG: %[[RESULT:.+]] = mhlo.convolution(%[[INPUT]], %[[NEW_FILTER]]) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = {{\[\[}}0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x256x256x2xf32>
%filter = mhlo.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]]> : tensor<1x1x3x2xf32>
%cst = mhlo.constant dense<[0.1, 0.2]> : tensor<2xf32>
%0 = mhlo.convolution(%input, %filter) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x256x256x2xf32>
%1 = "mhlo.broadcast_in_dim"(%cst) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>) -> tensor<1x256x256x2xf32>
%2 = mhlo.multiply %0, %1 : tensor<1x256x256x2xf32>
// CHECK-DAG: return %[[RESULT]]
func.return %2 : tensor<1x256x256x2xf32>
}
// -----
// CHECK-LABEL: @fuseMulAndConv2DDynamic
// CHECK-SAME: %[[INPUT:[^:[:space:]]+]]
func.func @fuseMulAndConv2DDynamic(%input: tensor<?x256x256x3xf32>) -> (tensor<?x256x256x2xf32>) {
// CHECK-DAG: %[[FILTER:.+]] = mhlo.constant dense<{{\[\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]]]> : tensor<1x1x3x2xf32>
// CHECK-DAG: %[[CST_0:.+]] = mhlo.constant dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>
// CHECK-DAG: %[[CST_1:.+]] = mhlo.constant dense<[3.000000e-01, 4.000000e-01]> : tensor<2xf32>
// CHECK: %[[CST_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[CST_0]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>) -> tensor<1x1x3x2xf32>
// CHECK: %[[NEW_FILTER:.+]] = mhlo.multiply %[[CST_BCAST]], %[[FILTER]] : tensor<1x1x3x2xf32>
// CHECK: %[[CONV:.+]] = mhlo.convolution(%[[INPUT]], %[[NEW_FILTER]]) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = {{\[\[}}0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<?x256x256x2xf32>
// CHECK: %[[SHAPE:.+]] = shape.shape_of %[[CONV]] : tensor<?x256x256x2xf32> -> tensor<4xindex>
// CHECK: %[[DYNAMIC_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[CST_1]], %[[SHAPE]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
// CHECK: %[[ADD:.+]] = mhlo.add %[[CONV]], %[[DYNAMIC_BCAST]] : tensor<?x256x256x2xf32>
%filter = mhlo.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]]> : tensor<1x1x3x2xf32>
%cst_0 = mhlo.constant dense<[0.1, 0.2]> : tensor<2xf32>
%cst_1 = mhlo.constant dense<[0.3, 0.4]> : tensor<2xf32>
%0 = mhlo.convolution(%input, %filter) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<?x256x256x2xf32>
%1 = shape.shape_of %0 : tensor<?x256x256x2xf32> -> tensor<4xindex>
%2 = "mhlo.dynamic_broadcast_in_dim"(%cst_0, %1) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
%3 = mhlo.multiply %0, %2 : tensor<?x256x256x2xf32>
%4 = "mhlo.dynamic_broadcast_in_dim"(%cst_1, %1) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
%5 = mhlo.add %3, %4 : tensor<?x256x256x2xf32>
// CHECK-DAG: return %[[ADD]]
func.return %5 : tensor<?x256x256x2xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/gelu.mlir
+49
View File
@@ -0,0 +1,49 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-legalize-chlo -split-input-file
//| FileCheck %s --dump-input=fail
module {
func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<2xf32>
%1 = stablehlo.constant dense<0.707106769> : tensor<2xf32>
%2 = stablehlo.constant dense<5.000000e-01> : tensor<2xf32>
%3 = stablehlo.multiply %arg0, %2 : tensor<2xf32>
%4 = stablehlo.multiply %arg0, %1 : tensor<2xf32>
%5 = stablehlo.custom_call @mhlo.erf(%4) {mhlo.attributes = {}, mhlo.version = 1 : i64} : (tensor<2xf32>) -> tensor<2xf32>
%6 = stablehlo.add %5, %0 : tensor<2xf32>
%7 = stablehlo.multiply %3, %6 : tensor<2xf32>
return %7 : tensor<2xf32>
}
}
// CHECK-LABEL: geluWithCustomCallErf
// -----
func.func @geluWithCHLOErf(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<2xf32>
%1 = stablehlo.constant dense<0.707106769> : tensor<2xf32>
%2 = stablehlo.constant dense<5.000000e-01> : tensor<2xf32>
%3 = stablehlo.multiply %arg0, %2 : tensor<2xf32>
%4 = stablehlo.multiply %arg0, %1 : tensor<2xf32>
%5 = chlo.erf %4 : tensor<2xf32> -> tensor<2xf32>
%6 = stablehlo.add %5, %0 : tensor<2xf32>
%7 = stablehlo.multiply %3, %6 : tensor<2xf32>
return %7 : tensor<2xf32>
}
// CHECK-LABEL: geluWithCHLOErf
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-inplaceupdate.mlir
+34
View File
@@ -0,0 +1,34 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tf-stablehlo | FileCheck %s
module attributes {tf.versions = {producer = 888 : i32}} {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
%1 = arith.constant dense<1> : tensor<1xi32>
%2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
%3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
: (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
func.return %3 : tensor<2x1x2xf32>
}
}
// CHECK-LABEL: @tfInplaceUpdate
// CHECK-DAG: %[[CST0:.*]] = stablehlo.constant dense<1> : tensor<i32>
// CHECK-DAG: %[[CST1:.*]] = stablehlo.constant dense<0> : tensor<i32>
// CHECK-DAG: %[[CST2:.*]] = stablehlo.constant dense<2.000000e+00> : tensor<1x1x2xf32>
// CHECK: %[[RES:.*]] = stablehlo.dynamic_update_slice %arg0, %[[CST2]], %[[CST0]], %[[CST1]], %[[CST1]] : (tensor<2x1x2xf32>, tensor<1x1x2xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2x1x2xf32>
// CHECK: return %[[RES]] : tensor<2x1x2xf32>
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-skip-partitioned-calls.mlir
+48
View File
@@ -0,0 +1,48 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s --tf-stablehlo=skip-partitioned-calls=true | FileCheck %s --check-prefix=CHECK-SKIP
// RUN: odml-to-stablehlo-opt %s --tf-stablehlo=skip-partitioned-calls=false | FileCheck %s --check-prefix=CHECK-NOSKIP
module {
func.func @partitioned_call(%arg0: tensor<1x2x2x3xf32>) -> (tensor<1x2x2x3xf32>) {
%0 = "tf.StatefulPartitionedCall"(%arg0) <{
config = "", config_proto = "", executor_type = "", f = @some_func
}> {
_collective_manager_ids = [], device = ""
} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
// CHECK-SKIP: tf.StatefulPartitionedCall
// CHECK-NOSKIP: call @some_func
// CHECK-NOSKIP-NOT: tf.StatefulPartitionedCall
%1 = "tf.PartitionedCall"(%0) <{
config = "", config_proto = "", executor_type = "", f = @some_other_func
}> {
_collective_manager_ids = [], device = ""
} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
// CHECK-SKIP: tf.PartitionedCall
// CHECK-NOSKIP: call @some_other_func
// CHECK-NOSKIP-NOT: tf.PartitionedCall
func.return %1: tensor<1x2x2x3xf32>
}
// CHECK-SKIP: func.func private @some_func
func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
return %arg0 : tensor<1x2x2x3xf32>
}
// CHECK-SKIP: func.func private @some_other_func
func.func private @some_other_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
return %arg0 : tensor<1x2x2x3xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-skip-quantization-ops.mlir
+23
View File
@@ -0,0 +1,23 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s --tf-stablehlo=skip-quantization-ops=true | FileCheck %s --check-prefix=CHECK-SKIP
// RUN: odml-to-stablehlo-opt %s --tf-stablehlo=skip-quantization-ops=false | FileCheck %s --check-prefix=CHECK-NOSKIP
func.func @fake_quant_with_min_max_vars(%arg0: tensor<1x1x28x48xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<1x1x28x48xf32> {
%0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {device = "", narrow_range = true, num_bits = 8 : i64} : (tensor<1x1x28x48xf32>, tensor<f32>, tensor<f32>) -> tensor<1x1x28x48xf32>
func.return %0 : tensor<1x1x28x48xf32>
// CHECK-SKIP: tf.FakeQuantWithMinMaxVars
// CHECK-NOSKIP-NOT: tf.
}
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-stablehlo-tfl-composite-quantize-dequantize.mlir
+47
View File
@@ -0,0 +1,47 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -stablehlo-composite-legalize-tfl-custom | FileCheck %s
module {
// CHECK-LABEL: func.func private @test_quantize_and_dequantize
func.func private @test_quantize_and_dequantize(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = "tfl.pseudo_const"() <{value = dense<1.000000e+00> : tensor<4x3xf32>}> : () -> tensor<4x3xf32>
// CHECK: %1:3 = "tfl.custom"(%0) <{custom_code = "odml.quantize_and_dequantize", custom_option = #tfl<const_bytes : "0x61786973006269747300020B0702010200040404042401">}> : (tensor<4x3xf32>) -> (tensor<4x3xf32>, tensor<4x3xf32>, tensor<1x3xf32>)
%1:3 = stablehlo.composite "odml.quantize_and_dequantize" %0 {composite_attributes = {axis = 0 : i64, bits = 4 : i64}, decomposition = @call_module_odml.quantize_and_dequantize.0} : (tensor<4x3xf32>) -> (tensor<4x3xf32>, tensor<4x3xf32>, tensor<1x3xf32>)
%2 = "tfl.batch_matmul"(%arg0, %1#0) <{adj_x = false, adj_y = false, asymmetric_quantize_inputs = false}> : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
func.func private @call_module_odml.quantize_and_dequantize.0(%arg0: tensor<4x3xf32>) -> (tensor<4x3xf32>, tensor<4x3xf32>, tensor<1x3xf32>) {
%0 = "tfl.abs"(%arg0) : (tensor<4x3xf32>) -> tensor<4x3xf32>
%1 = "tfl.pseudo_const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
%2 = "tfl.reduce_max"(%0, %1) <{keep_dims = false}> : (tensor<4x3xf32>, tensor<1xi32>) -> tensor<3xf32>
%3 = "tfl.pseudo_const"() <{value = dense<0.142857149> : tensor<f32>}> : () -> tensor<f32>
%4 = tfl.mul(%2, %3) <{fused_activation_function = "NONE"}> : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
%5 = "tfl.pseudo_const"() <{value = dense<9.99999993E-9> : tensor<f32>}> : () -> tensor<f32>
%6 = tfl.add(%4, %5) <{fused_activation_function = "NONE"}> : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
%7 = "tfl.pseudo_const"() <{value = dense<[1, 3]> : tensor<2xi32>}> : () -> tensor<2xi32>
%8 = "tfl.reshape"(%6, %7) : (tensor<3xf32>, tensor<2xi32>) -> tensor<1x3xf32>
%9 = tfl.div(%arg0, %8) <{fused_activation_function = "NONE"}> : (tensor<4x3xf32>, tensor<1x3xf32>) -> tensor<4x3xf32>
%10 = tfl.sub %9, %9 {fused_activation_function = "NONE"} : tensor<4x3xf32>
%11 = "tfl.round"(%9) : (tensor<4x3xf32>) -> tensor<4x3xf32>
%12 = tfl.add %10, %11 {fused_activation_function = "NONE"} : tensor<4x3xf32>
%13 = "tfl.pseudo_const"() <{value = dense<-8.000000e+00> : tensor<f32>}> : () -> tensor<f32>
%14 = "tfl.maximum"(%12, %13) : (tensor<4x3xf32>, tensor<f32>) -> tensor<4x3xf32>
%15 = "tfl.pseudo_const"() <{value = dense<7.000000e+00> : tensor<f32>}> : () -> tensor<f32>
%16 = "tfl.minimum"(%14, %15) : (tensor<4x3xf32>, tensor<f32>) -> tensor<4x3xf32>
%17 = tfl.mul(%16, %8) <{fused_activation_function = "NONE"}> : (tensor<4x3xf32>, tensor<1x3xf32>) -> tensor<4x3xf32>
return %17, %16, %8 : tensor<4x3xf32>, tensor<4x3xf32>, tensor<1x3xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-stablehlo-tfl-composite.mlir
+46
View File
@@ -0,0 +1,46 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -stablehlo-composite-legalize-tfl-custom | FileCheck %s
func.func private @odml.update_kv_cache.impl_0(%arg0: tensor<1x500x4x4xf32>, %arg1: tensor<1x500x4x4xf32>, %arg2: tensor<100xi64>, %arg3: tensor<1x100x4x4xf32>, %arg4: tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>)
// CHECK-LABEL: func.func private @test_multiple_kv_caches
func.func private @test_multiple_kv_caches(%arg0: tensor<1x500x4x4xf32>, %arg1: tensor<1x500x4x4xf32>, %arg2: tensor<100xi64>, %arg3: tensor<1x100x4x4xf32>, %arg4: tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>) {
// CHECK: %0:2 = "tfl.custom"(%arg2, %arg3, %arg4) <{custom_code = "odml.update_kv_cache", custom_option = #tfl<const_bytes : "0x6B765F63616368655F6D6178006C617965725F696E646578006E756D5F6C6179657273000325190E030001000300F40100000200050505092501">}> : (tensor<100xi64>, tensor<1x100x4x4xf32>, tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>)
// CHECK: %1:2 = "tfl.custom"(%arg2, %arg3, %arg4) <{custom_code = "odml.update_kv_cache", custom_option = #tfl<const_bytes : "0x6B765F63616368655F6D6178006C617965725F696E646578006E756D5F6C6179657273000325190E030001000300F40101000200050505092501">}> : (tensor<100xi64>, tensor<1x100x4x4xf32>, tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>)
%0:2 = stablehlo.composite "odml.update_kv_cache" %arg0, %arg1, %arg2, %arg3, %arg4 {composite_attributes = {kv_cache_max = 500 : i64}, decomposition = @odml.update_kv_cache.impl_0} : (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>, tensor<100xi64>, tensor<1x100x4x4xf32>, tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>)
%1:2 = stablehlo.composite "odml.update_kv_cache" %0#0, %0#1, %arg2, %arg3, %arg4 {composite_attributes = {kv_cache_max = 500 : i64}, decomposition = @odml.update_kv_cache.impl_0} : (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>, tensor<100xi64>, tensor<1x100x4x4xf32>, tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>)
return %1#0, %1#1 : tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>
}
// ---
func.func private @test_odml_detector.detector.impl_0(%arg0: tensor<2xf32>) -> tensor<2xf32>
// CHECK-LABEL: func.func private @test_odml_detector
func.func @test_odml_detector(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> (tensor<2xf32>) {
%0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<2xf32>
// CHECK %1 = "tfl.custom"(%0) <{custom_code = "odml.detector", custom_option = #tfl<const_bytes : "0x6E616D6500036F757400776F726B696E675F64697200082F746D702F7473740002211802010220101414042401">}> : (tensor<2xf32>) -> tensor<2xf32>
%1 = stablehlo.composite "odml.detector" %0 {composite_attributes = {name = "out", working_dir = "/tmp/tst"}, decomposition = @test_odml_detector.detector.impl_0} : (tensor<2xf32>) -> tensor<2xf32>
return %1 : tensor<2xf32>
}
// ---
func.func private @test_litert_custom_op.impl_0(%arg0: tensor<2xf32>) -> tensor<2xf32>
// CHECK-LABEL: func.func @test_litert_custom_op
func.func @test_litert_custom_op(%arg0: tensor<2xf32>) -> (tensor<2xf32>) {
// CHECK: "tfl.custom"(%arg0) <{custom_code = "litert_custom_op.my_op"
%0 = stablehlo.composite "litert_custom_op.my_op" %arg0 {composite_attributes = {}, decomposition = @test_litert_custom_op.impl_0} : (tensor<2xf32>) -> tensor<2xf32>
return %0 : tensor<2xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-stablehlo-vhlo.mlir
+142
View File
@@ -0,0 +1,142 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s --stablehlo-legalize-vhlo -reconcile-unrealized-casts -split-input-file | FileCheck %s
// RUN: odml-to-stablehlo-opt --stablehlo-legalize-vhlo -reconcile-unrealized-casts %s | odml-to-stablehlo-opt --vhlo-legalize-stablehlo -reconcile-unrealized-casts > %t.0
// RUN: odml-to-stablehlo-opt %s > %t.1
// RUN: diff %t.0 %t.1
// CHECK-LABEL: op_tfl
func.func @op_tfl(%arg0 : tensor<f32>) -> (tensor<f32>) {
// CHECK: %0 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : tensor<f32>
%0 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : tensor<f32>
return %0 : tensor<f32>
}
// -----
// CHECK-LABEL: op_shlo
func.func @op_shlo(%arg0 : tensor<f32>) -> (tensor<f32>) {
// CHECK: %0 = "vhlo.add_v1"(%arg0, %arg0) : (tensor<f32>, tensor<f32>) -> tensor<f32>
%0 = stablehlo.add %arg0, %arg0 : tensor<f32>
return %0 : tensor<f32>
}
// -----
// CHECK-LABEL: mixed_shlo_tfl_shlo
func.func @mixed_shlo_tfl_shlo(%arg0 : tensor<f32>) -> (tensor<f32>) {
// CHECK: %0 = "vhlo.abs_v1"(%arg0) : (tensor<f32>) -> tensor<f32>
// CHECK-NEXT: %1 = tfl.add %0, %arg0 {fused_activation_function = "NONE"} : tensor<f32>
// CHECK-NEXT: %2 = "vhlo.abs_v1"(%1) : (tensor<f32>) -> tensor<f32>
%0 = stablehlo.abs %arg0 : tensor<f32>
%1 = tfl.add %0, %arg0 {fused_activation_function = "NONE"} : tensor<f32>
%2 = stablehlo.abs %1 : tensor<f32>
return %2 : tensor<f32>
}
// -----
// CHECK-LABEL: mixed_tfl_shlo_tfl
func.func @mixed_tfl_shlo_tfl(%arg0 : tensor<f32>) -> (tensor<f32>) {
%0 = "tfl.abs"(%arg0) {fused_activation_function = "NONE"} : (tensor<f32>) -> tensor<f32>
// CHECK: %1 = "vhlo.add_v1"(%0, %arg0) : (tensor<f32>, tensor<f32>) -> tensor<f32>
%1 = stablehlo.add %0, %arg0 : tensor<f32>
%2 = "tfl.abs"(%1) {fused_activation_function = "NONE"} : (tensor<f32>) -> tensor<f32>
return %2 : tensor<f32>
}
// -----
// CHECK-LABEL: op_with_region
func.func @op_with_region(%arg0: tensor<1x16x16x320xf32>, %arg1: tensor<f32>) -> tensor<1x320xf32> {
// CHECK: %0 = "vhlo.reduce_v1"(%arg0, %arg1) <{{.*}}> ({
// CHECK-NEXT: ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
// CHECK-NEXT: %1 = "vhlo.add_v1"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
// CHECK-NEXT: "vhlo.return_v1"(%1) : (tensor<f32>) -> ()
// CHECK-NEXT: }) : (tensor<1x16x16x320xf32>, tensor<f32>) -> tensor<1x320xf32>
%0 = stablehlo.reduce(%arg0 init: %arg1) applies stablehlo.add across dimensions = [1, 2] : (tensor<1x16x16x320xf32>, tensor<f32>) -> tensor<1x320xf32>
return %0 : tensor<1x320xf32>
}
// -----
// CHECK-LABEL: op_with_region_mixed_tfl_shlo_tfl
func.func @op_with_region_mixed_tfl_shlo_tfl(%arg0: tensor<7x5xf32>, %arg1 : tensor<5xf32>) -> tensor<5xf32> {
%0 = "stablehlo.reduce"(%arg0, %arg1) ({
^bb0(%arg2: tensor<5xf32>, %arg3: tensor<5xf32>):
// CHECK: %1 = "tfl.abs"(%arg2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
// CHECK-NEXT: %2 = "vhlo.add_v1"(%1, %arg2) : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>
// CHECK-NEXT: %3 = "tfl.abs"(%2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
%1 = "tfl.abs"(%arg2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
%2 = stablehlo.add %1, %arg2 : tensor<5xf32>
%3 = "tfl.abs"(%2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
"stablehlo.return"(%3) : (tensor<5xf32>) -> ()
}) {dimensions = array<i64: 0>} : (tensor<7x5xf32>, tensor<5xf32>) -> tensor<5xf32>
func.return %0: tensor<5xf32>
}
// -----
// CHECK-LABEL: op_with_region_mixed_shlo_tfl_shlo
func.func @op_with_region_mixed_shlo_tfl_shlo(%arg0: tensor<7x5xf32>, %arg1 : tensor<5xf32>) -> tensor<5xf32> {
%0 = "stablehlo.reduce"(%arg0, %arg1) ({
^bb0(%arg2: tensor<5xf32>, %arg3: tensor<5xf32> ):
// CHECK: %1 = "vhlo.abs_v1"(%arg2) : (tensor<5xf32>) -> tensor<5xf32>
// CHECK-NEXT: %2 = tfl.add %1, %arg2 {fused_activation_function = "NONE"} : tensor<5xf32>
// CHECK-NEXT: %3 = "vhlo.abs_v1"(%2) : (tensor<5xf32>) -> tensor<5xf32>
%1 = stablehlo.abs %arg2 : tensor<5xf32>
%2 = tfl.add %1, %arg2 {fused_activation_function = "NONE"} : tensor<5xf32>
%3 = stablehlo.abs %2 : tensor<5xf32>
"stablehlo.return"(%3) : (tensor<5xf32>) -> ()
}) {dimensions = array<i64: 0>} : (tensor<7x5xf32>, tensor<5xf32>) -> tensor<5xf32>
func.return %0: tensor<5xf32>
}
// -----
// CHECK-LABEL: op_with_tfl_control_flow
func.func @op_with_tfl_control_flow() -> (tensor<1xf32>, !tfl.control) {
// CHECK: vhlo.constant_v1
%0 = stablehlo.constant dense<1.000000e+00> : tensor<1xf32>
// CHECK-NEXT: tfl.control_node
%outputs, %control = tfl.control_node {
%1 = "tfl.neg"(%0) : (tensor<1xf32>) -> tensor<1xf32>
"tfl.yield"(%1) : (tensor<1xf32>) -> ()
}
return %outputs, %control : tensor<1xf32>, !tfl.control
}
// -----
// CHECK-LABEL: func_with_tfl_attrs
func.func @func_with_tfl_attrs(%arg0: tensor<!tf_type.variant<tensor<2xi32>>>, %arg1: tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<2xi32>>> attributes {tf.entry_function = {inputs = "arg0,arg1", outputs = "arg0"}} {
return %arg0 : tensor<!tf_type.variant<tensor<2xi32>>>
}
// -----
// There are cases where ODML converter relies on constants not being folded or
// CSE'ed. This test ensures that StableHLO<->ODML conversion does not fold.
// CHECK-LABEL: mixed_no_constant_folding
func.func @mixed_no_constant_folding() -> (tensor<f32>) {
// CHECK: %[[CST0:.+]] = arith.constant dense<0.000000e+00>
// CHECK-NEXT: %[[CST1:.+]] = arith.constant dense<0.000000e+00>
// CHECK-NEXT: "vhlo.add_v1"(%[[CST0]], %[[CST1]]) : (tensor<f32>, tensor<f32>) -> tensor<f32>
%cst_0 = arith.constant dense<0.000000e+00> : tensor<f32>
%cst_1 = arith.constant dense<0.000000e+00> : tensor<f32>
%0 = stablehlo.add %cst_0, %cst_1 : tensor<f32>
return %0 : tensor<f32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tf.mlir
+2565
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-add.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
%0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.add", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
// CHECK-NEXT: %0 = stablehlo.add %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-broadcast.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<1x2xi32>) -> tensor<1x2x2xi32> {
%0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.broadcast_in_dim", custom_option = #tfl<const_bytes : "0x62726F6164636173745F64696D656E73696F6E73000201020119010101072C022401">} : (tensor<1x2xi32>) -> tensor<1x2x2xi32>
func.return %0 : tensor<1x2x2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<1x2xi32>) -> tensor<1x2x2xi32> {
// CHECK-NEXT: %0 = stablehlo.broadcast_in_dim %arg0, dims = [1, 2] : (tensor<1x2xi32>) -> tensor<1x2x2xi32>
// CHECK-NEXT: return %0 : tensor<1x2x2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-clamp.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>, %arg2: tensor<2xi32>) -> tensor<2xi32> {
%0 = "tfl.custom"(%arg0, %arg1, %arg2) {custom_code = "stablehlo.clamp", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>, %arg2: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %0 = stablehlo.clamp %arg0, %arg1, %arg2 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-concat.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
%0 = "tfl.custom"(%arg0, %arg1) {custom_code = "stablehlo.concatenate", custom_option = #tfl<const_bytes : "0x64696D656E73696F6E00010B0101010004022401">} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32>
func.return %0 : tensor<6x3xf32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
// CHECK-NEXT: %0 = stablehlo.concatenate %arg0, %arg1, dim = 0 : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32>
// CHECK-NEXT: return %0 : tensor<6x3xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-constant.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck -dump-input always %s
module {
func.func @main() -> tensor<1xi64> {
%0 = "tfl.custom"() {custom_code = "stablehlo.constant", custom_option = #tfl<const_bytes : "0x76616C75650001020109010101062C022401">} : () -> tensor<1xi64>
func.return %0 : tensor<1xi64>
}
}
// CHECK: module {
// CHECK-NEXT: func @main() -> tensor<1xi64> {
// CHECK-NEXT: %[[c0:.+]] = stablehlo.constant dense<2> : tensor<1xi64>
// CHECK-NEXT: return %[[c0]] : tensor<1xi64>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-conv.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck -dump-input always %s
module {
func.func @main(%arg0: tensor<8x8x1x207xf32>, %arg1: tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32> {
%0 = "tfl.custom"(%arg0, %arg1) {custom_code = "stablehlo.convolution", custom_option = #tfl<const_bytes : "0x62617463685F67726F75705F636F756E740064696D656E73696F6E5F6E756D62657273000200010404020001040402010204040902031103020F03000D040428040428040428666561747572655F67726F75705F636F756E74006C68735F64696C6174696F6E0002010170616464696E67000401010101707265636973696F6E5F636F6E666967000744454641554C54000744454641554C540002120A7268735F64696C6174696F6E0002010177696E646F775F726576657273616C0002010077696E646F775F737472696465730002010109D3C28F7C6D613C2D1B09010901AC017A70493A28170428042C2C3C2C902C122401">} : (tensor<8x8x1x207xf32>, tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32>
func.return %0 : tensor<16x8x8x1xf32>
}
}
// CHECK: module {
// CHECK-NEXT: func.func @main(%arg0: tensor<8x8x1x207xf32>, %arg1: tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32> {
// CHECK-NEXT: %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [0, 1, b, f]x[0, 1, o, i]->[f, 0, 1, b], window = {stride = [1, 1], pad = {{\[}}[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1], reverse = [true, false]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<8x8x1x207xf32>, tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32>
// CHECK-NEXT: return %0 : tensor<16x8x8x1xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-max.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
%0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.maximum", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
// CHECK-NEXT: %0 = stablehlo.maximum %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-mul.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
%0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.multiply", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-pad.mlir
+30
View File
@@ -0,0 +1,30 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck -dump-input always %s
module {
func.func @main(%arg0: tensor<8x128xf32>, %arg1: tensor<f32>) -> tensor<11x131xf32> {
%0 = "tfl.custom"(%arg0, %arg1) {custom_code = "stablehlo.pad", custom_option = #tfl<const_bytes : "0x656467655F70616464696E675F6869676800020203656467655F70616464696E675F6C6F7700020100696E746572696F725F70616464696E6700020000033E2A17030103311E0B2C2C2C062401">} : (tensor<8x128xf32>, tensor<f32>) -> tensor<11x131xf32>
func.return %0 : tensor<11x131xf32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<8x128xf32>, %arg1: tensor<f32>) -> tensor<11x131xf32> {
// CHECK-NEXT: %0 = stablehlo.pad %arg0, %arg1, low = [1, 0], high = [2, 3], interior = [0, 0] : (tensor<8x128xf32>, tensor<f32>) -> tensor<11x131xf32>
// CHECK-NEXT: return %0 : tensor<11x131xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-reshape.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck -dump-input always %s
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> {
%0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.reshape", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> {
// CHECK-NEXT: %0 = stablehlo.reshape %arg0 : (tensor<2xi32>) -> tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-rsqrt.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck -dump-input always %s
module {
func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.rsqrt", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
}
// CHECK: module
// CHECK-NEXT: func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK-NEXT: %0 = stablehlo.rsqrt %arg0 : tensor<2xf32>
// CHECK-NEXT: return %0 : tensor<2xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo-sub.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
%0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.subtract", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
// CHECK-NEXT: %0 = stablehlo.subtract %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %0 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-tfl-stablehlo.mlir
+31
View File
@@ -0,0 +1,31 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-parse-stablehlo-ops | FileCheck %s
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
%0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.add", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
%1 = "tfl.custom"(%0, %arg0) {custom_code = "stablehlo.subtract", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %1 : tensor<2xi32>
}
}
// CHECK: module {
// CHECK-NEXT: func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
// CHECK-NEXT: %0 = stablehlo.add %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT: %1 = stablehlo.subtract %0, %arg0 : tensor<2xi32>
// CHECK-NEXT: return %1 : tensor<2xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/legalize_hlo.mlir
+5260
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/stablehlo/tests/odml-to-stablehlo-allow-tf.mlir
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml_to_stablehlo %s --allow-tf=false -o /tmp/temp.mlir; [ -f /tmp/temp.mlir ]; [ -f /tmp/debug_stablehlo.mlir ]
// RUN: odml_to_stablehlo %s --allow-tf=true -o /tmp/temp2.mlir; [ -f /tmp/temp2.mlir ]
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 975 : i32}, tf_saved_model.semantics} {
func.func @serving_default(%arg0: tensor<1x20x20x28xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x40x40x28xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%c = stablehlo.constant dense<40> : tensor<2xi32>
%0 = "tf.UnconvertedOp"(%arg0, %c) {align_corners = false, half_pixel_centers = false} : (tensor<1x20x20x28xf32>, tensor<2xi32>) -> tensor<1x40x40x28xf32>
func.return %0 : tensor<1x40x40x28xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/odml-to-stablehlo-smuggle-resize.mlir
+29
View File
@@ -0,0 +1,29 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml_to_stablehlo %s -skip-resize -smuggle-disallowed-ops -o - | FileCheck %s
// RUN: odml-to-stablehlo-opt %s --smuggle-disallowed-ops-pass | FileCheck %s --check-prefix=CHECK-OPT
// CHECK-LABEL: @main
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 975 : i32}, tf_saved_model.semantics} {
func.func @serving_default(%arg0: tensor<1x32x32x128xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x64x64x128xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = "tf.Const"() {value = dense<[56, 904]> : tensor<2xi32>} : () -> tensor<2xi32>
// CHECK: %{{.*}} = stablehlo.custom_call @tf.ResizeBilinear(%arg0, %{{.*}}) {align_corners = false, device = "", half_pixel_centers = true} : (tensor<1x32x32x128xf32>, tensor<2xi32>) -> tensor<1x64x64x128xf32>
// CHECK-OPT: %{{.*}} = stablehlo.custom_call @tf.ResizeBilinear(%arg0, %cst) {align_corners = false, device = "", half_pixel_centers = true} : (tensor<1x32x32x128xf32>, tensor<2xi32>) -> tensor<1x64x64x128xf32>
%1 = "tf.ResizeBilinear"(%arg0, %0) {
align_corners = false, device = "", half_pixel_centers = true
} : (tensor<1x32x32x128xf32>, tensor<2xi32>) -> tensor<1x64x64x128xf32>
func.return %1 : tensor<1x64x64x128xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/optimize.mlir
+528
View File
@@ -0,0 +1,528 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -split-input-file -mhlo-optimize | FileCheck %s
// CHECK-LABEL: testDotToDotGeneralVectorVector
func.func @testDotToDotGeneralVectorVector(%arg0: tensor<3072xf32>, %arg1: tensor<3072xf32>) -> tensor<f32> {
%0 = "mhlo.dot"(%arg0, %arg1) : (tensor<3072xf32>, tensor<3072xf32>) -> tensor<f32>
func.return %0 : tensor<f32>
// CHECK: %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_contracting_dimensions = [0],
// CHECK-SAME: rhs_contracting_dimensions = [0]
// CHECK-SAME: >}> : (tensor<3072xf32>, tensor<3072xf32>) -> tensor<f32>
// CHECK: return %[[RES]] : tensor<f32>
}
// -----
// CHECK-LABEL: testDotToDotGeneralVectorMatrix
func.func @testDotToDotGeneralVectorMatrix(%arg0: tensor<3072xf32>, %arg1: tensor<3072x512xf32>) -> tensor<512xf32> {
%0 = "mhlo.dot"(%arg0, %arg1) : (tensor<3072xf32>, tensor<3072x512xf32>) -> tensor<512xf32>
func.return %0 : tensor<512xf32>
// CHECK: %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_contracting_dimensions = [0],
// CHECK-SAME: rhs_contracting_dimensions = [0]
// CHECK-SAME: >}> : (tensor<3072xf32>, tensor<3072x512xf32>) -> tensor<512xf32>
// CHECK: return %[[RES]] : tensor<512xf32>
}
// -----
// CHECK-LABEL: testDotToDotGeneralMatrixVector
func.func @testDotToDotGeneralMatrixVector(%arg0: tensor<2x3072xf32>, %arg1: tensor<3072xf32>) -> tensor<2xf32> {
%0 = "mhlo.dot"(%arg0, %arg1) : (tensor<2x3072xf32>, tensor<3072xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
// CHECK: %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_contracting_dimensions = [1],
// CHECK-SAME: rhs_contracting_dimensions = [0]
// CHECK-SAME: >}> : (tensor<2x3072xf32>, tensor<3072xf32>) -> tensor<2xf32>
// CHECK: return %[[RES]] : tensor<2xf32>
}
// -----
// CHECK-LABEL: testDotToDotGeneralMatrixMatrix
func.func @testDotToDotGeneralMatrixMatrix(%arg0: tensor<2x3072xf32>, %arg1: tensor<3072x512xf32>) -> tensor<2x512xf32> {
%0 = "mhlo.dot"(%arg0, %arg1) : (tensor<2x3072xf32>, tensor<3072x512xf32>) -> tensor<2x512xf32>
func.return %0 : tensor<2x512xf32>
// CHECK: %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_contracting_dimensions = [1],
// CHECK-SAME: rhs_contracting_dimensions = [0]
// CHECK-SAME: >}> : (tensor<2x3072xf32>, tensor<3072x512xf32>) -> tensor<2x512xf32>
// CHECK: return %[[RES]] : tensor<2x512xf32>
}
// -----
// CHECK-LABEL: testRemoveReshapeAroundDotGeneral
func.func @testRemoveReshapeAroundDotGeneral(%arg0: tensor<3x72x1x2048xf32>, %arg1: tensor<3x2048x512xf32>) -> tensor<3x72x1x512xf32> {
%0 = "mhlo.reshape"(%arg0) : (tensor<3x72x1x2048xf32>) -> tensor<3x72x2048xf32>
%1 = "mhlo.dot_general"(%0, %arg1) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [0],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [1]
>} : (tensor<3x72x2048xf32>, tensor<3x2048x512xf32>) -> tensor<3x72x512xf32>
%2 = "mhlo.reshape"(%1) : (tensor<3x72x512xf32>) -> tensor<3x72x1x512xf32>
func.return %2 : tensor<3x72x1x512xf32>
// CHECK: %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_batching_dimensions = [0],
// CHECK-SAME: rhs_batching_dimensions = [0],
// CHECK-SAME: lhs_contracting_dimensions = [3],
// CHECK-SAME: rhs_contracting_dimensions = [1]
// CHECK-SAME: >}> : (tensor<3x72x1x2048xf32>, tensor<3x2048x512xf32>) -> tensor<3x72x1x512xf32>
// CHECK: return %[[RES]] : tensor<3x72x1x512xf32>
}
// -----
// CHECK-LABEL: testRemoveReshapeAroundDot
func.func @testRemoveReshapeAroundDot(%arg0: tensor<1x1x512xf32>, %arg1: tensor<512x13x!quant.uniform<i8:f32, 0.00285>>) -> tensor<1x1x13xf32> {
%0 = "mhlo.reshape"(%arg0) : (tensor<1x1x512xf32>) -> tensor<1x512xf32>
%1 = "mhlo.dot"(%0, %arg1) : (tensor<1x512xf32>, tensor<512x13x!quant.uniform<i8:f32, 0.00285>>) -> tensor<1x13xf32>
%2 = "mhlo.reshape"(%1) : (tensor<1x13xf32>) -> tensor<1x1x13xf32>
func.return %2 : tensor<1x1x13xf32>
// CHECK: %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_contracting_dimensions = [2],
// CHECK-SAME: rhs_contracting_dimensions = [0]
// CHECK-SAME: >}> : (tensor<1x1x512xf32>, tensor<512x13x!quant.uniform<i8:f32, 2.850000e-03>>) -> tensor<1x1x13xf32>
// CHECK: return %[[RES]] : tensor<1x1x13xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePads
func.func @testTwoConsecutivePads(%arg0: tensor<10x10x10xf32>) -> (tensor<12x12x12xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<0> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<11x11x11xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<11x11x11xf32>, tensor<f32>) -> tensor<12x12x12xf32>
return %3 : tensor<12x12x12xf32>
// CHECK: %[[RES:.*]] = "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
// CHECK: return %[[RES]] : tensor<12x12x12xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsNegativeLowPad
func.func @testTwoConsecutivePadsNegativeLowPad(%arg0: tensor<10x10x10xf32>) -> (tensor<10x10x10xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<0> : tensor<3xi64>, edge_padding_low = dense<-1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<9x9x9xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<9x9x9xf32>, tensor<f32>) -> tensor<10x10x10xf32>
return %3 : tensor<10x10x10xf32>
// CHECK: %[[RES:.*]] = "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<-1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<10x10x10xf32>
// CHECK: return %[[RES]] : tensor<10x10x10xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsTwoNegativeHighPad
func.func @testTwoConsecutivePadsTwoNegativeHighPad(%arg0: tensor<10x10x10xf32>) -> (tensor<9x9x9xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<-1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<10x10x10xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<-1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<9x9x9xf32>
return %3 : tensor<9x9x9xf32>
// CHECK: %[[RES:.*]] = "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<-2> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<9x9x9xf32>
// CHECK: return %[[RES]] : tensor<9x9x9xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsPositiveNegativeHighPad
func.func @testTwoConsecutivePadsPositiveNegativeHighPad(%arg0: tensor<10x10x10xf32>) -> (tensor<11x11x11xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<-1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<12x12x12xf32>, tensor<f32>) -> tensor<11x11x11xf32>
return %3 : tensor<11x11x11xf32>
// CHECK: %[[RES:.*]] = "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<0> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<11x11x11xf32>
// CHECK: return %[[RES]] : tensor<11x11x11xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsNegativePositiveHighPad
func.func @testTwoConsecutivePadsNegativePositiveHighPad(%arg0: tensor<10x10x10xf32>) -> (tensor<11x11x11xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<-1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<10x10x10xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<11x11x11xf32>
return %3 : tensor<11x11x11xf32>
// CHECK: "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<-1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<10x10x10xf32>
// CHECK: "mhlo.pad"(%1, %0) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<0> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<11x11x11xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsDifferentPadVal
func.func @testTwoConsecutivePadsDifferentPadVal(%arg0: tensor<10x10x10xf32>) -> (tensor<14x14x14xf32>) {
%0 = mhlo.constant dense<1.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<12x12x12xf32>, tensor<f32>) -> tensor<14x14x14xf32>
return %3 : tensor<14x14x14xf32>
// CHECK: "mhlo.pad"(%arg0, %1) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
// CHECK: "mhlo.pad"(%2, %0) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<12x12x12xf32>, tensor<f32>) -> tensor<14x14x14xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsDifferentUsers
func.func @testTwoConsecutivePadsDifferentUsers(%arg0: tensor<10x10x10xf32>) -> (tensor<13x13x13xf32>, tensor<12x12x12xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
%2 = mhlo.exponential %1 : tensor<12x12x12xf32>
%3 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%4 = "mhlo.pad"(%1, %3) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<12x12x12xf32>, tensor<f32>) -> tensor<13x13x13xf32>
return %4, %2 : tensor<13x13x13xf32>, tensor<12x12x12xf32>
// CHECK: "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
// CHECK: "mhlo.pad"(%1, %0) <{
// CHECK-SAME: edge_padding_high = dense<1> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<0> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<12x12x12xf32>, tensor<f32>) -> tensor<13x13x13xf32>
}
// -----
// CHECK-LABEL: testTwoConsecutivePadsMultipleDownstreamUsers
func.func @testTwoConsecutivePadsMultipleDownstreamUsers(%arg0: tensor<10x10x10xf32>) -> (tensor<13x13x13xf32>, tensor<13x13x13xf32>) {
%0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%1 = "mhlo.pad"(%arg0, %0) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<1> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<12x12x12xf32>
%2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<12x12x12xf32>, tensor<f32>) -> tensor<13x13x13xf32>
%4 = mhlo.exponential %3 : tensor<13x13x13xf32>
%5 = mhlo.tanh %3 : tensor<13x13x13xf32>
return %4, %5 : tensor<13x13x13xf32>, tensor<13x13x13xf32>
// CHECK: "mhlo.pad"(%arg0, %0) <{
// CHECK-SAME: edge_padding_high = dense<2> : tensor<3xi64>,
// CHECK-SAME: edge_padding_low = dense<1> : tensor<3xi64>,
// CHECK-SAME: interior_padding = dense<0> : tensor<3xi64>
// CHECK-SAME: }> : (tensor<10x10x10xf32>, tensor<f32>) -> tensor<13x13x13xf32>
// CHECK: mhlo.exponential %1 : tensor<13x13x13xf32>
// CHECK: mhlo.tanh %1 : tensor<13x13x13xf32>
// CHECK: return %2, %3 : tensor<13x13x13xf32>, tensor<13x13x13xf32>
}
// -----
// CHECK-LABEL: testLiftDotConcatLHSSimple
func.func @testLiftDotConcatLHSSimple(%arg0: tensor<1x1x512xf32>, %arg1: tensor<2x1x512xf32>, %arg2: tensor<3x1x512xf32>, %arg3: tensor<512x13xf32>) -> tensor<6x1x13xf32> {
%0 = "mhlo.dot_general"(%arg0, %arg3) {
dot_dimension_numbers = #mhlo.dot<
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [0]
>} : (tensor<1x1x512xf32>, tensor<512x13xf32>) -> tensor<1x1x13xf32>
%1 = "mhlo.dot_general"(%arg1, %arg3) {
dot_dimension_numbers = #mhlo.dot<
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [0]
>} : (tensor<2x1x512xf32>, tensor<512x13xf32>) -> tensor<2x1x13xf32>
%2 = "mhlo.dot_general"(%arg2, %arg3) {
dot_dimension_numbers = #mhlo.dot<
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [0]
>} : (tensor<3x1x512xf32>, tensor<512x13xf32>) -> tensor<3x1x13xf32>
%r = "mhlo.concatenate"(%0, %1, %2) <{dimension = 0 : i64}> : (tensor<1x1x13xf32>, tensor<2x1x13xf32>, tensor<3x1x13xf32>) -> tensor<6x1x13xf32>
func.return %r : tensor<6x1x13xf32>
// CHECK: %[[R0:.*]] = "mhlo.concatenate"(%arg0, %arg1, %arg2) <{dimension = 0 : i64}> : (tensor<1x1x512xf32>, tensor<2x1x512xf32>, tensor<3x1x512xf32>) -> tensor<6x1x512xf32>
// CHECK: %[[R1:.*]] = "mhlo.dot_general"(%[[R0]], %arg3) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_contracting_dimensions = [2],
// CHECK-SAME: rhs_contracting_dimensions = [0]
// CHECK-SAME: >}> : (tensor<6x1x512xf32>, tensor<512x13xf32>) -> tensor<6x1x13xf32>
// CHECK: return %[[R1]] : tensor<6x1x13xf32>
}
// -----
// CHECK-LABEL: testLiftDotConcatLHSComplex
func.func @testLiftDotConcatLHSComplex(%arg0: tensor<1x9x2x3x8x4x10xf32>, %arg1: tensor<1x9x2x3x8x100x10xf32>, %arg2: tensor<9x2x1x5x10x5x8x7xf32>) -> tensor<1x2x3x104x5x5x7xf32> {
%0 = "mhlo.dot_general"(%arg0, %arg2) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0, 2],
rhs_batching_dimensions = [2, 1],
lhs_contracting_dimensions = [4, 1, 6],
rhs_contracting_dimensions = [6, 0, 4]
>} : (tensor<1x9x2x3x8x4x10xf32>, tensor<9x2x1x5x10x5x8x7xf32>) -> tensor<1x2x3x4x5x5x7xf32>
%1 = "mhlo.dot_general"(%arg1, %arg2) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0, 2],
rhs_batching_dimensions = [2, 1],
lhs_contracting_dimensions = [4, 1, 6],
rhs_contracting_dimensions = [6, 0, 4]
>} : (tensor<1x9x2x3x8x100x10xf32>, tensor<9x2x1x5x10x5x8x7xf32>) -> tensor<1x2x3x100x5x5x7xf32>
%r = "mhlo.concatenate"(%0, %1) <{dimension = 3 : i64}> : (tensor<1x2x3x4x5x5x7xf32>, tensor<1x2x3x100x5x5x7xf32>) -> tensor<1x2x3x104x5x5x7xf32>
func.return %r : tensor<1x2x3x104x5x5x7xf32>
// CHECK: %[[R0:.*]] = "mhlo.concatenate"(%arg0, %arg1) <{dimension = 5 : i64}> : (tensor<1x9x2x3x8x4x10xf32>, tensor<1x9x2x3x8x100x10xf32>) -> tensor<1x9x2x3x8x104x10xf32>
// CHECK: %[[R1:.*]] = "mhlo.dot_general"(%[[R0]], %arg2) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_batching_dimensions = [0, 2],
// CHECK-SAME: rhs_batching_dimensions = [2, 1],
// CHECK-SAME: lhs_contracting_dimensions = [4, 1, 6],
// CHECK-SAME: rhs_contracting_dimensions = [6, 0, 4]
// CHECK-SAME: >}> : (tensor<1x9x2x3x8x104x10xf32>, tensor<9x2x1x5x10x5x8x7xf32>) -> tensor<1x2x3x104x5x5x7xf32>
// CHECK: return %[[R1]] : tensor<1x2x3x104x5x5x7xf32>
}
// -----
// CHECK-LABEL: testLiftDotConcatLHSAndRHS
func.func @testLiftDotConcatLHSAndRHS(%arg0: tensor<1x72x128xf32>, %arg1: tensor<1x128x72xf32>, %arg2: tensor<1x72x128xf32>, %arg3: tensor<1x128x72xf32>, %arg4: tensor<1x72x128xf32>, %arg5: tensor<1x128x72xf32>, %arg6: tensor<1x72x128xf32>, %arg7: tensor<1x128x72xf32>) -> tensor<4x72x72xf32> {
%0 = "mhlo.dot_general"(%arg0, %arg1) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [0],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [1]
>} : (tensor<1x72x128xf32>, tensor<1x128x72xf32>) -> tensor<1x72x72xf32>
%1 = "mhlo.dot_general"(%arg2, %arg3) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [0],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [1]
>} : (tensor<1x72x128xf32>, tensor<1x128x72xf32>) -> tensor<1x72x72xf32>
%2 = "mhlo.dot_general"(%arg4, %arg5) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [0],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [1]
>} : (tensor<1x72x128xf32>, tensor<1x128x72xf32>) -> tensor<1x72x72xf32>
%3 = "mhlo.dot_general"(%arg6, %arg7) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [0],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [1]
>} : (tensor<1x72x128xf32>, tensor<1x128x72xf32>) -> tensor<1x72x72xf32>
%4 = "mhlo.concatenate"(%0, %1, %2, %3) <{dimension = 0 : i64}> : (tensor<1x72x72xf32>, tensor<1x72x72xf32>, tensor<1x72x72xf32>, tensor<1x72x72xf32>) -> tensor<4x72x72xf32>
func.return %4 : tensor<4x72x72xf32>
// CHECK: %[[R0:.*]] = "mhlo.concatenate"(%arg0, %arg2, %arg4, %arg6) <{dimension = 0 : i64}> : (tensor<1x72x128xf32>, tensor<1x72x128xf32>, tensor<1x72x128xf32>, tensor<1x72x128xf32>) -> tensor<4x72x128xf32>
// CHECK: %[[R1:.*]] = "mhlo.concatenate"(%arg1, %arg3, %arg5, %arg7) <{dimension = 0 : i64}> : (tensor<1x128x72xf32>, tensor<1x128x72xf32>, tensor<1x128x72xf32>, tensor<1x128x72xf32>) -> tensor<4x128x72xf32>
// CHECK: %[[R2:.*]] = "mhlo.dot_general"(%[[R0]], %[[R1]]) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_batching_dimensions = [0],
// CHECK-SAME: rhs_batching_dimensions = [0],
// CHECK-SAME: lhs_contracting_dimensions = [2],
// CHECK-SAME: rhs_contracting_dimensions = [1]
// CHECK-SAME: >}> : (tensor<4x72x128xf32>, tensor<4x128x72xf32>) -> tensor<4x72x72xf32>
// CHECK: return %[[R2]] : tensor<4x72x72xf32>
}
// -----
// CHECK-LABEL: testSliceConcat
func.func @testSliceConcat(%arg0: tensor<3x1x512xf32>) -> tensor<3x1x512xf32> {
%0 = "mhlo.slice"(%arg0) <{limit_indices = dense<[1, 1, 512]> : tensor<3xi64>, start_indices = dense<[0, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}> : (tensor<3x1x512xf32>) -> tensor<1x1x512xf32>
%1 = "mhlo.slice"(%arg0) <{limit_indices = dense<[2, 1, 512]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}> : (tensor<3x1x512xf32>) -> tensor<1x1x512xf32>
%2 = "mhlo.slice"(%arg0) <{limit_indices = dense<[3, 1, 512]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}> : (tensor<3x1x512xf32>) -> tensor<1x1x512xf32>
%r = "mhlo.concatenate"(%0, %1, %2) <{dimension = 0 : i64}> : (tensor<1x1x512xf32>, tensor<1x1x512xf32>, tensor<1x1x512xf32>) -> tensor<3x1x512xf32>
func.return %r : tensor<3x1x512xf32>
// CHECK: return %arg0 : tensor<3x1x512xf32>
}
// -----
// CHECK-LABEL: testConvertReshapeDotRhsToBatchedDot
func.func @testConvertReshapeDotRhsToBatchedDot(%arg0: tensor<1x72x72xf32>, %arg1: tensor<1x72x128xf32>) -> tensor<1x72x128xf32> {
%0 = mhlo.reshape %arg1 : (tensor<1x72x128xf32>) -> tensor<72x128xf32>
%1 = "mhlo.dot_general"(%arg0, %0) {
dot_dimension_numbers = #mhlo.dot<
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [0]
>} : (tensor<1x72x72xf32>, tensor<72x128xf32>) -> tensor<1x72x128xf32>
func.return %1 : tensor<1x72x128xf32>
// CHECK: %[[R:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME: dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME: lhs_batching_dimensions = [0],
// CHECK-SAME: rhs_batching_dimensions = [0],
// CHECK-SAME: lhs_contracting_dimensions = [2],
// CHECK-SAME: rhs_contracting_dimensions = [1]
// CHECK-SAME: >}> : (tensor<1x72x72xf32>, tensor<1x72x128xf32>) -> tensor<1x72x128xf32>
// CHECK: return %[[R]] : tensor<1x72x128xf32>
}
// -----
// CHECK-LABEL: broadcast_reshape_one_non_unit_dimnsion
func.func @broadcast_reshape_one_non_unit_dimnsion(%arg0: tensor<1x1x1x63xf32>) -> tensor<32x1x63xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x1x1x63xf32>) -> tensor<1x32x1x63xf32>
%1 = mhlo.reshape %0 : (tensor<1x32x1x63xf32>) -> tensor<32x1x63xf32>
return %1 : tensor<32x1x63xf32>
}
// CHECK: %0 = mhlo.reshape %arg0 : (tensor<1x1x1x63xf32>) -> tensor<63xf32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<63xf32>) -> tensor<32x1x63xf32>
// CHECK: return %1 : tensor<32x1x63xf32>
// -----
// CHECK-LABEL: broadcast_reshape_one_non_unit_dimnsion_trailing_zeros
func.func @broadcast_reshape_one_non_unit_dimnsion_trailing_zeros(%arg0: tensor<63x1x1x1xf32>) -> tensor<63x1x2xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<63x1x1x1xf32>) -> tensor<63x1x1x2xf32>
%1 = mhlo.reshape %0 : (tensor<63x1x1x2xf32>) -> tensor<63x1x2xf32>
return %1 : tensor<63x1x2xf32>
}
// CHECK: %0 = mhlo.reshape %arg0 : (tensor<63x1x1x1xf32>) -> tensor<63xf32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<63xf32>) -> tensor<63x1x2xf32>
// CHECK: return %1 : tensor<63x1x2xf32>
// -----
// CHECK-LABEL: broadcast_reshape_multiple_non_unit_dimension
func.func @broadcast_reshape_multiple_non_unit_dimension(%arg0: tensor<1x2x1x63xf32>) -> tensor<2x3x63xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x2x1x63xf32>) -> tensor<1x2x3x63xf32>
%1 = mhlo.reshape %0 : (tensor<1x2x3x63xf32>) -> tensor<2x3x63xf32>
return %1 : tensor<2x3x63xf32>
}
// CHECK: %0 = mhlo.reshape %arg0 : (tensor<1x2x1x63xf32>) -> tensor<2x63xf32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<[0, 2]> : tensor<2xi64>}> : (tensor<2x63xf32>) -> tensor<2x3x63xf32>
// CHECK: return %1 : tensor<2x3x63xf32>
// -----
// CHECK-LABEL: broadcast_reshape_multiple_non_unit_dimension_unsorted_broadcast_dims
func.func @broadcast_reshape_multiple_non_unit_dimension_unsorted_broadcast_dims(%arg0: tensor<1x2x1x63xf32>) -> tensor<3x2x63xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 2, 1, 3]> : tensor<4xi64>}> : (tensor<1x2x1x63xf32>) -> tensor<3x1x2x63xf32>
%1 = mhlo.reshape %0 : (tensor<3x1x2x63xf32>) -> tensor<3x2x63xf32>
return %1 : tensor<3x2x63xf32>
}
// CHECK: %0 = mhlo.reshape %arg0 : (tensor<1x2x1x63xf32>) -> tensor<2x63xf32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}> : (tensor<2x63xf32>) -> tensor<3x2x63xf32>
// CHECK: return %1 : tensor<3x2x63xf32>
// -----
// CHECK-LABEL: broadcast_reshape_broadcast_increases_rank
func.func @broadcast_reshape_broadcast_increases_rank(%arg0: tensor<1x2x1x63xf32>) -> tensor<2x3x63xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 4]> : tensor<4xi64>}> : (tensor<1x2x1x63xf32>) -> tensor<1x2x3x1x63xf32>
%1 = mhlo.reshape %0 : (tensor<1x2x3x1x63xf32>) -> tensor<2x3x63xf32>
return %1 : tensor<2x3x63xf32>
}
// CHECK: %0 = mhlo.reshape %arg0 : (tensor<1x2x1x63xf32>) -> tensor<2x63xf32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<[0, 2]> : tensor<2xi64>}> : (tensor<2x63xf32>) -> tensor<2x3x63xf32>
// CHECK: return %1 : tensor<2x3x63xf32>
// -----
// CHECK-LABEL: broadcast_reshape_not_same_non_unit_dims
func.func @broadcast_reshape_not_same_non_unit_dims(%arg0: tensor<63x1x1x1xf32>) -> tensor<2x1x63xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<63x1x1x1xf32>) -> tensor<63x1x1x2xf32>
%1 = mhlo.reshape %0 : (tensor<63x1x1x2xf32>) -> tensor<2x1x63xf32>
return %1 : tensor<2x1x63xf32>
}
// CHECK: %0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<63x1x1x1xf32>) -> tensor<63x1x1x2xf32>
// CHECK: %1 = mhlo.reshape %0 : (tensor<63x1x1x2xf32>) -> tensor<2x1x63xf32>
// CHECK: return %1 : tensor<2x1x63xf32>
// -----
// CHECK-LABEL: broadcast_reshape_multi_use
func.func @broadcast_reshape_multi_use(%arg0: tensor<1x1x1x63xf32>) -> (tensor<32x1x63xf32>, tensor<1x32x1x63xf32>) {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x1x1x63xf32>) -> tensor<1x32x1x63xf32>
%1 = mhlo.reshape %0 : (tensor<1x32x1x63xf32>) -> tensor<32x1x63xf32>
return %1, %0 : tensor<32x1x63xf32>, tensor<1x32x1x63xf32>
}
// CHECK: %0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x1x1x63xf32>) -> tensor<1x32x1x63xf32>
// CHECK: %1 = mhlo.reshape %0 : (tensor<1x32x1x63xf32>) -> tensor<32x1x63xf32>
// -----
// CHECK-LABEL: broadcast_reshape_rank_increase
func.func @broadcast_reshape_rank_increase(%arg0: tensor<1x1x1x63xf32>) -> tensor<32x1x1x1x1x63xf32> {
%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x1x1x63xf32>) -> tensor<1x32x1x63xf32>
%1 = mhlo.reshape %0 : (tensor<1x32x1x63xf32>) -> tensor<32x1x1x1x1x63xf32>
return %1 : tensor<32x1x1x1x1x63xf32>
}
// CHECK: %0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>}> : (tensor<1x1x1x63xf32>) -> tensor<1x32x1x63xf32>
// CHECK: %1 = mhlo.reshape %0 : (tensor<1x32x1x63xf32>) -> tensor<32x1x1x1x1x63xf32>
tensorflow/compiler/mlir/lite/stablehlo/tests/optimize_layout.mlir
+64
View File
@@ -0,0 +1,64 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s --transpose-commute-ops | FileCheck %s
// CHECK-LABEL: func.func @commute_transpose_pad(
// CHECK-SAME: %[[INPUT:.*]]: tensor<1x112x112x64xf32>,
// CHECK-SAME: %[[PAD_VAL:.*]]: tensor<f32>) -> tensor<1x64x114x114xf32> {
// CHECK: %[[PAD:.*]] = stablehlo.pad %[[INPUT]], %[[PAD_VAL]],
// CHECK: low = [0, 1, 1, 0], high = [0, 1, 1, 0], interior = [0, 0, 0, 0]
// CHECK: : (tensor<1x112x112x64xf32>, tensor<f32>) -> tensor<1x114x114x64xf32>
// CHECK: %[[TPOS:.*]] = stablehlo.transpose %[[PAD]], dims = [0, 3, 1, 2]
// CHECK: : (tensor<1x114x114x64xf32>) -> tensor<1x64x114x114xf32>
// CHECK: return %[[TPOS]] : tensor<1x64x114x114xf32>
func.func @commute_transpose_pad(
%arg0: tensor<1x112x112x64xf32>, %padding_val: tensor<f32>)
-> tensor<1x64x114x114xf32> {
%tspos = stablehlo.transpose %arg0, dims = [0, 3, 1, 2]
: (tensor<1x112x112x64xf32>) -> tensor<1x64x112x112xf32>
%ret = stablehlo.pad %tspos, %padding_val,
low = [0, 0, 1, 1], high = [0, 0, 1, 1], interior = [0, 0, 0, 0]
: (tensor<1x64x112x112xf32>, tensor<f32>) -> tensor<1x64x114x114xf32>
return %ret :tensor<1x64x114x114xf32>
}
// -----
// CHECK-LABEL: func.func @commute_transpose_reduce_window(
// CHECK-SAME: %[[INPUT:.*]]: tensor<1x114x114x64xf32>,
// CHECK-SAME: %[[PAD_VAL:.*]]: tensor<f32>) -> tensor<1x64x56x56xf32> {
// CHECK: %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[INPUT]], %[[PAD_VAL]])
// CHECK: <{window_dimensions = array<i64: 1, 3, 3, 1>,
// CHECK: window_strides = array<i64: 1, 2, 2, 1>}> ({
// CHECK: ^bb0(%[[ARG0:.*]]: tensor<f32>, %[[ARG1:.*]]: tensor<f32>):
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ARG0]], %[[ARG1]] : tensor<f32>
// CHECK: stablehlo.return %[[MAX]] : tensor<f32>
// CHECK: }) : (tensor<1x114x114x64xf32>, tensor<f32>) -> tensor<1x56x56x64xf32>
// CHECK: %[[TPOS:.*]] = stablehlo.transpose %[[REDUCE]], dims = [0, 3, 1, 2]
// CHECK: : (tensor<1x56x56x64xf32>) -> tensor<1x64x56x56xf32>
// CHECK: return %[[TPOS]] : tensor<1x64x56x56xf32>
func.func @commute_transpose_reduce_window(
%input: tensor<1x114x114x64xf32>,
%cst: tensor<f32>) -> tensor<1x64x56x56xf32> {
%tpos = stablehlo.transpose %input, dims = [0, 3, 1, 2]
: (tensor<1x114x114x64xf32>) -> tensor<1x64x114x114xf32>
%ret = "stablehlo.reduce_window"(%tpos, %cst)
<{window_dimensions = array<i64: 1, 1, 3, 3>, window_strides = array<i64: 1, 1, 2, 2>}> ({
^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
%max = stablehlo.maximum %arg0, %arg1 : tensor<f32>
stablehlo.return %max: tensor<f32>
}) : (tensor<1x64x114x114xf32>, tensor<f32>) -> tensor<1x64x56x56xf32>
return %ret : tensor<1x64x56x56xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/prepare_hlo.mlir
+909
View File
@@ -0,0 +1,909 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -prepare-hlo -split-input-file | FileCheck %s --dump-input=fail
// Just assert that pass is properly registered.
func.func @main(%arg0: tensor<f32>) -> tensor<f32> {
return %arg0: tensor<f32>
}
// CHECK-LABEL: main
// -----
//===----------------------------------------------------------------------===//
// mhlo.convolution
//===----------------------------------------------------------------------===//
// 2D
//=--
// CHECK-LABEL: transpose_conv2d_same_padding_nchw_ihwo
func.func @transpose_conv2d_same_padding_nchw_ihwo(%input: tensor<1x2x256x256xf32>, %filter:tensor<2x2x4x4xf32>) -> tensor<1x2x512x512xf32> {
%1 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[i, o, 0, 1]->[b, f, 0, 1],
window = {pad = [[2, 2], [2, 2]], lhs_dilate = [2, 2]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<1x2x256x256xf32>, tensor<2x2x4x4xf32>) -> tensor<1x2x512x512xf32>
func.return %1 : tensor<1x2x512x512xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: %[[TRANSPOSED_KERNEL:.*]] = "mhlo.transpose"(%arg1)
// CHECK-SAME: permutation
// CHECK-SAME: [1, 2, 3, 0]
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[TRANSPOSED_INPUT]], %[[TRANSPOSED_KERNEL]])
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "mhlo.transpose"(%[[CONV_OUT]])
// CHECK-SAME: permutation
// CHECK-SAME: [0, 3, 1, 2]
// -----
// CHECK-LABEL: transpose_conv2d_same_padding_nchw_oihw
func.func @transpose_conv2d_same_padding_nchw_oihw(%input: tensor<1x2x256x256xf32>, %filter:tensor<2x2x4x4xf32>) -> tensor<1x2x512x512xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1],
window = {pad = [[2, 2], [2, 2]], lhs_dilate = [2, 2]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<1x2x256x256xf32>, tensor<2x2x4x4xf32>) -> tensor<1x2x512x512xf32>
func.return %0 : tensor<1x2x512x512xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: %[[TRANSPOSED_KERNEL:.*]] = "mhlo.transpose"(%arg1)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[TRANSPOSED_INPUT]], %[[TRANSPOSED_KERNEL]])
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "mhlo.transpose"(%[[CONV_OUT]])
// CHECK-SAME: permutation
// CHECK-SAME: [0, 3, 1, 2]
// -----
// CHECK-LABEL: depthwise_transpose_conv2d_same_padding_nchw_hwoi
func.func @depthwise_transpose_conv2d_same_padding_nchw_hwoi(%input: tensor<1x2x20x20xf32>, %filter:tensor<8x8x2x1xf32>) -> tensor<1x2x80x80xf32> {
%1 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[0, 1, o, i]->[b, f, 0, 1],
window = {pad = [[5, 5], [5, 5]], lhs_dilate = [4, 4]}
{batch_group_count = 1 : i64, feature_group_count = 2 : i64}
: (tensor<1x2x20x20xf32>, tensor<8x8x2x1xf32>) -> tensor<1x2x80x80xf32>
func.return %1 : tensor<1x2x80x80xf32>
// CHECK: %0 = "mhlo.transpose"(%arg0) <{permutation = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : (tensor<1x2x20x20xf32>) -> tensor<1x20x20x2xf32>
// CHECK: %1 = "mhlo.transpose"(%arg1) <{permutation = dense<[2, 0, 1, 3]> : tensor<4xi64>}> : (tensor<8x8x2x1xf32>) -> tensor<2x8x8x1xf32>
// CHECK: %2 = "mhlo.slice"(%0) <{limit_indices = dense<[1, 20, 20, 1]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<1x20x20x2xf32>) -> tensor<1x20x20x1xf32>
// CHECK: %3 = "mhlo.slice"(%1) <{limit_indices = dense<[1, 8, 8, 1]> : tensor<4xi64>, start_indices = dense<0> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<2x8x8x1xf32>) -> tensor<1x8x8x1xf32>
// CHECK: %4 = mhlo.convolution(%2, %3) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f], window = {pad = {{\[\[}}5, 5], [5, 5]], lhs_dilate = [4, 4]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x20x20x1xf32>, tensor<1x8x8x1xf32>) -> tensor<1x80x80x1xf32>
// CHECK: %5 = "mhlo.slice"(%0) <{limit_indices = dense<[1, 20, 20, 2]> : tensor<4xi64>, start_indices = dense<[0, 0, 0, 1]> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<1x20x20x2xf32>) -> tensor<1x20x20x1xf32>
// CHECK: %6 = "mhlo.slice"(%1) <{limit_indices = dense<[2, 8, 8, 1]> : tensor<4xi64>, start_indices = dense<[1, 0, 0, 0]> : tensor<4xi64>, strides = dense<1> : tensor<4xi64>}> : (tensor<2x8x8x1xf32>) -> tensor<1x8x8x1xf32>
// CHECK: %7 = mhlo.convolution(%5, %6) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f], window = {pad = {{\[\[}}5, 5], [5, 5]], lhs_dilate = [4, 4]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x20x20x1xf32>, tensor<1x8x8x1xf32>) -> tensor<1x80x80x1xf32>
// CHECK: %8 = "mhlo.concatenate"(%4, %7) <{dimension = 3 : i64}> : (tensor<1x80x80x1xf32>, tensor<1x80x80x1xf32>) -> tensor<1x80x80x2xf32>
// CHECK: %9 = "mhlo.transpose"(%8) <{permutation = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : (tensor<1x80x80x2xf32>) -> tensor<1x2x80x80xf32>
// CHECK: return %9 : tensor<1x2x80x80xf32>
}
// -----
// CHECK-LABEL: conv2d_nhwc_ohwi_nhwc
func.func @conv2d_nhwc_ohwi_nhwc(%input: tensor<1x256x256x3xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
window_strides = dense<1> : tensor<2xi64>,
padding = dense<0> : tensor<2x2xi64>,
rhs_dilation = dense<[1, 1]> : tensor<2xi64>,
lhs_dilation = dense<[1, 1]> : tensor<2xi64>
} : (tensor<1x256x256x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
func.return %0 : tensor<1x256x256x2xf32>
}
// CHECK-NOT: transpose
// CHECK: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv2d_nhwc_ohwi_nhwc_dynamic
func.func @conv2d_nhwc_ohwi_nhwc_dynamic(%input: tensor<?x256x256x3xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<?x256x256x2xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
window_strides = dense<1> : tensor<2xi64>,
padding = dense<0> : tensor<2x2xi64>,
rhs_dilation = dense<[1, 1]> : tensor<2xi64>,
lhs_dilation = dense<[1, 1]> : tensor<2xi64>
} : (tensor<?x256x256x3xf32>, tensor<2x1x1x3xf32>) -> tensor<?x256x256x2xf32>
func.return %0 : tensor<?x256x256x2xf32>
}
// CHECK-NOT: transpose
// CHECK: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv2d_nchw_ohwi_nhwc
func.func @conv2d_nchw_ohwi_nhwc(%input: tensor<?x3x256x256xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<?x256x256x2xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[o, 0, 1, i]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<?x3x256x256xf32>, tensor<2x1x1x3xf32>) -> tensor<?x256x256x2xf32>
func.return %0 : tensor<?x256x256x2xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: mhlo.convolution(%[[TRANSPOSED_INPUT]], %arg1)
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv2d_nchw_ohwi_nhwc_dynamic_batch
func.func @conv2d_nchw_ohwi_nhwc_dynamic_batch(%input: tensor<?x3x256x256xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<?x256x256x2xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[o, 0, 1, i]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<?x3x256x256xf32>, tensor<2x1x1x3xf32>) -> tensor<?x256x256x2xf32>
func.return %0 : tensor<?x256x256x2xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: mhlo.convolution(%[[TRANSPOSED_INPUT]], %arg1)
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv2d_nhwc_hwio_nhwc
func.func @conv2d_nhwc_hwio_nhwc(%input: tensor<1x256x256x3xf32>, %filter: tensor<1x1x3x2xf32>) -> tensor<1x256x256x2xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<1x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x256x256x2xf32>
func.return %0 : tensor<1x256x256x2xf32>
}
// CHECK: %[[TRANSPOSED_KERNEL:.*]] = "mhlo.transpose"(%arg1)
// CHECK-SAME: permutation
// CHECK-SAME: [3, 0, 1, 2]
// CHECK: mhlo.convolution(%arg0, %[[TRANSPOSED_KERNEL]])
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv2d_nhwc_ohwi_nchw
func.func @conv2d_nhwc_ohwi_nchw(%input: tensor<1x256x256x3xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<1x2x256x256xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, f, 0, 1],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<1x256x256x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x2x256x256xf32>
func.return %0 : tensor<1x2x256x256xf32>
}
// CHECK-NOT: transpose
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "mhlo.transpose"(%[[CONV_OUT]])
// CHECK-SAME: permutation
// CHECK-SAME: [0, 3, 1, 2]
// -----
// CHECK-LABEL: conv2d_nchw_oihw_nchw
func.func @conv2d_nchw_oihw_nchw(%input: tensor<1x3x256x256xf32>, %filter: tensor<2x3x1x1xf32>) -> tensor<1x2x256x256xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<1x3x256x256xf32>, tensor<2x3x1x1xf32>) -> tensor<1x2x256x256xf32>
func.return %0 : tensor<1x2x256x256xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: %[[TRANSPOSED_KERNEL:.*]] = "mhlo.transpose"(%arg1)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[TRANSPOSED_INPUT]], %[[TRANSPOSED_KERNEL]])
// CHECK-SAME: [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "mhlo.transpose"(%[[CONV_OUT]])
// CHECK-SAME: permutation
// CHECK-SAME: [0, 3, 1, 2]
// -----
// CHECK-LABEL: conv2d_nhwc_ohwi_nhwc_padded
func.func @conv2d_nhwc_ohwi_nhwc_padded(%input: tensor<1x254x254x3xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32> {
%0 = "mhlo.convolution"(%input, %filter) {
dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]>,
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
window_strides = dense<1> : tensor<2xi64>,
padding = dense<1> : tensor<2x2xi64>,
rhs_dilation = dense<[1, 1]> : tensor<2xi64>,
lhs_dilation = dense<[1, 1]> : tensor<2xi64>
} : (tensor<1x254x254x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
func.return %0 : tensor<1x256x256x2xf32>
}
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"
// CHECK-SAME: edge_padding_high = dense<[0, 1, 1, 0]>
// CHECK-SAME: edge_padding_low = dense<[0, 1, 1, 0]>
// CHECK-SAME: interior_padding = dense<0>
// CHECK: mhlo.convolution(%[[PADDED_LHS]]
// CHECK-SAME: pad
// CHECK-SAME: [0, 0], [0, 0]
// CHECK-SAME: (tensor<1x256x256x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
// -----
// CHECK-LABEL: conv2d_nhwc_ohwi_nhwc_asymmetric_padded
func.func @conv2d_nhwc_ohwi_nhwc_asymmetric_padded(%input: tensor<1x255x255x3xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32> {
%0 = "mhlo.convolution"(%input, %filter) {
dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]>,
batch_group_count = 1 : i64,
feature_group_count = 1 : i64,
window_strides = dense<1> : tensor<2xi64>,
padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>,
rhs_dilation = dense<[1, 1]> : tensor<2xi64>,
lhs_dilation = dense<[1, 1]> : tensor<2xi64>
} : (tensor<1x255x255x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
func.return %0 : tensor<1x256x256x2xf32>
}
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"
// CHECK-SAME: edge_padding_high = dense<[0, 1, 1, 0]>
// CHECK-SAME: edge_padding_low = dense<0>
// CHECK-SAME: interior_padding = dense<0>
// CHECK: mhlo.convolution(%[[PADDED_LHS]]
// CHECK-SAME: pad
// CHECK-SAME: [0, 0], [0, 0]
// CHECK-SAME: (tensor<1x256x256x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
// -----
// CHECK-LABEL: conv2d_nchw_ohwi_nhwc_padded
func.func @conv2d_nchw_ohwi_nhwc_padded(%input: tensor<1x3x253x249xf32>, %filter: tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32> {
%0 = mhlo.convolution(%input, %filter)
dim_numbers = [b, f, 0, 1]x[o, 0, 1, i]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[1, 2], [3, 4]]} {
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<1x3x253x249xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
func.return %0 : tensor<1x256x256x2xf32>
}
// Want to ensure that we transpose before padding input (which this test does implicitly).
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"
// CHECK-SAME: edge_padding_high = dense<[0, 2, 4, 0]>
// CHECK-SAME: edge_padding_low = dense<[0, 1, 3, 0]>
// CHECK-SAME: interior_padding = dense<0>
// CHECK: mhlo.convolution(%[[PADDED_LHS]], %arg1)
// CHECK-SAME: pad
// CHECK-SAME: [0, 0], [0, 0]
// CHECK-SAME: (tensor<1x256x256x3xf32>, tensor<2x1x1x3xf32>) -> tensor<1x256x256x2xf32>
// -----
// CHECK-LABEL: conv2d_nchw_ohwi_nhwc_padded_dilated_lhs
func.func @conv2d_nchw_ohwi_nhwc_padded_dilated_lhs(%input: tensor<1x64x64x256xf32>, %filter: tensor<64x2x2x256xf32>) -> tensor<1x128x128x64xf32> {
%0 = "mhlo.convolution"(%input, %filter) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]>,
feature_group_count = 1 : i64,
lhs_dilation = dense<2> : tensor<2xi64>,
padding = dense<1> : tensor<2x2xi64>,
precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
rhs_dilation = dense<1> : tensor<2xi64>,
window_reversal = dense<false> : tensor<2xi1>,
window_strides = dense<1> : tensor<2xi64>} :
(tensor<1x64x64x256xf32>, tensor<64x2x2x256xf32>) -> tensor<1x128x128x64xf32>
func.return %0 : tensor<1x128x128x64xf32>
}
// CHECK-NOT: mhlo.pad
// CHECK: mhlo.convolution
// CHECK-SAME: pad
// CHECK-SAME: [1, 1], [1, 1]
// CHECK-SAME: lhs_dilate = [2, 2]
// -----
// CHECK-LABEL: depthwise_conv2d_nhwc_ohwi_nhwc
func.func @depthwise_conv2d_nhwc_ohwi_nhwc(%arg0: tensor<1x10x10x207xf32>, %arg1: tensor<3312x3x3x1xf32>) -> tensor<1x8x8x3312xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]>,
feature_group_count = 207 : i64,
lhs_dilation = dense<1> : tensor<2xi64>,
padding = dense<0> : tensor<2x2xi64>,
precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
rhs_dilation = dense<1> : tensor<2xi64>,
window_strides = dense<1> : tensor<2xi64>
} : (tensor<1x10x10x207xf32>, tensor<3312x3x3x1xf32>) -> tensor<1x8x8x3312xf32>
func.return %0 : tensor<1x8x8x3312xf32>
}
// CHECK: %[[TRANSPOSED_KERNEL:.*]] = "mhlo.transpose"(%arg1)
// CHECK-SAME: permutation
// CHECK-SAME: [3, 1, 2, 0]
// CHECK: mhlo.convolution(%arg0, %[[TRANSPOSED_KERNEL]])
// CHECK-SAME: [b, 0, 1, f]x[i, 0, 1, o]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: depthwise_conv2d_nchw_ihwo_nhwc
func.func @depthwise_conv2d_nchw_ihwo_nhwc(%arg0: tensor<1x207x10x10xf32>, %arg1: tensor<1x3x3x3312xf32>) -> tensor<1x8x8x3312xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, f, 0, 1]x[i, 0, 1, o]->[b, 0, 1, f]>,
feature_group_count = 207 : i64,
lhs_dilation = dense<1> : tensor<2xi64>,
padding = dense<0> : tensor<2x2xi64>,
precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
rhs_dilation = dense<1> : tensor<2xi64>,
window_strides = dense<1> : tensor<2xi64>
} : (tensor<1x207x10x10xf32>, tensor<1x3x3x3312xf32>) -> tensor<1x8x8x3312xf32>
func.return %0 : tensor<1x8x8x3312xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: mhlo.convolution(%[[TRANSPOSED_INPUT]], %arg1)
// CHECK-SAME: [b, 0, 1, f]x[i, 0, 1, o]->[b, 0, 1, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: depthwise_conv2d_nchw_ihwo_nhwc_padded
func.func @depthwise_conv2d_nchw_ihwo_nhwc_padded(%arg0: tensor<1x207x8x8xf32>, %arg1: tensor<1x3x3x3312xf32>) -> tensor<1x8x8x3312xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, f, 0, 1]x[i, 0, 1, o]->[b, 0, 1, f]>,
feature_group_count = 207 : i64,
lhs_dilation = dense<1> : tensor<2xi64>,
padding = dense<1> : tensor<2x2xi64>,
precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
rhs_dilation = dense<1> : tensor<2xi64>,
window_strides = dense<1> : tensor<2xi64>
} : (tensor<1x207x8x8xf32>, tensor<1x3x3x3312xf32>) -> tensor<1x8x8x3312xf32>
func.return %0 : tensor<1x8x8x3312xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 1]
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"(%[[TRANSPOSED_INPUT]]
// CHECK-SAME: edge_padding_high = dense<[0, 1, 1, 0]>
// CHECK-SAME: edge_padding_low = dense<[0, 1, 1, 0]>
// CHECK-SAME: interior_padding = dense<0>
// CHECK: mhlo.convolution(%[[PADDED_LHS]], %arg1)
// CHECK-SAME: [b, 0, 1, f]x[i, 0, 1, o]->[b, 0, 1, f]
// CHECK-SAME: pad =
// CHECK-SAME: [0, 0], [0, 0]
// -----
// 3D
//=--
// CHECK-LABEL: conv3d_ndhwc_dhwio_ndhwc
func.func @conv3d_ndhwc_dhwio_ndhwc(%arg0: tensor<1x8x8x32x207xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]>,
feature_group_count = 1 : i64} :
(tensor<1x8x8x32x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32>
func.return %0 : tensor<1x6x6x1x16xf32>
}
// CHECK-NOT: transpose
// CHECK: mhlo.convolution
// CHECK-SAME: [b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv3d_ncdhw_dhwio_ndhwc
func.func @conv3d_ncdhw_dhwio_ndhwc(%arg0: tensor<1x207x8x8x32xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, f, 0, 1, 2]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]>,
feature_group_count = 1 : i64} :
(tensor<1x207x8x8x32xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32>
func.return %0 : tensor<1x6x6x1x16xf32>
}
// CHECK: %[[TRANSPOSED_INPUT:.*]] = "mhlo.transpose"(%arg0)
// CHECK-SAME: permutation
// CHECK-SAME: [0, 2, 3, 4, 1]
// CHECK: mhlo.convolution(%[[TRANSPOSED_INPUT]], %arg1)
// CHECK-SAME: [b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv3d_ndhwc_odhwi_ndhwc
func.func @conv3d_ndhwc_odhwi_ndhwc(%arg0: tensor<1x8x8x32x207xf32>, %arg1: tensor<16x3x3x32x207xf32>) -> tensor<1x6x6x1x16xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, 1, 2, f]x[o, 0, 1, 2, i]->[b, 0, 1, 2, f]>,
feature_group_count = 1 : i64} :
(tensor<1x8x8x32x207xf32>, tensor<16x3x3x32x207xf32>) -> tensor<1x6x6x1x16xf32>
func.return %0 : tensor<1x6x6x1x16xf32>
}
// CHECK: %[[TRANSPOSED_KERNEL:.*]] = "mhlo.transpose"(%arg1)
// CHECK-SAME: permutation
// CHECK-SAME: [1, 2, 3, 4, 0]
// CHECK: mhlo.convolution(%arg0, %[[TRANSPOSED_KERNEL]])
// CHECK-SAME: [b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]
// CHECK-NOT: transpose
// -----
// CHECK-LABEL: conv3d_ndhwc_dhwio_ncdhw
func.func @conv3d_ndhwc_dhwio_ncdhw(%arg0: tensor<1x8x8x32x207xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<1x16x6x6x1xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, f, 0, 1, 2]>,
feature_group_count = 1 : i64} :
(tensor<1x8x8x32x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x16x6x6x1xf32>
func.return %0 : tensor<1x16x6x6x1xf32>
}
// CHECK-NOT: transpose
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution
// CHECK-SAME: [b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]
// CHECK: "mhlo.transpose"(%[[CONV_OUT]])
// CHECK-SAME: permutation
// CHECK-SAME: [0, 4, 1, 2, 3]
// -----
// CHECK-LABEL: conv3d_ndhwc_dhwio_ndhwc_padded
func.func @conv3d_ndhwc_dhwio_ndhwc_padded(%arg0: tensor<1x6x6x30x207xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]>,
feature_group_count = 1 : i64,
padding = dense<1> : tensor<3x2xi64>} :
(tensor<1x6x6x30x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32>
func.return %0 : tensor<1x6x6x1x16xf32>
}
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"
// CHECK-SAME: edge_padding_high = dense<[0, 1, 1, 1, 0]>
// CHECK-SAME: edge_padding_low = dense<[0, 1, 1, 1, 0]>
// CHECK-SAME: interior_padding = dense<0>
// CHECK: mhlo.convolution(%[[PADDED_LHS]], %arg1)
// CHECK-SAME: pad =
// CHECK-SAME: [0, 0], [0, 0], [0, 0]
// CHECK-SAME: (tensor<1x8x8x32x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32>
// -----
// CHECK-LABEL: conv3d_ncdhw_dhwio_ndhwc_padded
func.func @conv3d_ncdhw_dhwio_ndhwc_padded(%arg0: tensor<1x207x6x6x30xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, f, 0, 1, 2]x[0, 1, 2, i, o]->[b, 0, 1, 2, f]>,
feature_group_count = 1 : i64,
padding = dense<1> : tensor<3x2xi64>} :
(tensor<1x207x6x6x30xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32>
func.return %0 : tensor<1x6x6x1x16xf32>
}
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"
// CHECK-SAME: edge_padding_high = dense<[0, 1, 1, 1, 0]>
// CHECK-SAME: edge_padding_low = dense<[0, 1, 1, 1, 0]>
// CHECK-SAME: interior_padding = dense<0>
// CHECK: mhlo.convolution(%[[PADDED_LHS]], %arg1)
// CHECK-SAME: pad =
// CHECK-SAME: [0, 0], [0, 0], [0, 0]
// CHECK-SAME: (tensor<1x8x8x32x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<1x6x6x1x16xf32>
// -----
// 1D
//=--
// CHECK-LABEL: conv1d_nsc_osi_nsc
func.func @conv1d_nsc_osi_nsc(%arg0: tensor<16x32x256xf32>, %arg1: tensor<256x1x256xf32>) -> tensor<16x32x256xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, f]x[o, 0, i]->[b, 0, f]>,
feature_group_count = 1 : i64,
lhs_dilation = dense<1> : tensor<1xi64>,
padding = dense<0> : tensor<1x2xi64>,
precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
rhs_dilation = dense<1> : tensor<1xi64>,
window_strides = dense<1> : tensor<1xi64>
} : (tensor<16x32x256xf32>, tensor<256x1x256xf32>) -> tensor<16x32x256xf32>
func.return %0 : tensor<16x32x256xf32>
}
// CHECK: %[[RESHAPED_LHS:.*]] = "tfl.expand_dims"(%arg0
// CHECK: %[[RESHAPED_RHS:.*]] = "tfl.expand_dims"(%arg1
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[RESHAPED_LHS]], %[[RESHAPED_RHS]]) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "tfl.squeeze"(%[[CONV_OUT]]
// -----
// CHECK-LABEL: conv1d_nsc_sio_nsc
func.func @conv1d_nsc_sio_nsc(%arg0: tensor<16x32x256xf32>, %arg1: tensor<1x256x256xf32>) -> tensor<16x32x256xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, 0, f]x[0, i, o]->[b, 0, f]>,
feature_group_count = 1 : i64,
lhs_dilation = dense<1> : tensor<1xi64>,
padding = dense<0> : tensor<1x2xi64>,
precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
rhs_dilation = dense<1> : tensor<1xi64>,
window_strides = dense<1> : tensor<1xi64>
} : (tensor<16x32x256xf32>, tensor<1x256x256xf32>) -> tensor<16x32x256xf32>
func.return %0 : tensor<16x32x256xf32>
}
// CHECK: %[[RESHAPED_LHS:.*]] = "tfl.expand_dims"(%arg0
// CHECK: %[[RESHAPED_RHS:.*]] = "tfl.expand_dims"(%arg1
// CHECK: %[[TPOSED_RHS:.*]] = "mhlo.transpose"(%[[RESHAPED_RHS]]) <{permutation = dense<[3, 0, 1, 2]> : tensor<4xi64>}> : (tensor<1x1x256x256xf32>) -> tensor<256x1x1x256xf32>
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[RESHAPED_LHS]], %[[TPOSED_RHS]]) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "tfl.squeeze"(%[[CONV_OUT]]
// -----
// CHECK-LABEL: conv1d_ncs_osi_nsc_padded
func.func @conv1d_ncs_osi_nsc_padded(%arg0: tensor<16x256x30xf32>, %arg1: tensor<256x1x256xf32>) -> tensor<16x32x256xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, f, 0]x[o, 0, i]->[b, 0, f]>,
feature_group_count = 1 : i64,
padding = dense<1> : tensor<1x2xi64>
} : (tensor<16x256x30xf32>, tensor<256x1x256xf32>) -> tensor<16x32x256xf32>
func.return %0 : tensor<16x32x256xf32>
}
// CHECK: %[[RESHAPED_LHS:.*]] = "tfl.expand_dims"(%arg0{{.*}}-> tensor<16x256x30x1xf32>
// CHECK: %[[RESHAPED_RHS:.*]] = "tfl.expand_dims"(%arg1{{.*}}-> tensor<256x1x1x256xf32>
// CHECK: %[[TPOSED_LHS:.*]] = "mhlo.transpose"(%[[RESHAPED_LHS]]) <{permutation = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : (tensor<16x256x30x1xf32>) -> tensor<16x30x1x256xf32>
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"(%[[TPOSED_LHS]], %cst) <{edge_padding_high = dense<[0, 1, 0, 0]> : tensor<4xi64>, edge_padding_low = dense<[0, 1, 0, 0]> : tensor<4xi64>, interior_padding = dense<0> : tensor<4xi64>}> : (tensor<16x30x1x256xf32>, tensor<f32>) -> tensor<16x32x1x256xf32>
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[PADDED_LHS]], %[[RESHAPED_RHS]]) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "tfl.squeeze"(%[[CONV_OUT]]{{.*}}-> tensor<16x32x256xf32>
// -----
// CHECK-LABEL: conv1d_ncs_osi_nsc_padded_dynamic
func.func @conv1d_ncs_osi_nsc_padded_dynamic(%arg0: tensor<?x256x30xf32>, %arg1: tensor<256x1x256xf32>) -> tensor<?x32x256xf32> {
%0 = "mhlo.convolution"(%arg0, %arg1) {
batch_group_count = 1 : i64,
dimension_numbers = #mhlo.conv<[b, f, 0]x[o, 0, i]->[b, 0, f]>,
feature_group_count = 1 : i64,
padding = dense<1> : tensor<1x2xi64>
} : (tensor<?x256x30xf32>, tensor<256x1x256xf32>) -> tensor<?x32x256xf32>
func.return %0 : tensor<?x32x256xf32>
}
// CHECK: %[[RESHAPED_LHS:.*]] = "tfl.expand_dims"(%arg0{{.*}}-> tensor<?x256x30x1xf32>
// CHECK: %[[RESHAPED_RHS:.*]] = "tfl.expand_dims"(%arg1{{.*}}-> tensor<256x1x1x256xf32>
// CHECK: %[[TPOSED_LHS:.*]] = "mhlo.transpose"(%[[RESHAPED_LHS]]) <{permutation = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : (tensor<?x256x30x1xf32>) -> tensor<?x30x1x256xf32>
// CHECK: %[[PADDED_LHS:.*]] = "mhlo.pad"(%[[TPOSED_LHS]], %cst) <{edge_padding_high = dense<[0, 1, 0, 0]> : tensor<4xi64>, edge_padding_low = dense<[0, 1, 0, 0]> : tensor<4xi64>, interior_padding = dense<0> : tensor<4xi64>}> : (tensor<?x30x1x256xf32>, tensor<f32>) -> tensor<?x32x1x256xf32>
// CHECK: %[[CONV_OUT:.*]] = mhlo.convolution(%[[PADDED_LHS]], %[[RESHAPED_RHS]]) dim_numbers = [b, 0, 1, f]x[o, 0, 1, i]->[b, 0, 1, f]
// CHECK: "tfl.squeeze"(%[[CONV_OUT]]{{.*}}-> tensor<?x32x256xf32>
// -----
//===----------------------------------------------------------------------===//
// mhlo.pad
//===----------------------------------------------------------------------===//
// CHECK-LABEL: pad_2d
func.func @pad_2d(%arg0: tensor<3x3xf32>, %arg1: tensor<f32>) -> tensor<4x3xf32> {
%0 = "mhlo.pad"(%arg0, %arg1) {
edge_padding_low = dense<[0, -1]> : tensor<2xi64>,
edge_padding_high = dense<[1, 1]> : tensor<2xi64>,
interior_padding = dense<0> : tensor<2xi64>
} : (tensor<3x3xf32>, tensor<f32>) -> tensor<4x3xf32>
func.return %0 : tensor<4x3xf32>
}
// CHECK: mhlo.slice
// CHECK-SAME: limit_indices = dense<3>
// CHECK-SAME: start_indices = dense<[0, 1]>
// CHECK-SAME: (tensor<3x3xf32>) -> tensor<3x2xf32>
// CHECK: mhlo.pad
// CHECK-SAME: edge_padding_high = dense<1>
// CHECK-SAME: edge_padding_low = dense<0>
// CHECK-SAME: (tensor<3x2xf32>, tensor<f32>) -> tensor<4x3xf32>
// -----
// CHECK-LABEL: pad_2d_negative
func.func @pad_2d_negative(%arg0: tensor<3x3xf32>, %arg1: tensor<f32>) -> tensor<1x2xf32> {
%0 = "mhlo.pad"(%arg0, %arg1) {
edge_padding_low = dense<[-1, -1]> : tensor<2xi64>,
edge_padding_high = dense<[-1, 0]> : tensor<2xi64>,
interior_padding = dense<0> : tensor<2xi64>
} : (tensor<3x3xf32>, tensor<f32>) -> tensor<1x2xf32>
func.return %0 : tensor<1x2xf32>
}
// CHECK: mhlo.slice
// CHECK-SAME: limit_indices = dense<[2, 3]>
// CHECK-SAME: start_indices = dense<1>
// CHECK-SAME: (tensor<3x3xf32>) -> tensor<1x2xf32>
// CHECK-NOT: mhlo.pad
// -----
// CHECK-LABEL: pad_3d_mixed
func.func @pad_3d_mixed(%arg0: tensor<3x3x3xf32>, %arg1: tensor<f32>) -> tensor<3x3x3xf32> {
%0 = "mhlo.pad"(%arg0, %arg1) {
edge_padding_low = dense<[-1, 1, 0]> : tensor<3xi64>,
edge_padding_high = dense<[1, -1, 0]> : tensor<3xi64>,
interior_padding = dense<0> : tensor<3xi64>
} : (tensor<3x3x3xf32>, tensor<f32>) -> tensor<3x3x3xf32>
func.return %0 : tensor<3x3x3xf32>
}
// CHECK: mhlo.slice
// CHECK-SAME: limit_indices = dense<[3, 2, 3]>
// CHECK-SAME: start_indices = dense<[1, 0, 0]>
// CHECK-SAME: (tensor<3x3x3xf32>) -> tensor<2x2x3xf32>
// CHECK: mhlo.pad
// CHECK-SAME: edge_padding_high = dense<[1, 0, 0]>
// CHECK-SAME: edge_padding_low = dense<[0, 1, 0]>
// CHECK-SAME: (tensor<2x2x3xf32>, tensor<f32>) -> tensor<3x3x3xf32>
// -----
//===----------------------------------------------------------------------===//
// mhlo.reduce_window
//===----------------------------------------------------------------------===//
// CHECK-LABEL: reduce_window_valid_channel_first
func.func @reduce_window_valid_channel_first(%arg0: tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> {
// "0xFF800000" represents -INF for f32.
%0 = mhlo.constant dense<0xFF800000> : tensor<f32>
%1 = "mhlo.reduce_window"(%arg0, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%2 = mhlo.maximum %arg1, %arg2 : tensor<f32>
mhlo.return %2 : tensor<f32>
}) {
base_dilations = dense<1> : tensor<4xi64>,
padding = dense<0> : tensor<4x2xi64>,
window_dilations = dense<1> : tensor<4xi64>,
window_dimensions = dense<[1, 1, 3, 3]> : tensor<4xi64>,
window_strides = dense<[1, 1, 2, 2]> : tensor<4xi64>} : (tensor<4x3x16x16xf32>, tensor<f32>) -> tensor<4x3x7x7xf32>
func.return %1 : tensor<4x3x7x7xf32>
}
// CHECK: %[[INIT_CST:.*]] = mhlo.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[TPOSE_IN:.*]] = "mhlo.transpose"(%arg0) <{permutation = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : (tensor<4x3x16x16xf32>) -> tensor<4x16x16x3xf32>
// CHECK: %[[RW:.*]] = "mhlo.reduce_window"(%[[TPOSE_IN]], %[[INIT_CST]])
// CHECK-SAME: window_dimensions = dense<[1, 3, 3, 1]>
// CHECK-SAME: window_strides = dense<[1, 2, 2, 1]>
// CHECK: %3 = "mhlo.transpose"(%[[RW]]) <{permutation = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : (tensor<4x7x7x3xf32>) -> tensor<4x3x7x7xf32>
// -----
// CHECK-LABEL: reduce_window_same_channel_first
func.func @reduce_window_same_channel_first(%arg0: tensor<4x3x16x16xf32>) -> tensor<4x3x8x8xf32> {
// "0xFF800000" represents -INF for f32.
%0 = mhlo.constant dense<0xFF800000> : tensor<f32>
%1 = "mhlo.reduce_window"(%arg0, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%6 = mhlo.maximum %arg1, %arg2 : tensor<f32>
"mhlo.return"(%6) : (tensor<f32>) -> ()
}) {
base_dilations = dense<1> : tensor<4xi64>,
padding = dense<[[0, 0], [0, 0], [0, 1], [0, 1]]> : tensor<4x2xi64>,
window_dilations = dense<1> : tensor<4xi64>,
window_dimensions = dense<[1, 1, 3, 3]> : tensor<4xi64>,
window_strides = dense<[1, 1, 2, 2]> : tensor<4xi64>} : (tensor<4x3x16x16xf32>, tensor<f32>) -> tensor<4x3x8x8xf32>
func.return %1 : tensor<4x3x8x8xf32>
}
// CHECK: %[[INIT_CST:.*]] = mhlo.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[TPOSE_IN:.*]] = "mhlo.transpose"(%arg0) <{permutation = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : (tensor<4x3x16x16xf32>) -> tensor<4x16x16x3xf32>
// CHECK: %[[RW:.*]] = "mhlo.reduce_window"(%[[TPOSE_IN]], %[[INIT_CST]])
// CHECK-SAME: padding
// CHECK-SAME: [0, 0], [0, 1], [0, 1], [0, 0]
// CHECK-SAME: window_dimensions = dense<[1, 3, 3, 1]>
// CHECK-SAME: window_strides = dense<[1, 2, 2, 1]>
// CHECK: %3 = "mhlo.transpose"(%[[RW]]) <{permutation = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : (tensor<4x8x8x3xf32>) -> tensor<4x3x8x8xf32>
// -----
//===----------------------------------------------------------------------===//
// mhlo.dynamic_slice
//===----------------------------------------------------------------------===//
// CHECK-LABEL: dynamic_slice
func.func @dynamic_slice(%arg0: tensor<7x3xf32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<4x2xf32> {
%0 = "mhlo.dynamic_slice"(%arg0, %arg1, %arg2) <{slice_sizes = dense<[4, 2]> : tensor<2xi64>}> : (tensor<7x3xf32>, tensor<i32>, tensor<i32>) -> tensor<4x2xf32>
func.return %0 : tensor<4x2xf32>
}
// CHECK: mhlo.dynamic_slice
// CHECK-SAME: (tensor<7x3xf32>, tensor<i32>, tensor<i32>) -> tensor<4x2xf32>
// -----
// CHECK-LABEL: dynamic_slice_ui32
func.func @dynamic_slice_ui32(%arg0: tensor<7x3xf32>, %arg1: tensor<ui32>, %arg2: tensor<ui32>) -> tensor<4x2xf32> {
%0 = "mhlo.dynamic_slice"(%arg0, %arg1, %arg2) <{slice_sizes = dense<[4, 2]> : tensor<2xi64>}> : (tensor<7x3xf32>, tensor<ui32>, tensor<ui32>) -> tensor<4x2xf32>
func.return %0 : tensor<4x2xf32>
}
// CHECK: mhlo.dynamic_slice
// CHECK-SAME: (tensor<7x3xf32>, tensor<i32>, tensor<i32>) -> tensor<4x2xf32>
// CHECK-LABEL: dynamic_slice_ui64
func.func @dynamic_slice_ui64(%arg0: tensor<7x3xf32>, %arg1: tensor<ui64>, %arg2: tensor<ui64>) -> tensor<4x2xf32> {
%0 = "mhlo.dynamic_slice"(%arg0, %arg1, %arg2) <{slice_sizes = dense<[4, 2]> : tensor<2xi64>}> : (tensor<7x3xf32>, tensor<ui64>, tensor<ui64>) -> tensor<4x2xf32>
func.return %0 : tensor<4x2xf32>
}
// CHECK: mhlo.dynamic_slice
// CHECK-SAME: (tensor<7x3xf32>, tensor<i64>, tensor<i64>) -> tensor<4x2xf32>
// -----
// CHECK-LABEL: dynamic_slice_i64
func.func @dynamic_slice_i64(%arg0: tensor<7x3xf32>, %arg1: tensor<i64>, %arg2: tensor<i64>) -> tensor<4x2xf32> {
%0 = "mhlo.dynamic_slice"(%arg0, %arg1, %arg2) <{slice_sizes = dense<[4, 2]> : tensor<2xi64>}> : (tensor<7x3xf32>, tensor<i64>, tensor<i64>) -> tensor<4x2xf32>
func.return %0 : tensor<4x2xf32>
}
// CHECK: mhlo.dynamic_slice
// CHECK-SAME: (tensor<7x3xf32>, tensor<i64>, tensor<i64>) -> tensor<4x2xf32>
//===----------------------------------------------------------------------===//
// mhlo.custom_call
//===----------------------------------------------------------------------===//
// -----
// CHECK-LABEL: @shape_assertion_custom_call
func.func @shape_assertion_custom_call(%arg1: tensor<?x5xi32>) -> tensor<i32> {
%0 = mhlo.constant dense<3> : tensor<i32>
%1 = "mhlo.get_dimension_size"(%arg1) <{dimension = 0 : i64}> : (tensor<?x5xi32>) -> tensor<i32>
%ok = mhlo.compare EQ, %1, %0, SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>
mhlo.custom_call @shape_assertion(%ok) {
error_message = "The error message",
has_side_effect = true
} : (tensor<i1>) -> ()
return %1 : tensor<i32>
}
// CHECK-NOT: mhlo.custom_call
//===----------------------------------------------------------------------===//
// mhlo.fft
//===----------------------------------------------------------------------===//
// CHECK-LABEL: rfft_2d
func.func @rfft_2d(%arg0: tensor<1x512xf32>) -> tensor<1x257xcomplex<f32>> {
%0 = "mhlo.fft"(%arg0) <{fft_length = dense<512> : tensor<1xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<1x512xf32>) -> tensor<1x257xcomplex<f32>>
func.return %0 : tensor<1x257xcomplex<f32>>
}
// CHECK: %0 = "mhlo.fft"(%arg0) <{fft_length = dense<[1, 512]> : tensor<2xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<1x512xf32>) -> tensor<1x257xcomplex<f32>>
// CHECK: return %0 : tensor<1x257xcomplex<f32>>
// -----
// CHECK-LABEL: @fft
func.func @fft(%arg0: tensor<3x9xcomplex<f32>>) -> tensor<3x9xcomplex<f32>> {
%0 = "mhlo.fft"(%arg0) <{ fft_length = dense<9> : tensor<1xi64>, fft_type = #mhlo<fft_type FFT> }> : (tensor<3x9xcomplex<f32>>) -> tensor<3x9xcomplex<f32>>
func.return %0 : tensor<3x9xcomplex<f32>>
}
// CHECK: %0 = "mhlo.fft"(%arg0) <{fft_length = dense<9> : tensor<1xi64>, fft_type = #mhlo<fft_type FFT>}> : (tensor<3x9xcomplex<f32>>) -> tensor<3x9xcomplex<f32>>
// CHECK: return %0 : tensor<3x9xcomplex<f32>>
// -----
// CHECK-LABEL: @mhlo_nd_fft
func.func @mhlo_nd_fft(%arg0: tensor<2x3x345x256xf32>) -> tensor<2x3x345x129xcomplex<f32>> {
%43 = "mhlo.fft"(%arg0) <{fft_length = dense<256> : tensor<1xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<2x3x345x256xf32>) -> tensor<2x3x345x129xcomplex<f32>>
return %43 : tensor<2x3x345x129xcomplex<f32>>
}
// CHECK: %0 = mhlo.reshape %arg0 : (tensor<2x3x345x256xf32>) -> tensor<2x3x345x1x256xf32>
// CHECK: %1 = "mhlo.fft"(%0) <{fft_length = dense<[1, 256]> : tensor<2xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<2x3x345x1x256xf32>) -> tensor<2x3x345x1x129xcomplex<f32>>
// CHECK: %2 = mhlo.reshape %1 : (tensor<2x3x345x1x129xcomplex<f32>>) -> tensor<2x3x345x129xcomplex<f32>>
// CHECK: return %2 : tensor<2x3x345x129xcomplex<f32>>
// -----
// CHECK-LABEL: @mhlo_dynamic_fft_1
func.func @mhlo_dynamic_fft_1(%arg0: tensor<?x9x2560xf32>) -> tensor<?x9x1281xcomplex<f32>> {
%0 = "mhlo.fft"(%arg0) <{fft_length = dense<2560> : tensor<1xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<?x9x2560xf32>) -> tensor<?x9x1281xcomplex<f32>>
return %0 : tensor<?x9x1281xcomplex<f32>>
// CHECK: %4 = "mhlo.get_dimension_size"(%arg0) <{dimension = 0 : i64}> : (tensor<?x9x2560xf32>) -> tensor<i32>
// CHECK: %5 = mhlo.reshape %4 : (tensor<i32>) -> tensor<1xi32>
// CHECK: %6 = "mhlo.concatenate"(%5, %3, %2, %1) <{dimension = 0 : i64}> : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK: %7 = mhlo.dynamic_reshape %arg0, %6 : (tensor<?x9x2560xf32>, tensor<4xi32>) -> tensor<?x9x1x2560xf32>
// CHECK: %8 = "mhlo.fft"(%7) <{fft_length = dense<[1, 2560]> : tensor<2xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<?x9x1x2560xf32>) -> tensor<?x9x1x1281xcomplex<f32>>
// CHECK: %9 = "mhlo.get_dimension_size"(%8) <{dimension = 0 : i64}> : (tensor<?x9x1x1281xcomplex<f32>>) -> tensor<i32>
// CHECK: %10 = mhlo.reshape %9 : (tensor<i32>) -> tensor<1xi32>
// CHECK: %11 = "mhlo.concatenate"(%10, %3, %0) <{dimension = 0 : i64}> : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32>
// CHECK: %12 = mhlo.dynamic_reshape %8, %11 : (tensor<?x9x1x1281xcomplex<f32>>, tensor<3xi32>) -> tensor<?x9x1281xcomplex<f32>>
// CHECK: return %12 : tensor<?x9x1281xcomplex<f32>>
}
// -----
// CHECK-LABEL: @mhlo_dynamic_fft_2
func.func @mhlo_dynamic_fft_2(%arg0: tensor<?x?x2560xf32>) -> tensor<?x?x1281xcomplex<f32>> {
%0 = "mhlo.fft"(%arg0) <{fft_length = dense<2560> : tensor<1xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<?x?x2560xf32>) -> tensor<?x?x1281xcomplex<f32>>
return %0 : tensor<?x?x1281xcomplex<f32>>
// CHECK: %3 = "mhlo.get_dimension_size"(%arg0) <{dimension = 0 : i64}> : (tensor<?x?x2560xf32>) -> tensor<i32>
// CHECK: %4 = mhlo.reshape %3 : (tensor<i32>) -> tensor<1xi32>
// CHECK: %5 = "mhlo.get_dimension_size"(%arg0) <{dimension = 1 : i64}> : (tensor<?x?x2560xf32>) -> tensor<i32>
// CHECK: %6 = mhlo.reshape %5 : (tensor<i32>) -> tensor<1xi32>
// CHECK: %7 = "mhlo.concatenate"(%4, %6, %2, %1) <{dimension = 0 : i64}> : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK: %8 = mhlo.dynamic_reshape %arg0, %7 : (tensor<?x?x2560xf32>, tensor<4xi32>) -> tensor<?x?x1x2560xf32>
// CHECK: %9 = "mhlo.fft"(%8) <{fft_length = dense<[1, 2560]> : tensor<2xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<?x?x1x2560xf32>) -> tensor<?x?x1x1281xcomplex<f32>>
// CHECK: %10 = "mhlo.get_dimension_size"(%9) <{dimension = 0 : i64}> : (tensor<?x?x1x1281xcomplex<f32>>) -> tensor<i32>
// CHECK: %11 = mhlo.reshape %10 : (tensor<i32>) -> tensor<1xi32>
// CHECK: %12 = "mhlo.get_dimension_size"(%9) <{dimension = 1 : i64}> : (tensor<?x?x1x1281xcomplex<f32>>) -> tensor<i32>
// CHECK: %13 = mhlo.reshape %12 : (tensor<i32>) -> tensor<1xi32>
// CHECK: %14 = "mhlo.concatenate"(%11, %13, %0) <{dimension = 0 : i64}> : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32>
// CHECK: %15 = mhlo.dynamic_reshape %9, %14 : (tensor<?x?x1x1281xcomplex<f32>>, tensor<3xi32>) -> tensor<?x?x1281xcomplex<f32>>
// CHECK: return %15 : tensor<?x?x1281xcomplex<f32>>
}
// -----
// CHECK-LABEL: @mhlo_dynamic_fft_2_neg
func.func @mhlo_dynamic_fft_2_neg(%arg0: tensor<?x9x?xf32>) -> tensor<?x9x1281xcomplex<f32>> {
%0 = "mhlo.fft"(%arg0) <{fft_length = dense<2560> : tensor<1xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<?x9x?xf32>) -> tensor<?x9x1281xcomplex<f32>>
return %0 : tensor<?x9x1281xcomplex<f32>>
// CHECK: %0 = "mhlo.fft"(%arg0) <{fft_length = dense<2560> : tensor<1xi64>, fft_type = #mhlo<fft_type RFFT>}> : (tensor<?x9x?xf32>) -> tensor<?x9x1281xcomplex<f32>>
// CHECK: return %0 : tensor<?x9x1281xcomplex<f32>>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/stablehlo-custom-call-legalize-composite.mlir
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -stablehlo-custom-call-legalize-composite | FileCheck %s
// CHECK-LABEL: module
module {
// CHECK-LABEL: @main
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) {
// CHECK: stablehlo.custom_call @foo
stablehlo.custom_call @foo() : () -> ()
// CHECK-NOT: stablehlo.custom_call
// CHECK: stablehlo.composite "odml.foo" %arg0, %arg1 {composite_attributes = {bar = 500 : i64}, decomposition = @foo.impl} : (tensor<1xf32>, tensor<2xf32>) -> (tensor<2xf32>, tensor<1xf32>)
%1:2 = stablehlo.custom_call @stablehlo.composite(%arg0, %arg1) {called_computations = [@foo.impl], composite.backend_config = {attributes = {bar = 500 : i64}, name = "odml.foo"}} : (tensor<1xf32>, tensor<2xf32>) -> (tensor<2xf32>, tensor<1xf32>)
return
}
// CHECK-LABEL: func private @foo.impl
func.func private @foo.impl(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> (tensor<2xf32>, tensor<1xf32>) {
return %arg1, %arg0 : tensor<2xf32>, tensor<1xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/stablehlo_fuse_convolution.mlir
+59
View File
@@ -0,0 +1,59 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -stablehlo-fuse-convolution -cse | FileCheck %s
// CHECK-LABEL: @fuseMulAndConv2D
// CHECK-SAME: %[[INPUT:[^:[:space:]]+]]
func.func @fuseMulAndConv2D(%input: tensor<1x256x256x3xf32>) -> (tensor<1x256x256x2xf32>) {
// CHECK-DAG: %[[FILTER:.+]] = stablehlo.constant dense<{{\[\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]]]> : tensor<1x1x3x2xf32>
// CHECK-DAG: %[[CST:.+]] = stablehlo.constant dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>
// CHECK-DAG: %[[CST_BCAST:.+]] = stablehlo.broadcast_in_dim %[[CST]], dims = [3] : (tensor<2xf32>) -> tensor<1x1x3x2xf32>
// CHECK-DAG: %[[NEW_FILTER:.+]] = stablehlo.multiply %[[FILTER]], %[[CST_BCAST]] : tensor<1x1x3x2xf32>
// CHECK-DAG: %[[RESULT:.+]] = stablehlo.convolution(%[[INPUT]], %[[NEW_FILTER]]) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = {{\[\[}}0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x256x256x2xf32>
%filter = stablehlo.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]]> : tensor<1x1x3x2xf32>
%cst = stablehlo.constant dense<[0.1, 0.2]> : tensor<2xf32>
%0 = stablehlo.convolution(%input, %filter) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x256x256x2xf32>
%1 = stablehlo.broadcast_in_dim %cst, dims = [3] : (tensor<2xf32>) -> tensor<1x256x256x2xf32>
%2 = stablehlo.multiply %0, %1 : tensor<1x256x256x2xf32>
// CHECK-DAG: return %[[RESULT]]
func.return %2 : tensor<1x256x256x2xf32>
}
// -----
// CHECK-LABEL: @fuseMulAndConv2DDynamic
// CHECK-SAME: %[[INPUT:[^:[:space:]]+]]
func.func @fuseMulAndConv2DDynamic(%input: tensor<?x256x256x3xf32>) -> (tensor<?x256x256x2xf32>) {
// CHECK-DAG: %[[FILTER:.+]] = stablehlo.constant dense<{{\[\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]]]> : tensor<1x1x3x2xf32>
// CHECK-DAG: %[[CST_0:.+]] = stablehlo.constant dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>
// CHECK-DAG: %[[CST_1:.+]] = stablehlo.constant dense<[3.000000e-01, 4.000000e-01]> : tensor<2xf32>
// CHECK: %[[CST_BCAST:.+]] = stablehlo.broadcast_in_dim %[[CST_0]], dims = [3] : (tensor<2xf32>) -> tensor<1x1x3x2xf32>
// CHECK: %[[NEW_FILTER:.+]] = stablehlo.multiply %[[FILTER]], %[[CST_BCAST]] : tensor<1x1x3x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[INPUT]], %[[NEW_FILTER]]) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = {{\[\[}}0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<?x256x256x2xf32>
// CHECK: %[[SHAPE:.+]] = shape.shape_of %[[CONV]] : tensor<?x256x256x2xf32> -> tensor<4xindex>
// CHECK: %[[DYNAMIC_BCAST:.+]] = stablehlo.dynamic_broadcast_in_dim %[[CST_1]], %[[SHAPE]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BCAST]] : tensor<?x256x256x2xf32>
%filter = stablehlo.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]]> : tensor<1x1x3x2xf32>
%cst_0 = stablehlo.constant dense<[0.1, 0.2]> : tensor<2xf32>
%cst_1 = stablehlo.constant dense<[0.3, 0.4]> : tensor<2xf32>
%0 = stablehlo.convolution(%input, %filter) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x256x256x3xf32>, tensor<1x1x3x2xf32>) -> tensor<?x256x256x2xf32>
%1 = shape.shape_of %0 : tensor<?x256x256x2xf32> -> tensor<4xindex>
%2 = stablehlo.dynamic_broadcast_in_dim %cst_0, %1, dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
%3 = stablehlo.multiply %0, %2 : tensor<?x256x256x2xf32>
%4 = stablehlo.dynamic_broadcast_in_dim %cst_1, %1, dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
%5 = stablehlo.add %3, %4 : tensor<?x256x256x2xf32>
// CHECK-DAG: return %[[ADD]]
func.return %5 : tensor<?x256x256x2xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/stablehlo_unfuse_batch_norm.mlir
+186
View File
@@ -0,0 +1,186 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -stablehlo-unfuse-batch-norm -cse -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @batchNormInference_2D_inner_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_2D_inner_features(
%x: tensor<4x256xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>,
%mean: tensor<256xf32>, %variance: tensor<256xf32>)
-> (tensor<4x256xf32>) {
// CHECK: %[[CST:.+]] = stablehlo.constant dense<1.001000e-05> : tensor<f32>
// CHECK-NEXT: %[[EPS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[CST]], dims = [] : (tensor<f32>) -> tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = stablehlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = stablehlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = stablehlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = stablehlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[RHS:.+]] = stablehlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = stablehlo.broadcast_in_dim %[[MULTIPLIER]], dims = [1] : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[X_NORMED:.+]] = stablehlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<4x256xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[RHS]], dims = [1] : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[RESULT:.+]] = stablehlo.add %[[X_NORMED]], %[[RHS_BCAST]] : tensor<4x256xf32>
%0 = "stablehlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.001000e-05 : f32, feature_index = 1 : i64} :
(tensor<4x256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>,
tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: return %[[RESULT]]
func.return %0 : tensor<4x256xf32>
}
// -----
// CHECK-LABEL: @batchNormInference_4D_middle_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_4D_middle_features(
%x: tensor<3x4x256x6xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>,
%mean: tensor<256xf32>, %variance: tensor<256xf32>)
-> (tensor<3x4x256x6xf32>) {
// CHECK: %[[CST:.+]] = stablehlo.constant dense<1.001000e-05> : tensor<f32>
// CHECK-NEXT: %[[EPS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[CST]], dims = [] : (tensor<f32>) -> tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = stablehlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = stablehlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = stablehlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = stablehlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[RHS:.+]] = stablehlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = stablehlo.broadcast_in_dim %[[MULTIPLIER]], dims = [2] : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[RHS]], dims = [2] : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
%0 = "stablehlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.001000e-05 : f32, feature_index = 2 : i64} :
(tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>,
tensor<256xf32>) -> tensor<3x4x256x6xf32>
func.return %0 : tensor<3x4x256x6xf32>
}
// -----
// CHECK-LABEL: @batchNormInference_dynamic_shape
// Validate that dynamic shapes are handled properly.
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_dynamic_shape(
%x: tensor<?x?x?x?xf32>, %scale: tensor<?xf32>, %offset: tensor<?xf32>,
%mean: tensor<?xf32>, %variance: tensor<?xf32>)
-> tensor<?x?x?x?xf32> {
// CHECK-DAG: %[[EPS:.+]] = stablehlo.constant dense<1.000000e-03> : tensor<f32>
// CHECK-DAG: %[[VAR_SHAPE:.+]] = shape.shape_of %[[VARIANCE]] : tensor<?xf32> -> tensor<1xindex>
// CHECK-DAG: %[[EPS_BCAST:.+]] = stablehlo.dynamic_broadcast_in_dim %[[EPS]], %[[VAR_SHAPE]], dims = [] : (tensor<f32>, tensor<1xindex>) -> tensor<?xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = stablehlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<?xf32>
// CHECK-DAG: %[[R_STDDEV:.+]] = stablehlo.rsqrt %[[VARIANCE_EPS]] : tensor<?xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = stablehlo.multiply %[[R_STDDEV]], %[[SCALE]] : tensor<?xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = stablehlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<?xf32>
// CHECK-DAG: %[[RHS:.+]] = stablehlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<?xf32>
// CHECK-DAG: %[[X_SHAPE:.+]] = shape.shape_of %[[X]] : tensor<?x?x?x?xf32> -> tensor<4xindex>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = stablehlo.dynamic_broadcast_in_dim %[[MULTIPLIER]], %[[X_SHAPE]], dims = [1] : (tensor<?xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
// CHECK-DAG: %[[X_NORMED:.+]] = stablehlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = stablehlo.dynamic_broadcast_in_dim %[[RHS]], %[[X_SHAPE]], dims = [1] : (tensor<?xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
// CHECK-DAG: %[[RESULT:.+]] = stablehlo.add %[[X_NORMED]], %[[RHS_BCAST]] : tensor<?x?x?x?xf32>
%0 = "stablehlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 0.001 : f32, feature_index = 1 : i64} :
(tensor<?x?x?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>,
tensor<?xf32>) -> tensor<?x?x?x?xf32>
func.return %0 : tensor<?x?x?x?xf32>
}
// -----
// CHECK-LABEL: @batchNormInference_f64
// Validate that epsilon is properly promoted to f64
// CHECK: %[[CST:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<f64>
// CHECK-NEXT: [[EPS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[CST]], dims = [] : (tensor<f64>) -> tensor<256xf64>
func.func @batchNormInference_f64(
%x: tensor<4x256xf64>, %scale: tensor<256xf64>, %offset: tensor<256xf64>,
%mean: tensor<256xf64>, %variance: tensor<256xf64>)
-> (tensor<4x256xf64>) {
%0 = "stablehlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.0 : f32, feature_index = 1 : i64} :
(tensor<4x256xf64>, tensor<256xf64>, tensor<256xf64>, tensor<256xf64>,
tensor<256xf64>) -> tensor<4x256xf64>
func.return %0 : tensor<4x256xf64>
}
// -----
// CHECK-LABEL: @batchNormInference_f16
// Validate that epsilon is properly down to f16
// CHECK: %[[EPS:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<f16>
// CHECK-NEXT: %[[EPS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[EPS]], dims = [] : (tensor<f16>) -> tensor<256xf16>
func.func @batchNormInference_f16(
%x: tensor<4x256xf16>, %scale: tensor<256xf16>, %offset: tensor<256xf16>,
%mean: tensor<256xf16>, %variance: tensor<256xf16>)
-> (tensor<4x256xf16>) {
%0 = "stablehlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.0 : f32, feature_index = 1 : i64} :
(tensor<4x256xf16>, tensor<256xf16>, tensor<256xf16>, tensor<256xf16>,
tensor<256xf16>) -> tensor<4x256xf16>
func.return %0 : tensor<4x256xf16>
}
// -----
// Validate that epsilon is overflow
func.func @batchNormInference_f16_overflow(
%x: tensor<4x256xf16>, %scale: tensor<256xf16>, %offset: tensor<256xf16>,
%mean: tensor<256xf16>, %variance: tensor<256xf16>)
-> (tensor<4x256xf16>) {
// expected-warning @+1 {{Could not convert batch_norm epsilon to target fp type: opStatus = 24}}
%0 = "stablehlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 0.00000001 : f32, feature_index = 1 : i64} :
(tensor<4x256xf16>, tensor<256xf16>, tensor<256xf16>, tensor<256xf16>,
tensor<256xf16>) -> tensor<4x256xf16>
func.return %0 : tensor<4x256xf16>
}
// -----
// CHECK-LABEL: @batchNormTraining_4D_middle_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
func.func @batchNormTraining_4D_middle_features(
%x: tensor<3x4x256x6xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>)
-> (tensor<3x4x256x6xf32>) {
// CHECK: %[[CST:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
// CHECK-DAG: %[[CST_AXIS:.+]] = "tf.Const"() <{value = dense<[0, 1, 3]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK-DAG: %[[X_SHAPE:.+]] = shape.shape_of %[[X]] : tensor<3x4x256x6xf32> -> tensor<4xindex>
// CHECK-DAG: %[[MEAN:.+]] = "tf.Mean"(%arg0, %[[CST_AXIS]]) <{keep_dims = false}> : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32>
// CHECK-DAG: %[[MEAN_BCAST:.+]] = stablehlo.dynamic_broadcast_in_dim %[[MEAN]], %[[X_SHAPE]], dims = [2] : (tensor<256xf32>, tensor<4xindex>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[SQ_DIFF:.+]] = "tf.SquaredDifference"(%arg0, %[[MEAN_BCAST]]) : (tensor<3x4x256x6xf32>, tensor<3x4x256x6xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[VARIANCE:.+]] = "tf.Mean"(%[[SQ_DIFF]], %[[CST_AXIS]]) <{keep_dims = false}> : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32>
// CHECK-DAG: %[[EPS:.+]] = stablehlo.broadcast_in_dim %[[CST]], dims = [] : (tensor<f32>) -> tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = stablehlo.add %[[VARIANCE]], %[[EPS]] : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = stablehlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = stablehlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = stablehlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[RHS:.+]] = stablehlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = stablehlo.broadcast_in_dim %[[MULTIPLIER]], dims = [2] : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[X_NORMED:.+]] = stablehlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<3x4x256x6xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = stablehlo.broadcast_in_dim %[[RHS]], dims = [2] : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[RESULT:.+]] = stablehlo.add %[[X_NORMED]], %[[RHS_BCAST]] : tensor<3x4x256x6xf32>
%0:3 = "stablehlo.batch_norm_training"(%x, %scale, %offset)
{epsilon = 1.0 : f32, feature_index = 2 : i64} :
(tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>) -> (tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>)
func.return %0 : tensor<3x4x256x6xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-clamp.mlir
+23
View File
@@ -0,0 +1,23 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0 : tensor<5xf32>, %arg1 : tensor<f32>, %arg2 : tensor<f32>) -> tensor<5xf32> {
%0 = "tf.ClipByValue"(%arg0, %arg1, %arg2) : (tensor<5xf32>, tensor<f32>, tensor<f32>) -> tensor<5xf32>
func.return %0 : tensor<5xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-concat.mlir
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
%axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
%1 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor<i64>) -> tensor<6x3xf32>
func.return %1 : tensor<6x3xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-conv.mlir
+23
View File
@@ -0,0 +1,23 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0: tensor<4x68x68x3xf32>, %arg1: tensor<5x5x3x8xf32>) -> tensor<4x64x64x8xf32> {
%0 = "tf.Conv2D"(%arg0, %arg1) {padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<4x68x68x3xf32>, tensor<5x5x3x8xf32>) -> tensor<4x64x64x8xf32>
func.return %0 : tensor<4x64x64x8xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-division.mlir
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> {
%0 = "tf.Div"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-logistic.mlir
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = "tf.Sigmoid"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
func.return %0 : tensor<2xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-multiply.mlir
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> {
%0 = "tf.Mul"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
func.return %0 : tensor<2xi32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo-resize-bilinear.mlir
+24
View File
@@ -0,0 +1,24 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]
module {
func.func @main(%arg0: tensor<21x32x32x128xf32>, %arg1: tensor<2xi32>) -> tensor<1x64x64x128xf32> {
%0 = "tf.ResizeBilinear"(%arg0, %arg1) {align_corners = false, device = "", half_pixel_centers = true} : (tensor<21x32x32x128xf32>, tensor<2xi32>) -> tensor<1x64x64x128xf32>
func.return %0 : tensor<1x64x64x128xf32>
}
}
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo.mlir
+37
View File
@@ -0,0 +1,37 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s
module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
%1 = arith.constant dense<1> : tensor<1xi32>
%2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
%3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
: (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
func.return %3 : tensor<2x1x2xf32>
}
}
//CHECK: module attributes
//CHECK-SAME: keep_stablehlo_constant = "true"
//CHECK-NEXT: func.func @main(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> attributes {tf.entry_function = {inputs = "arg0", outputs = "vhlo.dynamic_update_slice_v1"}} {
//CHECK-DAG: %[[c0:.+]] = stablehlo.constant dense<2.000000e+00> : tensor<1x1x2xf32>
//CHECK-DAG: %[[c1:.+]] = stablehlo.constant dense<1> : tensor<i32>
//CHECK-DAG: %[[c2:.+]] = stablehlo.constant dense<0> : tensor<i32>
//CHECK-NEXT: %[[c3:.+]] = stablehlo.dynamic_update_slice %arg0, %[[c0]], %[[c1]], %[[c2]], %[[c2]] : (tensor<2x1x2xf32>, tensor<1x1x2xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2x1x2xf32>
//CHECK-NEXT: return %[[c3]] : tensor<2x1x2xf32>
//CHECK-NEXT: }
tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-tf-quantize.mlir
+37
View File
@@ -0,0 +1,37 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
//RUN: tf_tfl_translate --post-training-quantization --enable-stablehlo-conversion --input-mlir --output-mlir %s -o - | FileCheck %s
module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
%1 = arith.constant dense<1> : tensor<1xi32>
%2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
%3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
: (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
func.return %3 : tensor<2x1x2xf32>
}
}
//CHECK: module {
//CHECK-NEXT: func.func @main(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
//CHECK-DAG: %[[c0:.+]] = stablehlo.constant dense<2.000000e+00> : tensor<1x1x2xf32>
//CHECK-DAG: %[[c1:.+]] = stablehlo.constant dense<1> : tensor<i32>
//CHECK-DAG: %[[c2:.+]] = stablehlo.constant dense<0> : tensor<i32>
//CHECK-NEXT: %[[c3:.+]] = stablehlo.dynamic_update_slice %arg0, %[[c0]], %[[c1]], %[[c2]], %[[c2]] : (tensor<2x1x2xf32>, tensor<1x1x2xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2x1x2xf32>
//CHECK-NEXT: return %[[c3:.+]] : tensor<2x1x2xf32>
//CHECK-NEXT: }
//CHECK-NEXT:}
tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_chlo.mlir
+59
View File
@@ -0,0 +1,59 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-legalize-chlo -split-input-file | FileCheck %s --dump-input=fail
// Just assert that pass is properly registered.
func.func @main(%arg0: tensor<f32>) -> tensor<f32> {
return %arg0: tensor<f32>
}
// CHECK-LABEL: main
// -----
func.func @geluWithCustomCallErf(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<2xf32>
%1 = stablehlo.constant dense<0.707106769> : tensor<2xf32>
%2 = stablehlo.constant dense<5.000000e-01> : tensor<2xf32>
%3 = stablehlo.multiply %arg0, %2 : tensor<2xf32>
%4 = stablehlo.multiply %arg0, %1 : tensor<2xf32>
%5 = stablehlo.custom_call @mhlo.erf(%4) {mhlo.attributes = {}, mhlo.version = 1 : i64} : (tensor<2xf32>) -> tensor<2xf32>
%6 = stablehlo.add %5, %0 : tensor<2xf32>
%7 = stablehlo.multiply %3, %6 : tensor<2xf32>
return %7 : tensor<2xf32>
}
// CHECK-LABEL: geluWithCustomCallErf
// CHECK: "tfl.gelu"
// CHECK-NOT: stablehlo
// CHECK-NOT: chlo
// -----
func.func @geluWithCHLOErf(%arg0: tensor<2xf32>) -> tensor<2xf32> {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<2xf32>
%1 = stablehlo.constant dense<0.707106769> : tensor<2xf32>
%2 = stablehlo.constant dense<5.000000e-01> : tensor<2xf32>
%3 = stablehlo.multiply %arg0, %2 : tensor<2xf32>
%4 = stablehlo.multiply %arg0, %1 : tensor<2xf32>
%5 = chlo.erf %4 : tensor<2xf32> -> tensor<2xf32>
%6 = stablehlo.add %5, %0 : tensor<2xf32>
%7 = stablehlo.multiply %3, %6 : tensor<2xf32>
return %7 : tensor<2xf32>
}
// CHECK-LABEL: geluWithCHLOErf
// CHECK: "tfl.gelu"
// CHECK-NOT: stablehlo
// CHECK-NOT: chlo
tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_hlo.mlir
+3947
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_hlo_custom_call.mlir
+68
View File
@@ -0,0 +1,68 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -tfl-legalize-hlo -split-input-file | FileCheck %s --dump-input=fail
// CHECK-LABEL: mhlo_custom_call_test__legalize_string_backend_config
func.func @mhlo_custom_call_test__legalize_string_backend_config(%arg0: tensor<1x4xf32>) -> tensor<1x8xf32> {
%0 = mhlo.custom_call @custom_call.my_custom_op(%arg0) {
api_version = 1 : i32,
backend_config = "this_is_a_test_string"
} : (tensor<1x4xf32>) -> (tensor<1x8xf32>)
func.return %0 : tensor<1x8xf32>
// CHECK: %0 = "tfl.custom"(%arg0) <{
// CHECK-SAME: custom_code = "custom_call.my_custom_op",
// CHECK-SAME: custom_option = #tfl<const_bytes : "0x746869735F69735F615F746573745F737472696E67">
// CHECK-SAME: }> : (tensor<1x4xf32>) -> tensor<1x8xf32>
}
// CHECK-LABEL: mhlo_custom_call_test__dont_legalize_dict_backend_config
func.func @mhlo_custom_call_test__dont_legalize_dict_backend_config(%arg0: tensor<1x4xf32>) -> tensor<1x8xf32> {
%0 = mhlo.custom_call @custom_call.my_custom_op(%arg0) {
api_version = 4 : i32,
backend_config = {foo = "bar"}
} : (tensor<1x4xf32>) -> (tensor<1x8xf32>)
func.return %0 : tensor<1x8xf32>
// CHECK: %0 = mhlo.custom_call @custom_call.my_custom_op(%arg0) {
// CHECK-SAME: api_version = 4 : i32,
// CHECK-SAME: backend_config = {foo = "bar"}
// CHECK-SAME: } : (tensor<1x4xf32>) -> tensor<1x8xf32>
}
// CHECK-LABEL: mhlo_custom_call_test__api_version_4
func.func @mhlo_custom_call_test__api_version_4(%arg0: tensor<1x4xf32>) -> tensor<1x8xf32> {
%0 = mhlo.custom_call @custom_call.my_custom_op(%arg0) {
api_version = 4 : i32
} : (tensor<1x4xf32>) -> (tensor<1x8xf32>)
func.return %0 : tensor<1x8xf32>
// CHECK: %0 = "tfl.custom"(%arg0) <{
// CHECK-SAME: custom_code = "custom_call.my_custom_op",
// CHECK-SAME: custom_option = #tfl<const_bytes : "0x">
// CHECK-SAME: }> : (tensor<1x4xf32>) -> tensor<1x8xf32>
}
// CHECK-LABEL: mhlo_custom_call_does_not_legalize_tf_function
func.func @mhlo_custom_call_does_not_legalize_tf_function(%arg0: tensor<1x4xf32>) -> tensor<1x8xf32> {
%0 = mhlo.custom_call @tf.ResizeBilinear(%arg0) {
backend_config = "this_is_a_test_string"
} : (tensor<1x4xf32>) -> (tensor<1x8xf32>)
func.return %0 : tensor<1x8xf32>
// CHECK: %0 = mhlo.custom_call @tf.ResizeBilinear(%arg0) {
// CHECK-SAME: backend_config = "this_is_a_test_string"
// CHECK-SAME: } : (tensor<1x4xf32>) -> tensor<1x8xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/torch-build-stablehlo-composite.mlir
+46
View File
@@ -0,0 +1,46 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -build-stablehlo-composite -cse -canonicalize -cse | FileCheck %s --dump-input=fail
module {
func.func public @build_nested_composite(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
%c = stablehlo.constant dense<1> : tensor<i32>
%c_0 = stablehlo.constant dense<2> : tensor<i32>
%0 = stablehlo.custom_call @mark_tensor(%arg0) {backend_config = "{\22name\22: \22test.outer\22, \22pos\22: 0, \22id\22: \22f8cf5dba52ff4995b9f3810647aa69e2\22, \22is_input\22: true, \22attr\22: null}"} : (tensor<2x2xf32>) -> tensor<2x2xf32>
%1 = stablehlo.multiply %c_0, %c : tensor<i32>
%2 = stablehlo.convert %1 : (tensor<i32>) -> tensor<f32>
%3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<f32>) -> tensor<2x2xf32>
%4 = stablehlo.add %0, %3 : tensor<2x2xf32>
%5 = stablehlo.custom_call @mark_tensor(%4) {backend_config = "{\22name\22: \22test.inner\22, \22pos\22: 0, \22id\22: \22709cea3466314458963f3262d1deb27e\22, \22is_input\22: true, \22attr\22: null}"} : (tensor<2x2xf32>) -> tensor<2x2xf32>
%6 = stablehlo.multiply %c, %c : tensor<i32>
%7 = stablehlo.convert %6 : (tensor<i32>) -> tensor<f32>
%8 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor<f32>) -> tensor<2x2xf32>
%9 = stablehlo.add %5, %8 : tensor<2x2xf32>
%10 = stablehlo.custom_call @mark_tensor(%9) {backend_config = "{\22name\22: \22test.inner\22, \22pos\22: 0, \22id\22: \22709cea3466314458963f3262d1deb27e\22, \22is_input\22: false, \22attr\22: null}"} : (tensor<2x2xf32>) -> tensor<2x2xf32>
%11 = stablehlo.custom_call @mark_tensor(%10) {backend_config = "{\22name\22: \22test.outer\22, \22pos\22: 0, \22id\22: \22f8cf5dba52ff4995b9f3810647aa69e2\22, \22is_input\22: false, \22attr\22: null}"} : (tensor<2x2xf32>) -> tensor<2x2xf32>
return %11 : tensor<2x2xf32>
}
}
// CHECK-LABEL: build_nested_composite
// CHECK: stablehlo.composite "test.outer"
// CHECK: return
// CHECK: func.func private @test.inner.impl
// CHECK: stablehlo.add %arg0,
// CHECK: return
// CHECK: func.func private @test.outer.impl
// CHECK: stablehlo.add %arg0,
// CHECK: stablehlo.composite "test.inner"
// CHECK: return
tensorflow/compiler/mlir/lite/stablehlo/tests/unfold_splat_constant_pass.mlir
+43
View File
@@ -0,0 +1,43 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -unfold-splat-constant-pass -cse -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @unfold_splat_constant_float
func.func @unfold_splat_constant_float() -> tensor<1x750xf32> {
%cst = mhlo.constant dense<7.680000e+02> : tensor<1x750xf32>
func.return %cst : tensor<1x750xf32>
// CHECK-DAG: %0 = mhlo.constant dense<7.680000e+02> : tensor<f32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<1x750xf32>
// CHECK: return %1 : tensor<1x750xf32>
}
// CHECK-LABEL: @unfold_splat_constant_integer
func.func @unfold_splat_constant_integer() -> tensor<1x750xi32> {
%cst = mhlo.constant dense<1> : tensor<1x750xi32>
func.return %cst : tensor<1x750xi32>
// CHECK-DAG: %0 = mhlo.constant dense<1> : tensor<i32>
// CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<1x750xi32>
// CHECK: return %1 : tensor<1x750xi32>
}
// CHECK-LABEL: @splat_scalar_no_change
func.func @splat_scalar_no_change() -> (tensor<f32>, tensor<i32>) {
// CHECK-NOT: mhlo.broadcast_in_dim
%cst0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
%cst1 = mhlo.constant dense<0> : tensor<i32>
func.return %cst0, %cst1 : tensor<f32>, tensor<i32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/unfuse_mhlo_batch_norm.mlir
+169
View File
@@ -0,0 +1,169 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: odml-to-stablehlo-opt %s -unfuse-mhlo-batch-norm-pass -cse -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @batchNormInference_2D_inner_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_2D_inner_features(
%x: tensor<4x256xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>,
%mean: tensor<256xf32>, %variance: tensor<256xf32>)
-> (tensor<4x256xf32>) {
// CHECK-DAG: %[[EPS_BCAST:.+]] = mhlo.constant dense<1.001000e-05> : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = mhlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[RHS:.+]] = mhlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[MULTIPLIER]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[X_NORMED:.+]] = mhlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<4x256xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[RHS]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[RESULT:.+]] = mhlo.add %[[X_NORMED]], %[[RHS_BCAST]] : tensor<4x256xf32>
%0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.001000e-05 : f32, feature_index = 1 : i64} :
(tensor<4x256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>,
tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: return %[[RESULT]]
func.return %0 : tensor<4x256xf32>
}
// CHECK-LABEL: @batchNormInference_4D_middle_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_4D_middle_features(
%x: tensor<3x4x256x6xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>,
%mean: tensor<256xf32>, %variance: tensor<256xf32>)
-> (tensor<3x4x256x6xf32>) {
// CHECK-DAG: %[[EPS_BCAST:.+]] = mhlo.constant dense<1.001000e-05> : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = mhlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[RHS:.+]] = mhlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[MULTIPLIER]]) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[RHS]]) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
%0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.001000e-05 : f32, feature_index = 2 : i64} :
(tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>,
tensor<256xf32>) -> tensor<3x4x256x6xf32>
func.return %0 : tensor<3x4x256x6xf32>
}
// CHECK-LABEL: @batchNormInference_dynamic_shape
// Validate that dynamic shapes are handled properly.
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @batchNormInference_dynamic_shape(
%x: tensor<?x?x?x?xf32>, %scale: tensor<?xf32>, %offset: tensor<?xf32>,
%mean: tensor<?xf32>, %variance: tensor<?xf32>)
-> tensor<?x?x?x?xf32> {
// CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e-03> : tensor<f32>
// CHECK-DAG: %[[VAR_SHAPE:.+]] = shape.shape_of %[[VARIANCE]] : tensor<?xf32> -> tensor<1xindex>
// CHECK-DAG: %[[EPS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[EPS]], %[[VAR_SHAPE]]) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>, tensor<1xindex>) -> tensor<?xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<?xf32>
// CHECK-DAG: %[[R_STDDEV:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<?xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[R_STDDEV]], %[[SCALE]] : tensor<?xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = mhlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<?xf32>
// CHECK-DAG: %[[RHS:.+]] = mhlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<?xf32>
// CHECK-DAG: %[[X_SHAPE:.+]] = shape.shape_of %[[X]] : tensor<?x?x?x?xf32> -> tensor<4xindex>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[MULTIPLIER]], %[[X_SHAPE]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<?xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
// CHECK-DAG: %[[X_NORMED:.+]] = mhlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[RHS]], %[[X_SHAPE]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<?xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
// CHECK-DAG: %[[RESULT:.+]] = mhlo.add %[[X_NORMED]], %[[RHS_BCAST]] : tensor<?x?x?x?xf32>
%0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 0.001 : f32, feature_index = 1 : i64} :
(tensor<?x?x?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>,
tensor<?xf32>) -> tensor<?x?x?x?xf32>
func.return %0 : tensor<?x?x?x?xf32>
}
// CHECK-LABEL: @batchNormInference_f64
// Validate that epsilon is properly promoted to f64
// CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e+00> : tensor<256xf64>
func.func @batchNormInference_f64(
%x: tensor<4x256xf64>, %scale: tensor<256xf64>, %offset: tensor<256xf64>,
%mean: tensor<256xf64>, %variance: tensor<256xf64>)
-> (tensor<4x256xf64>) {
%0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.0 : f32, feature_index = 1 : i64} :
(tensor<4x256xf64>, tensor<256xf64>, tensor<256xf64>, tensor<256xf64>,
tensor<256xf64>) -> tensor<4x256xf64>
func.return %0 : tensor<4x256xf64>
}
// CHECK-LABEL: @batchNormInference_f16
// Validate that epsilon is properly down to f16
// CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e+00> : tensor<256xf16>
func.func @batchNormInference_f16(
%x: tensor<4x256xf16>, %scale: tensor<256xf16>, %offset: tensor<256xf16>,
%mean: tensor<256xf16>, %variance: tensor<256xf16>)
-> (tensor<4x256xf16>) {
%0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 1.0 : f32, feature_index = 1 : i64} :
(tensor<4x256xf16>, tensor<256xf16>, tensor<256xf16>, tensor<256xf16>,
tensor<256xf16>) -> tensor<4x256xf16>
func.return %0 : tensor<4x256xf16>
}
// Validate that epsilon is overflow
func.func @batchNormInference_f16_overflow(
%x: tensor<4x256xf16>, %scale: tensor<256xf16>, %offset: tensor<256xf16>,
%mean: tensor<256xf16>, %variance: tensor<256xf16>)
-> (tensor<4x256xf16>) {
// expected-warning @+1 {{Could not convert batch_norm epsilon to target fp type: opStatus = 24}}
%0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
{epsilon = 0.00000001 : f32, feature_index = 1 : i64} :
(tensor<4x256xf16>, tensor<256xf16>, tensor<256xf16>, tensor<256xf16>,
tensor<256xf16>) -> tensor<4x256xf16>
func.return %0 : tensor<4x256xf16>
}
// CHECK-LABEL: @batchNormTraining_4D_middle_features
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
func.func @batchNormTraining_4D_middle_features(
%x: tensor<3x4x256x6xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>)
-> (tensor<3x4x256x6xf32>) {
// CHECK-DAG: %[[CST_AXIS:.+]] = "tf.Const"() <{value = dense<[0, 1, 3]> : tensor<3xi32>}> : () -> tensor<3xi32>
// CHECK-DAG: %[[X_SHAPE:.+]] = shape.shape_of %[[X]] : tensor<3x4x256x6xf32> -> tensor<4xindex>
// CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e+00> : tensor<256xf32>
// CHECK-DAG: %[[MEAN:.+]] = "tf.Mean"(%arg0, %[[CST_AXIS]]) <{keep_dims = false}> : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32>
// CHECK-DAG: %[[MEAN_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[MEAN]], %[[X_SHAPE]]) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<256xf32>, tensor<4xindex>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[SQ_DIFF:.+]] = "tf.SquaredDifference"(%arg0, %[[MEAN_BCAST]]) : (tensor<3x4x256x6xf32>, tensor<3x4x256x6xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[VARIANCE:.+]] = "tf.Mean"(%[[SQ_DIFF]], %[[CST_AXIS]]) <{keep_dims = false}> : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS]] : tensor<256xf32>
// CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>
// CHECK-DAG: %[[MUL_MEAN:.+]] = mhlo.multiply %[[MULTIPLIER]], %[[MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[RHS:.+]] = mhlo.subtract %[[OFFSET]], %[[MUL_MEAN]] : tensor<256xf32>
// CHECK-DAG: %[[MULTIPLIER_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[MULTIPLIER]]) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[X_NORMED:.+]] = mhlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<3x4x256x6xf32>
// CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[RHS]]) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<3x4x256x6xf32>
// CHECK-DAG: %[[RESULT:.+]] = mhlo.add %[[X_NORMED]], %[[RHS_BCAST]] : tensor<3x4x256x6xf32>
%0:3 = "mhlo.batch_norm_training"(%x, %scale, %offset)
{epsilon = 1.0 : f32, feature_index = 2 : i64} :
(tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>) -> (tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>)
func.return %0 : tensor<3x4x256x6xf32>
}
tensorflow/compiler/mlir/lite/stablehlo/tests/uniform-quantized-stablehlo-to-tfl.mlir
+1636
View File
File diff suppressed because it is too large Load Diff