wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 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/quantization/stablehlo/tests/BUILD
+48
View File
@@ -0,0 +1,48 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
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"],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
# TODO: b/288344501 - Enable OSS tests again when stable-quant-opt works well.
default_tags = [
"no_oss",
"no_pip",
],
driver = "//tensorflow/compiler/mlir/quantization/stablehlo:run_lit.sh",
size_override = {
},
tags_override = {
},
test_file_exts = ["mlir"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir/quantization/stablehlo:stablehlo-quant-opt",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
"@llvm-project//mlir:run_lit.sh",
],
)
tf_cc_test(
name = "fill_quantization_options_test",
srcs = ["fill_quantization_options_test.cc"],
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:fill_quantization_options",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_options_proto_cc",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@tsl//tsl/platform:protobuf",
],
)
tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/BUILD
+35
View File
@@ -0,0 +1,35 @@
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"],
licenses = ["notice"],
)
glob_lit_tests(
name = "all_tests",
data = [":test_utilities"],
# TODO: b/288344501 - Enable OSS tests again when stable-quant-opt works well.
default_tags = [
"no_oss",
"no_pip",
],
driver = "//tensorflow/compiler/mlir/quantization/stablehlo:run_lit.sh",
size_override = {
},
tags_override = {
},
test_file_exts = ["mlir"],
)
# Bundle together all of the test utilities that are used by tests.
filegroup(
name = "test_utilities",
testonly = True,
data = [
"//tensorflow/compiler/mlir/quantization/stablehlo:stablehlo-quant-opt",
"@llvm-project//llvm:FileCheck",
"@llvm-project//llvm:not",
"@llvm-project//mlir:run_lit.sh",
],
)
tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-tf-quant-types.mlir
+517
View File
@@ -0,0 +1,517 @@
// 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: stablehlo-quant-opt %s -convert-tf-quant-types -split-input-file -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @relu_qint8
func.func @relu_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
// CHECK: %[[X:.*]] = "tf.Relu"(%arg0) : (tensor<1xi8>) -> tensor<1xi8>
%0 = "tf.Relu"(%arg0) : (tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8>
func.return %0: tensor<1x!tf_type.qint8>
}
// -----
// CHECK-LABEL: func @if_qint8(%arg0: tensor<i1>, %arg1: tensor<1xi8>, %arg2: tensor<1xi8>) -> tensor<1xi8>
func.func @if_qint8(%arg0: tensor<i1>, %arg1: tensor<1x!tf_type.qint8>, %arg2: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
// CHECK-NEXT: %0 = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({
// CHECK-NEXT: "tf.Yield"(%arg1) : (tensor<1xi8>) -> ()
// CHECK-NEXT: }, {
// CHECK-NEXT: "tf.Yield"(%arg2) : (tensor<1xi8>) -> ()
// CHECK-NEXT: }) : (tensor<i1>) -> tensor<1xi8>
// CHECK-NEXT: return %0 : tensor<1xi8>
%0 = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({
"tf.Yield"(%arg1) : (tensor<1x!tf_type.qint8>) -> ()
}, {
"tf.Yield"(%arg2) : (tensor<1x!tf_type.qint8>) -> ()
}) : (tensor<i1>) -> tensor<1x!tf_type.qint8>
func.return %0 : tensor<1x!tf_type.qint8>
}
// -----
// CHECK-LABEL: func @id_qint8(%arg0: tensor<1xi8>) -> tensor<1xi8> {
func.func @id_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
// CHECK-NEXT: return %arg0 : tensor<1xi8>
func.return %arg0: tensor<1x!tf_type.qint8>
}
// -----
// CHECK-LABEL: func @id_qint16(%arg0: tensor<1xi16>) -> tensor<1xi16> {
func.func @id_qint16(%arg0: tensor<1x!tf_type.qint16>) -> tensor<1x!tf_type.qint16> {
// CHECK-NEXT: return %arg0 : tensor<1xi16>
func.return %arg0: tensor<1x!tf_type.qint16>
}
// -----
// CHECK-LABEL: func @id_qint32(%arg0: tensor<1xi32>) -> tensor<1xi32> {
func.func @id_qint32(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1x!tf_type.qint32> {
// CHECK-NEXT: return %arg0 : tensor<1xi32>
func.return %arg0: tensor<1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @id_quint8(%arg0: tensor<1xui8>) -> tensor<1xui8> {
func.func @id_quint8(%arg0: tensor<1x!tf_type.quint8>) -> tensor<1x!tf_type.quint8> {
// CHECK-NEXT: return %arg0 : tensor<1xui8>
func.return %arg0: tensor<1x!tf_type.quint8>
}
// -----
// CHECK-LABEL: func @id_quint16(%arg0: tensor<1xui16>) -> tensor<1xui16> {
func.func @id_quint16(%arg0: tensor<1x!tf_type.quint16>) -> tensor<1x!tf_type.quint16> {
// CHECK-NEXT: return %arg0 : tensor<1xui16>
func.return %arg0: tensor<1x!tf_type.quint16>
}
// -----
// CHECK-LABEL: func @uniform_quantize
func.func @uniform_quantize(%arg0: tensor<1xf32>) -> tensor<1x!tf_type.qint8>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[qint:.*]] = "tf.UniformQuantize"
// CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
// CHECK: return %[[int]] : tensor<1xi8>
func.return %0 : tensor<1x!tf_type.qint8>
}
// -----
// CHECK-LABEL: func @uniform_quantize_no_return
func.func @uniform_quantize_no_return(%arg0: tensor<1xf32>) -> ()
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[qint:.*]] = "tf.UniformQuantize"
// CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
// CHECK: return
func.return
}
// -----
// CHECK-LABEL: func @uniform_dequantize
func.func @uniform_dequantize(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xf32>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[x:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1x!tf_type.qint8>
// CHECK: %[[y:.*]] = "tf.UniformDequantize"(%[[x]]
%0 = "tf.UniformDequantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
// CHECK: return %[[y]] : tensor<1xf32>
func.return %0 : tensor<1xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_dequantize
func.func @uniform_quantize_dequantize(%arg0: tensor<1xf32>) -> tensor<1xf32>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[qint0:.*]] = "tf.UniformQuantize"
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
// CHECK: %[[int:.*]] = "tf.Cast"(%[[qint0]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
// CHECK: %[[qint1:.*]] = "tf.Cast"(%[[int]]) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1x!tf_type.qint8>
// CHECK: %[[res:.*]] = "tf.UniformDequantize"(%[[qint1]]
%1 = "tf.UniformDequantize"(%0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
// CHECK: return %[[res]] : tensor<1xf32>
func.return %1 : tensor<1xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_add
func.func @uniform_quantized_add(%arg0: tensor<2x!tf_type.qint32>, %arg1: tensor<2x!tf_type.qint32>) -> tensor<2x!tf_type.qint32>
{
%input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[lhs:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2x!tf_type.qint32>
// CHECK: %[[rhs:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2x!tf_type.qint32>
// CHECK: %[[res_qint:.*]] = "tf.UniformQuantizedAdd"(%[[lhs]], %[[rhs]]
// CHECK: %[[res_int:.*]] = "tf.Cast"(%[[res_qint]]) <{Truncate = false}> : (tensor<2x!tf_type.qint32>) -> tensor<2xi32>
// CHECK: return %[[res_int]] : tensor<2xi32>
%1 = "tf.UniformQuantizedAdd"(
%arg0, %arg1,
%input_scales, %input_zps,
%bias_scales, %bias_zps,
%output_scales, %output_zps) {
lhs_quantization_axis = -1 : i64,
lhs_quantization_min_val = -2147483648 : i64,
lhs_quantization_max_val = 2147483647 : i64,
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -2147483648 : i64,
rhs_quantization_max_val = 2147483647 : i64,
output_quantization_axis = -1 : i64,
output_quantization_min_val = -2147483648 : i64,
output_quantization_max_val = 2147483647 : i64} : (
tensor<2x!tf_type.qint32>, tensor<2x!tf_type.qint32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>) -> tensor<2x!tf_type.qint32>
func.return %1 : tensor<2x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @while_region_qint
func.func @while_region_qint(%arg0: tensor<2x2xf32>) -> (tensor<2x?xf32>, tensor<?x2xf32>) {
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps2 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
%zps4 = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[qint_0:.*]] = "tf.UniformQuantize"
// CHECK: %[[int_0:.*]] = "tf.Cast"(%[[qint_0]]) <{Truncate = false}> : (tensor<2x2x!tf_type.qint8>) -> tensor<2x2xi8>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps2) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint8>
// CHECK: %[[qint_1:.*]] = "tf.UniformQuantize"
// CHECK: %[[int_1:.*]] = "tf.Cast"(%[[qint_1]]) <{Truncate = false}> : (tensor<2x2x!tf_type.qint8>) -> tensor<2x2xi8>
%1 = "tf.UniformQuantize"(%arg0, %scales, %zps4) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<2x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint8>
// CHECK: %[[while_result:.*]]:2 = "tf.WhileRegion"(%[[int_0]], %[[int_1]])
%2:2 = "tf.WhileRegion"(%0, %1) ({
^bb0(%carg0: tensor<2x?x!tf_type.qint8>, %carg1: tensor<?x2x!tf_type.qint8>):
%cst = "tf.Const"() {value = dense<1> : tensor<i1>} : () -> tensor<i1>
"tf.Yield"(%cst) : (tensor<i1>) -> ()
}, {
^bb0(%barg0: tensor<2x?x!tf_type.qint8>, %barg1: tensor<?x2x!tf_type.qint8>):
%id = "tf.Identity"(%barg0) : (tensor<2x?x!tf_type.qint8>) -> tensor<2x?x!tf_type.qint8>
"tf.Yield"(%id, %barg1) : (tensor<2x?x!tf_type.qint8>, tensor<?x2x!tf_type.qint8>) -> ()
}) {is_stateless = false} : (tensor<2x2x!tf_type.qint8>, tensor<2x2x!tf_type.qint8>) -> (tensor<2x?x!tf_type.qint8>, tensor<?x2x!tf_type.qint8>)
// CHECK: %[[out_qint_0:.*]] = "tf.Cast"(%[[while_result]]#0) <{Truncate = false}> : (tensor<2x?xi8>) -> tensor<2x?x!tf_type.qint8>
// CHECK: %[[out_f_0:.*]] = "tf.UniformDequantize"(%[[out_qint_0]]
%3 = "tf.UniformDequantize"(%2#0, %scales, %zps2) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<2x?x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<2x?xf32>
// CHECK: %[[out_qint_1:.*]] = "tf.Cast"(%[[while_result]]#1) <{Truncate = false}> : (tensor<?x2xi8>) -> tensor<?x2x!tf_type.qint8>
// CHECK: %[[out_f_1:.*]] = "tf.UniformDequantize"(%[[out_qint_1]]
%4 = "tf.UniformDequantize"(%2#1, %scales, %zps4) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<?x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?x2xf32>
// CHECK: return %[[out_f_0]], %[[out_f_1]]
func.return %3, %4 : tensor<2x?xf32>, tensor<?x2xf32>
}
// -----
// CHECK-LABEL: func @concat_uniform_quantize
func.func @concat_uniform_quantize(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3x!tf_type.qint8> {
%axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[input:.*]] = "tf.ConcatV2"(%arg0, %arg1
// CHECK: %[[output_qint:.*]] = "tf.UniformQuantize"(%[[input]]
// CHECK: %[[output:.*]] = "tf.Cast"(%[[output_qint]]) <{Truncate = false}> : (tensor<6x3x!tf_type.qint8>) -> tensor<6x3xi8>
// CHECK: return %[[output]] : tensor<6x3xi8>
%0 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor<i64>) -> tensor<6x3xf32>
%1 = "tf.UniformQuantize"(%0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<6x3xf32>, tensor<f32>, tensor<i32>) -> tensor<6x3x!tf_type.qint8>
func.return %1 : tensor<6x3x!tf_type.qint8>
}
// -----
// CHECK-LABEL: func @concat_uniform_dequantize
func.func @concat_uniform_dequantize(%arg0: tensor<3x3x!tf_type.qint8>, %arg1: tensor<3x3x!tf_type.qint8>) -> tensor<6x3xf32> {
%axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[input:.*]] = "tf.ConcatV2"(%arg0, %arg1, %[[VAL:.*]]) : (tensor<3x3xi8>, tensor<3x3xi8>, tensor<i64>) -> tensor<6x3xi8>
// CHECK: %[[input_qint:.*]] = "tf.Cast"(%[[input]]) <{Truncate = false}> : (tensor<6x3xi8>) -> tensor<6x3x!tf_type.qint8>
// CHECK: %[[output:.*]] = "tf.UniformDequantize"(%[[input_qint]]
// CHECK: return %[[output]] : tensor<6x3xf32>
%0 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3x!tf_type.qint8>, tensor<3x3x!tf_type.qint8>, tensor<i64>) -> tensor<6x3x!tf_type.qint8>
%1 = "tf.UniformDequantize"(%0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<6x3x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<6x3xf32>
func.return %1 : tensor<6x3xf32>
}
// -----
// CHECK-LABEL: func @tf_const_qint32
func.func @tf_const_qint32() -> tensor<1x!tf_type.qint32> {
// CHECK: %[[result:.*]] = "tf.Const"() <{value = dense<127> : tensor<1xi32>}> : () -> tensor<1xi32>
%0 = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<1x!tf_type.qint32> } : () -> tensor<1x!tf_type.qint32>
// CHECK: return %[[result]] : tensor<1xi32>
func.return %0 : tensor<1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @tf_const_qint8
func.func @tf_const_qint8() -> tensor<2x!tf_type.qint8> {
// CHECK: %[[result:.*]] = "tf.Const"() <{value = dense<[127, 18]> : tensor<2xi8>}> : () -> tensor<2xi8>
%0 = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656e736f722464747970653a2044545f51494e54382074656e736f725f7368617065207b2064696d207b2073697a653a2032207d207d2074656e736f725f636f6e74656e743a20225c3137375c30323222"> : tensor<2x!tf_type.qint8> } : () -> tensor<2x!tf_type.qint8>
// CHECK: return %[[result]] : tensor<2xi8>
func.return %0 : tensor<2x!tf_type.qint8>
}
// -----
func.func @tf_const_invalid_proto() -> tensor<2x!tf_type.qint32> {
// expected-error@+2 {{failed to get DenseElementAttr}}
// expected-error@+1 {{failed to legalize operation 'tf.Const'}}
%0 = "tf.Const"() { value = #tf_type<tensor_proto : "0x2532"> : tensor<2x!tf_type.qint32> } : () -> tensor<2x!tf_type.qint32>
func.return %0 : tensor<2x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @cast_op_qint32_int32
func.func @cast_op_qint32_int32(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xi32> {
// CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi32>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint32>) -> tensor<1xi32>
func.return %0: tensor<1xi32>
}
// -----
// CHECK-LABEL: func @cast_op_int32_qint32
func.func @cast_op_int32_qint32(%arg0: tensor<1xi32>) -> tensor<1x!tf_type.qint32> {
// CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi32>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1x!tf_type.qint32>
func.return %0: tensor<1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @cast_op_qint8_int8
func.func @cast_op_qint8_int8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xi8> {
// CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi8>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
func.return %0: tensor<1xi8>
}
// -----
// CHECK-LABEL: func @cast_op_int8_qint8
func.func @cast_op_int8_qint8(%arg0: tensor<1xi8>) -> tensor<1x!tf_type.qint8> {
// CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi8>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint8>
func.return %0: tensor<1x!tf_type.qint8>
}
// -----
// CHECK-LABEL: func @cast_op_qint32_int8
func.func @cast_op_qint32_int8(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xi8> {
// CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi8>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint32>) -> tensor<1xi8>
func.return %0: tensor<1xi8>
}
// -----
// CHECK-LABEL: func @cast_op_int8_qint32
func.func @cast_op_int8_qint32(%arg0: tensor<1xi8>) -> tensor<1x!tf_type.qint32> {
// CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi32>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint32>
func.return %0: tensor<1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @cast_uniform_dequantize
func.func @cast_uniform_dequantize(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xf32>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[x:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi8>
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint32>) -> tensor<1x!tf_type.qint8>
// CHECK: %[[y:.*]] = "tf.Cast"(%[[x]]) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1x!tf_type.qint8>
// CHECK: %[[z:.*]] = "tf.UniformDequantize"(%[[y]]
%1 = "tf.UniformDequantize"(%0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
// CHECK: return %[[z]] : tensor<1xf32>
func.return %1 : tensor<1xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_cast
func.func @uniform_quantize_cast(%arg0: tensor<1xf32>) -> tensor<1x!tf_type.qint32>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: tf.UniformQuantize
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
// CHECK: %1 = "tf.Cast"(%0) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
// CHECK: %2 = "tf.Cast"(%1) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi32>
%1 = "tf.Cast"(%0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint32>
// CHECK: return %2 : tensor<1xi32>
func.return %1 : tensor<1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_cast_dequantize
func.func @uniform_quantize_cast_dequantize(%arg0: tensor<1xf32>) -> tensor<1xf32>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
%scales1 = "tf.Const"() { value = dense<3.0> : tensor<f32> } : () -> tensor<f32>
%zps1 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
// CHECK: %[[qint_1:.*]] = "tf.UniformQuantize"
// CHECK: %[[int_1:.*]] = "tf.Cast"(%[[qint_1]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
// CHECK: %[[int_2:.*]] = "tf.Cast"(%[[int_1]]) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi32>
%1 = "tf.Cast"(%0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint32>
// CHECK: %[[qint_2:.*]] = "tf.Cast"(%[[int_2]]) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1x!tf_type.qint32>
// CHECK: %[[int_3:.*]] = "tf.UniformDequantize"(%[[qint_2]]
%2 = "tf.UniformDequantize"(%1, %scales1, %zps1) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
// CHECK: return %[[int_3]] : tensor<1xf32>
func.return %2 : tensor<1xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_clip_min_cast
func.func @uniform_quantize_clip_min_cast(%arg0: tensor<1x2x2x1x!tf_type.qint32>, %arg1: tensor<i32>) -> tensor<1x2x2x1x!tf_type.qint32> {
%scale = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
%q_min = "tf.Cast"(%arg1) {Truncate = false} : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK: "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MIN_QINT]]
%output = "tf.UniformQuantizedClipByValue"(%arg0, %q_min, %q_min, %scale, %zp)
{quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} :
(tensor<1x2x2x1x!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x1x!tf_type.qint32>
return %output : tensor<1x2x2x1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_clip_input_cast
func.func @uniform_quantize_clip_input_cast(%arg0: tensor<1x2x2x1xi32>, %arg1: tensor<!tf_type.qint32>) -> tensor<1x2x2x1x!tf_type.qint32> {
%scale = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>
%q_input = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK: "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MAX_QINT]]
%output = "tf.UniformQuantizedClipByValue"(%q_input, %arg1, %arg1, %scale, %zp)
{quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} :
(tensor<1x2x2x1x!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x1x!tf_type.qint32>
return %output : tensor<1x2x2x1x!tf_type.qint32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_clip_output_cast
func.func @uniform_quantize_clip_output_cast(%arg0: tensor<1x2x2x1x!tf_type.qint32>, %arg1: tensor<!tf_type.qint32>) -> tensor<1x2x2x1xi32> {
%scale = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK: %[[OUTPUT_QINT:.*]] = "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MAX_QINT]]
%q_output = "tf.UniformQuantizedClipByValue"(%arg0, %arg1, %arg1, %scale, %zp)
{quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} :
(tensor<1x2x2x1x!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK: %[[OUTPUT:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>
%output = "tf.Cast"(%q_output) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>
return %output : tensor<1x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @uniform_quantize_clip_output_cast_multiple_uses
func.func @uniform_quantize_clip_output_cast_multiple_uses(%arg0: tensor<1x2x2x1x!tf_type.qint32>, %arg1: tensor<!tf_type.qint32>) -> tensor<1x2x2x1xi32> {
%scale = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<i32>) -> tensor<!tf_type.qint32>
// CHECK: %[[OUTPUT_QINT:.*]] = "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MAX_QINT]]
%q_output = "tf.UniformQuantizedClipByValue"(%arg0, %arg1, %arg1, %scale, %zp)
{quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} :
(tensor<1x2x2x1x!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK-DAG: %[[OUTPUT_1:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>
%output = "tf.Cast"(%q_output) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>
// CHECK-DAG: %[[OUTPUT_2:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>
// CHECK-DAG: %[[OUTPUT_QINT_1:.*]] = "tf.Cast"(%[[OUTPUT_1]]) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>
// CHECK: "tf.UniformDequantize"(%[[OUTPUT_QINT_1:.*]]
%dq = "tf.UniformDequantize"(%q_output, %scale, %zp) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x2x2x1x!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<1x2x2x1xf32>
return %output : tensor<1x2x2x1xi32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert_tf_quant_ops_to_mhlo.mlir
+129
View File
@@ -0,0 +1,129 @@
// 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: stablehlo-quant-opt %s -quant-convert-tf-quant-ops-to-mhlo | FileCheck %s
// CHECK-LABEL: func @quantized_matmul_fn
func.func @quantized_matmul_fn(%input: tensor<?x?xf32>) -> tensor<?x?xf32> {
%weight = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E54382074656E736F725F7368617065207B2064696D207B2073697A653A2032207D2064696D207B2073697A653A2032207D207D2074656E736F725F636F6E74656E743A20225C3030315C3030325C3030335C30303422"> : tensor<2x2x!tf_type.qint8> } : () -> tensor<2x2x!tf_type.qint8>
%weight_scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// CHECK: "tf.AddV2"
// CHECK: mhlo.constant
// CHECK-SAME{LITERAL}: dense<[[1, 2], [3, 4]]> : tensor<2x2xi8>
%0 = "tf.AddV2"(%input, %input) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: "mhlo.dot"
// CHECK-SAME: (tensor<?x?xf32>, tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?x?xf32>
%1 = "tf.UniformQuantizedDotHybrid"(%0, %weight, %weight_scales, %weight_zps) {rhs_quantization_axis = -1 : i64, rhs_quantization_min_val = -128 : i64, rhs_quantization_max_val = 127 : i64} : (tensor<?x?xf32>, tensor<2x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?x?xf32>
func.return %1 : tensor<?x?xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_add
func.func @uniform_quantized_add(%input: tensor<3x2xf32>) -> tensor<3x2xf32> {
%input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// tensor_proto that points to dense<127> of type !tf_type.qint32.
// CHECK-DAG: %[[RHS:.*]] = mhlo.constant() <{value = dense<127> : tensor<2xi32>}> : () -> tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>
%bias = "tf.Const"() { value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F51494E5433322074656E736F725F7368617065207B207D2074656E736F725F636F6E74656E743A20225C3137375C3030305C3030305C30303022"> : tensor<2x!tf_type.qint32> } : () -> tensor<2x!tf_type.qint32>
%bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[LHS:.*]] = mhlo.uniform_quantize %arg0 : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
// CHECK-DAG: %[[LHS1:.*]] = mhlo.bitcast_convert %[[LHS]] : (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>) -> tensor<3x2xi32>
// CHECK-DAG: %[[LHS2:.*]] = mhlo.bitcast_convert %[[LHS1]] : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
%0 = "tf.UniformQuantize"(%input, %input_scales, %input_zps) {
quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
} : (tensor<3x2xf32>, tensor<f32>, tensor<i32>) -> tensor<3x2x!tf_type.qint32>
// CHECK: chlo.broadcast_add %[[LHS2]], %[[RHS]] {broadcast_dimensions = array<i64: 1>} :
// CHECK-SAME: (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>, tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>)
// CHECK-SAME: -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
%1 = "tf.UniformQuantizedAdd"(
%0, %bias,
%input_scales, %input_zps,
%bias_scales, %bias_zps,
%output_scales, %output_zps) {
lhs_quantization_axis = -1 : i64,
lhs_quantization_min_val = -2147483648 : i64,
lhs_quantization_max_val = 2147483647 : i64,
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -2147483648 : i64,
rhs_quantization_max_val = 2147483647 : i64,
output_quantization_axis = -1 : i64,
output_quantization_min_val = -2147483648 : i64,
output_quantization_max_val = 2147483647 : i64} : (
tensor<3x2x!tf_type.qint32>, tensor<2x!tf_type.qint32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>) -> tensor<3x2x!tf_type.qint32>
%2 = "tf.UniformDequantize"(%1, %output_scales, %output_zps) {
quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
} : (tensor<3x2x!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<3x2xf32>
func.return %2 : tensor<3x2xf32>
}
// -----
// CHECK-LABEL: func @uniform_quantized_add_bias_not_const
func.func @uniform_quantized_add_bias_not_const(%input1: tensor<3x2xi32>, %input2: tensor<2xi32>) -> tensor<3x2xi32> {
%input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
%output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
%output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
// CHECK-DAG: %[[LHS_1:.*]] = mhlo.convert %arg0 : tensor<3x2xi32>
// CHECK-DAG: %[[LHS_2:.*]] = mhlo.bitcast_convert %[[LHS_1]] : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
// CHECK-DAG: %[[RHS_1:.*]] = mhlo.convert %arg1 : tensor<2xi32>
// CHECK-DAG: %[[RHS_2:.*]] = mhlo.bitcast_convert %[[RHS_1]] : (tensor<2xi32>) -> tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>
%input1_qint = "tf.Cast"(%input1) {Truncate = false} : (tensor<3x2xi32>) -> tensor<3x2x!tf_type.qint32>
%input2_qint = "tf.Cast"(%input2) {Truncate = false} : (tensor<2xi32>) -> tensor<2x!tf_type.qint32>
// CHECK: %[[RES:.*]] = chlo.broadcast_add %[[LHS_2]], %[[RHS_2]] {broadcast_dimensions = array<i64: 1>} :
// CHECK-SAME: (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>, tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>)
// CHECK-SAME: -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
%result = "tf.UniformQuantizedAdd"(
%input1_qint, %input2_qint,
%input_scales, %input_zps,
%bias_scales, %bias_zps,
%output_scales, %output_zps) {
lhs_quantization_axis = -1 : i64,
lhs_quantization_min_val = -2147483648 : i64,
lhs_quantization_max_val = 2147483647 : i64,
rhs_quantization_axis = -1 : i64,
rhs_quantization_min_val = -2147483648 : i64,
rhs_quantization_max_val = 2147483647 : i64,
output_quantization_axis = -1 : i64,
output_quantization_min_val = -2147483648 : i64,
output_quantization_max_val = 2147483647 : i64} : (
tensor<3x2x!tf_type.qint32>, tensor<2x!tf_type.qint32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>,
tensor<f32>, tensor<i32>) -> tensor<3x2x!tf_type.qint32>
// CHECK: %[[RES_INT_1:.*]] = mhlo.bitcast_convert %[[RES]] : (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>) -> tensor<3x2xi32>
// CHECK: %[[RES_INT_2:.*]] = mhlo.convert %[[RES_INT_1]] : tensor<3x2xi32>
%result_int = "tf.Cast"(%result) {Truncate = false} : (tensor<3x2x!tf_type.qint32>) -> tensor<3x2xi32>
// CHECK: return %[[RES_INT_2]] : tensor<3x2xi32>
func.return %result_int : tensor<3x2xi32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/optimize.mlir
+262
View File
@@ -0,0 +1,262 @@
// 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: stablehlo-quant-opt -optimize-int-graph -split-input-file %s -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @convolution_add_add
func.func @convolution_add_add(
%lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<?x2x2x1xi32>, %bias: tensor<1xi32>
) -> tensor<?x2x2x1xi32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
return %2 : tensor<?x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @convolution_add_add_static
func.func @convolution_add_add_static(
%lhs: tensor<2x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<2x2x2x1xi32>, %bias: tensor<1xi32>
) -> tensor<2x2x2x1xi32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<2x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<2x2x2x1xi32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>) -> tensor<2x2x2x1xi32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<2x2x2x1xi32>, tensor<1xi32>) ->tensor<2x2x2x1xi32>
return %2 : tensor<2x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @convolution_add_subtract
func.func @convolution_add_subtract(
%lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<?x2x2x1xi32>, %bias: tensor<1xi32>
) -> tensor<?x2x2x1xi32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_subtract %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>
%2 = chlo.broadcast_subtract %1, %bias : (
tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
return %2 : tensor<?x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @convolution_subtract_subtract
func.func @convolution_subtract_subtract(
%lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<?x2x2x1xi32>, %bias: tensor<1xi32>
) -> tensor<?x2x2x1xi32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_subtract %[[conv]], %[[combined]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
%1 = chlo.broadcast_subtract %0, %zp_offset : (
tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>
%2 = chlo.broadcast_subtract %1, %bias : (
tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
return %2 : tensor<?x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @convolution_subtract_add
func.func @convolution_subtract_add(
%lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<?x2x2x1xi32>, %bias: tensor<1xi32>
) -> tensor<?x2x2x1xi32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_subtract %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_subtract %[[conv]], %[[combined]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
%1 = chlo.broadcast_subtract %0, %zp_offset : (
tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
return %2 : tensor<?x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @convolution_add_add_add
func.func @convolution_add_add_add(
%lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<?x2x2x1xi32>, %bias: tensor<1xi32>
) -> tensor<?x2x2x1xi32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined1:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[combined2:.*]] = chlo.broadcast_add %[[combined1]], %[[bias]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined2]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
%3 = chlo.broadcast_add %2, %bias : (
tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
return %3 : tensor<?x2x2x1xi32>
}
// -----
// CHECK-LABEL: func @convolution_add_add_i8
func.func @convolution_add_add_i8(
%lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
%zp_offset: tensor<?x2x2x1xi8>, %bias: tensor<1xi8>
) -> tensor<?x2x2x1xi8> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi8>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<?x2x2x1xi8>, tensor<?x2x2x1xi8>) -> tensor<?x2x2x1xi8>
%2 = chlo.broadcast_add %1, %bias : (
tensor<?x2x2x1xi8>, tensor<1xi8>) ->tensor<?x2x2x1xi8>
return %2 : tensor<?x2x2x1xi8>
}
// -----
// CHECK-LABEL: func @convolution_add_add_f32
func.func @convolution_add_add_f32(
%lhs: tensor<?x3x2x1xf32>, %rhs: tensor<2x1x1x1xf32>,
%zp_offset: tensor<?x2x2x1xf32>, %bias: tensor<1xf32>
) -> tensor<?x2x2x1xf32> {
// CHECK-DAG: %[[conv:.*]] = mhlo.convolution
// CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[conv:.*]], %[[zp_offset:.*]]
// CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[combined:.*]], %[[bias:.*]]
// CHECK: return %[[result]]
%0 = mhlo.convolution(%lhs, %rhs)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
{batch_group_count = 1 : i64, feature_group_count = 1 : i64}
: (tensor<?x3x2x1xf32>, tensor<2x1x1x1xf32>) -> tensor<?x2x2x1xf32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<?x2x2x1xf32>, tensor<?x2x2x1xf32>) -> tensor<?x2x2x1xf32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<?x2x2x1xf32>, tensor<1xf32>) ->tensor<?x2x2x1xf32>
return %2 : tensor<?x2x2x1xf32>
}
// -----
// CHECK-LABEL: func @dot_general_add_add
func.func @dot_general_add_add(
%lhs: tensor<?x?x6xi8>, %rhs: tensor<6x8xi8>,
%zp_offset: tensor<8xi32>, %bias: tensor<8xi32>
) -> tensor<?x?x8xi32> {
// CHECK: %[[dot:.*]] = "mhlo.dot_general"
// CHECK: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK: %[[result:.*]] = chlo.broadcast_add %[[dot]], %[[combined]]
// CHECK: return %[[result]]
%0 = "mhlo.dot_general" (%lhs, %rhs) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [],
rhs_batching_dimensions = [],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [0]
>} : (
tensor<?x?x6xi8>, tensor<6x8xi8>
) -> tensor<?x?x8xi32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>
return %2 : tensor<?x?x8xi32>
}
// -----
// CHECK-LABEL: func @dot_general_add_add_static
func.func @dot_general_add_add_static(
%lhs: tensor<2x5x6xi8>, %rhs: tensor<6x8x2xi8>,
%zp_offset: tensor<2x5x8xi32>, %bias: tensor<2x5x8xi32>
) -> tensor<2x5x8xi32> {
// CHECK: %[[dot:.*]] = "mhlo.dot_general"
// CHECK: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
// CHECK: %[[result:.*]] = chlo.broadcast_add %[[dot]], %[[combined]]
// CHECK: return %[[result]]
%0 = "mhlo.dot_general" (%lhs, %rhs) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [2],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [0]
>} : (
tensor<2x5x6xi8>, tensor<6x8x2xi8>
) -> tensor<2x5x8xi32>
%1 = chlo.broadcast_add %0, %zp_offset : (
tensor<2x5x8xi32>, tensor<2x5x8xi32>) -> tensor<2x5x8xi32>
%2 = chlo.broadcast_add %1, %bias : (
tensor<2x5x8xi32>, tensor<2x5x8xi32>) -> tensor<2x5x8xi32>
return %2 : tensor<2x5x8xi32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/verify-quant-legalization.mlir
+81
View File
@@ -0,0 +1,81 @@
// 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: stablehlo-quant-opt %s -verify-quant-legalization -split-input-file -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @legal_i8
func.func @legal_i8(%arg0: tensor<1xi8>) -> tensor<1xi8> {
func.return %arg0: tensor<1xi8>
}
// -----
func.func @illegal_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
// expected-error@+1 {{'func.return' op is illegal as it is a UQ op or contains uq/qint types}}
func.return %arg0: tensor<1x!tf_type.qint8>
}
// -----
func.func @illegal_cast(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xi8> {
// expected-error@+1 {{'tf.Cast' op is illegal as it is a UQ op or contains uq/qint types}}
%0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
func.return %0: tensor<1xi8>
}
// -----
func.func @illegal_mhlo_uniform_quantize(%arg0: tensor<?x?xf32>) -> tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>> {
// expected-error@+1 {{'mhlo.uniform_quantize' op is illegal as it is a UQ op or contains uq/qint types}}
%0 = mhlo.uniform_quantize %arg0 : (tensor<?x?xf32>) -> tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>
return %0 : tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>
}
// -----
func.func @illegal_mhlo_uniform_dequantize(%arg0: tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?x?xf32> {
// expected-error@+1 {{'mhlo.uniform_dequantize' op is illegal as it is a UQ op or contains uq/qint types}}
%0 = mhlo.uniform_dequantize %arg0 : (tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32>
}
// -----
func.func @illegal_tf_uniform_quantize(%arg0 : tensor<1xf32>) -> tensor<1xf32> {
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// expected-error@+1 {{'tf.UniformQuantize' op is illegal as it is a UQ op or contains uq/qint types}}
%0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
%1 = "tf.UniformDequantize"(%0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
func.return %1 : tensor<1xf32>
}
// -----
func.func @illegal_tf_uniform_dequantize(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xf32>
{
%scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
// expected-error@+1 {{'tf.UniformDequantize' op is illegal as it is a UQ op or contains uq/qint types}}
%0 = "tf.UniformDequantize"(%arg0, %scales, %zps) {
quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
} : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
func.return %0 : tensor<1xf32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/components/post_calibration_component.mlir
+88
View File
@@ -0,0 +1,88 @@
// 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: stablehlo-quant-opt %s -stablehlo-test-post-calibration-component \
// RUN: -split-input-file | FileCheck %s
// RUN: stablehlo-quant-opt %s -stablehlo-test-post-calibration-component='unpack-quantized-types=false' \
// RUN: -split-input-file | FileCheck %s --check-prefix=CHECK-NO-UNPACK
// Tests that a simple dot_general (lifted as a function) with CustomAggregators
// around it is quantized. The resulting graph has quantized types unpacked into
// int ops.
func.func @main(%arg0: tensor<1x1024xf32>) -> tensor<1x3xf32> {
%0 = "tf.Const"() <{value = dense<0.5> : tensor<1024x3xf32>}> : () -> tensor<1024x3xf32>
%1:4 = "tf.CustomAggregator"(%arg0) <{id = "1", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> {min = 7.547870e-07 : f32, max = 0.999992311 : f32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%2 = "tf.XlaCallModule"(%1#0, %0) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
%3:4 = "tf.CustomAggregator"(%2) <{id = "2", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> {min = -17.5216827 : f32, max = 18.3033524 : f32} : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
return %3#0 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func @main
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x1024xf32>) -> tensor<1x3xf32>
// Tests that the dot_general accepts i8 tensors and outputs an i32 tensor.
// Note: Argument quantization sequence omitted.
// CHECK: stablehlo.dot_general %{{.+}}, %{{.+}}, contracting_dims = [1] x [0] : (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>
// Note: Result dequantization sequence omitted.
// CHECK: return %{{.+}} : tensor<1x3xf32>
// -----
// Tests that a simple dot_general (lifted as a function) with CustomAggregators
// around it is quantized, when the 'unpack-quantized-types' option is set to
// false. This test case inputs the same graph as the test above. Note that now
// the uniform quantized types are directly expressed within the graph.
func.func @main_no_unpack(%arg0: tensor<1x1024xf32>) -> tensor<1x3xf32> {
%0 = "tf.Const"() <{value = dense<0.5> : tensor<1024x3xf32>}> : () -> tensor<1024x3xf32>
%1:4 = "tf.CustomAggregator"(%arg0) <{id = "1", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> {device = "", max = 0.999992311 : f32, min = 7.547870e-07 : f32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%2 = "tf.XlaCallModule"(%1#0, %0) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
%3:4 = "tf.CustomAggregator"(%2) <{id = "2", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> {device = "", max = 18.3033524 : f32, min = -17.5216827 : f32} : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
return %3#0 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-NO-UNPACK-LABEL: func.func @main_no_unpack
// CHECK-NO-UNPACK-SAME: (%[[ARG_0:.+]]: tensor<1x1024xf32>) -> tensor<1x3xf32>
// CHECK-NO-UNPACK-DAG: %[[CONST:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1024x3xi8>}> : () -> tensor<1024x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>
// CHECK-NO-UNPACK: %[[QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x1024xf32>) -> tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-NO-UNPACK: %[[DOT:.+]] = stablehlo.dot_general %[[QUANTIZE_0]], %[[CONST]]
// CHECK-NO-UNPACK: %[[QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[DOT]] : (tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-NO-UNPACK: %[[DEQUANTIZE:.+]] = stablehlo.uniform_dequantize %[[QUANTIZE_1]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3xf32>
// CHECK-NO-UNPACK: return %[[DEQUANTIZE]] : tensor<1x3xf32>
// -----
// Tests that a simple dot_general without CustomAggregators is not quantized.
func.func @main(%arg0: tensor<1x1024xf32>) -> tensor<1x3xf32> {
%0 = "tf.Const"() <{value = dense<0.5> : tensor<1024x3xf32>}> : () -> tensor<1024x3xf32>
%2 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: func.func @main
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x1024xf32>) -> tensor<1x3xf32>
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<{{.*}}> : tensor<1024x3xf32>
// CHECK: stablehlo.dot_general %[[ARG_0]], %[[CONST_0]]
// CHECK-NOT: tf.XlaCallModule
tensorflow/compiler/mlir/quantization/stablehlo/tests/components/pre_calibration_component.mlir
+76
View File
@@ -0,0 +1,76 @@
// 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: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \
// RUN: -stablehlo-test-pre-calibration-component | FileCheck %s
func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
// CHECK: @main(%[[ARG_0:.+]]: tensor<1x4xf32>) -> tensor<1x3xf32>
// CHECK-DAG: %[[CST:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<4x3xf32>
// CHECK: %[[CUSTOM_AGGREGATOR_0:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[ARG_0]]) <{calibration_method = 1 : i32, id = "composite_dot_general_fn_1_arg_0_calibration_method_1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
// CHECK: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[CST]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE]]) <{calibration_method = 1 : i32, id = "composite_dot_general_fn_1_calibration_method_1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
// CHECK: return %[[CUSTOM_AGGREGATOR_1]] : tensor<1x3xf32>
// CHECK: }
// CHECK: }
// -----
// Tests that stablehlo op serialization also works for the "serving_default"
// function.
func.func @serving_default(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
// CHECK: @serving_default(%[[ARG_0:.+]]: tensor<1x4xf32>) -> tensor<1x3xf32>
// CHECK-DAG: %[[CST:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<4x3xf32>
// CHECK: %[[CUSTOM_AGGREGATOR_0:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[ARG_0]]) <{calibration_method = 1 : i32, id = "composite_dot_general_fn_1_arg_0_calibration_method_1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
// CHECK: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[CST]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE]]) <{calibration_method = 1 : i32, id = "composite_dot_general_fn_1_calibration_method_1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
// CHECK: return %[[CUSTOM_AGGREGATOR_1]] : tensor<1x3xf32>
// CHECK: }
// CHECK: }
// -----
// Tests that `stablehlo.convolution` with NCHW format is converted to NHWC.
func.func @main(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32> {
%0 = stablehlo.constant() {value = dense<3.000000e+00> : tensor<8x8x3x3xf32>} : () -> tensor<8x8x3x3xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x4x4xf32>, tensor<8x8x3x3xf32>) -> tensor<1x8x4x4xf32>
return %2 : tensor<1x8x4x4xf32>
}
// CHECK: @main(%[[ARG:.+]]: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32>
// Contains the `stablehlo.transpose` op of the arg (e.g. [b, f, 0, 1] to
// [b, 0, 1, f]). The weight constant is folded into [0, 1, i, o] format.
// CHECK-DAG: %[[CST:.+]] = stablehlo.constant dense<3.000000e+00> : tensor<3x3x8x8xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %arg0, dims = [0, 2, 3, 1] : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>
// CHECK: %[[CUSTOM_AGGREGATOR_0:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[TRANSPOSE_1]]) {{.*}} : (tensor<1x4x4x8xf32>) -> (tensor<1x4x4x8xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
// Corresponds to the converted `stablehlo.convolution`. Note that the shapes
// correspond to the dimension numbers of: [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[CST]]) {{.*}} : (tensor<1x4x4x8xf32>, tensor<3x3x8x8xf32>) -> tensor<1x4x4x8xf32>
// CHECK: %[[CUSTOM_AGGREGATOR_1:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE]]) {{.*}} : (tensor<1x4x4x8xf32>) -> (tensor<1x4x4x8xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
// CHECK: %[[TRANSPOSE_2:.+]] = stablehlo.transpose %[[CUSTOM_AGGREGATOR_1]], dims = [0, 3, 1, 2] : (tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32>
// CHECK: return %[[TRANSPOSE_2]] : tensor<1x8x4x4xf32>
tensorflow/compiler/mlir/quantization/stablehlo/tests/components/tf_to_stablehlo.mlir
+172
View File
@@ -0,0 +1,172 @@
// 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: stablehlo-quant-opt %s -split-input-file -verify-diagnostics -stablehlo-test-tf-to-stablehlo | FileCheck %s
// TODO(b/330759552): Fix the msan issue and enable this test.
// func.func @fused_batchnorm_no_training() -> tensor<1x1x2x8xf32> {
// %cst_0 = "tf.Const"() {value = dense<[[[[0.1, 0.2, 0.1, 0.2, 0.1, 0.2, 0.1, 0.2], [0.3, 0.4, 0.3, 0.4, 0.3, 0.4, 0.3, 0.4]]]]> : tensor<1x1x2x8xf32>} : () -> tensor<1x1x2x8xf32>
// %cst_1 = "tf.Const"() {value = dense<[0.1, 0.2, 0.1, 0.2, 0.1, 0.2, 0.1, 0.2]> : tensor<8xf32>} : () -> tensor<8xf32>
// %cst_2 = "tf.Const"() {value = dense<[0.3, 0.4, 0.3, 0.4, 0.3, 0.4, 0.3, 0.4]> : tensor<8xf32>} : () -> tensor<8xf32>
// %0:6 = "tf.FusedBatchNormV3"(%cst_0, %cst_1, %cst_2, %cst_1, %cst_2) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<1x1x2x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<1x1x2x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)
// func.return %0#0 : tensor<1x1x2x8xf32>
// }
// COM: CHECK: func.func @main() -> tensor<1x1x2x8xf32>
// COM: CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<{{.*}}> : tensor<1x1x2x8xf32>
// COM: CHECK: return %[[CONST]] : tensor<1x1x2x8xf32>
func.func @fused_batchnorm_no_training_arg_input(%arg_0: tensor<1x1x2x8xf32>) -> (tensor<1x1x2x8xf32>) {
%cst_0 = "tf.Const"() {value = dense<[0.1, 0.2, 0.1, 0.2, 0.1, 0.2, 0.1, 0.2]> : tensor<8xf32>} : () -> tensor<8xf32>
%cst_1 = "tf.Const"() {value = dense<[0.3, 0.4, 0.3, 0.4, 0.3, 0.4, 0.3, 0.4]> : tensor<8xf32>} : () -> tensor<8xf32>
%0:6 = "tf.FusedBatchNormV3"(%arg_0, %cst_0, %cst_1, %cst_0, %cst_1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<1x1x2x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<1x1x2x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)
func.return %0#0 : tensor<1x1x2x8xf32>
}
// CHECK: func.func @main(%[[ARG:.+]]: tensor<1x1x2x8xf32>) -> tensor<1x1x2x8xf32>
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<{{.*}}> : tensor<8xf32>
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<{{.*}}> : tensor<8xf32>
// CHECK: %[[BROADCAST_0:.*]] = stablehlo.broadcast_in_dim %[[CONST_0]], dims = [3] : (tensor<8xf32>) -> tensor<1x1x2x8xf32>
// CHECK: %[[MUL:.*]] = stablehlo.multiply %[[ARG]], %[[BROADCAST_0]] : tensor<1x1x2x8xf32>
// CHECK: %[[BROADCAST_1:.*]] = stablehlo.broadcast_in_dim %[[CONST_1]], dims = [3] : (tensor<8xf32>) -> tensor<1x1x2x8xf32>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[MUL]], %[[BROADCAST_1]] : tensor<1x1x2x8xf32>
// CHECK: return %[[ADD]] : tensor<1x1x2x8xf32>
// -----
func.func @fuse_conv_batchnorm(%arg_0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf32>) {
%cst_0 = "tf.Const"() {value = dense<[[[[-6.30731344, 5.4962182], [1.80364347, -7.64542675], [-2.11145878, -7.08605719]], [[-9.54062747, -6.14013147], [6.12640238, -4.18223286], [5.05738974, 8.99269962]], [[3.3535192, 0.84816426], [-6.64676809, -7.95477629], [5.81315517, 9.21566581]]], [[[1.38622558, 4.63866329], [4.54742622, -1.43770897], [-3.96835279, 2.99996852]], [[0.989735424, -4.83384752], [-7.27702999, 1.17216611], [1.33735656, 0.728900194]], [[5.1286211, 8.98645591], [1.55008793, -3.85491467], [3.7003777, 9.26594448]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_1 = "tf.Const"() {value = dense<[0.1, 0.2]> : tensor<2xf32>} : () -> tensor<2xf32>
%cst_2 = "tf.Const"() {value = dense<[0.3, 0.4]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.Conv2D"(%arg_0, %cst_0) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32>
%1:6 = "tf.FusedBatchNormV3"(%0, %cst_1, %cst_2, %cst_1, %cst_2) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<1x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<1x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>)
func.return %1#0 : tensor<1x3x2x2xf32>
}
// CHECK: func.func @main(%[[ARG:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x2x2xf32> {
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2xf32>
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2x3x3x2xf32>
// CHECK-DAG: %[[CONV:.*]] = stablehlo.convolution(%[[ARG]], %[[CONST_1]]) {{.*}} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.broadcast_in_dim %[[CONST_0]], dims = [3] : (tensor<2xf32>) -> tensor<1x3x2x2xf32>
// CHECK-DAG: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST]] : tensor<1x3x2x2xf32>
// CHECK: return %[[ADD]] : tensor<1x3x2x2xf32>
// -----
func.func @fuse_conv_batchnorm_dynamic(%arg_0: tensor<?x3x4x3xf32>) -> (tensor<?x3x2x2xf32>) {
%cst_0 = "tf.Const"() {value = dense<[[[[-6.30731344, 5.4962182], [1.80364347, -7.64542675], [-2.11145878, -7.08605719]], [[-9.54062747, -6.14013147], [6.12640238, -4.18223286], [5.05738974, 8.99269962]], [[3.3535192, 0.84816426], [-6.64676809, -7.95477629], [5.81315517, 9.21566581]]], [[[1.38622558, 4.63866329], [4.54742622, -1.43770897], [-3.96835279, 2.99996852]], [[0.989735424, -4.83384752], [-7.27702999, 1.17216611], [1.33735656, 0.728900194]], [[5.1286211, 8.98645591], [1.55008793, -3.85491467], [3.7003777, 9.26594448]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_1 = "tf.Const"() {value = dense<[0.1, 0.2]> : tensor<2xf32>} : () -> tensor<2xf32>
%cst_2 = "tf.Const"() {value = dense<[0.3, 0.4]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.Conv2D"(%arg_0, %cst_0) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x2x2xf32>
%1:6 = "tf.FusedBatchNormV3"(%0, %cst_1, %cst_2, %cst_1, %cst_2) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<?x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<?x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>)
func.return %1#0 : tensor<?x3x2x2xf32>
}
// CHECK: func.func @main(%[[ARG:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x2x2xf32> {
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2xf32>
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2x3x3x2xf32>
// CHECK-DAG: %[[CONV:.*]] = stablehlo.convolution(%[[ARG]], %[[CONST_1]]) {{.*}} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x2x2xf32>
// CHECK-DAG: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] : tensor<?x3x2x2xf32> -> tensor<4xindex>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST_0]], %[[SHAPE_OF]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x3x2x2xf32>
// CHECK-DAG: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST]] : tensor<?x3x2x2xf32>
// CHECK: return %[[ADD]] : tensor<?x3x2x2xf32>
// -----
func.func @func_conv_batchnorm_relu6(%arg_0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf32>) {
%cst_0 = "tf.Const"() {value = dense<[[[[-6.30731344, 5.4962182], [1.80364347, -7.64542675], [-2.11145878, -7.08605719]], [[-9.54062747, -6.14013147], [6.12640238, -4.18223286], [5.05738974, 8.99269962]], [[3.3535192, 0.84816426], [-6.64676809, -7.95477629], [5.81315517, 9.21566581]]], [[[1.38622558, 4.63866329], [4.54742622, -1.43770897], [-3.96835279, 2.99996852]], [[0.989735424, -4.83384752], [-7.27702999, 1.17216611], [1.33735656, 0.728900194]], [[5.1286211, 8.98645591], [1.55008793, -3.85491467], [3.7003777, 9.26594448]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_1 = "tf.Const"() {value = dense<[0.1, 0.2]> : tensor<2xf32>} : () -> tensor<2xf32>
%cst_2 = "tf.Const"() {value = dense<[0.3, 0.4]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.Conv2D"(%arg_0, %cst_0) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32>
%1:6 = "tf.FusedBatchNormV3"(%0, %cst_1, %cst_2, %cst_1, %cst_2) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<1x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<1x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>)
%2 = "tf.Relu6"(%1#0) : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
func.return %2 : tensor<1x3x2x2xf32>
}
// CHECK: func.func @main(%[[ARG:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x2x2xf32> {
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2xf32>
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00> : tensor<f32>
// CHECK-DAG: %[[CONST_2:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-DAG: %[[CONST_3:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2x3x3x2xf32>
// CHECK-DAG: %[[CONV:.*]] = stablehlo.convolution(%[[ARG]], %[[CONST_3]]) {{.*}} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.broadcast_in_dim %[[CONST_0]], dims = [3] : (tensor<2xf32>) -> tensor<1x3x2x2xf32>
// CHECK-DAG: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST]] : tensor<1x3x2x2xf32>
// CHECK-DAG: %[[RELU6:.*]] = stablehlo.clamp %[[CONST_2]], %[[ADD]], %[[CONST_1]] : (tensor<f32>, tensor<1x3x2x2xf32>, tensor<f32>) -> tensor<1x3x2x2xf32>
// CHECK: return %[[RELU6]] : tensor<1x3x2x2xf32>
// -----
func.func @func_conv_batchnorm_relu6_dynamic(%arg_0: tensor<?x3x4x3xf32>) -> (tensor<?x3x2x2xf32>) {
%cst_0 = "tf.Const"() {value = dense<[[[[-6.30731344, 5.4962182], [1.80364347, -7.64542675], [-2.11145878, -7.08605719]], [[-9.54062747, -6.14013147], [6.12640238, -4.18223286], [5.05738974, 8.99269962]], [[3.3535192, 0.84816426], [-6.64676809, -7.95477629], [5.81315517, 9.21566581]]], [[[1.38622558, 4.63866329], [4.54742622, -1.43770897], [-3.96835279, 2.99996852]], [[0.989735424, -4.83384752], [-7.27702999, 1.17216611], [1.33735656, 0.728900194]], [[5.1286211, 8.98645591], [1.55008793, -3.85491467], [3.7003777, 9.26594448]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_1 = "tf.Const"() {value = dense<[0.1, 0.2]> : tensor<2xf32>} : () -> tensor<2xf32>
%cst_2 = "tf.Const"() {value = dense<[0.3, 0.4]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.Conv2D"(%arg_0, %cst_0) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x2x2xf32>
%1:6 = "tf.FusedBatchNormV3"(%0, %cst_1, %cst_2, %cst_1, %cst_2) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<?x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<?x3x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>)
%2 = "tf.Relu6"(%1#0) : (tensor<?x3x2x2xf32>) -> tensor<?x3x2x2xf32>
func.return %2 : tensor<?x3x2x2xf32>
}
// CHECK: func.func @main(%[[ARG:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x2x2xf32> {
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2xf32>
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00> : tensor<f32>
// CHECK-DAG: %[[CONST_2:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-DAG: %[[CONST_3:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2x3x3x2xf32>
// CHECK-DAG: %[[CONV:.*]] = stablehlo.convolution(%[[ARG]], %[[CONST_3]]) {{.*}} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x2x2xf32>
// CHECK-DAG: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] : tensor<?x3x2x2xf32> -> tensor<4xindex>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST_0]], %[[SHAPE_OF]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x3x2x2xf32>
// CHECK-DAG: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST]] : tensor<?x3x2x2xf32>
// CHECK-DAG: %[[RELU6:.*]] = stablehlo.clamp %[[CONST_2]], %[[ADD]], %[[CONST_1]] : (tensor<f32>, tensor<?x3x2x2xf32>, tensor<f32>) -> tensor<?x3x2x2xf32>
// CHECK: return %[[RELU6]] : tensor<?x3x2x2xf32>
// -----
// This test makes sure functions with tf._noinline=true is not inlined.
module {
func.func @stateful_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>
func.return %0: tensor<1x2x2x3xf32>
}
func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
return %arg0 : tensor<1x2x2x3xf32>
}
}
// CHECK: module
// CHECK: tf.StatefulPartitionedCall
// CHECK: func.func private @some_func
// CHECK-NOT: func.call
// -----
// This test makes sure functions without tf._noinline=true is inlined.
module {
func.func @partitioned_call(%arg0: tensor<1x2x2x3xf32>) -> (tensor<1x2x2x3xf32>) {
%0 = "tf.PartitionedCall"(%arg0) <{
config = "", config_proto = "", executor_type = "", f = @some_func
}> {
_collective_manager_ids = [], device = ""
} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
func.return %0: tensor<1x2x2x3xf32>
}
func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> {
return %arg0 : tensor<1x2x2x3xf32>
}
}
// CHECK: module
// CHECK-NOT: tf.PartitionedCall
// CHECK-NOT: some_func
// CHECK-NOT: func.call
tensorflow/compiler/mlir/quantization/stablehlo/tests/fill_quantization_options_test.cc
+110
View File
@@ -0,0 +1,110 @@
/* Copyright 2023 The StableHLO 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/fill_quantization_options.h"
#include <ostream>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
#include "tsl/platform/protobuf.h" // IWYU pragma: keep
namespace mlir::quant::stablehlo {
namespace {
using ::stablehlo::quantization::PresetQuantizationMethod;
using ::stablehlo::quantization::QuantizationComponentSpec;
using ::stablehlo::quantization::QuantizationOptions;
// Simple implementation of ::testing::EqualsProto equivalent until open source
// b/135192747 is fixed. Originally from type_to_shape_test.cc.
class ProtoStringMatcher {
public:
explicit ProtoStringMatcher(const tsl::protobuf::Message& expected)
: expected_(expected.SerializeAsString()) {}
template <typename Message>
bool MatchAndExplain(const Message& p, testing::MatchResultListener*) const {
return p.SerializeAsString() == expected_;
}
void DescribeTo(::std::ostream* os) const { *os << expected_; }
void DescribeNegationTo(::std::ostream* os) const {
*os << "not equal to expected message: " << expected_;
}
private:
const std::string expected_;
};
inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
const tsl::protobuf::Message& x) {
return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
}
void FillPresetQuantizationOptionsTestHelper(
const PresetQuantizationMethod::PresetMethod preset_quantization_options,
const QuantizationComponentSpec expected_activation_component,
const QuantizationComponentSpec expected_weight_component,
const QuantizationComponentSpec expected_bias_component) {
QuantizationOptions quantization_options;
quantization_options.mutable_quantization_method()
->mutable_preset_quantization_method()
->set_preset_method(preset_quantization_options);
QuantizationOptions filled_quantization_options =
quant::stablehlo::FillPresetQuantizationOptions(quantization_options);
for (QuantizationComponentSpec component :
filled_quantization_options.quantization_method()
.custom_quantization_method()
.quantization_component_spec()) {
switch (component.quantization_component()) {
case (QuantizationComponentSpec::COMPONENT_ACTIVATION):
EXPECT_THAT(component, EqualsProto(expected_activation_component));
break;
case (QuantizationComponentSpec::COMPONENT_WEIGHT):
EXPECT_THAT(component, EqualsProto(expected_weight_component));
break;
case (QuantizationComponentSpec::COMPONENT_BIAS):
EXPECT_THAT(component, EqualsProto(expected_bias_component));
break;
default:
break;
}
}
}
TEST(FillQuantizationOptionsTest, PresetFloat16) {
QuantizationComponentSpec activation_component, weight_component,
bias_component;
weight_component.set_quantization_component(
QuantizationComponentSpec::COMPONENT_WEIGHT);
weight_component.set_bit_width(QuantizationComponentSpec::BIT_WIDTH_16);
weight_component.set_bit_type(QuantizationComponentSpec::BIT_TYPE_FLOAT);
bias_component.set_quantization_component(
QuantizationComponentSpec::COMPONENT_BIAS);
bias_component.set_bit_width(QuantizationComponentSpec::BIT_WIDTH_16);
bias_component.set_bit_type(QuantizationComponentSpec::BIT_TYPE_FLOAT);
FillPresetQuantizationOptionsTestHelper(
/*preset_quantization_options=*/PresetQuantizationMethod::FLOAT16,
/*expected_activation_component=*/activation_component,
/*expected_weight_component*/ weight_component,
/*expected_bias_component*/ bias_component);
}
} // namespace
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/convert_func_to_bfloat16.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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-convert-func-to-bfloat16 -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @add_f32(%arg0: tensor<3x3xbf16>, %arg1: tensor<3x3xbf16>) -> tensor<3x3xbf16>
func.func @add_f32(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> {
// CHECK-NOT: f32
// CHECK: stablehlo.add
%0 = stablehlo.add %arg0, %arg1: (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
// -----
// CHECK-LABEL: @add_f64(%arg0: tensor<3x3xbf16>, %arg1: tensor<3x3xbf16>) -> tensor<3x3xbf16>
func.func @add_f64(%arg0: tensor<3x3xf64>, %arg1: tensor<3x3xf64>) -> tensor<3x3xf64> {
// CHECK-NOT: f64
// CHECK: stablehlo.add
%0 = stablehlo.add %arg0, %arg1: (tensor<3x3xf64>, tensor<3x3xf64>) -> tensor<3x3xf64>
return %0 : tensor<3x3xf64>
}
// -----
// CHECK-LABEL: @constant_f32() -> tensor<2x2xbf16>
func.func @constant_f32() -> tensor<2x2xf32> {
// CHECK-NOT: f32
// CHECK{LITERAL}: stablehlo.constant dense<[[1.398440e+00, 0.000000e+00], [3.093750e+00, -2.001950e-01]]> : tensor<2x2xbf16>
%0 = stablehlo.constant dense<[[1.4, 0.0], [3.1, -0.2]]> : tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
// -----
func.func @constant_elided() -> tensor<2x2xf32> {
// expected-error @+1 {{failed to legalize operation 'stablehlo.constant' that was explicitly marked illegal}}
%0 = stablehlo.constant dense_resource<__elided__> : tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
// -----
// CHECK-LABEL: @reduce_window_f32(%arg0: tensor<2x3x1x3xbf16>) -> tensor<2x3x1x3xbf16>
func.func @reduce_window_f32(%arg0: tensor<2x3x1x3xf32>) -> tensor<2x3x1x3xf32> {
// CHECK-NOT: f32
// CHECK: stablehlo.reduce_window
%0 = stablehlo.constant dense<0.0> : tensor<f32>
%1 = "stablehlo.reduce_window"(%arg0, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%2 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %2 : tensor<f32>
}) {padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>} : (tensor<2x3x1x3xf32>, tensor<f32>) -> tensor<2x3x1x3xf32>
return %1 : tensor<2x3x1x3xf32>
}
// -----
// CHECK-LABEL: @bitcast_convert_i32_f32(%arg0: tensor<1x256128xi32>) -> tensor<1x256128xbf16>
func.func @bitcast_convert_i32_f32(%arg0: tensor<1x256128xi32>) -> tensor<1x256128xf32> {
// CHECK: %[[BITCAST:.*]] = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xi32>) -> tensor<1x256128xf32>
// CHECK: %[[CONVERT:.*]] = stablehlo.convert %[[BITCAST]] : (tensor<1x256128xf32>) -> tensor<1x256128xbf16>
// CHECK: return %[[CONVERT]] : tensor<1x256128xbf16>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xi32>) -> tensor<1x256128xf32>
return %20 : tensor<1x256128xf32>
}
// -----
// CHECK-LABEL: @bitcast_convert_f32_i32(%arg0: tensor<1x256128xbf16>) -> tensor<1x256128xi32>
func.func @bitcast_convert_f32_i32(%arg0: tensor<1x256128xf32>) -> tensor<1x256128xi32> {
// CHECK: %[[CONVERT:.*]] = stablehlo.convert %arg0 : (tensor<1x256128xbf16>) -> tensor<1x256128xf32>
// CHECK: %[[BITCAST:.*]] = stablehlo.bitcast_convert %[[CONVERT]] : (tensor<1x256128xf32>) -> tensor<1x256128xi32>
// CHECK: return %[[BITCAST]] : tensor<1x256128xi32>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xf32>) -> tensor<1x256128xi32>
return %20 : tensor<1x256128xi32>
}
// -----
// CHECK-LABEL: @bitcast_convert_ui32_f32(%arg0: tensor<1x256128xui32>) -> tensor<1x256128xbf16>
func.func @bitcast_convert_ui32_f32(%arg0: tensor<1x256128xui32>) -> tensor<1x256128xf32> {
// CHECK: %[[BITCAST:.*]] = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xui32>) -> tensor<1x256128xf32>
// CHECK: %[[CONVERT:.*]] = stablehlo.convert %[[BITCAST]] : (tensor<1x256128xf32>) -> tensor<1x256128xbf16>
// CHECK: return %[[CONVERT]] : tensor<1x256128xbf16>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xui32>) -> tensor<1x256128xf32>
return %20 : tensor<1x256128xf32>
}
// -----
// CHECK-LABEL: @bitcast_convert_f32_ui32(%arg0: tensor<1x256128xbf16>) -> tensor<1x256128xui32>
func.func @bitcast_convert_f32_ui32(%arg0: tensor<1x256128xf32>) -> tensor<1x256128xui32> {
// CHECK: %[[CONVERT:.*]] = stablehlo.convert %arg0 : (tensor<1x256128xbf16>) -> tensor<1x256128xf32>
// CHECK: %[[BITCAST:.*]] = stablehlo.bitcast_convert %[[CONVERT]] : (tensor<1x256128xf32>) -> tensor<1x256128xui32>
// CHECK: return %[[BITCAST]] : tensor<1x256128xui32>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xf32>) -> tensor<1x256128xui32>
return %20 : tensor<1x256128xui32>
}
// -----
// CHECK-LABEL: @bitcast_convert_f32_f32(%arg0: tensor<1x256128xbf16>) -> tensor<1x256128xbf16>
func.func @bitcast_convert_f32_f32(%arg0: tensor<1x256128xf32>) -> tensor<1x256128xf32> {
// Convert bitcast_convert to no-op for f32->f32.
// CHECK: return %arg0 : tensor<1x256128xbf16>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xf32>) -> tensor<1x256128xf32>
return %20 : tensor<1x256128xf32>
}
// -----
// CHECK-LABEL: @bitcast_convert_i32_ui32(%arg0: tensor<1x256128xi32>) -> tensor<1x256128xui32>
func.func @bitcast_convert_i32_ui32(%arg0: tensor<1x256128xi32>) -> tensor<1x256128xui32> {
// Do not convert bitcast_convert for legal types.
// CHECK: %[[BITCAST:.*]] = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xi32>) -> tensor<1x256128xui32>
// CHECK: return %[[BITCAST]] : tensor<1x256128xui32>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xi32>) -> tensor<1x256128xui32>
return %20 : tensor<1x256128xui32>
}
// -----
// CHECK-LABEL: @bitcast_convert_bf16_bf16(%arg0: tensor<1x256128xbf16>) -> tensor<1x256128xbf16>
func.func @bitcast_convert_bf16_bf16(%arg0: tensor<1x256128xbf16>) -> tensor<1x256128xbf16> {
// Do not convert bitcast_convert for legal types.
// CHECK: %[[BITCAST:.*]] = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xbf16>) -> tensor<1x256128xbf16>
// CHECK: return %[[BITCAST]] : tensor<1x256128xbf16>
%20 = stablehlo.bitcast_convert %arg0 : (tensor<1x256128xbf16>) -> tensor<1x256128xbf16>
return %20 : tensor<1x256128xbf16>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/convert_xla_call_module_op_to_bfloat16.mlir
+56
View File
@@ -0,0 +1,56 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-xla-call-module-serialization -tf-stablehlo-convert-xla-call-module-op-to-bfloat16 -tf-xla-call-module-deserialization | FileCheck %s
// ConvertXlaCallModuleOpToBfloat16Pass works on XlaCallModuleOps with
// serialized modules. Which makes verification difficult. Therefore we add
// (de)serialization passes so that the input and output are deserializated
// StableHLO functions.
// CHECK-LABEL: module
module {
// CHECK-LABEL: func @main
// CHECK-SAME: %[[ARG_0:.*]]: tensor<10xf32>, %[[ARG_1:.*]]: tensor<10xf32>, %[[ARG_2:.*]]: tensor<6xi32>
func.func @main(
%arg0: tensor<10xf32>, %arg1: tensor<10xf32>, %arg2: tensor<6xi32>
) -> (tensor<10xf32>, tensor<6xi32>) {
// CHECK: %[[CAST_0:.*]] = "tf.Cast"(%[[ARG_0]]) <{Truncate = false}> : (tensor<10xf32>) -> tensor<10xbf16>
// CHECK: %[[CAST_1:.*]] = "tf.Cast"(%[[ARG_1]]) <{Truncate = false}> : (tensor<10xf32>) -> tensor<10xbf16>
// CHECK: %[[RESULT:.*]]:2 = "tf.XlaCallModule"(%[[CAST_0]], %[[CAST_1]], %[[ARG_2]])
// CHECK-SAME: _stablehlo_version = "1.0.0"
// CHECK-SAME: (tensor<10xbf16>, tensor<10xbf16>, tensor<6xi32>) -> (tensor<10xbf16>, tensor<6xi32>)
// CHECK: %[[RESULT_CAST:.*]] = "tf.Cast"(%[[RESULT]]#0) <{Truncate = false}> : (tensor<10xbf16>) -> tensor<10xf32>
%0:2 = "tf.XlaCallModule"(%arg0, %arg1, %arg2) {
Sout = [#tf_type.shape<10>], dim_args_spec = [],
_entry_function = @main_0,
_stablehlo_version = "1.0.0",
_stablehlo_module_attrs = { mhlo.num_partitions = 1 }, module = "",
platforms = [], version = 5 : i64
} : (tensor<10xf32>, tensor<10xf32>, tensor<6xi32>) -> (tensor<10xf32>, tensor<6xi32>)
// CHECK: return %[[RESULT_CAST]], %[[RESULT]]#1 : tensor<10xf32>, tensor<6xi32>
func.return %0#0, %0#1 : tensor<10xf32>, tensor<6xi32>
}
// CHECK-LABEL: func private @main_0
// CHECK-SAME: %[[ARG_0:.*]]: tensor<10xbf16>, %[[ARG_1:.*]]: tensor<10xbf16>, %[[ARG_2:.*]]: tensor<6xi32>
func.func private @main_0(
%arg0: tensor<10xf32>, %arg1: tensor<10xf32>, %arg2: tensor<6xi32>
) -> (tensor<10xf32>, tensor<6xi32>) attributes {_from_xla_call_module} {
// CHECK: %[[ADD:.*]] = stablehlo.add %[[ARG_0]], %[[ARG_1]] : tensor<10xbf16>
%0 = stablehlo.add %arg0, %arg1 : tensor<10xf32>
// CHECK: return %[[ADD]], %[[ARG_2]] : tensor<10xbf16>, tensor<6xi32>
return %0, %arg2 : tensor<10xf32>, tensor<6xi32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/defer_activation_transpose.mlir
+321
View File
@@ -0,0 +1,321 @@
// 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: stablehlo-quant-opt %s -tf-stablehlo-defer-activation-transpose \
// RUN: -split-input-file -verify-diagnostics | FileCheck %s
// Tests that an `add(transpose(arg0), arg1)` pattern is converted to
// `transpose(add(arg0, transpose(arg1)))`. The transpose in the activation is
// deferred to the output of `stablehlo.add` and an extra transpose op is
// inserted to the RHS to match the shape of the operand.
// CHECK-LABEL: add_with_activation_transpose
func.func @add_with_activation_transpose(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<1x4x3x3xf32>
%1 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
%2 = stablehlo.add %1, %0 : tensor<1x4x3x3xf32>
return %2 : tensor<1x4x3x3xf32>
}
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[CONST_0]], dims = [0, 2, 3, 1] : (tensor<1x4x3x3xf32>) -> tensor<1x3x3x4xf32>
// Check that the shape of the add is changed to reflect the deferred transpose.
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[ARG_0]], %[[TRANSPOSE_0]] : tensor<1x3x3x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose
// CHECK: return %[[TRANSPOSE_1]]
// -----
// Tests that an `add(transpose(arg0), broadcast_in_dim(arg1))` pattern is
// converted to `transpose(add(arg0, transpose(broadcast_in_dim(arg1))))`.
// The transpose in the activation is deferred to the output of `stablehlo.add`
// and an extra transpose op is inserted to the RHS to match the shape of the
// operand.
// CHECK-LABEL: add_with_activation_transpose_broadcasted_rhs
func.func @add_with_activation_transpose_broadcasted_rhs(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<4xf32>
%1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<4xf32>) -> tensor<1x4x3x3xf32>
%2 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
%3 = stablehlo.add %2, %1 : tensor<1x4x3x3xf32>
return %3 : tensor<1x4x3x3xf32>
}
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant
// CHECK-DAG: %[[BROADCAST:.+]] = stablehlo.broadcast_in_dim %[[CONST_0]], dims = [1] : (tensor<4xf32>) -> tensor<1x4x3x3xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[BROADCAST]], dims = [0, 2, 3, 1] : (tensor<1x4x3x3xf32>) -> tensor<1x3x3x4xf32>
// Check that the shape of the add is changed to reflect the deferred transpose.
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[ARG_0]], %[[TRANSPOSE_0]] : tensor<1x3x3x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose
// CHECK: return %[[TRANSPOSE_1]]
// -----
// [No change] Tests that the activation transpose whose permutation is not
// `[0, 3, 1, 2]` is not deferred.
// CHECK-LABEL: add_with_activation_transpose_permutation_mismatch
func.func @add_with_activation_transpose_permutation_mismatch(
%arg0: tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x2x4xf32>
%1 = stablehlo.transpose %arg0, dims = [0, 2, 1, 3] : (tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32>
%2 = stablehlo.add %1, %0 : tensor<1x3x2x4xf32>
return %2 : tensor<1x3x2x4xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[TRANSPOSE_0]], {{.*}}
// CHECK: return %[[ADD_0]]
// -----
// [No change] Tests that the activation transpose whose rank is not 4 is not
// deferred.
// CHECK-LABEL: add_with_activation_transpose_rank_two
func.func @add_with_activation_transpose_rank_two(%arg0: tensor<1x2xf32>) -> tensor<2x1xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<2x1xf32>
%1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
%2 = stablehlo.add %1, %0 : tensor<2x1xf32>
return %2 : tensor<2x1xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[TRANSPOSE_0]], {{.*}}
// CHECK: return %[[ADD_0]]
// -----
// [No change] Tests that the right-hand side that is not a constant is not
// deferred.
// CHECK-LABEL: add_with_activation_transpose_nonconst_rhs
func.func @add_with_activation_transpose_nonconst_rhs(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<1x4x3x3xf32>) -> tensor<1x4x3x3xf32> {
%0 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
%1 = stablehlo.add %0, %arg1 : tensor<1x4x3x3xf32>
return %1 : tensor<1x4x3x3xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[TRANSPOSE_0]], {{.*}}
// CHECK: return %[[ADD_0]]
// -----
// Tests that the transpose of the input of `stablehlo.reduce_window` is
// deferred to the result. The attributes are permutated according to the new
// input shape.
// CHECK-LABEL: reduce_window_max_activation_transpose
func.func @reduce_window_max_activation_transpose(%arg0: tensor<1x16x16x4xf32>) -> tensor<1x4x8x8xf32> {
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32> // -inf
%1 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x16x16x4xf32>) -> tensor<1x4x16x16xf32>
%2 = "stablehlo.reduce_window"(%1, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %3 : tensor<f32>
}) {window_dimensions = array<i64: 1, 1, 2, 2>, window_strides = array<i64: 1, 1, 2, 2>} : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x8x8xf32>
return %2 : tensor<1x4x8x8xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<1x16x16x4xf32>
// CHECK-DAG: %[[INIT_VALUE_CONST:.+]] = stablehlo.constant dense<0xFF800000>
// Check that the body is not modified.
// CHECK: %[[REDUCE_WINDOW:.+]] = "stablehlo.reduce_window"(%[[ARG]], %[[INIT_VALUE_CONST]])
// CHECK: <{window_dimensions = array<i64: 1, 2, 2, 1>, window_strides = array<i64: 1, 2, 2, 1>}>
// CHECK: ^bb0(%[[REDUCE_ARG_0:.+]]: tensor<f32>, %[[REDUCE_ARG_1:.+]]: tensor<f32>):
// CHECK: %[[MAX:.+]] = stablehlo.maximum %[[REDUCE_ARG_0]], %[[REDUCE_ARG_1]]
// CHECK: stablehlo.return %[[MAX]]
// Check that the attributes window_dimensions & window_strides are also
// permutated to match the new input shape.
// CHECK: (tensor<1x16x16x4xf32>, tensor<f32>) -> tensor<1x8x8x4xf32>
// Check that a `stablehlo.transpose` is added to the result to match the shape
// of the users.
// CHECK: %[[TRANSPOSE:.+]] = stablehlo.transpose %[[REDUCE_WINDOW]], dims = [0, 3, 1, 2] : (tensor<1x8x8x4xf32>) -> tensor<1x4x8x8xf32>
// CHECK: return %[[TRANSPOSE]]
// -----
// Tests that the transpose of the input of `stablehlo.reduce_window` is
// deferred to the result. The attributes are permutated according to the new
// input shape. This test is similar to the test above with the difference that
// the `stablehlo.reduce_window` has explicit optional attributes:
// `base_dilations` and `window_dilations`.
// CHECK-LABEL: reduce_window_max_activation_transpose_explicit_optional_attrs
func.func @reduce_window_max_activation_transpose_explicit_optional_attrs(
%arg0: tensor<1x16x16x4xf32>) -> tensor<1x4x15x15xf32> {
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32> // -inf
%1 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x16x16x4xf32>) -> tensor<1x4x16x16xf32>
%2 = "stablehlo.reduce_window"(%1, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %3 : tensor<f32>
}) {
window_dimensions = array<i64: 1, 1, 2, 2>,
window_strides = array<i64: 1, 1, 2, 2>,
base_dilations = array<i64: 1, 1, 2, 2>,
window_dilations = array<i64: 1, 1, 2, 2>
} : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x15x15xf32>
return %2 : tensor<1x4x15x15xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<1x16x16x4xf32>
// CHECK-DAG: %[[INIT_VALUE_CONST:.+]] = stablehlo.constant dense<0xFF800000>
// Check that the body is not modified.
// CHECK: %[[REDUCE_WINDOW:.+]] = "stablehlo.reduce_window"(%[[ARG]], %[[INIT_VALUE_CONST]])
// CHECK: <{base_dilations = array<i64: 1, 2, 2, 1>, window_dilations = array<i64: 1, 2, 2, 1>, window_dimensions = array<i64: 1, 2, 2, 1>, window_strides = array<i64: 1, 2, 2, 1>}>
// CHECK: ^bb0(%[[REDUCE_ARG_0:.+]]: tensor<f32>, %[[REDUCE_ARG_1:.+]]: tensor<f32>):
// CHECK: %[[MAX:.+]] = stablehlo.maximum %[[REDUCE_ARG_0]], %[[REDUCE_ARG_1]]
// CHECK: stablehlo.return %[[MAX]]
// Check that the attributes window_dimensions & window_strides along with
// optional attributes base_dilations and window_dilations are also permutated
// to match the new input shape.
// CHECK: (tensor<1x16x16x4xf32>, tensor<f32>) -> tensor<1x15x15x4xf32>
// Check that a `stablehlo.transpose` is added to the result to match the shape
// of the users.
// CHECK: %[[TRANSPOSE:.+]] = stablehlo.transpose %[[REDUCE_WINDOW]], dims = [0, 3, 1, 2] : (tensor<1x15x15x4xf32>) -> tensor<1x4x15x15xf32>
// CHECK: return %[[TRANSPOSE]]
// -----
// [No change] Tests that the transpose of the input of
// `stablehlo.reduce_window` is NOT deferred to the result, when the input
// tensor does not have rank 4.
// CHECK-LABEL: reduce_window_max_activation_transpose
// CHECK-SAME: (%[[ARG:.+]]: tensor<16x8xf32>) -> tensor<4x8xf32>
func.func @reduce_window_max_activation_transpose_rank2(%arg0: tensor<16x8xf32>) -> tensor<4x8xf32> {
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32> // -inf
%1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<16x8xf32>) -> tensor<8x16xf32>
%2 = "stablehlo.reduce_window"(%1, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %3 : tensor<f32>
}) {window_dimensions = array<i64: 2, 2>, window_strides = array<i64: 2, 2>} : (tensor<8x16xf32>, tensor<f32>) -> tensor<4x8xf32>
return %2 : tensor<4x8xf32>
}
// CHECK-DAG: stablehlo.constant
// CHECK: stablehlo.transpose %[[ARG]]
// CHECK: stablehlo.reduce_window
// -----
// [No change] Tests that the transpose of the input of
// `stablehlo.reduce_window` is NOT deferred to the result, when it has an
// explicit `padding` attribute.
// CHECK-LABEL: reduce_window_max_activation_transpose_with_padding
func.func @reduce_window_max_activation_transpose_with_padding(%arg0: tensor<1x16x16x4xf32>) -> tensor<1x4x9x9xf32> {
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32> // -inf
%1 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x16x16x4xf32>) -> tensor<1x4x16x16xf32>
%2 = "stablehlo.reduce_window"(%1, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %3 : tensor<f32>
}) {
window_dimensions = array<i64: 1, 1, 2, 2>,
window_strides = array<i64: 1, 1, 2, 2>,
padding = dense<[[0, 0], [0, 0], [1, 1], [1, 1]]> : tensor<4x2xi64>
} : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x9x9xf32>
return %2 : tensor<1x4x9x9xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<1x16x16x4xf32>
// CHECK-DAG: stablehlo.constant
// CHECK: stablehlo.transpose %[[ARG]]
// CHECK: stablehlo.reduce_window
// -----
// [No change] Tests that the transpose of the input of
// `stablehlo.reduce_window` is NOT deferred to the result, when the transpose
// isn't `[0, 3, 1, 2]` (i.e. NCHW->NHWC).
// CHECK-LABEL: reduce_window_max_activation_transpose_with_padding
func.func @reduce_window_max_activation_transpose_with_padding(%arg0: tensor<16x16x4x1xf32>) -> tensor<1x4x8x8xf32> {
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32> // -inf
%1 = stablehlo.transpose %arg0, dims = [3, 2, 1, 0] : (tensor<16x16x4x1xf32>) -> tensor<1x4x16x16xf32>
%2 = "stablehlo.reduce_window"(%1, %0) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %3 : tensor<f32>
}) {
window_dimensions = array<i64: 1, 1, 2, 2>,
window_strides = array<i64: 1, 1, 2, 2>
} : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x8x8xf32>
return %2 : tensor<1x4x8x8xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<16x16x4x1xf32>
// CHECK-DAG: stablehlo.constant
// CHECK: stablehlo.transpose %[[ARG]]
// CHECK: stablehlo.reduce_window
// -----
// Tests that an `max(transpose(arg0), arg1)` pattern is converted to
// `transpose(max(arg0, transpose(arg1)))`. The transpose in the activation is
// deferred to the output of `stablehlo.max` and an extra transpose op is
// inserted to the RHS to match the shape of the operand.
// CHECK-LABEL: max_with_activation_transpose
func.func @max_with_activation_transpose(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<1x4x3x3xf32>
%1 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
%2 = stablehlo.maximum %1, %0 : tensor<1x4x3x3xf32>
return %2 : tensor<1x4x3x3xf32>
}
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[CONST_0]], dims = [0, 2, 3, 1] : (tensor<1x4x3x3xf32>) -> tensor<1x3x3x4xf32>
// Check that the shape of the add is changed to reflect the deferred transpose.
// CHECK: %[[MAX_0:.+]] = stablehlo.maximum %[[ARG_0]], %[[TRANSPOSE_0]] : tensor<1x3x3x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose
// CHECK: return %[[TRANSPOSE_1]]
// -----
// [No change] Tests that the activation transpose of `stablehlo.maximum` whose
// permutation is not `[0, 3, 1, 2]` is not deferred.
// CHECK-LABEL: max_with_activation_transpose_permutation_mismatch
func.func @max_with_activation_transpose_permutation_mismatch(
%arg0: tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x2x4xf32>
%1 = stablehlo.transpose %arg0, dims = [0, 2, 1, 3] : (tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32>
%2 = stablehlo.maximum %1, %0 : tensor<1x3x2x4xf32>
return %2 : tensor<1x3x2x4xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose
// CHECK: %[[MAX_0:.+]] = stablehlo.maximum %[[TRANSPOSE_0]], {{.*}}
// CHECK: return %[[MAX_0]]
// -----
// [No change] Tests that the activation transpose of `stablehlo.maximum` whose
// rank is not 4 is not deferred.
// CHECK-LABEL: max_with_activation_transpose_rank_two
func.func @max_with_activation_transpose_rank_two(%arg0: tensor<1x2xf32>) -> tensor<2x1xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<2x1xf32>
%1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
%2 = stablehlo.maximum %1, %0 : tensor<2x1xf32>
return %2 : tensor<2x1xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose
// CHECK: %[[MAX_0:.+]] = stablehlo.maximum %[[TRANSPOSE_0]], {{.*}}
// CHECK: return %[[MAX_0]]
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/fold_constant_transpose.mlir
+73
View File
@@ -0,0 +1,73 @@
// 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: stablehlo-quant-opt %s -tf-stablehlo-fold-constant-transpose \
// RUN: -split-input-file | FileCheck %s
// CHECK-LABEL: transpose_simple_1d
func.func @transpose_simple_1d() -> tensor<2xf32> {
%0 = stablehlo.constant dense<[0.000000e+0, 1.000000e+0]> : tensor<2xf32>
%1 = stablehlo.transpose %0, dims = [0] : (tensor<2xf32>) -> tensor<2xf32>
return %1 : tensor<2xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>
// CHECK-NOT: transpose
// CHECK: return %[[CONST_0]] : tensor<2xf32>
// -----
// CHECK-LABEL: transpose_simple_2d
func.func @transpose_simple_2d() -> tensor<3x2xf32> {
%0 = stablehlo.constant dense<[[0.000000e+0, 1.000000e+0, 2.000000e+0], [3.000000e+0, 4.000000e+0, 5.000000e+0]]> : tensor<2x3xf32>
%1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3xf32>) -> tensor<3x2xf32>
return %1 : tensor<3x2xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<{{\[\[}}0.000000e+00, 3.000000e+00], [1.000000e+00, 4.000000e+00], [2.000000e+00, 5.000000e+00]]> : tensor<3x2xf32>
// CHECK-NOT: transpose
// CHECK: return %[[CONST_0]] : tensor<3x2xf32>
// -----
// CHECK-LABEL: transpose_simple_4d
func.func @transpose_simple_4d() -> tensor<5x2x3x4xf32> {
%0 = stablehlo.constant dense<1.000000e+0> : tensor<2x3x4x5xf32>
%1 = stablehlo.transpose %0, dims = [3, 0, 1, 2] : (tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32>
return %1 : tensor<5x2x3x4xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<5x2x3x4xf32>
// CHECK-NOT: transpose
// CHECK: return %[[CONST_0]] : tensor<5x2x3x4xf32>
// -----
// Tests that int constants are not folded.
// CHECK-LABEL: transpose_int
func.func @transpose_int() -> tensor<3x2xi32> {
%0 = stablehlo.constant dense<0> : tensor<2x3xi32>
%1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3xi32>) -> tensor<3x2xi32>
return %1 : tensor<3x2xi32>
}
// CHECK: transpose
// -----
// Tests that transposing an argument cannot be folded.
// CHECK-LABEL: transpose_arg
func.func @transpose_arg(%arg0: tensor<2x3xf32>) -> tensor<3x2xf32> {
%0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x3xf32>) -> tensor<3x2xf32>
return %0 : tensor<3x2xf32>
}
// CHECK: transpose
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_calibration_statistics_saver.mlir
+233
View File
@@ -0,0 +1,233 @@
// 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: stablehlo-quant-opt %s -split-input-file -mlir-disable-threading -tf-stablehlo-insert-calibration-statistics-saver | FileCheck %s
func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {
%cst = "tf.Const"() <{value = dense<[[[[-0.891899645, 0.392044574], [0.77720493, 1.31188095], [0.255048186, 2.700150e+00]], [[-1.08111858, -0.406604826], [-0.298575521, -2.25356531], [-1.00201964, 2.54532099]], [[-1.34911358, 0.279911458], [-0.868258893, -1.36708188], [0.866317451, -2.05804896]]], [[[-0.591397941, 0.331505477], [0.715151429, 2.64073896], [1.27163255, 0.206143498]], [[0.474211812, 1.45044816], [0.119936548, 2.54149938], [-0.939900994, 0.438387245]], [[-1.12486279, -1.09022558], [0.82202208, 1.04652023], [1.30316162, 2.62054276]]]]> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 5 : i32, id = "0", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
%0 = "tf.Conv2D"(%output, %cst) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
%output_1, %min_2, %max_3, %histogram_4 = "tf.CustomAggregator"(%0) <{calibration_method = 5 : i32, id = "1", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x2x2x2xf32>) -> (tensor<1x2x2x2xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
%1 = "tf.Identity"(%output_1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
return %1 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: @serving_default
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "0", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "1", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: "tf.CalibrationStatisticsSaver"(%[[MIN_O]], %[[MAX_O]], %[[HISTOGRAM_0]], %[[MIN_1]], %[[MAX_1]], %[[HISTOGRAM_1]])
// CHECK-SAME: <{calibration_methods = [5 : i32, 5 : i32], ids = ["0", "1"], output_file_path = "serving_default_0.pb"}> : (tensor<f32>, tensor<f32>, tensor<512xi64>, tensor<f32>, tensor<f32>, tensor<512xi64>) -> ()
// CHECK: return
// -----
// No CustomAggregator ops exist.
func.func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x2x2x2xf32> attributes {tf_quant.composite_function} {
%0 = "tf.Conv2D"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
%1 = "tf.BiasAdd"(%0, %arg2) <{data_format = "NHWC"}> : (tensor<1x2x2x2xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
%2 = "tf.Relu6"(%1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
return %2 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: @composite_conv2d_with_bias_and_relu6_fn_1
// CHECK-NOT: "tf.CalibrationStatisticsSaver"
// -----
// Check the IfOp is set to stateful.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1833 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: func.func @serving_default
// CHECK: "tf.If"
// CHECK-SAME: is_stateless = false
func.func @serving_default(%arg0: tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_x:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%cst = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi32>}> {device = ""} : () -> tensor<2xi32>
%cst_0 = "tf.Const"() <{value = dense<1.000000e+01> : tensor<f32>}> {device = ""} : () -> tensor<f32>
%0 = "tf.Sum"(%arg0, %cst) <{keep_dims = false}> {device = ""} : (tensor<1x4xf32>, tensor<2xi32>) -> tensor<f32>
%1 = "tf.Greater"(%0, %cst_0) {device = ""} : (tensor<f32>, tensor<f32>) -> tensor<i1>
%2:2 = "tf.If"(%1, %arg0) <{else_branch = @cond_false_80, is_stateless = true, then_branch = @cond_true_70}> {Tcond = i1, Tin = [f32], Tout = [i1, f32], _lower_using_switch_merge = true, _read_only_resource_inputs = [], device = ""} : (tensor<i1>, tensor<1x4xf32>) -> (tensor<i1>, tensor<1x3xf32>)
%3 = "tf.Identity"(%2#1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @cond_false_80
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "cond_false_80_0.pb"
func.func private @cond_false_80(%arg0: tensor<1x4xf32> {tf._user_specified_name = "x"}) -> (tensor<i1>, tensor<1x3xf32>) attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x4>], tf._original_func_name = "cond_false_8"} {
%cst = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = ""} : () -> tensor<i1>
%cst_0 = "tf.Const"() <{value = dense<[0.117216609, 0.933735609, 0.0728900209]> : tensor<3xf32>}> {device = ""} : () -> tensor<3xf32>
%cst_1 = "tf.Const"() <{value = dense<[[-0.795477629, 0.581315517, 0.921566545], [0.138622552, 0.463866323, 0.95474267], [-0.143770888, -0.796835303, 0.899996876], [0.0989735424, -0.483384758, -7.277030e-01]]> : tensor<4x3xf32>}> {device = ""} : () -> tensor<4x3xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%0 = "tf.Identity"(%cst) {device = ""} : (tensor<i1>) -> tensor<i1>
%1 = "tf.PartitionedCall"(%output, %cst_1, %cst_0) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_with_bias_fn_1}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x4xf32>, tensor<4x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
%output_2, %min_3, %max_4, %histogram_5 = "tf.CustomAggregator"(%1) <{calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%2 = "tf.Identity"(%output_2) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %0, %2 : tensor<i1>, tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @cond_true_70
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "cond_true_70_0.pb"
func.func private @cond_true_70(%arg0: tensor<1x4xf32> {tf._user_specified_name = "x"}) -> (tensor<i1>, tensor<1x3xf32>) attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x4>], tf._original_func_name = "cond_true_7"} {
%cst = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = ""} : () -> tensor<i1>
%cst_0 = "tf.Const"() <{value = dense<[0.335351914, 0.084816426, -0.664676845]> : tensor<3xf32>}> {device = ""} : () -> tensor<3xf32>
%cst_1 = "tf.Const"() <{value = dense<[[-0.630731344, 0.54962182, 0.180364341], [-0.764542698, -0.211145893, -0.708605706], [-0.954062759, -0.614013135, 0.612640202], [-0.418223292, 5.057390e-01, 0.899269938]]> : tensor<4x3xf32>}> {device = ""} : () -> tensor<4x3xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "2", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%0 = "tf.Identity"(%cst) {device = ""} : (tensor<i1>) -> tensor<i1>
%1 = "tf.PartitionedCall"(%output, %cst_1, %cst_0) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_with_bias_fn_2}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x4xf32>, tensor<4x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
%output_2, %min_3, %max_4, %histogram_5 = "tf.CustomAggregator"(%1) <{calibration_method = 1 : i32, id = "3", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%2 = "tf.Identity"(%output_2) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %0, %2 : tensor<i1>, tensor<1x3xf32>
}
func.func private @composite_matmul_with_bias_fn_1(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x3xf32> attributes {tf_quant.composite_function} {
%0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%1 = "tf.BiasAdd"(%0, %arg2) <{data_format = "NHWC"}> {device = ""} : (tensor<1x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
func.func private @composite_matmul_with_bias_fn_2(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x3xf32> attributes {tf_quant.composite_function} {
%0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%1 = "tf.BiasAdd"(%0, %arg2) <{data_format = "NHWC"}> {device = ""} : (tensor<1x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
}
// -----
// Check the IfRegion is set to stateful.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1833 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: func.func @serving_default
// CHECK: "tf.IfRegion"
// CHECK-SAME: is_stateless = false
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "serving_default_0.pb"
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "serving_default_1.pb"
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "serving_default_2.pb"
func.func @serving_default(%arg0: tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_x:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%cst = "tf.Const"() <{value = dense<1.000000e+01> : tensor<f32>}> {device = ""} : () -> tensor<f32>
%cst_0 = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi32>}> {device = ""} : () -> tensor<2xi32>
%cst_1 = "tf.Const"() <{value = dense<[[-0.630731344, 0.54962182, 0.180364341], [-0.764542698, -0.211145893, -0.708605706], [-0.954062759, -0.614013135, 0.612640202], [-0.418223292, 5.057390e-01, 0.899269938]]> : tensor<4x3xf32>}> {device = ""} : () -> tensor<4x3xf32>
%cst_2 = "tf.Const"() <{value = dense<[0.335351914, 0.084816426, -0.664676845]> : tensor<3xf32>}> {device = ""} : () -> tensor<3xf32>
%cst_3 = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = ""} : () -> tensor<i1>
%cst_4 = "tf.Const"() <{value = dense<[[-0.795477629, 0.581315517, 0.921566545], [0.138622552, 0.463866323, 0.95474267], [-0.143770888, -0.796835303, 0.899996876], [0.0989735424, -0.483384758, -7.277030e-01]]> : tensor<4x3xf32>}> {device = ""} : () -> tensor<4x3xf32>
%cst_5 = "tf.Const"() <{value = dense<[0.117216609, 0.933735609, 0.0728900209]> : tensor<3xf32>}> {device = ""} : () -> tensor<3xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%0 = "tf.Sum"(%output, %cst_0) <{keep_dims = false}> {device = ""} : (tensor<1x4xf32>, tensor<2xi32>) -> tensor<f32>
%1 = "tf.Greater"(%0, %cst) {device = ""} : (tensor<f32>, tensor<f32>) -> tensor<i1>
%2:2 = "tf.IfRegion"(%1) <{_else_func_name = "cond_false_80", _then_func_name = "cond_true_70", is_stateless = true}> ({
%4 = "tf.Identity"(%cst_3) {device = ""} : (tensor<i1>) -> tensor<i1>
%5 = "tf.PartitionedCall"(%output, %cst_1, %cst_2) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_with_bias_fn_2}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x4xf32>, tensor<4x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
%output_6, %min_7, %max_8, %histogram_9 = "tf.CustomAggregator"(%5) <{calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%6 = "tf.Identity"(%output_6) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
"tf.Yield"(%4, %6) {device = ""} : (tensor<i1>, tensor<1x3xf32>) -> ()
}, {
%4 = "tf.Identity"(%cst_3) {device = ""} : (tensor<i1>) -> tensor<i1>
%5 = "tf.PartitionedCall"(%output, %cst_4, %cst_5) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_with_bias_fn_1}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x4xf32>, tensor<4x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
%output_6, %min_7, %max_8, %histogram_9 = "tf.CustomAggregator"(%5) <{calibration_method = 1 : i32, id = "2", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%6 = "tf.Identity"(%output_6) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
"tf.Yield"(%4, %6) {device = ""} : (tensor<i1>, tensor<1x3xf32>) -> ()
}) {_lower_using_switch_merge = true, _read_only_resource_inputs = [], device = ""} : (tensor<i1>) -> (tensor<i1>, tensor<1x3xf32>)
%3 = "tf.Identity"(%2#1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
func.func private @composite_matmul_with_bias_fn_2(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x3xf32> attributes {tf_quant.composite_function} {
%0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%1 = "tf.BiasAdd"(%0, %arg2) <{data_format = "NHWC"}> {device = ""} : (tensor<1x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
func.func private @composite_matmul_with_bias_fn_1(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x3xf32> attributes {tf_quant.composite_function} {
%0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%1 = "tf.BiasAdd"(%0, %arg2) <{data_format = "NHWC"}> {device = ""} : (tensor<1x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
}
// -----
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1836 : i32}, tf_saved_model.semantics} {
func.func @main(%arg0: tensor<10x1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<10x1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%cst = stablehlo.constant dense<0.000000e+00>: tensor<10x1024x3xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<10x1x1024xf32>) -> (tensor<10x1x1024xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%0 = "tf.XlaCallModule"(%output, %cst) <{Sout = [#tf_type.shape<10x1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
%output_0, %min_1, %max_2, %histogram_3 = "tf.CustomAggregator"(%0) <{calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<10x1x3xf32>) -> (tensor<10x1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
return %output_0 : tensor<10x1x3xf32>
}
// CHECK-LABEL: @main
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: "tf.CalibrationStatisticsSaver"(%[[MIN_O]], %[[MAX_O]], %[[HISTOGRAM_0]], %[[MIN_1]], %[[MAX_1]], %[[HISTOGRAM_1]])
// CHECK-SAME: <{calibration_methods = [1 : i32, 1 : i32], ids = ["0", "1"], output_file_path = "main_0.pb"}> : (tensor<f32>, tensor<f32>, tensor<0xi64>, tensor<f32>, tensor<f32>, tensor<0xi64>) -> ()
// CHECK: return
func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<10x1x1024xf32>, %arg1: tensor<10x1024x3xf32>) -> tensor<10x1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<10x1x3xf32>
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
%1 = stablehlo.maximum %0, %cst : tensor<10x1x3xf32>
return %1 : tensor<10x1x3xf32>
}
// CHECK-LABEL: func.func private @composite_dot_general_with_relu_fn_1
// CHECK-NOT: "tf.CalibrationStatisticsSaver"
}
// -----
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1836 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: func.func @main
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "main_0.pb"
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "main_1.pb"
// CHECK: "tf.CalibrationStatisticsSaver"
// CHECK-SAME: output_file_path = "main_2.pb"
func.func @main(%arg0: tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_x:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%cst = stablehlo.constant dense<1.000000e+01> : tensor<f32>
%cst_0 = stablehlo.constant dense<[[-0.630731344, 0.54962182, 0.180364341], [-0.764542698, -0.211145893, -0.708605706], [-0.954062759, -0.614013135, 0.612640202], [-0.418223292, 5.057390e-01, 0.899269938]]> : tensor<4x3xf32>
%c = stablehlo.constant dense<true> : tensor<i1>
%cst_1 = stablehlo.constant dense<[[-0.795477629, 0.581315517, 0.921566545], [0.138622552, 0.463866323, 0.95474267], [-0.143770888, -0.796835303, 0.899996876], [0.0989735424, -0.483384758, -7.277030e-01]]> : tensor<4x3xf32>
%cst_2 = stablehlo.constant dense<-0.000000e+00> : tensor<f32>
%cst_3 = stablehlo.constant dense<[[0.335351914, 0.084816426, -0.664676845]]> : tensor<1x3xf32>
%cst_4 = stablehlo.constant dense<[[0.117216609, 0.933735609, 0.0728900209]]> : tensor<1x3xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%0 = stablehlo.reduce(%output init: %cst_2) applies stablehlo.add across dimensions = [0, 1] : (tensor<1x4xf32>, tensor<f32>) -> tensor<f32>
%1 = stablehlo.compare GT, %0, %cst : (tensor<f32>, tensor<f32>) -> tensor<i1>
%2:2 = "stablehlo.if"(%1) ({
%3 = "tf.XlaCallModule"(%output, %cst_0, %cst_3) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_bias_same_shape_fn_2, _original_entry_function = "composite_dot_general_with_bias_same_shape_fn_2", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4xf32>, tensor<4x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%output_5, %min_6, %max_7, %histogram_8 = "tf.CustomAggregator"(%3) <{calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
stablehlo.return %c, %output_5 : tensor<i1>, tensor<1x3xf32>
}, {
%3 = "tf.XlaCallModule"(%output, %cst_1, %cst_4) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_bias_same_shape_fn_1, _original_entry_function = "composite_dot_general_with_bias_same_shape_fn_1", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4xf32>, tensor<4x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%output_5, %min_6, %max_7, %histogram_8 = "tf.CustomAggregator"(%3) <{calibration_method = 1 : i32, id = "2", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
stablehlo.return %c, %output_5 : tensor<i1>, tensor<1x3xf32>
}) : (tensor<i1>) -> (tensor<i1>, tensor<1x3xf32>)
return %2#1 : tensor<1x3xf32>
}
func.func private @composite_dot_general_with_bias_same_shape_fn_2(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%1 = stablehlo.add %0, %arg2 : tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
func.func private @composite_dot_general_with_bias_same_shape_fn_1(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%1 = stablehlo.add %0, %arg2 : tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_calibration_statistics_saver_with_skipping.mlir
+61
View File
@@ -0,0 +1,61 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-insert-calibration-statistics-saver='aggregator-ops-to-ignore=skipping_id' | FileCheck %s
func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {
%cst = "tf.Const"() <{value = dense<[[[[-0.891899645, 0.392044574], [0.77720493, 1.31188095], [0.255048186, 2.700150e+00]], [[-1.08111858, -0.406604826], [-0.298575521, -2.25356531], [-1.00201964, 2.54532099]], [[-1.34911358, 0.279911458], [-0.868258893, -1.36708188], [0.866317451, -2.05804896]]], [[[-0.591397941, 0.331505477], [0.715151429, 2.64073896], [1.27163255, 0.206143498]], [[0.474211812, 1.45044816], [0.119936548, 2.54149938], [-0.939900994, 0.438387245]], [[-1.12486279, -1.09022558], [0.82202208, 1.04652023], [1.30316162, 2.62054276]]]]> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
%0 = "tf.Conv2D"(%output, %cst) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
%output_1, %min_2, %max_3, %histogram_4 = "tf.CustomAggregator"(%0) <{calibration_method = 5 : i32, id = "keeping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x2x2x2xf32>) -> (tensor<1x2x2x2xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
%1 = "tf.Identity"(%output_1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
return %1 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: @serving_default
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "keeping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: "tf.CalibrationStatisticsSaver"(%[[MIN_1]], %[[MAX_1]], %[[HISTOGRAM_1]])
// CHECK-SAME: <{calibration_methods = [5 : i32], ids = ["keeping_id"], output_file_path = "serving_default_0.pb"}> : (tensor<f32>, tensor<f32>, tensor<512xi64>) -> ()
// CHECK: return
// -----
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1836 : i32}, tf_saved_model.semantics} {
func.func @main(%arg0: tensor<10x1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<10x1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%cst = stablehlo.constant dense<0.000000e+00>: tensor<10x1024x3xf32>
%output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "skipping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<10x1x1024xf32>) -> (tensor<10x1x1024xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
%0 = "tf.XlaCallModule"(%output, %cst) <{Sout = [#tf_type.shape<10x1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
%output_0, %min_1, %max_2, %histogram_3 = "tf.CustomAggregator"(%0) <{calibration_method = 1 : i32, id = "keeping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<10x1x3xf32>) -> (tensor<10x1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
return %output_0 : tensor<10x1x3xf32>
}
// CHECK-LABEL: @main
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 1 : i32, id = "skipping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 1 : i32, id = "keeping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: "tf.CalibrationStatisticsSaver"(%[[MIN_1]], %[[MAX_1]], %[[HISTOGRAM_1]])
// CHECK-SAME: <{calibration_methods = [1 : i32], ids = ["keeping_id"], output_file_path = "main_0.pb"}> : (tensor<f32>, tensor<f32>, tensor<0xi64>) -> ()
// CHECK: return
func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<10x1x1024xf32>, %arg1: tensor<10x1024x3xf32>) -> tensor<10x1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<10x1x3xf32>
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
%1 = stablehlo.maximum %0, %cst : tensor<10x1x3xf32>
return %1 : tensor<10x1x3xf32>
}
// CHECK-LABEL: func.func private @composite_dot_general_with_relu_fn_1
// CHECK-NOT: "tf.CalibrationStatisticsSaver"
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_weight_param.mlir
+388
View File
@@ -0,0 +1,388 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-insert-weight-param | FileCheck %s
// Test that q/dq pair with per-tensor quantization parameter is inserted
// between constant and XlaCallModule op with empty `weight_only_ptq` method
// and function name containing conv.
func.func @qdq_for_conv_weight_empty(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x2x2x2>], _entry_function = @composite_conv_fn,
_original_entry_function = "composite_conv_fn",
_stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq { }",
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
return %0 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: func.func @qdq_for_conv_weight_empty
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<2x3x3x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG_0]], %[[DQ]])
// CHECK-SAME: _entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn", _quantization_method = "weight_only_ptq { }"
// CHECK-SAME: (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
// CHECK: return %[[CALL]] : tensor<1x2x2x2xf32>
// -----
// Test that q/dq pair with per-tensor quantization parameter is inserted
// between constant and XlaCallModule op with empty `weight_only_ptq` method and
// function name containing dot_general.
func.func @qdq_for_dot_general_weight_empty(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {},
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func @qdq_for_dot_general_weight_empty
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3xf32>}> : () -> tensor<2x3xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<2x3xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG_0]], %[[DQ]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "weight_only_ptq { }"
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
// CHECK: return %[[CALL]] : tensor<1x3xf32>
// -----
// Test that q/dq pair with per-tensor quantization parameter is inserted
// between constant and XlaCallModule op with `weight_only_ptq` method of
// `per_tensor` and function name containing conv.
func.func @qdq_for_conv_weight_per_tensor(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x2x2x2>], _entry_function = @composite_conv_fn,
_original_entry_function = "composite_conv_fn",
_stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {per_tensor {}}}}",
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
return %0 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: func.func @qdq_for_conv_weight_per_tensor
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<2x3x3x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG_0]], %[[DQ]])
// CHECK-SAME: _entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn", _quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {per_tensor {}}}}"
// CHECK-SAME: (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
// CHECK: return %[[CALL]] : tensor<1x2x2x2xf32>
// -----
// Test that q/dq pair with per-tensor quantization parameter is inserted
// between constant and XlaCallModule op with `weight_only_ptq` method of
// `per_tensor` and function name containing dot_general.
func.func @qdq_for_dot_general_weight_per_tensor(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {per_tensor {}}}}", _stablehlo_module_attrs = {},
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func @qdq_for_dot_general_weight_per_tensor
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3xf32>}> : () -> tensor<2x3xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<2x3xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG_0]], %[[DQ]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {per_tensor {}}}}"
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
// CHECK: return %[[CALL]] : tensor<1x3xf32>
// -----
// Test that q/dq pair with per-channel quantization parameter is inserted
// between constant and XlaCallModule op with `weight_only_ptq` method of
// `quatized_type` without specified quantization dimension and function name
// containing conv.
module attributes {tf_saved_model.semantics} {
func.func private @qdq_for_conv_weight_per_channel_default(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3x4x2>], dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [], version = 5 : i64,
_entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn",
_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {}}}}",
_stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",
device = ""
} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
// CHECK: func.func private @qdq_for_conv_weight_per_channel_default(%[[ARG0:.+]]: tensor<1x3x4x3xf32>)
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<2x3x3x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[DQ]])
// CHECK-SAME: (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CALL]]
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
// CHECK: func private @composite_conv_fn
// CHECK: %[[CONV:.+]] = stablehlo.convolution
// CHECK: return %[[CONV]]
}
// -----
// Test that q/dq pair with per-channel quantization parameter is inserted
// between constant and XlaCallModule op with `weight_only_ptq` method of
// `quatized_type` without specified quantization dimension and function name
// containing dot_general.
module attributes {tf_saved_model.semantics} {
func.func private @qdq_for_dot_general_weight_per_channel_default(%arg0: tensor<4x3x6x5xf32>) -> tensor<4x3x6x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<4x3x5x2xf32>} : () -> tensor<4x3x5x2xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<4x3x6x2>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {}}}}",
_stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
return %0 : tensor<4x3x6x2xf32>
}
// CHECK: func.func private @qdq_for_dot_general_weight_per_channel_default(%[[ARG0:.+]]: tensor<4x3x6x5xf32>)
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<4x3x5x2xf32>}> : () -> tensor<4x3x5x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<4x3x5x2xf32>) -> tensor<4x3x5x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<4x3x5x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<4x3x5x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[DQ]])
// CHECK-SAME: (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
// CHECK: return %[[CALL]]
func.func private @composite_dot_general_fn(%arg0: tensor<4x3x6x5xf32>, %arg1: tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] : (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
return %0 : tensor<4x3x6x2xf32>
}
// CHECK: func private @composite_dot_general_fn
// CHECK: %[[DOT:.+]] = stablehlo.dot_general
// CHECK: return %[[DOT]]
}
// -----
// Test that q/dq pair with per-channel quantization parameter is inserted
// between constant and XlaCallModule op with `weight_only_ptq` method of
// `quatized_type` with specified quantization dimension and function name
// containing conv.
module attributes {tf_saved_model.semantics} {
func.func private @qdq_for_conv_weight_per_channel(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3x4x2>], dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [], version = 5 : i64,
_entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn",
_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",
device = ""
} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
// CHECK: func.func private @qdq_for_conv_weight_per_channel(%[[ARG0:.+]]: tensor<1x3x4x3xf32>)
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<2x3x3x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[DQ]])
// CHECK-SAME: (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CALL]]
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
// CHECK: func private @composite_conv_fn
// CHECK: %[[CONV:.+]] = stablehlo.convolution
// CHECK: return %[[CONV]]
}
// -----
// Test that q/dq pair with per-channel quantization parameter is inserted
// between constant and XlaCallModule op with `weight_only_ptq` method of
// `quatized_type` with specified quantization dimension and function name
// containing dot_general.
module attributes {tf_saved_model.semantics} {
func.func private @qdq_for_dot_general_weight_per_channel(%arg0: tensor<4x3x6x5xf32>) -> tensor<4x3x6x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<4x3x5x2xf32>} : () -> tensor<4x3x5x2xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<4x3x6x2>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
return %0 : tensor<4x3x6x2xf32>
}
// CHECK: func.func private @qdq_for_dot_general_weight_per_channel(%[[ARG0:.+]]: tensor<4x3x6x5xf32>)
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<4x3x5x2xf32>}> : () -> tensor<4x3x5x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<4x3x5x2xf32>) -> tensor<4x3x5x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>
// CHECK: %[[DQ:.+]] = "quantization.dcast"(%[[Q]]) : (tensor<4x3x5x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<4x3x5x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[DQ]])
// CHECK-SAME: (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
// CHECK: return %[[CALL]]
func.func private @composite_dot_general_fn(%arg0: tensor<4x3x6x5xf32>, %arg1: tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] : (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
return %0 : tensor<4x3x6x2xf32>
}
// CHECK: func private @composite_dot_general_fn
// CHECK: %[[DOT:.+]] = stablehlo.dot_general
// CHECK: return %[[DOT]]
}
// -----
// Test that q/dq pair is not inserted between constant and XlaCallModule op
// whose entry function name does not include conv nor dot_general.
func.func @no_qdq_except_conv_and_dot_general(%arg0: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<3x4x2xf32>} : () -> tensor<3x4x2xf32>
%0 = "tf.XlaCallModule"(%cst, %arg0) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_gather_fn,
_original_entry_function = "composite_gather_fn", _quantization_method = "weight_only_ptq { }",
_stablehlo_module_attrs = {}, device = "", dim_args_spec = [],
disabled_checks = [], has_token_input_output = false, module = "",
platforms = [], version = 5 : i64
} : (tensor<3x4x2xf32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32>
return %0 : tensor<2x3x2x2xf32>
}
// CHECK-LABEL: func.func @no_qdq_except_conv_and_dot_general
// CHECK-NOT: quantization.qcast
// CHECK-NOT: quantization.dcast
// -----
// Test that q/dq pair is not inserted for constant whose operand number is
// not 1.
func.func @no_qdq_for_non_weight_constant(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<4.000000e-02> : tensor<3xf32>} : () -> tensor<3xf32>
%0 = "tf.XlaCallModule"(%arg0, %arg1, %cst) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_with_bias_fn,
_original_entry_function = "composite_dot_general_with_bias_fn",
_stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq { }",
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<1x2xf32>, tensor<2x3xf32>, tensor<3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func @no_qdq_for_non_weight_constant
// CHECK-NOT: quantization.qcast
// CHECK-NOT: quantization.dcast
// -----
// Test that q/dq pair is not inserted between constant and XlaCallModule op
// without `weight_only_ptq` method.
func.func @no_qdq_for_not_quantizable_call(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_stablehlo_module_attrs = {}, device = "", dim_args_spec = [],
disabled_checks = [], has_token_input_output = false, module = "",
platforms = [], version = 5 : i64
} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func @no_qdq_for_not_quantizable_call
// CHECK-NOT: quantization.qcast
// CHECK-NOT: quantization.dcast
// -----
// Test that q/dq pair is not inserted between constant and XlaCallModule op
// with different method.
func.func @no_qdq_for_not_quantizable_call(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_stablehlo_module_attrs = {}, device = "", dim_args_spec = [],
disabled_checks = [], has_token_input_output = false, module = "",
platforms = [], _quantization_method = "static_range_ptq { }", version = 5 : i64
} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func @no_qdq_for_not_quantizable_call
// CHECK-NOT: quantization.qcast
// CHECK-NOT: quantization.dcast
// -----
// Test that q/dq pair is not inserted when constant has multiple users.
func.func @no_qdq_for_multiple_users(%arg0: tensor<2x2xf32>) -> tensor<2x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst) {
Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn,
_original_entry_function = "composite_dot_general_fn",
_stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq { }",
device = "", dim_args_spec = [], disabled_checks = [],
has_token_input_output = false, module = "", platforms = [],
version = 5 : i64
} : (tensor<2x2xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
%2 = stablehlo.add %cst, %0 : tensor<2x3xf32>
return %2 : tensor<2x3xf32>
}
// CHECK-LABEL: func.func @no_qdq_for_multiple_users
// CHECK-NOT: quantization.qcast
// CHECK-NOT: quantization.dcast
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/lift_quantizable_spots_as_functions.mlir
+875
View File
@@ -0,0 +1,875 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-lift-quantizable-spots-as-functions | FileCheck %s
// CHECK-LABEL: @conv_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.convolution(%arg0, %0) 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>
func.return %1: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_fn_1
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: return %[[CONV]] : tensor<1x3x3x4xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
return %1 : tensor<1x1x64xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: return %[[DOT_GENERAL:.*]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_bias_same_shape_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x2xf32>
func.func @dot_general_with_bias_same_shape_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<2x3xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3xf32>
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%3 = stablehlo.add %2, %1 : tensor<1x3xf32>
func.return %3: tensor<1x3xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_same_shape_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %arg2
// CHECK: return %[[ADD]] : tensor<1x3xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_bias_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_with_bias_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<4xf32>
%2 = stablehlo.convolution(%arg0, %0) 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>
%3 = stablehlo.broadcast_in_dim %1, dims = [3] : (tensor<4xf32>) -> tensor<1x3x3x4xf32>
%4 = stablehlo.add %2, %3 : tensor<1x3x3x4xf32>
func.return %4: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_bias_fn_1
// CHECK: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg2
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST_IN_DIM]]
// CHECK: return %[[ADD]] : tensor<1x3x3x4xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_bias_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_with_bias_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<64xf32>
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%3 = stablehlo.broadcast_in_dim %1, dims = [2] : (tensor<64xf32>) -> tensor<1x1x64xf32>
%4 = stablehlo.add %2, %3 : tensor<1x1x64xf32>
func.return %4: tensor<1x1x64xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_fn_1
// CHECK: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg2
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[BROADCAST_IN_DIM]]
// CHECK: return %[[ADD]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_bias_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_dynamic_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%3 = shape.shape_of %2 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%5 = stablehlo.add %2, %4 : tensor<?x28x28x16xf32>
func.return %5: tensor<?x28x28x16xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x28x28x16xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_bias_dynamic_fn_1
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM]]
// CHECK: return %[[ADD]] : tensor<?x28x28x16xf32>
// CHECK: }
// -----
// Because the operand of shape_of is other than the target conv,
// should not match conv bias pattern.
// CHECK-LABEL: @conv_with_bias_dynamic_shape_not_same_op_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_dynamic_shape_not_same_op_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%4 = shape.shape_of %3 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%5 = stablehlo.dynamic_broadcast_in_dim %1, %4, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%6 = stablehlo.add %2, %5 : tensor<?x28x28x16xf32>
func.return %6: tensor<?x28x28x16xf32>
}
// CHECK-NOT: @composite_conv_with_bias_dynamic_fn_1
// -----
// CHECK-LABEL: @dot_general_with_bias_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x12544xf32>
func.func @dot_general_with_bias_dynamic_fn(%arg0: tensor<?x12544xf32>) -> tensor<?x10xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<10xf32>
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<?x12544xf32>, tensor<12544x10xf32>) -> tensor<?x10xf32>
%3 = shape.shape_of %2 : tensor<?x10xf32> -> tensor<2xindex>
%4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [1] : (tensor<10xf32>, tensor<2xindex>) -> tensor<?x10xf32>
%5 = stablehlo.add %2, %4 : tensor<?x10xf32>
func.return %5: tensor<?x10xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x10xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_dynamic_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[DOT_GENERAL]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK: return %[[ADD]] : tensor<?x10xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_relu_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_with_relu_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32>
%2 = stablehlo.convolution(%arg0, %0) 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>
%3 = stablehlo.maximum %2, %1 : tensor<1x3x3x4xf32>
func.return %3: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_relu_fn_1
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[CONV]], %[[CONST]]
// CHECK: return %[[MAX]] : tensor<1x3x3x4xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_relu_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>,
func.func @dot_general_with_relu_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%3 = stablehlo.maximum %2, %1 : tensor<1x1x64xf32>
return %3 : tensor<1x1x64xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_relu_fn_1
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[DOT_GENERAL]], %[[CONST]]
// CHECK: return %[[MAX:.*]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_relu_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_relu_dynamic_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%3 = shape.shape_of %2 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [] : (tensor<f32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%5 = stablehlo.maximum %2, %4 : tensor<?x28x28x16xf32>
func.return %5: tensor<?x28x28x16xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x28x28x16xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_relu_dynamic_fn_1
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]]
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF]]
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM]]
// CHECK: return %[[MAX]] : tensor<?x28x28x16xf32>
// CHECK: }
// -----
// Because the operand of shape_of is other than the target conv,
// should not match conv relu dynamic pattern.
// CHECK-LABEL: @conv_with_relu_dynamic_shape_not_same_op_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_relu_dynamic_shape_not_same_op_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%4 = shape.shape_of %3 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%5 = stablehlo.dynamic_broadcast_in_dim %1, %4, dims = [] : (tensor<f32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%6 = stablehlo.maximum %2, %5 : tensor<?x28x28x16xf32>
func.return %6: tensor<?x28x28x16xf32>
}
// CHECK-NOT: private @composite_conv_with_relu_dynamic_fn_1
// -----
// CHECK-LABEL: @dot_general_with_relu_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x12544xf32>
func.func @dot_general_with_relu_dynamic_fn(%arg0: tensor<?x12544xf32>) -> tensor<?x10xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<?x12544xf32>, tensor<12544x10xf32>) -> tensor<?x10xf32>
%3 = shape.shape_of %2 : tensor<?x10xf32> -> tensor<2xindex>
%4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [] : (tensor<f32>, tensor<2xindex>) -> tensor<?x10xf32>
%5 = stablehlo.maximum %2, %4 : tensor<?x10xf32>
func.return %5: tensor<?x10xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x10xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_relu_dynamic_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[DOT_GENERAL]]
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF]]
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM]]
// CHECK: return %[[MAX]] : tensor<?x10xf32>
// CHECK: }
// -----
// The pattern should not match when the const value for relu is not 0.
// CHECK-LABEL: @conv_with_relu_wrong_const_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_with_relu_wrong_const_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x3x4xf32>
%2 = stablehlo.convolution(%arg0, %0) 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>
%3 = stablehlo.maximum %2, %1 : tensor<1x3x3x4xf32>
func.return %3: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]])
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[XLA_CALL_MODULE]], %[[CONST_1]]
// CHECK: return %[[MAX]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_fn_1
// CHECK-NOT: private @composite_conv_with_relu_fn_1
// -----
// CHECK-LABEL: @conv_with_relu6_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_with_relu6_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32>
%2 = stablehlo.constant dense<6.000000e+00> : tensor<1x3x3x4xf32>
%3 = stablehlo.convolution(%arg0, %0) 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>
%4 = stablehlo.clamp %1, %3, %2 : tensor<1x3x3x4xf32>
func.return %4: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_relu6_fn_1
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[CONV]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<1x3x3x4xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_relu6_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_with_relu6_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%2 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32>
%3 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%4 = stablehlo.clamp %1, %3, %2 : tensor<1x1x64xf32>
return %4 : tensor<1x1x64xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_relu6_fn_1
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[DOT_GENERAL]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_relu6_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_relu6_dynamic_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%2 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%3 = stablehlo.convolution(%arg0, %0) 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<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%4 = stablehlo.clamp %1, %3, %2 : (tensor<f32>, tensor<?x28x28x16xf32>, tensor<f32>) -> tensor<?x28x28x16xf32>
func.return %4: tensor<?x28x28x16xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x28x28x16xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_relu6_fn_1
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[CONV]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<?x28x28x16xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_relu6_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x12544xf32>
func.func @dot_general_with_relu6_dynamic_fn(%arg0: tensor<?x12544xf32>) -> tensor<?x10xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32>
%1 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%2 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%3 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<?x12544xf32>, tensor<12544x10xf32>) -> tensor<?x10xf32>
%4 = stablehlo.clamp %1, %3, %2 : (tensor<f32>, tensor<?x10xf32>, tensor<f32>) -> tensor<?x10xf32>
func.return %4: tensor<?x10xf32>
}
// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x10xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_relu6_fn_1
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[DOT_GENERAL]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<?x10xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_bias_same_shape_and_relu_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_with_bias_same_shape_and_relu_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<1x1x64xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%3 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%4 = stablehlo.add %3, %1 : tensor<1x1x64xf32>
%5 = stablehlo.maximum %4, %2 : tensor<1x1x64xf32>
func.return %5: tensor<1x1x64xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_same_shape_and_relu_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %arg2
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[CONST]]
// CHECK: return %[[MAX]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_bias_and_relu_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_with_bias_and_relu_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<4xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32>
%3 = stablehlo.convolution(%arg0, %0) 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>
%4 = stablehlo.broadcast_in_dim %1, dims = [3] : (tensor<4xf32>) -> tensor<1x3x3x4xf32>
%5 = stablehlo.add %3, %4 : tensor<1x3x3x4xf32>
%6 = stablehlo.maximum %5, %2 : tensor<1x3x3x4xf32>
func.return %6: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_bias_and_relu_fn_1
// CHECK: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg2
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST_IN_DIM]]
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[CONST]]
// CHECK: return %[[MAX]] : tensor<1x3x3x4xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_bias_and_relu_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_with_bias_and_relu_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<64xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%3 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%4 = stablehlo.broadcast_in_dim %1, dims = [2] : (tensor<64xf32>) -> tensor<1x1x64xf32>
%5 = stablehlo.add %3, %4 : tensor<1x1x64xf32>
%6 = stablehlo.maximum %5, %2 : tensor<1x1x64xf32>
func.return %6: tensor<1x1x64xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_and_relu_fn_1
// CHECK: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg2
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[BROADCAST_IN_DIM]]
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[CONST]]
// CHECK: return %[[MAX]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_bias_and_relu_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_and_relu_dynamic_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%4 = shape.shape_of %3 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%5 = stablehlo.dynamic_broadcast_in_dim %1, %4, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%6 = stablehlo.add %3, %5 : tensor<?x28x28x16xf32>
%7 = shape.shape_of %6 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%8 = stablehlo.dynamic_broadcast_in_dim %2, %7, dims = [] : (tensor<f32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%9 = stablehlo.maximum %6, %8 : tensor<?x28x28x16xf32>
func.return %9: tensor<?x28x28x16xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x28x28x16xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_bias_and_relu_dynamic_fn_1
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[CONV]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK: %[[SHAPE_OF_1:.*]] = shape.shape_of %[[ADD]]
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_1:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF_1]]
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[DYNAMIC_BROADCAST_IN_DIM_1]]
// CHECK: return %[[MAX]] : tensor<?x28x28x16xf32>
// CHECK: }
// -----
// Because the operand of shape_of is other than the target conv,
// should not match conv bias relu dynamic pattern.
// CHECK-LABEL: @conv_with_bias_and_relu_dynamic_shape_not_same_op_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_and_relu_dynamic_shape_not_same_op_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%4 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%5 = shape.shape_of %4 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%6 = stablehlo.dynamic_broadcast_in_dim %1, %5, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%7 = stablehlo.add %3, %6 : tensor<?x28x28x16xf32>
%8 = shape.shape_of %7 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%9 = stablehlo.dynamic_broadcast_in_dim %2, %8, dims = [] : (tensor<f32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%10 = stablehlo.maximum %7, %9 : tensor<?x28x28x16xf32>
func.return %10: tensor<?x28x28x16xf32>
}
// CHECK-NOT: private @composite_conv_with_bias_and_relu_dynamic_fn_1
// -----
// CHECK-LABEL: @dot_general_with_bias_and_relu_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x12544xf32>
func.func @dot_general_with_bias_and_relu_dynamic_fn(%arg0: tensor<?x12544xf32>) -> tensor<?x10xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<10xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%3 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<?x12544xf32>, tensor<12544x10xf32>) -> tensor<?x10xf32>
%4 = shape.shape_of %3 : tensor<?x10xf32> -> tensor<2xindex>
%5 = stablehlo.dynamic_broadcast_in_dim %1, %4, dims = [1] : (tensor<10xf32>, tensor<2xindex>) -> tensor<?x10xf32>
%6 = stablehlo.add %3, %5 : tensor<?x10xf32>
%7 = shape.shape_of %6 : tensor<?x10xf32> -> tensor<2xindex>
%8 = stablehlo.dynamic_broadcast_in_dim %2, %7, dims = [] : (tensor<f32>, tensor<2xindex>) -> tensor<?x10xf32>
%9 = stablehlo.maximum %6, %8 : tensor<?x10xf32>
func.return %9: tensor<?x10xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x10xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_and_relu_dynamic_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[DOT_GENERAL]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK: %[[SHAPE_OF_1:.*]] = shape.shape_of %[[ADD]]
// CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_1:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF_1]]
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[DYNAMIC_BROADCAST_IN_DIM_1]]
// CHECK: return %[[MAX]] : tensor<?x10xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_bias_same_shape_and_relu6_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_with_bias_same_shape_and_relu6_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<1x1x64xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%3 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32>
%4 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%5 = stablehlo.add %4, %1 : tensor<1x1x64xf32>
%6 = stablehlo.clamp %2, %5, %3 : tensor<1x1x64xf32>
func.return %6: tensor<1x1x64xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_same_shape_and_relu6_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %arg2
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[ADD]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_bias_and_relu6_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>
func.func @conv_with_bias_and_relu6_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<4xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32>
%3 = stablehlo.constant dense<6.000000e+00> : tensor<1x3x3x4xf32>
%4 = stablehlo.convolution(%arg0, %0) 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>
%5 = stablehlo.broadcast_in_dim %1, dims = [3] : (tensor<4xf32>) -> tensor<1x3x3x4xf32>
%6 = stablehlo.add %4, %5 : tensor<1x3x3x4xf32>
%7 = stablehlo.clamp %2, %6, %3 : tensor<1x3x3x4xf32>
func.return %7: tensor<1x3x3x4xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_bias_and_relu6_fn_1
// CHECK: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg2
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST_IN_DIM]]
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[ADD]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<1x3x3x4xf32>
// CHECK: }
// -----
// CHECK-LABEL: @dot_general_with_bias_and_relu6_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>
func.func @dot_general_with_bias_and_relu6_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<64xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%3 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32>
%4 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
%5 = stablehlo.broadcast_in_dim %1, dims = [2] : (tensor<64xf32>) -> tensor<1x1x64xf32>
%6 = stablehlo.add %4, %5 : tensor<1x1x64xf32>
%7 = stablehlo.clamp %2, %6, %3 : tensor<1x1x64xf32>
func.return %7: tensor<1x1x64xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_and_relu6_fn_1
// CHECK: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg2
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[BROADCAST_IN_DIM]]
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[ADD]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<1x1x64xf32>
// CHECK: }
// -----
// CHECK-LABEL: @conv_with_bias_and_relu6_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_and_relu6_dynamic_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%3 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%4 = stablehlo.convolution(%arg0, %0) 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<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%5 = shape.shape_of %4 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%6 = stablehlo.dynamic_broadcast_in_dim %1, %5, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%7 = stablehlo.add %4, %6 : tensor<?x28x28x16xf32>
%8 = stablehlo.clamp %2, %7, %3 : (tensor<f32>, tensor<?x28x28x16xf32>, tensor<f32>) -> tensor<?x28x28x16xf32>
func.return %8: tensor<?x28x28x16xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x28x28x16xf32>
// CHECK: }
// CHECK-LABEL: private @composite_conv_with_bias_and_relu6_dynamic_fn_1
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1)
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[CONV]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[ADD]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<?x28x28x16xf32>
// CHECK: }
// -----
// Because the operand of shape_of is other than the target conv,
// should not match conv bias relu6 dynamic pattern.
// CHECK-LABEL: @conv_with_bias_and_relu6_dynamic_shape_not_same_op_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_and_relu6_dynamic_shape_not_same_op_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%3 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%4 = stablehlo.convolution(%arg0, %0) 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<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%5 = stablehlo.convolution(%arg0, %0) 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<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>
%6 = shape.shape_of %5 : tensor<?x28x28x16xf32> -> tensor<4xindex>
%7 = stablehlo.dynamic_broadcast_in_dim %1, %6, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor<?x28x28x16xf32>
%8 = stablehlo.add %4, %7 : tensor<?x28x28x16xf32>
%9 = stablehlo.clamp %2, %8, %3 : (tensor<f32>, tensor<?x28x28x16xf32>, tensor<f32>) -> tensor<?x28x28x16xf32>
func.return %9: tensor<?x28x28x16xf32>
}
// CHECK-NOT: private @composite_conv_with_bias_and_relu6_dynamic_fn_1
// -----
// CHECK-LABEL: @dot_general_with_bias_and_relu6_dynamic_fn(
// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x12544xf32>
func.func @dot_general_with_bias_and_relu6_dynamic_fn(%arg0: tensor<?x12544xf32>) -> tensor<?x10xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32>
%1 = stablehlo.constant dense<2.000000e+00> : tensor<10xf32>
%2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%3 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%4 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<?x12544xf32>, tensor<12544x10xf32>) -> tensor<?x10xf32>
%5 = shape.shape_of %4 : tensor<?x10xf32> -> tensor<2xindex>
%6 = stablehlo.dynamic_broadcast_in_dim %1, %5, dims = [1] : (tensor<10xf32>, tensor<2xindex>) -> tensor<?x10xf32>
%7 = stablehlo.add %4, %6 : tensor<?x10xf32>
%8 = stablehlo.clamp %2, %7, %3 : (tensor<f32>, tensor<?x10xf32>, tensor<f32>) -> tensor<?x10xf32>
func.return %8: tensor<?x10xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<?x10xf32>
// CHECK: }
// CHECK-LABEL: private @composite_dot_general_with_bias_and_relu6_dynamic_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[DOT_GENERAL]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<6.000000e+00>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<0.000000e+00>
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[CONST_0]], %[[ADD]], %[[CONST_1]]
// CHECK: return %[[CLAMP]] : tensor<?x10xf32>
// CHECK: }
// -----
// CHECK-LABEL: @gather_fn(
func.func @gather_fn() -> tensor<2x3x2x2xi32> {
%0 = stablehlo.constant dense<1> : tensor<3x4x2xi32>
%1 = stablehlo.constant dense<1> : tensor<2x3x2xi64>
%2 = "stablehlo.gather"(%0, %1) {
dimension_numbers = #stablehlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [1, 0],
index_vector_dim = 2>,
slice_sizes = array<i64: 1, 2, 2>,
indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
func.return %2: tensor<2x3x2x2xi32>
}
// CHECK: %[[OPERAND:.*]] = stablehlo.constant
// CHECK: %[[INDICES:.*]] = stablehlo.constant
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[OPERAND]], %[[INDICES]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<2x3x2x2xi32>
// CHECK: }
// CHECK-LABEL: private @composite_gather_fn_1
// CHECK: %[[GATHER:.*]] = "stablehlo.gather"(%arg0, %arg1)
// CHECK: return %[[GATHER]] : tensor<2x3x2x2xi32>
// CHECK: }
// -----
// Test that the name of composite functions are deterministic. There are 3
// unsorted functions in this module and each function has 2 quantizable ops.
module {
func.func @conv_3_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.convolution(%arg0, %0) 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>
%2 = stablehlo.convolution(%1, %0) 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>
func.return %2: tensor<1x3x3x4xf32>
}
func.func @conv_1_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.convolution(%arg0, %0) 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>
%2 = stablehlo.convolution(%1, %0) 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>
func.return %2: tensor<1x3x3x4xf32>
}
func.func @conv_2_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
%1 = stablehlo.convolution(%arg0, %0) 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>
%2 = stablehlo.convolution(%1, %0) 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>
func.return %2: tensor<1x3x3x4xf32>
}
}
// CHECK-LABEL: @conv_3_fn
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_6, _original_entry_function = "composite_conv_fn_6"
// CHECK-SAME: _stablehlo_version = "{{.*}}"
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_5, _original_entry_function = "composite_conv_fn_5"
// CHECK-SAME: _stablehlo_version = "{{.*}}"
// CHECK-LABEL: @conv_1_fn
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_2, _original_entry_function = "composite_conv_fn_2"
// CHECK-SAME: _stablehlo_version = "{{.*}}"
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_1, _original_entry_function = "composite_conv_fn_1"
// CHECK-SAME: _stablehlo_version = "{{.*}}"
// CHECK-LABEL: @conv_2_fn
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_4, _original_entry_function = "composite_conv_fn_4"
// CHECK-SAME: _stablehlo_version = "{{.*}}"
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @composite_conv_fn_3, _original_entry_function = "composite_conv_fn_3"
// CHECK-SAME: _stablehlo_version = "{{.*}}"
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/lift_quantizable_spots_as_functions_with_quantization_specs.mlir
+171
View File
@@ -0,0 +1,171 @@
// 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: stablehlo-quant-opt %s -stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs="quantization-specs=disable-all-dot-general" \
// RUN: -split-input-file | FileCheck %s --check-prefix=DISABLE-ALL-DOT-GENERAL
// Tests that `composite_dot_general_fn_1` and its corresponding XlaCallModuleOp
// contains attributes required for quantization, including the
// `_quantization_method` attribute that contains textpb of `Method`.
// DISABLE-ALL-DOT-GENERAL: @main
func.func @main(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
return %1 : tensor<1x1x64xf32>
}
// DISABLE-ALL-DOT-GENERAL: %[[CONST:.+]] = stablehlo.constant dense<2.000000e+00>
// DISABLE-ALL-DOT-GENERAL: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// Check that the `_quantization_method` attribute contains the quantization
// method in textproto format. The dot_general op quantization is explicitly
// disabled by having `_quantization_method = "no_quantization { }"`.
// DISABLE-ALL-DOT-GENERAL-SAME: _entry_function = @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: _original_entry_function
// DISABLE-ALL-DOT-GENERAL-SAME: _quantization_method = "no_quantization { }"
// DISABLE-ALL-DOT-GENERAL-SAME: _tfl_quant_trait = "fully_quantizable"
// DISABLE-ALL-DOT-GENERAL: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>
// DISABLE-ALL-DOT-GENERAL: }
// DISABLE-ALL-DOT-GENERAL-LABEL: private @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: tf_quant.composite_function
// DISABLE-ALL-DOT-GENERAL: %[[DOT_GENERAL:.+]] = stablehlo.dot_general %arg0, %arg1
// DISABLE-ALL-DOT-GENERAL: return %[[DOT_GENERAL:.+]] : tensor<1x1x64xf32>
// DISABLE-ALL-DOT-GENERAL: }
// -----
// RUN: stablehlo-quant-opt %s -stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs="quantization-specs=empty" \
// RUN: -split-input-file | FileCheck %s --check-prefix=EMPTY
// Tests that `composite_dot_general_fn_1` and its corresponding XlaCallModuleOp
// contains attributes required for quantization. `_quantization_method` is not
// set, as it is implicitly disabled.
// EMPTY: @main
func.func @main(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
return %1 : tensor<1x1x64xf32>
}
// EMPTY: %[[CONST:.+]] = stablehlo.constant dense<2.000000e+00>
// EMPTY: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// Check that the `_quantization_method` attribute doesn't contain the
// quantization method, implying "no_quantization".
// EMPTY-SAME: _entry_function = @composite_dot_general_fn_1
// EMPTY-SAME: _original_entry_function
// EMPTY-NOT: _quantization_method
// EMPTY-SAME: _tfl_quant_trait = "fully_quantizable"
// EMPTY: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>
// EMPTY: }
// EMPTY-LABEL: private @composite_dot_general_fn_1
// EMPTY-SAME: tf_quant.composite_function
// EMPTY: %[[DOT_GENERAL:.+]] = stablehlo.dot_general %arg0, %arg1
// EMPTY: return %[[DOT_GENERAL:.+]] : tensor<1x1x64xf32>
// EMPTY: }
// -----
// RUN: stablehlo-quant-opt %s -stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs="quantization-specs=static-range-ptq-to-all" \
// RUN: -split-input-file | FileCheck %s --check-prefix=STATIC-RANGE-PTQ-TO-ALL
// STATIC-RANGE-PTQ-TO-ALL: @main
func.func @main(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
return %1 : tensor<1x1x64xf32>
}
// Tests that `composite_dot_general_fn_1` and its corresponding XlaCallModuleOp
// contains attributes required for quantization, including the
// `_quantization_method` attribute that contains textpb of `Method`.
// STATIC-RANGE-PTQ-TO-ALL: %[[CONST:.+]] = stablehlo.constant dense<2.000000e+00>
// STATIC-RANGE-PTQ-TO-ALL: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// Check that the `_quantization_method` attribute contains the quantization
// method in textproto format, enabling static-range PTQ.
// STATIC-RANGE-PTQ-TO-ALL-SAME: _entry_function = @composite_dot_general_fn_1
// STATIC-RANGE-PTQ-TO-ALL-SAME: _original_entry_function
// STATIC-RANGE-PTQ-TO-ALL-SAME: _quantization_method = "static_range_ptq { }"
// STATIC-RANGE-PTQ-TO-ALL-SAME: _tfl_quant_trait = "fully_quantizable"
// STATIC-RANGE-PTQ-TO-ALL: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>
// STATIC-RANGE-PTQ-TO-ALL: }
// STATIC-RANGE-PTQ-TO-ALL-LABEL: private @composite_dot_general_fn_1
// STATIC-RANGE-PTQ-TO-ALL-SAME: tf_quant.composite_function
// STATIC-RANGE-PTQ-TO-ALL: %[[DOT_GENERAL:.+]] = stablehlo.dot_general %arg0, %arg1
// STATIC-RANGE-PTQ-TO-ALL: return %[[DOT_GENERAL:.+]] : tensor<1x1x64xf32>
// STATIC-RANGE-PTQ-TO-ALL: }
// -----
// RUN: stablehlo-quant-opt %s -stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs="quantization-specs=static-range-ptq-to-compute-heavy" \
// RUN: -split-input-file | FileCheck %s --check-prefix=STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: @main
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
%0 = stablehlo.add %arg0, %arg0 : tensor<1x2xf32>
return %0 : tensor<1x2xf32>
}
// Tests that `composite_add_fn_1` does not quantize when quantizing
// only compute-heavy ops.
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: %[[CONST:.+]] = stablehlo.constant dense<2.000000e+00>
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%arg0, %arg0)
// Check that the `_quantization_method` attribute contains the quantization
// method in textproto format, enabling static-range PTQ.
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: _entry_function = @composite_add_fn_1
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: _original_entry_function
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY-NOT: _quantization_method
// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: _tfl_quant_trait = "fully_quantizable"
// -----
// RUN: stablehlo-quant-opt %s -stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs="quantization-specs=static-range-ptq-to-all" \
// RUN: -split-input-file | FileCheck %s --check-prefix=STATIC-RANGE-PTQ-TO-ALL
// STATIC-RANGE-PTQ-TO-ALL-LABEL: @some_func
func.func @some_func(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
return %1 : tensor<1x1x64xf32>
}
// Tests that XlaCallModuleOp in non-main function has attributes set correctly.
// STATIC-RANGE-PTQ-TO-ALL: %[[CONST:.+]] = stablehlo.constant dense<2.000000e+00>
// STATIC-RANGE-PTQ-TO-ALL: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%arg0, %[[CONST]])
// Check that the `_quantization_method` attribute contains the quantization
// method in textproto format, enabling static-range PTQ.
// STATIC-RANGE-PTQ-TO-ALL-SAME: _entry_function = @composite_dot_general_fn_1
// STATIC-RANGE-PTQ-TO-ALL-SAME: _original_entry_function
// STATIC-RANGE-PTQ-TO-ALL-SAME: _quantization_method = "static_range_ptq { }"
// STATIC-RANGE-PTQ-TO-ALL-SAME: _tfl_quant_trait = "fully_quantizable"
// STATIC-RANGE-PTQ-TO-ALL: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>
// STATIC-RANGE-PTQ-TO-ALL: }
// STATIC-RANGE-PTQ-TO-ALL-LABEL: private @composite_dot_general_fn_1
// STATIC-RANGE-PTQ-TO-ALL-SAME: tf_quant.composite_function
// STATIC-RANGE-PTQ-TO-ALL: %[[DOT_GENERAL:.+]] = stablehlo.dot_general %arg0, %arg1
// STATIC-RANGE-PTQ-TO-ALL: return %[[DOT_GENERAL:.+]] : tensor<1x1x64xf32>
// STATIC-RANGE-PTQ-TO-ALL: }
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/merge-fusion-with-dequantize.mlir
+212
View File
@@ -0,0 +1,212 @@
// 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: stablehlo-quant-opt %s -tf-stablehlo-merge-fusion-with-dequantize -split-input-file -verify-diagnostics | FileCheck %s
// Merge fusion with dequantize for relu case.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @merge_relu_fusion
func.func private @merge_relu_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @quantized_dot_general_relu_fn
// CHECK-SAME: -> tensor<1x3xf32>
%2 = call @quantized_dot_general_relu_fn(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @quantized_dot_general_relu_fn
func.func private @quantized_dot_general_relu_fn(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>> attributes {_from_xla_call_module} {
// CHECK: %[[MIN:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]]
// CHECK: %[[MAX:.*]] = chlo.broadcast_maximum %[[DQ]], %[[MIN]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
}
// -----
// Merge fusion with dequantize for relu6 case.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @merge_relu6_fusion
func.func private @merge_relu6_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @quantized_dot_general_relu6_fn
// CHECK-SAME: -> tensor<1x3xf32>
%2 = call @quantized_dot_general_relu6_fn(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @quantized_dot_general_relu6_fn
func.func private @quantized_dot_general_relu6_fn(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>> attributes {_from_xla_call_module} {
// CHECK-DAG: %[[MIN:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-DAG: %[[MAX:.*]] = stablehlo.constant dense<6.000000e+00> : tensor<f32>
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]]
// CHECK: %[[CLAMP:.*]] = stablehlo.clamp %[[MIN]], %[[DQ]], %[[MAX]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
}
// -----
// Merge fusion with dequantize for no activation case.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @merge_no_act_fusion
func.func private @merge_no_act_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @quantized_dot_general_fn
// CHECK-SAME: -> tensor<1x3xf32>
%2 = call @quantized_dot_general_fn(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @quantized_dot_general_fn
func.func private @quantized_dot_general_fn(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]]
// CHECK: return %[[DQ]] : tensor<1x3xf32>
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
}
// -----
// Do not merge when quant.uniform result is used directly.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @no_merge_fusion_direct_usage
func.func private @no_merge_fusion_direct_usage(%arg0: tensor<1x4xf32>) -> (tensor<1x3xf32>, tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @quantized_dot_general_relu_fn
// CHECK-SAME: -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%2 = call @quantized_dot_general_relu_fn(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %3, %2 : tensor<1x3xf32>, tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
// CHECK-LABEL: func.func private @quantized_dot_general_relu_fn
func.func private @quantized_dot_general_relu_fn(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
}
// -----
// Do not merge when fusion and dequantize is already merged.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @no_merge_fusion_already_merged
func.func private @no_merge_fusion_already_merged(%arg0: tensor<1x4xf32>) -> (tensor<1x3xf32>) {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @quantized_dot_general_fn
// CHECK-SAME: -> tensor<1x3xf32>
%2 = call @quantized_dot_general_fn(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @quantized_dot_general_fn
func.func private @quantized_dot_general_fn(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
%1 = stablehlo.uniform_dequantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
}
// -----
// Do not merge when function is not quantized function.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @merge_relu_fusion
func.func private @merge_relu_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @some_func
// CHECK-SAME: -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%2 = call @some_func(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @some_func
func.func private @some_func(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
}
// -----
// Do not merge when the quantized fusion is invalid.
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: func.func private @merge_relu_fusion
func.func private @merge_relu_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: call @quantized_dot_general_relu_fn
// CHECK-SAME: -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%2 = call @quantized_dot_general_relu_fn(%1, %0) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @quantized_dot_general_relu_fn
func.func private @quantized_dot_general_relu_fn(
%arg0: tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>,
%arg1: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>> attributes {_from_xla_call_module} {
%0 = stablehlo.constant() {value = dense<2> : tensor<1x3xi8>} : () -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
return %0 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/nchw_convolution_to_nhwc.mlir
+110
View File
@@ -0,0 +1,110 @@
// 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: stablehlo-quant-opt %s -tf-stablehlo-nchw-convolution-to-nhwc \
// RUN: -split-input-file -verify-diagnostics | FileCheck %s
// Tests that `stablehlo.transpose` ops are inserted for each of input, filter,
// and output.
// Output dimension numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
// CHECK-LABEL: nchw_conv
// CHECK-SAME: %[[ARG:.+]]: tensor<1x8x4x4xf32>
func.func @nchw_conv(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32> {
%0 = stablehlo.constant() {value = dense<7.000000e+00> : tensor<8x8x3x3xf32>} : () -> tensor<8x8x3x3xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x4x4xf32>, tensor<8x8x3x3xf32>) -> tensor<1x8x4x4xf32>
return %2 : tensor<1x8x4x4xf32>
}
// CHECK-DAG: %[[CONST:.+]] = stablehlo.constant {{.*}} : tensor<8x8x3x3xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>
// CHECK-DAG: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[CONST]], dims = [2, 3, 1, 0] : (tensor<8x8x3x3xf32>) -> tensor<3x3x8x8xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[TRANSPOSE_1]]) 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<1x4x4x8xf32>, tensor<3x3x8x8xf32>) -> tensor<1x4x4x8xf32>
// CHECK: %[[TRANSPOSE_2:.+]] = stablehlo.transpose %[[CONV]], dims = [0, 3, 1, 2] : (tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32>
// -----
// Tests that the conversion doesn't happen when the input dimension numbers
// are not [b, f, 0, 1].
// CHECK-LABEL: conv_input_dim_numbers_mismatch
func.func @conv_input_dim_numbers_mismatch(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> {
%0 = stablehlo.constant() {value = dense<7.000000e+00> : tensor<8x8x3x3xf32>} : () -> tensor<8x8x3x3xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x4x8xf32>, tensor<8x8x3x3xf32>) -> tensor<1x8x4x4xf32>
return %2 : tensor<1x8x4x4xf32>
}
// CHECK-NOT: stablehlo.transpose
// CHECK: %[[CONV:.+]] = stablehlo.convolution
// CHECK-SAME{LITERAL}: [b, 0, 1, f]x[o, i, 0, 1]->[b, f, 0, 1]
// CHECK-NOT: stablehlo.transpose
// -----
// Tests that the conversion doesn't happen when the feature dimension numbers
// are not [i, 0, 1, o].
// CHECK-LABEL: conv_feature_dim_numbers_mismatch
func.func @conv_feature_dim_numbers_mismatch(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32> {
%0 = stablehlo.constant() {value = dense<7.000000e+00> : tensor<8x3x3x8xf32>} : () -> tensor<8x3x3x8xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[i, 0, 1, o]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x4x4xf32>, tensor<8x3x3x8xf32>) -> tensor<1x8x4x4xf32>
return %2 : tensor<1x8x4x4xf32>
}
// CHECK-NOT: stablehlo.transpose
// CHECK: %[[CONV:.+]] = stablehlo.convolution
// CHECK-SAME{LITERAL}: [b, f, 0, 1]x[i, 0, 1, o]->[b, f, 0, 1]
// CHECK-NOT: stablehlo.transpose
// -----
// Tests that the conversion doesn't happen when the output dimension numbers
// are not [b, 0, 1, f].
// CHECK-LABEL: conv_output_dim_numbers_mismatch
func.func @conv_output_dim_numbers_mismatch(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32> {
%0 = stablehlo.constant() {value = dense<7.000000e+00> : tensor<8x8x3x3xf32>} : () -> tensor<8x8x3x3xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x4x4xf32>, tensor<8x8x3x3xf32>) -> tensor<1x4x4x8xf32>
return %2 : tensor<1x4x4x8xf32>
}
// CHECK-NOT: stablehlo.transpose
// CHECK: %[[CONV:.+]] = stablehlo.convolution
// CHECK-SAME{LITERAL}: [b, f, 0, 1]x[o, i, 0, 1]->[b, 0, 1, f]
// CHECK-NOT: stablehlo.transpose
// -----
// Tests that a quantized convolution does not match. No conversion occurs.
// CHECK-LABEL: quantized_convolution
func.func @quantized_convolution(%arg0: tensor<1x4x3x3x!quant.uniform<i8:f32, 1.000000e+0:-100>>, %arg1: tensor<2x4x3x3x!quant.uniform<i8:f32:0, {2.000000e+2, 3.000000e+3}>>) -> tensor<1x2x3x3x!quant.uniform<i8:f32, 4.000000e+0>> {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x3x3x!quant.uniform<i8:f32, 1.000000e+0:-100>>, tensor<2x4x3x3x!quant.uniform<i8:f32:0, {2.000000e+2, 3.000000e+3}>>) -> tensor<1x2x3x3x!quant.uniform<i8:f32, 4.000000e+0>>
return %0 : tensor<1x2x3x3x!quant.uniform<i8:f32, 4.000000e+0>>
}
// CHECK-NOT: stablehlo.transpose
// -----
// Tests that a quantized convolution with rank > 4 does not match.
// No conversion occurs.
// CHECK-LABEL: convolution_3d
func.func @convolution_3d(%arg0: tensor<1x4x28x28x1xf32>, %arg1: tensor<2x3x3x1x16xf32>) -> tensor<1x3x26x26x16xf32> {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f], window = {} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x28x28x1xf32>, tensor<2x3x3x1x16xf32>) -> tensor<1x3x26x26x16xf32>
return %0 : tensor<1x3x26x26x16xf32>
}
// CHECK-NOT: stablehlo.transpose
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/optimize_graph.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: stablehlo-quant-opt %s -split-input-file -tf-optimize-graph | FileCheck %s
// CHECK-LABEL: @merge_requantization_followed_by_dequantization
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x4x3xf32>
func.func @merge_requantization_followed_by_dequantization(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> {
// CHECK: %[[CST:.*]] = stablehlo.constant dense<4.000000e-01> : tensor<2x3x3x2xf32>
// CHECK: %[[QUANT_CST:.*]] = stablehlo.uniform_quantize %[[CST]]
// CHECK: %[[QUANT_ARG_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]]
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%[[QUANT_ARG_0]], %[[QUANT_CST]])
// CHECK-NOT: stablehlo.uniform_quantize
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]]
// CHECK: return %[[DEQUANT]]
%cst = stablehlo.constant dense<0.4> : tensor<2x3x3x2xf32>
%quant_cst = stablehlo.uniform_quantize %cst : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.015>>
%quant_arg = stablehlo.uniform_quantize %arg0 : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0039207626791561354:-128>>
%conv = stablehlo.convolution(%quant_arg, %quant_cst) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0039207626791561354:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.015>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32, 5.8949912267181218E-5>>
%requant = stablehlo.uniform_quantize %conv : (tensor<1x3x4x2x!quant.uniform<i32:f32, 5.8949912267181218E-5>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, 0.045673100153605144:-62>>
%dequant = stablehlo.uniform_dequantize %requant : (tensor<1x3x4x2x!quant.uniform<i8:f32, 0.045673100153605144:-62>>) -> tensor<1x3x4x2xf32>
func.return %dequant : tensor<1x3x4x2xf32>
}
// -----
// CHECK-LABEL: @dont_merge_quantization_followed_by_quantization
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x4x3xf32>
func.func @dont_merge_quantization_followed_by_quantization(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32> {
// CHECK: %[[QUANT_ARG_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]]
// CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[QUANT_ARG_0]]
// CHECK: return %[[DEQUANT]]
%quant_arg = stablehlo.uniform_quantize %arg0 : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0039207626791561354:-128>>
%dequant = stablehlo.uniform_dequantize %quant_arg : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0039207626791561354:-128>>) -> tensor<1x3x4x3xf32>
func.return %dequant : tensor<1x3x4x3xf32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/post_quantize.mlir
+86
View File
@@ -0,0 +1,86 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-post-quantize | FileCheck %s
// CHECK-LABEL: @remove_volatile_qdq
func.func @remove_volatile_qdq() -> tensor<3x2xf32> {
// CHECK: %[[CST:.*]] = stablehlo.constant
// CHECK-NOT: "quantization.qcast"
// CHECK-NOT: "quantization.dcast"
// CHECK: return %[[CST]]
%cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32>
%q = "quantization.qcast"(%cst) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>
%dq = "quantization.dcast"(%q) : (tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>) -> tensor<3x2xf32>
func.return %dq : tensor<3x2xf32>
}
// -----
// CHECK-LABEL: @remove_volatile_qdq_with_requantization
// CHECK-SAME: %[[ARG0:.*]]: tensor<3x2xf32>
func.func @remove_volatile_qdq_with_requantization(%arg0: tensor<3x2xf32>) -> tensor<3x2xf32> {
// CHECK: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]]
// CHECK: %[[Q2:.*]] = stablehlo.uniform_quantize %[[Q1]]
// CHECK: %[[ABS:.*]] = stablehlo.abs %[[Q2]]
// CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[ABS]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[ARG0]], %[[DQ]]
// CHECK: return %[[ADD]]
%q1 = "quantization.qcast"(%arg0) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>
%q2 = "quantization.qcast"(%q1) {volatile} : (tensor<3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>
%dq1 = "quantization.dcast"(%q2) : (tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>) -> tensor<3x2xf32>
%abs = stablehlo.abs %q2 : (tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>) -> tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>
%dq2 = "quantization.dcast"(%abs) : (tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>) -> tensor<3x2xf32>
%add = stablehlo.add %dq1, %dq2 : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32>
func.return %add : tensor<3x2xf32>
}
// -----
// CHECK-LABEL: @quantize_constant
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3xf32>
func.func @quantize_constant(%arg0: tensor<1x3xf32>) -> tensor<1x2xf32> {
// CHECK-DAG: %[[QCST:.*]] = stablehlo.constant() <{value = dense<-78> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>
// CHECK-DAG: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]]
// CHECK-NOT: "quantization.qcast"
// CHECK: %[[DOT:.*]] = stablehlo.dot %[[Q1]], %[[QCST]]
// CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]]
// CHECK: return %[[DQ]]
%cst = stablehlo.constant dense<-0.390246302> : tensor<3x2xf32>
%q1 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>
%q2 = "quantization.qcast"(%cst) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>
%dot = stablehlo.dot %q1, %q2 : (tensor<1x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>) -> tensor<1x2x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%dq = "quantization.dcast"(%dot) : (tensor<1x2x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x2xf32>
func.return %dq : tensor<1x2xf32>
}
// -----
// CHECK-LABEL: @convert_quantization_qdq_to_stablehlo_uniform_qdq
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<3x2xf32>
func.func @convert_quantization_qdq_to_stablehlo_uniform_qdq(%arg0: tensor<1x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<1x2xf32> {
// CHECK: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]]
// CHECK-NOT: "quantization.qcast"
// CHECK: %[[Q2:.*]] = stablehlo.uniform_quantize %[[ARG1]]
// CHECK-NOT: "quantization.qcast"
// CHECK: %[[DOT:.*]] = stablehlo.dot %[[Q1]], %[[Q2]]
// CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]]
// CHECK: return %[[DQ]]
%q1 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>
%q2 = "quantization.qcast"(%arg1) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>
%dot = stablehlo.dot %q1, %q2 : (tensor<1x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>) -> tensor<1x2x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%dq = "quantization.dcast"(%dot) : (tensor<1x2x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x2xf32>
func.return %dq : tensor<1x2xf32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/prepare_quantize/prepare_quantize.mlir
+154
View File
@@ -0,0 +1,154 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-prepare-quantize=enable-per-channel-quantized-weight=false -verify-diagnostics | FileCheck %s
// -----
// CHECK-LABEL: func @dot
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
// CHECK: %[[cst:.*]] = stablehlo.constant
// CHECK: %[[q1:.*]] = "quantization.qcast"(%[[cst]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
// CHECK: %[[dq1:.*]] = "quantization.dcast"(%[[q1]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
%cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32>
// CHECK: %[[q2:.*]] = "quantization.qcast"(%[[ARG_0]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0078408040252386357:-1>
// CHECK: %[[dq2:.*]] = "quantization.dcast"(%[[q2]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0078408040252386357:-1>
%0 = "quantization.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<?x3xf32>) -> tensor<?x3xf32>
// CHECK: %[[dot:.*]] = stablehlo.dot %[[dq2]], %[[dq1]]
%1 = stablehlo.dot %0, %cst : (tensor<?x3xf32>, tensor<3x2xf32>) -> tensor<?x2xf32>
// CHECK: %[[q3:.*]] = "quantization.qcast"(%[[dot]])
// CHECK-SAME: quant.uniform<i8:f32, 0.036254937041039562:-28>>
// CHECK: %[[dq3:.*]] = "quantization.dcast"(%[[q3]])
// CHECK-SAME: quant.uniform<i8:f32, 0.036254937041039562:-28>>
%2 = "quantization.stats"(%1) {bitsNum = 8 : i64, layerStats = dense<[-3.6289506, 5.61605835]> : tensor<2xf32>, narrowRange = false} : (tensor<?x2xf32>) -> tensor<?x2xf32>
// CHECK: return %[[dq3]]
func.return %2 : tensor<?x2xf32>
}
// -----
// CHECK-LABEL: func @duplicate_stats
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<2x3xf32>) -> tensor<2x3xf32>
func.func @duplicate_stats(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
// CHECK: %[[q1:.*]] = "quantization.qcast"(%[[ARG_0]])
// CHECK: %[[dq1:.*]] = "quantization.dcast"(%[[q1]])
// CHECK: %[[q2:.*]] = "quantization.qcast"(%[[dq1]])
// CHECK: %[[dq2:.*]] = "quantization.dcast"(%[[q2]])
// CHECK: stablehlo.convert %[[dq2]]
%0 = "quantization.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<2x3xf32>) -> tensor<2x3xf32>
%1 = "quantization.stats"(%0) {bitsNum = 8 : i64, layerStats = dense<[-2.0, 2.0]> : tensor<2xf32>, narrowRange = false} : (tensor<2x3xf32>) -> tensor<2x3xf32>
%2 = stablehlo.convert %1 : (tensor<2x3xf32>) -> (tensor<2x3xf32>)
func.return %2 : tensor<2x3xf32>
}
// -----
// CHECK-LABEL: func @dot_redundant_stats
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot_redundant_stats(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
// CHECK: %[[cst:.*]] = stablehlo.constant
// CHECK: %[[q1:.*]] = "quantization.qcast"(%[[cst]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
// CHECK: %[[dq1:.*]] = "quantization.dcast"(%[[q1]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
%cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32>
// CHECK: %[[q2:.*]] = "quantization.qcast"(%[[ARG_0]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0078408040252386357:-1>
// CHECK: %[[dq2:.*]] = "quantization.dcast"(%[[q2]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0078408040252386357:-1>
%0 = "quantization.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-100.2, 212.4]> : tensor<2xf32>, narrowRange = false} : (tensor<?x3xf32>) -> tensor<?x3xf32>
%1 = "quantization.qcast"(%0) {volatile} : (tensor<?x3xf32>) -> tensor<?x3x!quant.uniform<i8:f32, 0.0078408040252386357:-1>>
%2 = "quantization.dcast"(%1) : (tensor<?x3x!quant.uniform<i8:f32, 0.0078408040252386357:-1>>) -> tensor<?x3xf32>
// CHECK: %[[dot:.*]] = stablehlo.dot %[[dq2]], %[[dq1]]
%3 = stablehlo.dot %2, %cst : (tensor<?x3xf32>, tensor<3x2xf32>) -> tensor<?x2xf32>
// CHECK: %[[q3:.*]] = "quantization.qcast"(%[[dot]])
// CHECK-SAME: quant.uniform<i8:f32, 0.036254937041039562:-28>>
// CHECK: %[[dq3:.*]] = "quantization.dcast"(%[[q3]])
// CHECK-SAME: quant.uniform<i8:f32, 0.036254937041039562:-28>>
%4 = "quantization.stats"(%3) {bitsNum = 8 : i64, layerStats = dense<[-3.6289506, 5.61605835]> : tensor<2xf32>, narrowRange = false} : (tensor<?x2xf32>) -> tensor<?x2xf32>
// CHECK: return %[[dq3]]
func.return %4 : tensor<?x2xf32>
}
// -----
// CHECK-LABEL: func @reshape_same_scale_propagate
func.func @reshape_same_scale_propagate(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {
// CHECK: %[[dq:.*]] = "quantization.dcast"
// CHECK-SAME: (tensor<2x3x!quant.uniform<i8:f32, 0.0078408040252386357:-1>>)
%0 = "quantization.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<2x3xf32>) -> tensor<2x3xf32>
// CHECK: %[[reshape:.*]] = stablehlo.reshape %[[dq]]
%1 = stablehlo.reshape %0 : (tensor<2x3xf32>) -> (tensor<6xf32>)
// CHECK: %[[q:.*]] = "quantization.qcast"(%[[reshape]])
// CHECK-SAME: -> tensor<6x!quant.uniform<i8:f32, 0.0078408040252386357:-1>>
%2 = "quantization.stats"(%1) {bitsNum = 8 : i64, layerStats = dense<[-2.0, 2.0]> : tensor<2xf32>, narrowRange = false} : (tensor<6xf32>) -> tensor<6xf32>
func.return %2 : tensor<6xf32>
}
// -----
// CHECK-LABEL: func @merge_consecutive_qcast
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?xf32>, %[[ARG_1:.*]]: tensor<?xf32>, %[[ARG_2:.*]]: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>)
func.func @merge_consecutive_qcast(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>) {
// CHECK: "quantization.qcast"(%[[ARG_1]])
// CHECK-SAME: -> tensor<?x!quant.uniform<i8:f32, 0.02454993117089365:-64>>
// CHECK: "quantization.qcast"(%[[ARG_1]])
// CHECK-SAME: -> tensor<?x!quant.uniform<i8:f32, 0.013075299590241675:-64>>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[-0.83811146, 2.4960899]> : tensor<2xf32>} : (tensor<?xf32>) -> tensor<?xf32>
%1 = "quantization.stats"(%arg1) {layerStats = dense<[-0.835039615, 1.000000e+00]> : tensor<2xf32>} : (tensor<?xf32>) -> tensor<?xf32>
%2 = "stablehlo.concatenate"(%0, %1) {dimension = 0 : i64} : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
%3 = "quantization.stats"(%2) {layerStats = dense<[-0.83811146, 2.4960899]> : tensor<2xf32>} : (tensor<?xf32>) -> tensor<?xf32>
%4 = "quantization.stats"(%arg2) {layerStats = dense<[-1.5726943, 1.07351148]> : tensor<2xf32>} : (tensor<?xf32>) -> tensor<?xf32>
%5 = "stablehlo.concatenate"(%4, %1) {dimension = 0 : i64} : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
%6 = "quantization.stats"(%5) {layerStats = dense<[-1.5726943, 4.6875381]> : tensor<2xf32>} : (tensor<?xf32>) -> tensor<?xf32>
func.return %3, %6 : tensor<?xf32>, tensor<?xf32>
}
// -----
// CHECK-LABEL: func @skip_nan_inf_constant
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x112x112x64xf32>) -> tensor<?x56x56x64xf32>
func.func @skip_nan_inf_constant(%arg0: tensor<?x112x112x64xf32>) -> tensor<?x56x56x64xf32> {
// CHECK-DAG: %[[cst0:.*]] = stablehlo.constant dense<0xFF800000> : tensor<f32
// CHECK-DAG: %[[cst1:.*]] = stablehlo.constant dense<0x7FC00000> : tensor<f32>
// CHECK-DAG: %[[cst2:.*]] = stablehlo.constant dense<6.000000e+00> : tensor<f32>
// CHECK-DAG: %[[cst3:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
// CHECK-NOT: %[[q0:.*]] = "quantization.qcast"(%[[cst0]])
// CHECK-NOT: %[[q1:.*]] = "quantization.qcast"(%[[cst1]])
// CHECK: %[[q2:.*]] = "quantization.qcast"(%[[cst2]])
// CHECK-SAME: quant.uniform<i8:f32, 0.023529411764705882:-128>
// CHECK: %[[dq2:.*]] = "quantization.dcast"(%[[q2]])
// CHECK-SAME: quant.uniform<i8:f32, 0.023529411764705882:-128>
// CHECK: %[[q3:.*]] = "quantization.qcast"(%[[cst3]])
// CHECK-SAME: quant.uniform<i8:f32, 3.9215686274509805E-9>
// CHECK: %[[dq3:.*]] = "quantization.dcast"(%[[q3]])
// CHECK-SAME: quant.uniform<i8:f32, 3.9215686274509805E-9>
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%1 = stablehlo.constant dense<0x7FC00000> : tensor<f32>
%2 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%3 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%4 = "stablehlo.add"(%0, %1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
%5 = stablehlo.clamp %3, %arg0, %2 : (tensor<f32>, tensor<?x112x112x64xf32>, tensor<f32>) -> tensor<?x112x112x64xf32>
%6 = "stablehlo.reduce_window"(%5, %4) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%7 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %7 : tensor<f32>
}) {padding = dense<[[0, 0], [0, 1], [0, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>, window_strides = array<i64: 1, 2, 2, 1>} : (tensor<?x112x112x64xf32>, tensor<f32>) -> tensor<?x56x56x64xf32>
return %6 : tensor<?x56x56x64xf32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/prepare_quantize/prepare_quantize_int4.mlir
+40
View File
@@ -0,0 +1,40 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-prepare-quantize=bit-width=4 -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func @dot_int4
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot_int4(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
// CHECK: %[[cst:.*]] = stablehlo.constant
// CHECK: %[[q1:.*]] = "quantization.qcast"(%[[cst]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
// CHECK: %[[dq1:.*]] = "quantization.dcast"(%[[q1]])
// CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
%cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32>
// CHECK: %[[q2:.*]] = "quantization.qcast"(%[[ARG_0]])
// CHECK-SAME: quant.uniform<i4:f32, 0.13329366842905679:-1>
// CHECK: %[[dq2:.*]] = "quantization.dcast"(%[[q2]])
// CHECK-SAME: quant.uniform<i4:f32, 0.13329366842905679:-1>
%0 = "quantization.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<?x3xf32>) -> tensor<?x3xf32>
// CHECK: %[[dot:.*]] = stablehlo.dot %[[dq2]], %[[dq1]]
%1 = stablehlo.dot %0, %cst : (tensor<?x3xf32>, tensor<3x2xf32>) -> tensor<?x2xf32>
// CHECK: %[[q3:.*]] = "quantization.qcast"(%[[dot]])
// CHECK-SAME: quant.uniform<i4:f32, 0.61633392969767253:-2>>
// CHECK: %[[dq3:.*]] = "quantization.dcast"(%[[q3]])
// CHECK-SAME: quant.uniform<i4:f32, 0.61633392969767253:-2>>
%2 = "quantization.stats"(%1) {bitsNum = 8 : i64, layerStats = dense<[-3.6289506, 5.61605835]> : tensor<2xf32>, narrowRange = false} : (tensor<?x2xf32>) -> tensor<?x2xf32>
// CHECK: return %[[dq3]]
func.return %2 : tensor<?x2xf32>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/prepare_quantize/prepare_quantize_per_channel.mlir
+144
View File
@@ -0,0 +1,144 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-prepare-quantize=enable-per-channel-quantized-weight=true -verify-diagnostics | FileCheck %s
// -----
module {
// CHECK-LABEL: conv_with_bias_and_relu
func.func private @conv_with_bias_and_relu(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> {
%cst = "tf.Const"() {device = "", value = dense<[7.11401462, 7.05456924]> : tensor<2xf32>} : () -> tensor<2xf32>
// CHECK: %[[q_weight_per_channel:.*]] = "quantization.qcast"
// CHECK-SAME: -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.075123051020104109,0.072960192762960605}>>
// CHECK: %[[dq_weight:.*]] = "quantization.dcast"(%[[q_weight_per_channel]])
%cst_0 = "tf.Const"() {device = "", value = dense<[[[[-6.30731344, 5.4962182], [1.80364347, -7.64542675], [-2.11145878, -7.08605719]], [[-9.54062747, -6.14013147], [6.12640238, -4.18223286], [5.05738974, 8.99269962]], [[3.3535192, 0.84816426], [-6.64676809, -7.95477629], [5.81315517, 9.21566581]]], [[[1.38622558, 4.63866329], [4.54742622, -1.43770897], [-3.96835279, 2.99996852]], [[0.989735424, -4.83384752], [-7.27702999, 1.17216611], [1.33735656, 0.728900194]], [[5.1286211, 8.98645591], [1.55008793, -3.85491467], [3.7003777, 9.26594448]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
// CHECK: %[[q_act:.*]] = "quantization.qcast"(%arg0)
// CHECK-SAME: -> tensor<1x3x2x3x!quant.uniform<i8:f32, 0.018920717052384919:-128>>
// CHECK: %[[dq_act:.*]] = "quantization.dcast"(%[[q_act]])
%0 = "quantization.stats"(%arg0) {layerStats = dense<[1.27501142, 4.824783]> : tensor<2xf32>} : (tensor<1x3x2x3xf32>) -> tensor<1x3x2x3xf32>
// CHECK: "tf.XlaCallModule"(%[[dq_act]], %[[dq_weight]]
%1 = "tf.XlaCallModule"(%0, %cst_0, %cst) {
Sout = [#tf_type.shape<1x2x2x2>], config = "",
module = "composite_conv2d_with_bias_and_relu6_fn_10",
_entry_function = @composite_conv2d_with_bias_and_relu6_fn_10,
// Represents a per-channel quantization for the operand index 1 with
// quantization dimension of 3
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
platforms = [], version = 4 : i64
} : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[0.000000e+00, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
return %2 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: composite_conv2d_with_bias_and_relu6_fn_10
func.func private @composite_conv2d_with_bias_and_relu6_fn_10(%arg0: tensor<1x3x2x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x2x2x2xf32> attributes {tf.tf_quant.composite_function} {
%0 = "quantization.stats"(%arg1) {layerStats = dense<[-3.54062747, 0.54742622]> : tensor<2xf32>} : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2xf32>
%1 = "quantization.stats"(%arg0) {layerStats = dense<[1.27501142, 2.824783]> : tensor<2xf32>} : (tensor<1x3x2x3xf32>) -> tensor<1x3x2x3xf32>
%2 = stablehlo.convolution(%1, %0)
dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
window = {
stride = [1, 1], pad = [[0, 0], [1, 1]],
lhs_dilate = [1, 1],
rhs_dilate = [1, 1]
}
{
batch_group_count = 1 : i64,
feature_group_count = 1 : i64
} : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>)
-> tensor<1x2x2x2xf32>
%3 = "quantization.stats"(%arg2) {layerStats = dense<[7.05456924, 7.11401462]> : tensor<2xf32>} : (tensor<2xf32>) -> tensor<2xf32>
%4 = "quantization.stats"(%2) {layerStats = dense<[-1.36523, 3.57373]> : tensor<2xf32>} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
%5 = "chlo.broadcast_add"(%4, %3) : (tensor<1x2x2x2xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
%6 = "quantization.stats"(%5) {layerStats = dense<[-1.31055, 2.62842]> : tensor<2xf32>} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
%cst_min = stablehlo.constant dense<0.0> : tensor<f32>
%cst_max = stablehlo.constant dense<6.0> : tensor<f32>
%7 = "stablehlo.clamp"(%cst_min, %6, %cst_max) {device = ""} : (tensor<f32>, tensor<1x2x2x2xf32>, tensor<f32>) -> tensor<1x2x2x2xf32>
%8 = "quantization.stats"(%7) {layerStats = dense<[0.000000e+00, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
return %8 : tensor<1x2x2x2xf32>
}
}
// -----
module {
// CHECK-LABEL: dot_general
func.func private @dot_general(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
// CHECK: %[[q_weight:.*]] = "quantization.qcast"
// CHECK-SAME: -> tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {0.049663885371891529,0.060200210631363035}>>
// CHECK: %[[dq_weight:.*]] = "quantization.dcast"(%[[q_weight]])
%cst = "tf.Const"() {device = "", value = dense<[[-6.30731344, 5.4962182], [1.80364347, -7.64542675]]> : tensor<2x2xf32>} : () -> tensor<2x2xf32>
// CHECK: %[[q_act:.*]] = "quantization.qcast"(%arg0)
// CHECK-SAME: -> tensor<2x2x!quant.uniform<i8:f32, 0.018920717052384919:-128>>
// CHECK: %[[dq_act:.*]] = "quantization.dcast"(%[[q_act]])
%0 = "quantization.stats"(%arg0) {layerStats = dense<[1.27501142, 4.824783]> : tensor<2xf32>} : (tensor<2x2xf32>) -> tensor<2x2xf32>
// CHECK: "tf.XlaCallModule"(%[[dq_act]], %[[dq_weight]]
%1 = "tf.XlaCallModule"(%0, %cst) {
Sout = [#tf_type.shape<2x2>], config = "",
_entry_function = @composite_dot_general,
module = "composite_dot_general",
platforms = [], version = 4 : i64
} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[0.000000e+00, 6.000000e+00]> : tensor<2xf32>} : (tensor<2x2xf32>) -> tensor<2x2xf32>
return %2 : tensor<2x2xf32>
}
// CHECK-LABEL: composite_dot_general
func.func private @composite_dot_general(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
%0 = "stablehlo.dot_general"(%arg0, %arg1) {
dot_dimension_numbers = #stablehlo.dot<
lhs_batching_dimensions = [],
rhs_batching_dimensions = [],
lhs_contracting_dimensions = [1],
rhs_contracting_dimensions = [0]
>
} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
}
// -----
// Tests that the `PrepareQuantizePass` prepares for per-tensor quantization for
// the weight of convolution. This is based on the `_quantization_method` that
// does not have a `input_quantized_types` with a specified `dimension_specs`.
// CHECK-LABEL: conv_per_tensor_quantized_method
func.func private @conv_per_tensor_quantized_method(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> {
%cst = "tf.Const"() {device = "", value = dense<[7.11401462, 7.05456924]> : tensor<2xf32>} : () -> tensor<2xf32>
%cst_0 = "tf.Const"() {device = "", value = dense<[[[[-6.30731344, 5.4962182], [1.80364347, -7.64542675], [-2.11145878, -7.08605719]], [[-9.54062747, -6.14013147], [6.12640238, -4.18223286], [5.05738974, 8.99269962]], [[3.3535192, 0.84816426], [-6.64676809, -7.95477629], [5.81315517, 9.21566581]]], [[[1.38622558, 4.63866329], [4.54742622, -1.43770897], [-3.96835279, 2.99996852]], [[0.989735424, -4.83384752], [-7.27702999, 1.17216611], [1.33735656, 0.728900194]], [[5.1286211, 8.98645591], [1.55008793, -3.85491467], [3.7003777, 9.26594448]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[1.27501142, 4.824783]> : tensor<2xf32>} : (tensor<1x3x2x3xf32>) -> tensor<1x3x2x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst_0, %cst) {
Sout = [#tf_type.shape<1x2x2x2>], config = "",
module = "composite_conv_fn_1",
_entry_function = @composite_conv_fn_1,
_quantization_method = "static_range_ptq {}",
platforms = [], version = 4 : i64
} : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[0.000000e+00, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
return %2 : tensor<1x2x2x2xf32>
}
// CHECK-SAME: %[[ARG_0:.+]]: tensor<1x3x2x3xf32>
// Test that the weight is prepared for per-tensor quantization, based on the
// `_quantization_method` attribute without a `dimension_specs` field in
// `QuantizedType`.
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} tensor<2x3x3x2xf32>
// CHECK: %[[Q_WEIGHT_PER_TENSOR:.*]] = "quantization.qcast"(%[[WEIGHT_CONST]]) {{.*}} (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[DQ_WEIGHT:.*]] = "quantization.dcast"(%[[Q_WEIGHT_PER_TENSOR]])
// CHECK: %[[Q_ACTIVATION:.*]] = "quantization.qcast"(%[[ARG_0]])
// CHECK-SAME: -> tensor<1x3x2x3x!quant.uniform<i8:f32, 0.018920717052384919:-128>>
// CHECK: %[[DQ_ACTIVATION:.*]] = "quantization.dcast"(%[[Q_ACTIVATION]])
// CHECK: "tf.XlaCallModule"(%[[DQ_ACTIVATION]], %[[DQ_WEIGHT]]
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize.mlir
+88
View File
@@ -0,0 +1,88 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-quantize -verify-each=false | FileCheck %s
// Tests for PopulateFusedGemmStylePatterns are handled in
// quantize_composite_functions for module-level evaluation of functions.
module attributes {tf_saved_model.semantics} {
// CHECK: quantize_simple_xla_call_module(%[[ARG_0:.+]]: tensor<1x4xf32>)
func.func private @quantize_simple_xla_call_module(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<4x3xf32>
%1 = "quantization.qcast"(%0) {volatile} : (tensor<4x3xf32>) -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03, 5.000000e-03, 5.000000e-03}>>
%2 = "quantization.dcast"(%1) : (tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03, 5.000000e-03, 5.000000e-03}>>) -> tensor<4x3xf32>
%3 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
%4 = "quantization.dcast"(%3) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x4xf32>
%5 = "tf.XlaCallModule"(%4, %2) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
%6 = "quantization.qcast"(%5) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%7 = "quantization.dcast"(%6) : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %7 : tensor<1x3xf32>
}
// Test that the inputs and output of the tf.XlaCallModule op has been replaced
// by quantized types, and the corresponding quantization.dcast ops that turned
// those quantized types back to float types are removed.
// CHECK: %[[CONST_0:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<4x3xf32>
// CHECK-DAG: %[[QCAST_0:.+]] = "quantization.qcast"(%[[CONST_0]]) {volatile} : (tensor<4x3xf32>) -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
// CHECK-DAG: %[[QCAST_1:.+]] = "quantization.qcast"(%[[ARG_0]]) {volatile} : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: %[[CALL_0:.+]] = call @quantized_dot_general_fn(%[[QCAST_1]], %[[QCAST_0]])
// Test that the `Method` has been copied over.
// CHECK-SAME: {_quantization_method = "static_range_ptq { }"}
// CHECK-SAME: : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
// CHECK: %[[DCAST_0:.+]] = "quantization.dcast"(%[[CALL_0]]) : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
// CHECK: return
func.func private @composite_dot_general_fn(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
// Tests that the output of the tf.XlaCallModule op has been replaced by
// a quantized type, and the corresponding quantization.qcast ops that turned
// the float output to a quantized type is removed.
// CHECK-LABEL: quantize_simple_xla_call_module_no_operand
func.func private @quantize_simple_xla_call_module_no_operand() -> tensor<1x3xf32> {
%0 = "tf.XlaCallModule"() {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : () -> tensor<1x3xf32>
%1 = "quantization.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%2 = "quantization.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
// CHECK: %[[XLA_CALL_MODULE_0:.+]] = "tf.XlaCallModule"() <{{{.*}}}> {{{.*}}} : () -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
// CHECK: %[[DCAST_0:.+]] = "quantization.dcast"(%[[XLA_CALL_MODULE_0]]) : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
// CHECK: "func.return"(%[[DCAST_0]]) : (tensor<1x3xf32>) -> ()
// -----
// Tests for emitting an error when there is no corresponding entry
// function to quantize (@composite_dot_general_fn).
module attributes {tf_saved_model.semantics} {
func.func private @error_when_no_entry_function(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%0 = stablehlo.constant dense<1.000000e+00> : tensor<2x3xf32>
%1 = "quantization.qcast"(%0) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>
%2 = "quantization.dcast"(%1) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>) -> tensor<2x3xf32>
%3 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>
%4 = "quantization.dcast"(%3) : (tensor<1x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x2xf32>
// expected-error @+2 {{Failed to find a valid entry function}}
// expected-error @+1 {{'tf.XlaCallModule' op operand #0 must be variadic of tensor of tf.dtype values}}
%5 = "tf.XlaCallModule"(%4, %2) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%6 = "quantization.qcast"(%5) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%7 = "quantization.dcast"(%6) : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
return %7 : tensor<1x3xf32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_op_with_region.mlir
+255
View File
@@ -0,0 +1,255 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-quantize -verify-each=false | FileCheck %s
// Tests if reduce_window op following quantized function is quantized.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1722 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: main_00
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x1x1024xf32>
func.func private @main_00(%arg0: tensor<2x3x1x1024xf32>) -> tensor<2x3x1x3xf32> attributes {tf._original_func_name = "main_0"} {
// CHECK: %[[CST0:.*]] = stablehlo.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[CST1:.*]] = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
// CHECK: %[[Q0:.*]] = "quantization.qcast"(%[[CST0]])
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[CST1]])
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG0]])
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q2]], %[[Q1]])
// CHECK: %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[CALL]], %[[Q0]])
// CHECK{LITERAL}: padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>
// CHECK-SAME: window_dimensions = array<i64: 1, 3, 3, 1>
// CHECK: %[[ARG1:.*]]: tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>, %[[ARG2:.*]]: tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ARG1]], %[[ARG2]] : tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: stablehlo.return %[[MAX]] : tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>, tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[REDUCE]])
// CHECK: return %[[DQ]]
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%1 = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
%2 = "quantization.qcast"(%0) {volatile} : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
%3 = "quantization.dcast"(%2) : (tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<f32>
%4 = "quantization.qcast"(%1) {volatile} : (tensor<2x3x1024x3xf32>) -> tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>
%5 = "quantization.dcast"(%4) : (tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>) -> tensor<2x3x1024x3xf32>
%6 = "quantization.qcast"(%arg0) {volatile} : (tensor<2x3x1x1024xf32>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%7 = "quantization.dcast"(%6) : (tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1x1024xf32>
%8 = "tf.XlaCallModule"(%7, %5) <{Sout = [#tf_type.shape<2x3x1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
%9 = "quantization.qcast"(%8) {volatile} : (tensor<2x3x1x3xf32>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%10 = "quantization.dcast"(%9) : (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3xf32>
%11 = "stablehlo.reduce_window"(%10, %3) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%14 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %14 : tensor<f32>
}) {padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>} : (tensor<2x3x1x3xf32>, tensor<f32>) -> tensor<2x3x1x3xf32>
%12 = "quantization.qcast"(%11) {volatile} : (tensor<2x3x1x3xf32>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%13 = "quantization.dcast"(%12) : (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3xf32>
return %13 : tensor<2x3x1x3xf32>
}
// CHECK: quantized_dot_general_fn_1
func.func private @composite_dot_general_fn_1(%arg0: tensor<2x3x1x1024xf32>, %arg1: tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general
// CHECK: %[[RQ:.*]] = stablehlo.uniform_quantize %[[DOT]]
// CHECK: return %[[RQ]]
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
return %0 : tensor<2x3x1x3xf32>
}
}
// -----
// Tests if reduce_window op preceding quantized function is quantized.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1722 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: main_00
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x1x1024xf32>
func.func private @main_00(%arg0: tensor<2x3x1x1024xf32>) -> tensor<2x3x1x3xf32> attributes {tf._original_func_name = "main_0"} {
// CHECK: %[[CST0:.*]] = stablehlo.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[CST1:.*]] = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
// CHECK: %[[Q0:.*]] = "quantization.qcast"(%[[CST0]])
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]])
// CHECK: %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[Q1]], %[[Q0]])
// CHECK{LITERAL}: padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>
// CHECK-SAME: window_dimensions = array<i64: 1, 3, 3, 1>
// CHECK: %[[ARG1:.*]]: tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>, %[[ARG2:.*]]: tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ARG1]], %[[ARG2]] : tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: stablehlo.return %[[MAX]] : tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: (tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>, tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[CST1]])
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[REDUCE]], %[[Q2]])
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[CALL]])
// CHECK: return %[[DQ]]
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%1 = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
%2 = "quantization.qcast"(%0) {volatile} : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
%3 = "quantization.dcast"(%2) : (tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<f32>
%4 = "quantization.qcast"(%arg0) {volatile} : (tensor<2x3x1x1024xf32>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%5 = "quantization.dcast"(%4) : (tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1x1024xf32>
%6 = "stablehlo.reduce_window"(%5, %3) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%14 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %14 : tensor<f32>
}) {padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array<i64: 1, 3, 3, 1>} : (tensor<2x3x1x1024xf32>, tensor<f32>) -> tensor<2x3x1x1024xf32>
%7 = "quantization.qcast"(%6) {volatile} : (tensor<2x3x1x1024xf32>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%8 = "quantization.dcast"(%7) : (tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1x1024xf32>
%9 = "quantization.qcast"(%1) {volatile} : (tensor<2x3x1024x3xf32>) -> tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>
%10 = "quantization.dcast"(%9) : (tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>) -> tensor<2x3x1024x3xf32>
%11 = "tf.XlaCallModule"(%8, %10) <{Sout = [#tf_type.shape<2x3x1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
%12 = "quantization.qcast"(%11) {volatile} : (tensor<2x3x1x3xf32>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%13 = "quantization.dcast"(%12) : (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3xf32>
return %13 : tensor<2x3x1x3xf32>
}
// CHECK: quantized_dot_general_fn_1
func.func private @composite_dot_general_fn_1(%arg0: tensor<2x3x1x1024xf32>, %arg1: tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general
// CHECK: %[[RQ:.*]] = stablehlo.uniform_quantize %[[DOT]]
// CHECK: return %[[RQ]]
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
return %0 : tensor<2x3x1x3xf32>
}
}
// -----
// Tests if reduce_window op following quantized same-scale op is quantized.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1722 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: main_00
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x1x1024xf32>
func.func private @main_00(%arg0: tensor<2x3x1x1024xf32>) -> tensor<2x3x3xf32> attributes {tf._original_func_name = "main_0"} {
// CHECK: %[[CST0:.*]] = stablehlo.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[CST1:.*]] = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
// CHECK: %[[Q0:.*]] = "quantization.qcast"(%[[CST0]])
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[CST1]])
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG0]])
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q2]], %[[Q1]])
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[CALL]]
// CHECK: %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[RESHAPE]], %[[Q0]])
// CHECK{LITERAL}: padding = dense<[[0, 0], [1, 1], [0, 0]]> : tensor<3x2xi64>
// CHECK-SAME: window_dimensions = array<i64: 1, 3, 1>
// CHECK: %[[ARG1:.*]]: tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>, %[[ARG2:.*]]: tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ARG1]], %[[ARG2]] : tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: stablehlo.return %[[MAX]] : tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: (tensor<2x3x3x!quant.uniform<i8:f32, 3.000000e-01:1>>, tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[REDUCE]])
// CHECK: return %[[DQ]]
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%1 = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
%2 = "quantization.qcast"(%0) {volatile} : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>
%3 = "quantization.dcast"(%2) : (tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<f32>
%4 = "quantization.qcast"(%1) {volatile} : (tensor<2x3x1024x3xf32>) -> tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>
%5 = "quantization.dcast"(%4) : (tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>) -> tensor<2x3x1024x3xf32>
%6 = "quantization.qcast"(%arg0) {volatile} : (tensor<2x3x1x1024xf32>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%7 = "quantization.dcast"(%6) : (tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1x1024xf32>
%8 = "tf.XlaCallModule"(%7, %5) <{Sout = [#tf_type.shape<2x3x1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
%9 = "quantization.qcast"(%8) {volatile} : (tensor<2x3x1x3xf32>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%10 = "quantization.dcast"(%9) : (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3xf32>
%11 = stablehlo.reshape %10 : (tensor<2x3x1x3xf32>) -> tensor<2x3x3xf32>
%12 = "quantization.qcast"(%11) {volatile} : (tensor<2x3x3xf32>) -> tensor<2x3x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%13 = "quantization.dcast"(%12) : (tensor<2x3x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x3xf32>
%14 = "stablehlo.reduce_window"(%13, %3) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%17 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %17 : tensor<f32>
}) {padding = dense<[[0, 0], [1, 1], [0, 0]]> : tensor<3x2xi64>, window_dimensions = array<i64: 1, 3, 1>} : (tensor<2x3x3xf32>, tensor<f32>) -> tensor<2x3x3xf32>
%15 = "quantization.qcast"(%14) {volatile} : (tensor<2x3x3xf32>) -> tensor<2x3x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%16 = "quantization.dcast"(%15) : (tensor<2x3x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x3xf32>
return %16 : tensor<2x3x3xf32>
}
// CHECK: quantized_dot_general_fn_1
func.func private @composite_dot_general_fn_1(%arg0: tensor<2x3x1x1024xf32>, %arg1: tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general
// CHECK: %[[RQ:.*]] = stablehlo.uniform_quantize %[[DOT]]
// CHECK: return %[[RQ]]
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
return %0 : tensor<2x3x1x3xf32>
}
}
// -----
// Tests if reduce_window op preceding quantized same-scale op is quantized.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1722 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: main_00
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x1024xf32>
func.func private @main_00(%arg0: tensor<2x3x1024xf32>) -> tensor<2x3x1x3xf32> attributes {tf._original_func_name = "main_0"} {
// CHECK: %[[CST0:.*]] = stablehlo.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[CST1:.*]] = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
// CHECK: %[[Q0:.*]] = "quantization.qcast"(%[[CST0]])
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]])
// CHECK: %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[Q1]], %[[Q0]])
// CHECK{LITERAL}: padding = dense<[[0, 0], [1, 1], [0, 0]]> : tensor<3x2xi64>
// CHECK-SAME: window_dimensions = array<i64: 1, 3, 1>
// CHECK: %[[ARG1:.*]]: tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>, %[[ARG2:.*]]: tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ARG1]], %[[ARG2]] : tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: stablehlo.return %[[MAX]] : tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: (tensor<2x3x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>, tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[REDUCE]]
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[CST1]])
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[RESHAPE]], %[[Q2]])
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[CALL]])
// CHECK: return %[[DQ]]
%0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%1 = stablehlo.constant dense<0xFF80000E> : tensor<2x3x1024x3xf32>
%2 = "quantization.qcast"(%0) {volatile} : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>
%3 = "quantization.dcast"(%2) : (tensor<!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<f32>
%4 = "quantization.qcast"(%arg0) {volatile} : (tensor<2x3x1024xf32>) -> tensor<2x3x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%5 = "quantization.dcast"(%4) : (tensor<2x3x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1024xf32>
%6 = "stablehlo.reduce_window"(%5, %3) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%17 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %17 : tensor<f32>
}) {padding = dense<[[0, 0], [1, 1], [0, 0]]> : tensor<3x2xi64>, window_dimensions = array<i64: 1, 3, 1>} : (tensor<2x3x1024xf32>, tensor<f32>) -> tensor<2x3x1024xf32>
%7 = "quantization.qcast"(%6) {volatile} : (tensor<2x3x1024xf32>) -> tensor<2x3x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%8 = "quantization.dcast"(%7) : (tensor<2x3x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1024xf32>
%9 = stablehlo.reshape %8 : (tensor<2x3x1024xf32>) -> tensor<2x3x1x1024xf32>
%10 = "quantization.qcast"(%9) {volatile} : (tensor<2x3x1x1024xf32>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>
%11 = "quantization.dcast"(%10) : (tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>) -> tensor<2x3x1x1024xf32>
%12 = "quantization.qcast"(%1) {volatile} : (tensor<2x3x1024x3xf32>) -> tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>
%13 = "quantization.dcast"(%12) : (tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>) -> tensor<2x3x1024x3xf32>
%14 = "tf.XlaCallModule"(%11, %13) <{Sout = [#tf_type.shape<2x3x1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
%15 = "quantization.qcast"(%14) {volatile} : (tensor<2x3x1x3xf32>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
%16 = "quantization.dcast"(%15) : (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3xf32>
return %16 : tensor<2x3x1x3xf32>
}
// CHECK: quantized_dot_general_fn_1
func.func private @composite_dot_general_fn_1(%arg0: tensor<2x3x1x1024xf32>, %arg1: tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general
// CHECK: %[[RQ:.*]] = stablehlo.uniform_quantize %[[DOT]]
// CHECK: return %[[RQ]]
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<2x3x1x1024xf32>, tensor<2x3x1024x3xf32>) -> tensor<2x3x1x3xf32>
return %0 : tensor<2x3x1x3xf32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_same_scale.mlir
+387
View File
@@ -0,0 +1,387 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-quantize -verify-each=false | FileCheck %s
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: same_scale_after_composite
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
func.func private @same_scale_after_composite(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<3x1xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[CALL]] : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[RESHAPE]]) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<1x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
%3 = "quantization.dcast"(%2) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x3xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%7 = stablehlo.reshape %6 : (tensor<1x3xf32>) -> tensor<3x1xf32>
%8 = "quantization.qcast"(%7) {volatile} : (tensor<3x1xf32>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%9 = "quantization.dcast"(%8) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
return %9 : tensor<3x1xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: same_scale_indirectly_connected
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
func.func private @same_scale_indirectly_connected(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[CALL]] : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[TRANSPOSE:.*]] = stablehlo.transpose %[[RESHAPE]], dims = [1, 0] : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[TRANSPOSE]]) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<1x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
%3 = "quantization.dcast"(%2) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x3xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%7 = stablehlo.reshape %6 : (tensor<1x3xf32>) -> tensor<3x1xf32>
%8 = "quantization.qcast"(%7) {volatile} : (tensor<3x1xf32>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%9 = "quantization.dcast"(%8) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
%10 = stablehlo.transpose %9, dims = [1, 0] : (tensor<3x1xf32>) -> tensor<1x3xf32>
%11 = "quantization.qcast"(%10) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%12 = "quantization.dcast"(%11) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
return %12 : tensor<1x3xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
// CHECK-LABEL: same_scale_not_connected_to_composite
func.func @same_scale_not_connected_to_composite() -> tensor<3x1xf32> {
// CHECK: %[[CST:.*]] = stablehlo.constant
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[CST]]) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ1:.*]] = "quantization.dcast"(%[[Q1]]) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[DQ1]]
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[RESHAPE]]) {volatile} : (tensor<3x1xf32>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ2:.*]] = "quantization.dcast"(%[[Q2]]) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
// CHECK: return %[[DQ2]]
%0 = stablehlo.constant dense<1.000000e+00> : tensor<1x3xf32>
%1 = "quantization.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%2 = "quantization.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%3 = stablehlo.reshape %2 : (tensor<1x3xf32>) -> tensor<3x1xf32>
%4 = "quantization.qcast"(%3) {volatile} : (tensor<3x1xf32>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%5 = "quantization.dcast"(%4) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
return %5 : tensor<3x1xf32>
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: concatenate_and_composite
// CHECK-SAME: %[[ARG0:.*]]: tensor<3x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<1x2xf32>
// CHECK-SAME: %[[ARG2:.*]]: tensor<2x5xf32>
func.func private @concatenate_and_composite(%arg0: tensor<3x2xf32>, %arg1: tensor<1x2xf32>, %arg2: tensor<2x5xf32>) -> tensor<4x5xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
// CHECK: %[[CONCAT:.*]] = stablehlo.concatenate %[[Q1]], %[[Q2]], dim = 0
// CHECK-SAME: (tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>, tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>) -> tensor<4x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
// CHECK: %[[Q3:.*]] = "quantization.qcast"(%[[ARG2]]) {volatile} : (tensor<2x5xf32>) -> tensor<2x5x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[CONCAT]], %[[Q3]])
// CHECK-SAME: (tensor<4x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>, tensor<2x5x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<4x5x!quant.uniform<i8:f32, 1.000000e-03:-3>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[CALL]]) : (tensor<4x5x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<4x5xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
%1 = "quantization.dcast"(%0) : (tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>) -> tensor<3x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
%3 = "quantization.dcast"(%2) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>) -> tensor<1x2xf32>
%4 = "stablehlo.concatenate"(%1, %3) {
dimension = 0 : i64
} : (tensor<3x2xf32>, tensor<1x2xf32>) -> tensor<4x2xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<4x2xf32>) -> tensor<4x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
%6 = "quantization.dcast"(%5) : (tensor<4x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>) -> tensor<4x2xf32>
%7 = "quantization.qcast"(%arg2) {volatile} : (tensor<2x5xf32>) -> tensor<2x5x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>
%8 = "quantization.dcast"(%7) : (tensor<2x5x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x5xf32>
%9 = "tf.XlaCallModule"(%6, %8) {Sout = [#tf_type.shape<4x5>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<4x2xf32>, tensor<2x5xf32>) -> tensor<4x5xf32>
%10 = "quantization.qcast"(%9) {volatile} : (tensor<4x5xf32>) -> tensor<4x5x!quant.uniform<i8:f32, 1.000000e-03:-3>>
%11 = "quantization.dcast"(%10) : (tensor<4x5x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<4x5xf32>
return %11 : tensor<4x5xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG3:.*]]: tensor<4x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>
// CHECK-SAME: %[[ARG4:.*]]: tensor<2x5x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<4x2xf32>, %arg1: tensor<2x5xf32>) -> tensor<4x5xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG3]], %[[ARG4]]
// CHECK-SAME: (tensor<4x2x!quant.uniform<i8:f32, 5.000000e-03:-1>>, tensor<2x5x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<4x5x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<4x5x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<4x5x!quant.uniform<i8:f32, 1.000000e-03:-3>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<4x2xf32>, tensor<2x5xf32>) -> tensor<4x5xf32>
return %0 : tensor<4x5xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: composite_and_pad
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
// CHECK-SAME: %[[ARG2:.*]]: tensor<f32>
func.func private @composite_and_pad(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<f32>) -> tensor<3x9xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[Q3:.*]] = "quantization.qcast"(%arg2) {volatile} : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[PAD:.*]] = stablehlo.pad %[[CALL]], %[[Q3]]
// CHECK-SAME: (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>, tensor<!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x9x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[PAD]]) : (tensor<3x9x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x9xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<1x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
%3 = "quantization.dcast"(%2) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x3xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%7 = "quantization.qcast"(%arg2) {volatile} : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%8 = "quantization.dcast"(%7) : (tensor<!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<f32>
%9 = stablehlo.pad %6, %8, low = [0, 1], high = [2, 1], interior = [0, 2] : (tensor<1x3xf32>, tensor<f32>) -> tensor<3x9xf32>
%10 = "quantization.qcast"(%9) {volatile} : (tensor<3x9xf32>) -> tensor<3x9x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%11 = "quantization.dcast"(%10) : (tensor<3x9x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x9xf32>
return %11 : tensor<3x9xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: composite_and_select
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
// CHECK-SAME: %[[ARG2:.*]]: tensor<1x3xi1>
// CHECK-SAME: %[[ARG3:.*]]: tensor<1x3xf32>
func.func private @composite_and_select(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<1x3xi1>, %arg3: tensor<1x3xf32>) -> tensor<1x3xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[Q3:.*]] = "quantization.qcast"(%[[ARG3]]) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[SELECT:.*]] = stablehlo.select %[[ARG2]], %[[CALL]], %[[Q3]] : tensor<1x3xi1>, tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[SELECT]]) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<1x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
%3 = "quantization.dcast"(%2) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x3xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%7 = "quantization.qcast"(%arg3) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%8 = "quantization.dcast"(%7) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%9 = stablehlo.select %arg2, %6, %8 : (tensor<1x3xi1>, tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%10 = "quantization.qcast"(%9) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%11 = "quantization.dcast"(%10) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
return %11 : tensor<1x3xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: composite_and_broadcast_in_dim
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
func.func private @composite_and_broadcast_in_dim(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<2x3x2xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[BROADCAST:.*]] = stablehlo.broadcast_in_dim %[[CALL]], dims = [2, 1] : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[BROADCAST]]) : (tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<1x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
%3 = "quantization.dcast"(%2) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x3xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
%7 = stablehlo.broadcast_in_dim %6, dims = [2, 1] : (tensor<1x3xf32>) -> tensor<2x3x2xf32>
%8 = "quantization.qcast"(%7) {volatile} : (tensor<2x3x2xf32>) -> tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%9 = "quantization.dcast"(%8) : (tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2xf32>
return %9 : tensor<2x3x2xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: composite_and_gather
// CHECK-SAME: %[[ARG0:.*]]: tensor<3x4x5xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<3x5x2xf32>
// CHECK-SAME: %[[ARG2:.*]]: tensor<2x3x2xi64>
func.func private @composite_and_gather(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x2xf32>, %arg2: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<3x4x5xf32>) -> tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<3x5x2xf32>) -> tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[GATHER:.*]] = "stablehlo.gather"(%[[CALL]], %[[ARG2]])
// CHECK-SAME: (tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>, tensor<2x3x2xi64>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[GATHER]]) : (tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2x2xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<3x4x5xf32>) -> tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<3x4x5xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<3x5x2xf32>) -> tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>
%3 = "quantization.dcast"(%2) : (tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>) -> tensor<3x5x2xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<3x4x5xf32>, tensor<3x5x2xf32>) -> tensor<3x4x2xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<3x4x2xf32>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x4x2xf32>
%7 = "stablehlo.gather"(%6, %arg2) {
dimension_numbers = #stablehlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [1, 0],
index_vector_dim = 2>,
slice_sizes = array<i64: 1, 2, 2>,
indices_are_sorted = false
} : (tensor<3x4x2xf32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32>
%8 = "quantization.qcast"(%7) {volatile} : (tensor<2x3x2x2xf32>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%9 = "quantization.dcast"(%8) : (tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2x2xf32>
return %9 : tensor<2x3x2x2xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x2xf32>) -> tensor<3x4x2xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>) -> tensor<3x4x2x!quant.uniform<i32:f32, 3.000000e-05>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<3x4x2x!quant.uniform<i32:f32, 3.000000e-05>>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<3x4x5xf32>, tensor<3x5x2xf32>) -> tensor<3x4x2xf32>
return %0 : tensor<3x4x2xf32>
}
}
// -----
module attributes {tf_saved_model.semantics} {
// CHECK-LABEL: composite_and_slice
// CHECK-SAME: %[[ARG0:.*]]: tensor<3x2xf32>
// CHECK-SAME: %[[ARG1:.*]]: tensor<2x4xf32>
func.func private @composite_and_slice(%arg0: tensor<3x2xf32>, %arg1: tensor<2x4xf32>) -> tensor<2x2xf32> {
// CHECK: %[[Q1:.*]] = "quantization.qcast"(%[[ARG0]]) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK: %[[Q2:.*]] = "quantization.qcast"(%[[ARG1]]) {volatile} : (tensor<2x4xf32>) -> tensor<2x4x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>
// CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
// CHECK-SAME: (tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x4x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<3x4x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[SLICE:.*]] = stablehlo.slice %[[CALL]] [1:3, 2:4] : (tensor<3x4x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: %[[DQ:.*]] = "quantization.dcast"(%[[SLICE]]) : (tensor<2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x2xf32>
// CHECK: return %[[DQ]]
%0 = "quantization.qcast"(%arg0) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<3x2xf32>
%2 = "quantization.qcast"(%arg1) {volatile} : (tensor<2x4xf32>) -> tensor<2x4x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>
%3 = "quantization.dcast"(%2) : (tensor<2x4x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<2x4xf32>
%4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<3x2xf32>, tensor<2x4xf32>) -> tensor<3x4xf32>
%5 = "quantization.qcast"(%4) {volatile} : (tensor<3x4xf32>) -> tensor<3x4x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%6 = "quantization.dcast"(%5) : (tensor<3x4x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x4xf32>
%7 = stablehlo.slice %6 [1:3, 2:4] : (tensor<3x4xf32>) -> tensor<2x2xf32>
%8 = "quantization.qcast"(%7) {volatile} : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
%9 = "quantization.dcast"(%8) : (tensor<2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x2xf32>
return %9 : tensor<2x2xf32>
}
// CHECK: quantized_dot_general_fn_1
// CHECK-SAME: %[[ARG2:.*]]: tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03>>
// CHECK-SAME: %[[ARG3:.*]]: tensor<2x4x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>
func.func private @composite_dot_general_fn_1(%arg0: tensor<3x2xf32>, %arg1: tensor<2x4xf32>) -> tensor<3x4xf32> attributes {_from_xla_call_module} {
// CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK-SAME: (tensor<3x2x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<2x4x!quant.uniform<i8<-127:127>:f32:1, {6.000000e-03,6.000000e-03,6.000000e-03,6.000000e-03}>>) -> tensor<3x4x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05,3.000000e-05}>>
// CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<3x4x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<3x4x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
// CHECK: return %[[Q3]]
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<3x2xf32>, tensor<2x4xf32>) -> tensor<3x4xf32>
return %0 : tensor<3x4xf32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_weight_only.mlir
+80
View File
@@ -0,0 +1,80 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-quantize | FileCheck %s
// Test that hybrid quantized dot_general is produced when q/dq pair only exists
// for weight.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
%0 = "quantization.qcast"(%cst) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>
%1 = "quantization.dcast"(%0) : (tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<2x3xf32>
%2 = "tf.XlaCallModule"(%arg0, %1) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// CHECK-LABEL: quantize_dot_general_fn
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>
// CHECK: %[[CALL:.+]] = call @quantized_dot_general_fn(%[[ARG0]], %[[Q]])
// CHECK-SAME: {_quantization_method = "weight_only_ptq { }"} : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
// CHECK: return %[[CALL]]
// CHECK: quantized_dot_general_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x2xf32>, %[[ARG2:.+]]: tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
// CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG1]], %[[ARG2]]
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
// CHECK: return %[[DOT]]
// -----
// Test that hybrid quantized convolution is produced when q/dq pair only exists
// for weight.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_fn(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = stablehlo.constant dense<3.000000e-01> : tensor<2x3x3x2xf32>
%0 = "quantization.qcast"(%cst) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>
%1 = "quantization.dcast"(%0) : (tensor<2x3x3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<2x3x3x2xf32>
%2 = "tf.XlaCallModule"(%arg0, %1) <{Sout = [#tf_type.shape<1x3x4x2>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn", _quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
}
// CHECK-LABEL: quantize_conv_fn
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x3x4x3xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant dense<3.000000e-01> : tensor<2x3x3x2xf32>
// CHECK: %[[Q:.+]] = "quantization.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>
// CHECK: %[[CALL:.+]] = call @quantized_conv_fn(%[[ARG0]], %[[Q]]) {_quantization_method = "weight_only_ptq { }"} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CALL]]
// CHECK: quantized_conv_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x3x4x3xf32>, %[[ARG2:.+]]: tensor<2x3x3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3x4x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[ARG1]], %[[ARG2]])
// CHECK-SAME: (tensor<1x3x4x3xf32>, tensor<2x3x3x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CONV]]
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize_composite_functions.mlir
+910
View File
@@ -0,0 +1,910 @@
// 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: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \
// RUN: -tf-stablehlo-quantize-composite-functions | FileCheck %s
// RUN: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \
// RUN: -tf-stablehlo-quantize-composite-functions=enable-per-channel-quantized-weight=false | FileCheck --check-prefix=CHECK-PER-TENSOR %s
// Tests that basic dot_general is properly quantized.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
// Checks that the quantized XlaCallModule has been replaced by a CallOp, which
// calls the quantized entry function.
// CHECK: func.func private @quantize_dot_general_fn(%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_dot_general_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3xf32>
// CHECK-PER-TENSOR: func.func private @quantize_dot_general_fn(%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_dot_general_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3xf32>
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// Checks that the entry function is quantized for dot_general. Quantized
// dot_general outputs an i32 quantized tensor, followed by requantization to
// i8 quantized tensor.
// CHECK: func.func private @quantized_dot_general_fn(%[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]], contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[DOT_GENERAL_0]] : (tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_dot_general_fn(%[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]], contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[DOT_GENERAL_0]] : (tensor<1x3x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Tests that `stablehlo.dot_general` with `batching_dim` is quantized.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_batch_per_tensor_quantized_fn(%arg0: tensor<2x2x2xf32>) -> tensor<2x2x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x2x3xf32>} : () -> tensor<2x2x3xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<2x2x2xf32>, tensor<2x2x3xf32>) -> tensor<2x2x3xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<2x2x3xf32>) -> tensor<2x2x3xf32>
return %2 : tensor<2x2x3xf32>
}
// CHECK: func.func private @quantize_dot_general_batch_per_tensor_quantized_fn(%[[ARG_0:.+]]: tensor<2x2x2xf32>) -> tensor<2x2x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x2x3xi8>}> : () -> tensor<2x2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<2x2x2xf32>) -> tensor<2x2x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_dot_general_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]]) {_quantization_method = "static_range_ptq { }"} : (tensor<2x2x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<2x2x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<2x2x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<2x2x3xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<2x2x3xf32>
func.func private @composite_dot_general_fn(%arg0: tensor<2x2x2xf32>, %arg1: tensor<2x2x3xf32>) -> tensor<2x2x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<2x2x2xf32>, tensor<2x2x3xf32>) -> tensor<2x2x3xf32>
return %0 : tensor<2x2x3xf32>
}
}
// -----
// Tests that fused pattern for dot_general + bias is properly quantized.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_with_bias_same_shape_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<1x3xf32>} : () -> tensor<1x3xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_with_bias_same_shape_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_with_bias_same_shape_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
// CHECK: func.func private @quantize_dot_general_with_bias_same_shape_fn(%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>
// CHECK: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x3xi32>}> : () -> tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_dot_general_with_bias_same_shape_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>, tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3xf32>
// CHECK-PER-TENSOR: func.func private @quantize_dot_general_with_bias_same_shape_fn(%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x3xi32>}> : () -> tensor<1x3x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_dot_general_with_bias_same_shape_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>, tensor<1x3x!quant.uniform<i32:f32, {{.*}}>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3xf32>
func.func private @composite_dot_general_with_bias_same_shape_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%1 = stablehlo.add %0, %arg2 : tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
// CHECK: func.func private @quantized_dot_general_with_bias_same_shape_fn(%[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]], contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>) -> tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[DOT_GENERAL_0]], %[[ARG_3]] : tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_dot_general_with_bias_same_shape_fn(%[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x3x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]], contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<1x3x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[DOT_GENERAL_0]], %[[ARG_3]] : tensor<1x3x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<1x3x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Tests that fused pattern for dot_general + bias with dynamic batch dimension
// is properly quantized.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_with_bias_dynamic_fn(%arg0: tensor<?x2xf32>) -> tensor<?x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<3xf32>} : () -> tensor<3xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<?x2xf32>) -> tensor<?x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {Sout = [#tf_type.shape<?x3>], _entry_function = @composite_dot_general_with_bias_dynamic_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_with_bias_dynamic_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<?x2xf32>, tensor<2x3xf32>, tensor<3xf32>) -> tensor<?x3xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<?x3xf32>) -> tensor<?x3xf32>
return %2 : tensor<?x3xf32>
}
// CHECK: func.func private @quantize_dot_general_with_bias_dynamic_fn(%[[ARG_0:.+]]: tensor<?x2xf32>) -> tensor<?x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>
// CHECK-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<3xi32>}> : () -> tensor<3x!quant.uniform<i32:f32:0, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x2xf32>) -> tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_dot_general_with_bias_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]]) {_quantization_method = "static_range_ptq { }"} : (tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>, tensor<3x!quant.uniform<i32:f32:0, {{.*}}>) -> tensor<?x3x!quant.uniform<i8:f32, {{.*}}>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3xf32>
// CHECK-PER-TENSOR: func.func private @quantize_dot_general_with_bias_dynamic_fn(%[[ARG_0:.+]]: tensor<?x2xf32>) -> tensor<?x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<3xi32>}> : () -> tensor<3x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x2xf32>) -> tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_dot_general_with_bias_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]]) {_quantization_method = "static_range_ptq { }"} : (tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>, tensor<3x!quant.uniform<i32:f32, {{.*}}>) -> tensor<?x3x!quant.uniform<i8:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3xf32>
func.func private @composite_dot_general_with_bias_dynamic_fn(%arg0: tensor<?x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<3xf32>) -> tensor<?x3xf32> attributes {_from_xla_call_module} {
%cst_0 = stablehlo.constant dense<2> : tensor<1xi32>
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<?x2xf32>, tensor<2x3xf32>) -> tensor<?x3xf32>
%1 = stablehlo.get_dimension_size %0, dim = 0 : (tensor<?x3xf32>) -> tensor<i32>
%2 = stablehlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
%3 = stablehlo.concatenate %2, %cst_0, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
%4 = stablehlo.dynamic_broadcast_in_dim %arg2, %3, dims = [1] : (tensor<3xf32>, tensor<2xi32>) -> tensor<?x3xf32>
%5 = stablehlo.add %0, %4 : tensor<?x3xf32>
return %5 : tensor<?x3xf32>
}
}
// CHECK: func.func private @quantized_dot_general_with_bias_dynamic_fn(%[[ARG_1:.+]]: tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>, %[[ARG_3:.+]]: tensor<3x!quant.uniform<i32:f32:0, {{.*}}>>) -> tensor<?x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[CONST_2:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]], contracting_dims = [1] x [0] : (tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>) -> tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[DOT_GENERAL_0]], dim = 0 : (tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>)
// CHECK: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [1] : (tensor<3x!quant.uniform<i32:f32:0, {{.*}}>>, tensor<2xi32>) -> tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[DOT_GENERAL_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x!quant.uniform<i32:f32:1, {{.*}}>>)
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<?x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_dot_general_with_bias_dynamic_fn(%[[ARG_1:.+]]: tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<3x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<?x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR: %[[CONST_2:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK-PER-TENSOR: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]], contracting_dims = [1] x [0] : (tensor<?x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<?x3x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[DOT_GENERAL_0]], dim = 0 : (tensor<?x3x!quant.uniform<i32:f32, {{.*}}>)
// CHECK-PER-TENSOR: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
// CHECK-PER-TENSOR: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [1] : (tensor<3x!quant.uniform<i32:f32, {{.*}}>>, tensor<2xi32>) -> tensor<?x3x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[DOT_GENERAL_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x!quant.uniform<i32:f32, {{.*}}>>)
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_1]] : tensor<?x3x!quant.uniform<i8:f32, {{.*}}>>
// -----
// Tests that basic convolution is properly quantized. It is per-channel
// quantized unless `enable-per-channel-quantized-weight=false`, according to
// `_quantization_method` with an `input_quantized_types` and explicit
// `dimension_specs`.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_fn(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {
Sout = [#tf_type.shape<1x3x4x2>],
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64,
_entry_function = @composite_conv_fn,
_stableghlo_version = "1.0.0",
_original_entry_function = "composite_conv_fn",
// Per-channel quantization at dimension 3 for input index 1.
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_stablehlo_module_attrs = {},
_tfl_quant_trait = "fully_quantizable",
device = ""
} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
// Check that the quantized XlaCallModule has been replaced by a CallOp, which
// calls the quantized entry function.
// CHECK: func.func private @quantize_conv_fn(%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]])
// CHECK-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
// CHECK-PER-TENSOR: func.func private @quantize_conv_fn(%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]])
// CHECK-PER-TENSOR-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
// Checks that the entry function is quantized for convolution. Quantized
// convolution outputs an i32 quantized tensor, followed by requantization to
// i8 quantized tensor.
// CHECK: func.func private @quantized_conv_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[CONVOLUTION_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_conv_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[CONVOLUTION_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Tests that basic convolution is properly quantized. In this example, the
// convolution is always per-tensor quantized (even if
// enable-per-channel-quantized-weights=true), according to
// `_quantization_method`.
// CHECK-LABEL: quantize_conv_fn_per_tensor
func.func @quantize_conv_fn_per_tensor(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {
Sout = [#tf_type.shape<1x3x4x2>],
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64,
_entry_function = @composite_conv_fn,
_stablehlo_version = "1.0.0",
_original_entry_function = "composite_conv_fn",
_quantization_method = "static_range_ptq { }",
_stablehlo_module_attrs = {},
_tfl_quant_trait = "fully_quantizable",
device = ""
} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
// Check that the quantized XlaCallModule has been replaced by a CallOp, which
// calls the quantized entry function.
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32>
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8:f32, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
// Checks that the entry function is quantized for convolution. Quantized
// convolution outputs an i32 quantized tensor, followed by requantization to
// i8 quantized tensor.
// CHECK: func.func private @quantized_conv_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[CONVOLUTION_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// -----
// Tests that fused pattern for convolution + bias is properly quantized.
// Checks that fused functions with 1D bias is properly quantized.
// The 1D bias should be broadcasted in dims [3], where it initially has
// `quantizedDimension=0`, but has `quantizedDimension=3` after broadcasting.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_with_bias_1d_fn(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {
Sout = [#tf_type.shape<1x3x4x2>],
_entry_function = @composite_conv_with_bias_1d_fn,
_stablehlo_version = "1.0.0",
_original_entry_function = "composite_conv_with_bias_1d_fn",
_stablehlo_module_attrs = {},
// Per-channel quantization at dimension 3 for input index 1.
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_tfl_quant_trait = "fully_quantizable",
device = "",
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64
} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
// CHECK: func.func private @quantize_conv_with_bias_1d_fn(%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>
// CHECK-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<47978> : tensor<2xi32>}> : () -> tensor<2x!quant.uniform<i32:f32:0, {8.3371932554046126E-6,8.3371932554046126E-6}>>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_with_bias_1d_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>, tensor<2x!quant.uniform<i32:f32:0, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, 0.0027450979924669452:-128>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
// CHECK-PER-TENSOR: func.func private @quantize_conv_with_bias_1d_fn(%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2xi32>}> : () -> tensor<2x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_with_bias_1d_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-PER-TENSOR-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>, tensor<2x!quant.uniform<i32:f32, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
func.func private @composite_conv_with_bias_1d_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.broadcast_in_dim %arg2, dims = [3] : (tensor<2xf32>) -> tensor<1x3x4x2xf32>
%1 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
%2 = stablehlo.add %1, %0 : tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
// CHECK: func.func private @quantized_conv_with_bias_1d_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>, %[[ARG_3:.+]]: tensor<2x!quant.uniform<i32:f32:0, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[BROADCAST_IN_DIM:.+]] = stablehlo.broadcast_in_dim %arg2, dims = [3] : (tensor<2x!quant.uniform<i32:f32:0, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[BROADCAST_IN_DIM]] : tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_conv_with_bias_1d_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR: %[[BROADCAST_IN_DIM:.+]] = stablehlo.broadcast_in_dim %[[ARG_3]]
// CHECK-PER-TENSOR: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[BROADCAST_IN_DIM]] : tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Checks that fused functions with 4D bias is properly quantized.
// The 4D bias should be braoadcasted in dims [0, 1, 2, 3], where it
// already has `quantizedDimension=3`.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_with_bias_fn(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {
Sout = [#tf_type.shape<1x3x4x2>],
_entry_function = @composite_conv_with_bias_fn,
_stablehlo_version = "1.0.0",
_original_entry_function = "composite_conv_with_bias_fn",
_stablehlo_module_attrs = {},
// Per-channel quantization at dimension 3 for input index 1.
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_tfl_quant_trait = "fully_quantizable",
device = "",
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64
} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<1x1x1x2xf32>) -> tensor<1x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
// CHECK: func.func private @quantize_conv_with_bias_fn(%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>
// CHECK-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_with_bias_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>, tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
// CHECK-PER-TENSOR: func.func private @quantize_conv_with_bias_fn(%[[ARG_0:.+]]: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_with_bias_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-PER-TENSOR-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>, tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x4x2xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3x4x2xf32>
func.func private @composite_conv_with_bias_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<1x1x1x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.broadcast_in_dim %arg2, dims = [0, 1, 2, 3] : (tensor<1x1x1x2xf32>) -> tensor<1x3x4x2xf32>
%1 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
%2 = stablehlo.add %1, %0 : tensor<1x3x4x2xf32>
return %2 : tensor<1x3x4x2xf32>
}
// CHECK: func.func private @quantized_conv_with_bias_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK: %[[BROADCAST_IN_DIM:.+]] = stablehlo.broadcast_in_dim %arg2, dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[BROADCAST_IN_DIM]] : tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_conv_with_bias_fn(%[[ARG_1:.+]]: tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR: %[[BROADCAST_IN_DIM:.+]] = stablehlo.broadcast_in_dim %arg2
// CHECK-PER-TENSOR: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<1x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[BROADCAST_IN_DIM]] : tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<1x3x4x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Tests that fused pattern for convolution + bias with dynamic batch dimension
// is properly quantized.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_with_bias_dynamic_fn(%arg0: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {
Sout = [#tf_type.shape<1x3x4x2>],
_entry_function = @composite_conv_with_bias_dynamic_fn,
_stablehlo_version = "1.0.0",
_original_entry_function = "composite_conv_with_bias_dynamic_fn",
_stablehlo_module_attrs = {},
// Per-channel quantization at dimension 3 for input index 1.
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_tfl_quant_trait = "fully_quantizable",
device = "",
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64
} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<1x1x1x2xf32>) -> tensor<?x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<?x3x4x2xf32>) -> tensor<?x3x4x2xf32>
return %2 : tensor<?x3x4x2xf32>
}
// CHECK: func.func private @quantize_conv_with_bias_dynamic_fn(%[[ARG_0:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>
// CHECK-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_with_bias_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>, tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3x4x2xf32>
// CHECK-PER-TENSOR: func.func private @quantize_conv_with_bias_dynamic_fn(%[[ARG_0:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_with_bias_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-PER_TENSOR-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>, tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3x4x2xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3x4x2xf32>
func.func private @composite_conv_with_bias_dynamic_fn(%arg0: tensor<?x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<1x1x1x2xf32>) -> tensor<?x3x4x2xf32> attributes {_from_xla_call_module} {
%cst_0 = stablehlo.constant dense<3> : tensor<1xi32>
%cst_1 = stablehlo.constant dense<4> : tensor<1xi32>
%cst_2 = stablehlo.constant dense<2> : tensor<1xi32>
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x4x2xf32>
%1 = stablehlo.get_dimension_size %0, dim = 0 : (tensor<?x3x4x2xf32>) -> tensor<i32>
%2 = stablehlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
%3 = stablehlo.concatenate %2, %cst_0, %cst_1, %cst_2, dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
%4 = stablehlo.dynamic_broadcast_in_dim %arg2, %3, dims = [0, 1, 2, 3] : (tensor<1x1x1x2xf32>, tensor<4xi32>) -> tensor<?x3x4x2xf32>
%5 = stablehlo.add %0, %4 : tensor<?x3x4x2xf32>
return %5 : tensor<?x3x4x2xf32>
}
}
// CHECK: func.func private @quantized_conv_with_bias_dynamic_fn(%[[ARG_1:.+]]: tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>) -> tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[CONVOLUTION_0]], dim = 0 : (tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>)
// CHECK: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>, tensor<4xi32>) -> tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>)
// CHECK: return %[[UNIFORM_QUANTIZE_0]] : tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: func.func private @quantized_conv_with_bias_dynamic_fn(%[[ARG_1:.+]]: tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-PER-TENSOR-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-PER-TENSOR-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>) -> tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[CONVOLUTION_0]], dim = 0 : (tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>)
// CHECK-PER-TENSOR: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK-PER-TENSOR: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>, tensor<4xi32>) -> tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>)
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_0]] : tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// -----
// Tests that fused pattern for convolution + bias + relu with
// dynamic batch dimension is properly quantized.
// Note that this checks for identical condition as
// quantize_conv_with_bias_dynamic_fn, omitting stablehlo.maximum.
// This is because activation clipping which includes 0.0f can be simply
// omitted from the graph as the lifted function's out_scale and out_zp are
// already calculated based on the clipped distribution.
// Note that the resulting scale and zero point should be calculated based on
// clipped range [0, r_max].
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_with_bias_and_relu_dynamic_fn(%arg0: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {
Sout = [#tf_type.shape<1x3x4x2>],
_entry_function = @composite_conv_with_bias_and_relu_dynamic_fn,
_stablehlo_version = "1.0.0",
_original_entry_function = "composite_conv_with_bias_and_relu_dynamic_fn",
_stablehlo_module_attrs = {},
// Per-channel quantization at dimension 3 for input index 1.
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_tfl_quant_trait = "fully_quantizable",
device = "",
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64
} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<1x1x1x2xf32>) -> tensor<?x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[0.00000000e-6, 8.00000000e-1]> : tensor<2xf32>} : (tensor<?x3x4x2xf32>) -> tensor<?x3x4x2xf32>
return %2 : tensor<?x3x4x2xf32>
}
// CHECK: func.func private @quantize_conv_with_bias_and_relu_dynamic_fn(%[[ARG_0:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>
// CHECK-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_with_bias_and_relu_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>, tensor<1x1x1x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0031372549487095253:-128>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3x4x2xf32>
// CHECK-PER-TENSOR: func.func private @quantize_conv_with_bias_and_relu_dynamic_fn(%[[ARG_0:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_with_bias_and_relu_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-PER-TENSOR-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>, tensor<1x1x1x2x!quant.uniform<i32:f32, 8.3371932554046126E-6>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0031372549487095253:-128>>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3x4x2xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3x4x2xf32>
func.func private @composite_conv_with_bias_and_relu_dynamic_fn(%arg0: tensor<?x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<1x1x1x2xf32>) -> tensor<?x3x4x2xf32> attributes {_from_xla_call_module} {
%cst_0 = stablehlo.constant dense<3> : tensor<1xi32>
%cst_1 = stablehlo.constant dense<4> : tensor<1xi32>
%cst_2 = stablehlo.constant dense<2> : tensor<1xi32>
%cst_3 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%cst_4 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x4x2xf32>
%1 = stablehlo.get_dimension_size %0, dim = 0 : (tensor<?x3x4x2xf32>) -> tensor<i32>
%2 = stablehlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
%3 = stablehlo.concatenate %2, %cst_0, %cst_1, %cst_2, dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
%4 = stablehlo.dynamic_broadcast_in_dim %arg2, %3, dims = [0, 1, 2, 3] : (tensor<1x1x1x2xf32>, tensor<4xi32>) -> tensor<?x3x4x2xf32>
%5 = stablehlo.add %0, %4 : tensor<?x3x4x2xf32>
%6 = stablehlo.clamp %cst_3, %5, %cst_4 : (tensor<f32>, tensor<?x3x4x2xf32>, tensor<f32>) -> tensor<?x3x4x2xf32>
return %6 : tensor<?x3x4x2xf32>
}
}
// CHECK: func.func private @quantized_conv_with_bias_and_relu_dynamic_fn(%[[ARG_1:.+]]: tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<?x3x4x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>
// CHECK: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[CONVOLUTION_0]], dim = 0 : (tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>)
// CHECK: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>, tensor<4xi32>) -> tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x4x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0031372549487095253:-128>>
// CHECK: return %[[UNIFORM_QUANTIZE_0]] : tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0031372549487095253:-128>>
// CHECK-PER-TENSOR: func.func private @quantized_conv_with_bias_and_relu_dynamic_fn(%[[ARG_1:.+]]: tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-PER-TENSOR-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-PER-TENSOR-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<?x3x4x2x!quant.uniform<i32:f32, 8.3371932554046126E-6>>
// CHECK-PER-TENSOR: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[CONVOLUTION_0]], dim = 0 : (tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>)
// CHECK-PER-TENSOR: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK-PER-TENSOR: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>, tensor<4xi32>) -> tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x4x2x!quant.uniform<i32:f32, 8.3371932554046126E-6>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0031372549487095253:-128>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_0]] : tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0031372549487095253:-128>>
// -----
// Tests that fused pattern for convolution + bias + relu6 with
// dynamic batch dimension is properly quantized.
// Note that this checks for identical condition as
// quantize_conv_with_bias_dynamic_fn, omitting stablehlo.clamp.
// This is because activation clipping which includes 0.0f can be simply
// omitted from the graph as the lifted function's out_scale and out_zp are
// already calculated based on the clipped distribution.
// Note that the resulting scale and zero point should be calculated based on
// clipped range [0, r_max].
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_with_bias_and_relu6_dynamic_fn(%arg0: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
%cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
%0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {
Sout = [#tf_type.shape<1x3x4x2>],
_entry_function = @composite_conv_with_bias_and_relu6_dynamic_fn,
_stablehlo_version = "1.0.0",
_original_entry_function = "composite_conv_with_bias_and_relu6_dynamic_fn",
_stablehlo_module_attrs = {},
// Per-channel quantization at dimension 3 for input index 1.
_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
_tfl_quant_trait = "fully_quantizable",
device = "",
dim_args_spec = [],
disabled_checks = [],
has_token_input_output = false,
module = "",
platforms = [],
version = 5 : i64
} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<1x1x1x2xf32>) -> tensor<?x3x4x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 6.00000000e-1]> : tensor<2xf32>} : (tensor<?x3x4x2xf32>) -> tensor<?x3x4x2xf32>
return %2 : tensor<?x3x4x2xf32>
}
// CHECK: func.func private @quantize_conv_with_bias_and_relu6_dynamic_fn(%[[ARG_0:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>
// CHECK-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_conv_with_bias_and_relu6_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-SAME: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>, tensor<1x1x1x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0023529412699680704:-128>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3x4x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3x4x2xf32>
// CHECK-PER-TENSOR: func.func private @quantize_conv_with_bias_and_relu6_dynamic_fn(%[[ARG_0:.+]]: tensor<?x3x4x3xf32>) -> tensor<?x3x4x2xf32> attributes {tf._original_func_name = "main_0"}
// CHECK-PER-TENSOR-DAG: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>
// CHECK-PER-TENSOR-DAG: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<?x3x4x3xf32>) -> tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[CALL_0:.+]] = call @quantized_conv_with_bias_and_relu6_dynamic_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]], %[[CONST_1]])
// CHECK-PER-TENSOR: {_quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}"} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>, tensor<1x1x1x2x!quant.uniform<i32:f32, 8.3371932554046126E-6>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0023529412699680704:-128>>
// CHECK-PER-TENSOR: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<?x3x4x2xf32>
// CHECK-PER-TENSOR: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<?x3x4x2xf32>
func.func private @composite_conv_with_bias_and_relu6_dynamic_fn(%arg0: tensor<?x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<1x1x1x2xf32>) -> tensor<?x3x4x2xf32> attributes {_from_xla_call_module} {
%cst_0 = stablehlo.constant dense<3> : tensor<1xi32>
%cst_1 = stablehlo.constant dense<4> : tensor<1xi32>
%cst_2 = stablehlo.constant dense<2> : tensor<1xi32>
%cst_3 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%cst_4 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x4x2xf32>
%1 = stablehlo.get_dimension_size %0, dim = 0 : (tensor<?x3x4x2xf32>) -> tensor<i32>
%2 = stablehlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
%3 = stablehlo.concatenate %2, %cst_0, %cst_1, %cst_2, dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
%4 = stablehlo.dynamic_broadcast_in_dim %arg2, %3, dims = [0, 1, 2, 3] : (tensor<1x1x1x2xf32>, tensor<4xi32>) -> tensor<?x3x4x2xf32>
%5 = stablehlo.add %0, %4 : tensor<?x3x4x2xf32>
%6 = stablehlo.clamp %cst_3, %5, %cst_4 : (tensor<f32>, tensor<?x3x4x2xf32>, tensor<f32>) -> tensor<?x3x4x2xf32>
return %6 : tensor<?x3x4x2xf32>
}
}
// CHECK: func.func private @quantized_conv_with_bias_and_relu6_dynamic_fn(%[[ARG_1:.+]]: tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<?x3x4x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>
// CHECK: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[CONVOLUTION_0]], dim = 0 : (tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>)
// CHECK: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32:3, {{.*}}>>, tensor<4xi32>) -> tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x4x2x!quant.uniform<i32:f32:3, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x4x2x!quant.uniform<i32:f32:3, {8.3371932554046126E-6,8.3371932554046126E-6}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0023529412699680704:-128>>
// CHECK: return %[[UNIFORM_QUANTIZE_0]] : tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0023529412699680704:-128>>
// CHECK-PER-TENSOR: func.func private @quantized_conv_with_bias_and_relu6_dynamic_fn(%[[ARG_1:.+]]: tensor<?x3x4x3x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, {{.*}}>>, %[[ARG_3:.+]]: tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
// CHECK-PER-TENSOR-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-PER-TENSOR-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-PER-TENSOR-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONVOLUTION_0:.+]] = stablehlo.convolution(%[[ARG_1]], %[[ARG_2]]) {{.*}} : (tensor<?x3x4x3x!quant.uniform<i8:f32, 0.0035294116712084002:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<?x3x4x2x!quant.uniform<i32:f32, 8.3371932554046126E-6>>
// CHECK-PER-TENSOR: %[[GET_DIMENSION_SIZE_0:.+]] = stablehlo.get_dimension_size %[[CONVOLUTION_0]], dim = 0 : (tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>)
// CHECK-PER-TENSOR: %[[RESHAPE_0:.+]] = stablehlo.reshape %[[GET_DIMENSION_SIZE_0]] : (tensor<i32>) -> tensor<1xi32>
// CHECK-PER-TENSOR: %[[CONCATENATE_0:.+]] = stablehlo.concatenate %[[RESHAPE_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK-PER-TENSOR: %[[DYNAMIC_BROADCAST_IN_DIM_0:.+]] = stablehlo.dynamic_broadcast_in_dim %[[ARG_3]], %[[CONCATENATE_0]], dims = [0, 1, 2, 3] : (tensor<1x1x1x2x!quant.uniform<i32:f32, {{.*}}>>, tensor<4xi32>) -> tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[ADD_0:.+]] = stablehlo.add %[[CONVOLUTION_0]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] : tensor<?x3x4x2x!quant.uniform<i32:f32, {{.*}}>>
// CHECK-PER-TENSOR: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ADD_0]] : (tensor<?x3x4x2x!quant.uniform<i32:f32, 8.3371932554046126E-6>>) -> tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0023529412699680704:-128>>
// CHECK-PER-TENSOR: return %[[UNIFORM_QUANTIZE_0]] : tensor<?x3x4x2x!quant.uniform<i8:f32, 0.0023529412699680704:-128>>
// -----
// Tests that XlaCallModule op is not quantized and converted to func.call without the quantization.stats ops.
module attributes {tf_saved_model.semantics} {
func.func private @not_quantized_without_stats_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%1 = "tf.XlaCallModule"(%arg0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stableghlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
// Check that "tf.Const" is converted to stablehlo.constant. XlaCallModule is
// not quantized.
// CHECK: func.func private @not_quantized_without_stats_fn(%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"}
// CHECK: %[[CONST_0:.+]] = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
// CHECK: %[[CALL:.+]] = call @composite_dot_general_fn(%[[ARG_0]], %[[CONST_0]]) : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
// CHECK: return %[[CALL]]
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: func.func private @composite_dot_general_fn(%[[ARG_1:.+]]: tensor<1x2xf32>, %[[ARG_2:.+]]: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module}
// Check that the composite_dot_general_fn is untouched.
// CHECK: %[[DOT_GENERAL_0:.+]] = stablehlo.dot_general %[[ARG_1]], %[[ARG_2]]
// CHECK: return %[[DOT_GENERAL_0]]
}
// -----
// Tests that basic `stablehlo.gather` is properly quantized.
module attributes {tf_saved_model.semantics} {
// CHECK: func.func private @quantize_gather_fn(%[[ARG:.+]]: tensor<3x4x2xf32>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"}
func.func private @quantize_gather_fn(%arg: tensor<3x4x2xf32>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<1> : tensor<2x3x2xi32>} : () -> tensor<2x3x2xi32>
%0 = "quantization.stats"(%arg) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<3x4x2xf32>) -> tensor<3x4x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<2x3x2x2>], _entry_function = @composite_gather_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_gather_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<3x4x2xf32>, tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<2x3x2x2xf32>) -> tensor<2x3x2x2xf32>
return %2 : tensor<2x3x2x2xf32>
}
// Checks that the quantized XlaCallModule has been replaced by a CallOp, which
// calls the quantized entry function.
// CHECK: %[[CONST:.+]] = stablehlo.constant dense<{{.*}}> : tensor<2x3x2xi32>
// CHECK: %[[UNIFORM_QUANTIZE:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<3x4x2xf32>) -> tensor<3x4x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL:.+]] = call @quantized_gather_fn(%[[UNIFORM_QUANTIZE]], %[[CONST]]) {_quantization_method = "static_range_ptq { }"} : (tensor<3x4x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x2xi32>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE:.+]] = stablehlo.uniform_dequantize %[[CALL]] : (tensor<2x3x2x2x!quant.uniform<i8:f32, {{.*}}>) -> tensor<2x3x2x2xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE]] : tensor<2x3x2x2xf32>
// CHECK: func.func private @quantized_gather_fn(%[[ARG_0:.+]]: tensor<3x4x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_1:.+]]: tensor<2x3x2xi32>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
func.func private @composite_gather_fn(%arg0: tensor<3x4x2xf32>, %arg1: tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32> attributes {_from_xla_call_module} {
%0 = "stablehlo.gather"(%arg0, %arg1) {
dimension_numbers = #stablehlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [1, 0],
index_vector_dim = 2>,
slice_sizes = array<i64: 1, 2, 2>,
indices_are_sorted = false
} : (tensor<3x4x2xf32>, tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32>
return %0 : tensor<2x3x2x2xf32>
}
// CHECK: %[[GATHER:.+]] = "stablehlo.gather"(%[[ARG_0]], %[[ARG_1]]) {{.*}} : (tensor<3x4x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x2xi32>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[GATHER]] : tensor<2x3x2x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE_0]] : tensor<2x3x2x2x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Tests that a basic `stablehlo.add` and a fused `stablehlo.dot_general`
// are properly quantized.
module attributes {tf_saved_model.semantics} {
// CHECK: func.func private @quantize_add_fn(%[[ARG:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"}
func.func private @quantize_add_fn(%arg: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%cst_0 = "tf.Const"() {value = dense<1.00000000e-1> : tensor<1x2xf32>} : () -> tensor<1x2xf32>
%cst_1 = "tf.Const"() {value = dense<1.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
%0 = "quantization.stats"(%arg) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst_0) {Sout = [#tf_type.shape<1x2>], _entry_function = @composite_add_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_add_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x2xf32>
%2 = "quantization.stats"(%1) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
%3 = "quantization.stats"(%2) {layerStats = dense<[5.00000000e-6, 6.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
%4 = "tf.XlaCallModule"(%3, %cst_1) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
%5 = "quantization.stats"(%4) {layerStats = dense<[5.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
return %5 : tensor<1x3xf32>
}
// CHECK: %[[CONST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<1x2xi8>}> : () -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE:.+]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL:.+]] = call @quantized_add_fn(%[[UNIFORM_QUANTIZE]], %[[CONST]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE:.+]] = stablehlo.uniform_dequantize %[[CALL]] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x2xf32>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[UNIFORM_DEQUANTIZE]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CALL_0:.+]] = call @quantized_dot_general_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]]) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[UNIFORM_DEQUANTIZE_0:.+]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3xf32>
// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3xf32>
// CHECK: func.func private @quantized_add_fn(%[[ARG_0:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
func.func private @composite_add_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<1x2xf32>) -> tensor<1x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.add %arg0, %arg1 : tensor<1x2xf32>
return %0 : tensor<1x2xf32>
}
// CHECK: %[[ADD:.+]] = stablehlo.add %arg0, %arg1 : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[ADD]] : tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: func.func private @quantized_dot_general_fn(%[[ARG_0:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_1:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: %[[DOT_GENERAL:.+]] = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1,{{.*}}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>
// CHECK: %[[UNIFORM_QUANTIZE:.+]] = stablehlo.uniform_quantize %[[DOT_GENERAL]] : (tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: return %[[UNIFORM_QUANTIZE]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
}
// -----
// Tests that `stablehlo.add` is not quantized and emits error when the function
// does not include two ops.
module attributes {tf_saved_model.semantics} {
func.func private @not_quantize_fn_when_not_singular(%arg: tensor<1x2xf32>) -> tensor<1x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<1.00000000e-1> : tensor<1x2xf32>} : () -> tensor<1x2xf32>
%0 = "quantization.stats"(%arg) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x2>], _entry_function = @composite_add_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_add_fn", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x2xf32>
// expected-error@+1 {{'stablehlo.uniform_dequantize' op operand #0 must be ranked tensor of per-tensor integer quantized or per-axis integer quantized values, but got 'tensor<1x2xf32>'}}
%2 = "quantization.stats"(%1) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
return %2 : tensor<1x2xf32>
}
func.func private @composite_add_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<1x2xf32>) -> tensor<1x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.add %arg0, %arg1 : tensor<1x2xf32>
%1 = stablehlo.add %0, %arg1 : tensor<1x2xf32>
return %1 : tensor<1x2xf32>
}
}
// -----
// Tests that `stablehlo.gather` without `static_range_ptq` is not quantized.
module attributes {tf_saved_model.semantics} {
func.func private @not_quantize_singular_op_without_static_range_ptq(%arg: tensor<3x4x2xf32>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"} {
%cst = "tf.Const"() {value = dense<1> : tensor<2x3x2xi32>} : () -> tensor<2x3x2xi32>
%0 = "quantization.stats"(%arg) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<3x4x2xf32>) -> tensor<3x4x2xf32>
%1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<2x3x2x2>], _entry_function = @composite_gather_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_gather_fn", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<3x4x2xf32>, tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32>
// expected-error@+1 {{'stablehlo.uniform_dequantize' op operand #0 must be ranked tensor of per-tensor integer quantized or per-axis integer quantized values, but got 'tensor<2x3x2x2xf32>'}}
%2 = "quantization.stats"(%1) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<2x3x2x2xf32>) -> tensor<2x3x2x2xf32>
return %2 : tensor<2x3x2x2xf32>
}
func.func private @composite_gather_fn(%arg0: tensor<3x4x2xf32>, %arg1: tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32> attributes {_from_xla_call_module} {
%0 = "stablehlo.gather"(%arg0, %arg1) {
dimension_numbers = #stablehlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [1, 0],
index_vector_dim = 2>,
slice_sizes = array<i64: 1, 2, 2>,
indices_are_sorted = false
} : (tensor<3x4x2xf32>, tensor<2x3x2xi32>) -> tensor<2x3x2x2xf32>
return %0 : tensor<2x3x2x2xf32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize_composite_functions_weight_only.mlir
+136
View File
@@ -0,0 +1,136 @@
// 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: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \
// RUN: -tf-stablehlo-quantize-composite-functions | FileCheck --check-prefix=CHECK %s
// Test that per-tensor weight-only quantized dot_general op is produced when
// empty `weight_only_ptq` is provided.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_per_tensor(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%0 = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// CHECK-LABEL: quantize_dot_general_per_tensor
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[CALL:.+]] = call @quantized_dot_general_fn(%[[ARG0]], %[[CST]]) {_quantization_method = "weight_only_ptq { }"} : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3xf32>
// CHECK: return %[[CALL]]
// CHECK: quantized_dot_general_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x2xf32>, %[[ARG2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3xf32>
// CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG1]], %[[ARG2]]
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3xf32>
// CHECK: return %[[DOT]]
// -----
// Test that per-tensor weight-only quantized convolution op is produced when
// empty `weight_only_ptq` is provided.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_per_tensor(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%0 = stablehlo.constant dense<3.000000e-01> : tensor<2x3x3x2xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<1x3x4x2>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn", _quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %1 : tensor<1x3x4x2xf32>
}
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
}
// CHECK-LABEL: quantize_conv_per_tensor
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x3x4x3xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[CALL:.+]] = call @quantized_conv_fn(%[[ARG0]], %[[CST]]) {_quantization_method = "weight_only_ptq { }"} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CALL]]
// CHECK: quantized_conv_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x3x4x3xf32>, %[[ARG2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3x4x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[ARG1]], %[[ARG2]])
// CHECK-SAME: (tensor<1x3x4x3xf32>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CONV]]
// -----
// Test that per-channel weight-only quantized dot_general op is produced when
// `weight_only_ptq` with `dimension_specs` is provided.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_dot_general_per_channel(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
%0 = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {}}}}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// CHECK-LABEL: quantize_dot_general_per_channel
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {0.0023622048182750312,0.0023622048182750312,0.0023622048182750312}>>
// CHECK: %[[CALL:.+]] = call @quantized_dot_general_fn(%[[ARG0]], %[[CST]]) {_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {}}}}"}
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {0.0023622048182750312,0.0023622048182750312,0.0023622048182750312}>>) -> tensor<1x3xf32>
// CHECK: return %[[CALL]]
// CHECK: quantized_dot_general_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x2xf32>, %[[ARG2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {0.0023622048182750312,0.0023622048182750312,0.0023622048182750312}>>) -> tensor<1x3xf32>
// CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG1]], %[[ARG2]]
// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {0.0023622048182750312,0.0023622048182750312,0.0023622048182750312}>>) -> tensor<1x3xf32>
// CHECK: return %[[DOT]]
// -----
// Test that per-channel weight-only quantized convolution op is produced when
// `weight_only_ptq` with `dimension_specs` is provided.
module attributes {tf_saved_model.semantics} {
func.func private @quantize_conv_per_channel(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "main_0"} {
%0 = stablehlo.constant dense<3.000000e-01> : tensor<2x3x3x2xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<1x3x4x2>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_conv_fn, _original_entry_function = "composite_conv_fn", _quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {}}}}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %1 : tensor<1x3x4x2xf32>
}
func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
return %0 : tensor<1x3x4x2xf32>
}
}
// CHECK-LABEL: quantize_conv_per_channel
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x3x4x3xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>
// CHECK: %[[CALL:.+]] = call @quantized_conv_fn(%[[ARG0]], %[[CST]]) {_quantization_method = "weight_only_ptq {input_quantized_types {key: 1, value {dimension_specs {}}}}"}
// CHECK-SAME: (tensor<1x3x4x3xf32>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CALL]]
// CHECK: quantized_conv_fn
// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x3x4x3xf32>, %[[ARG2:.+]]: tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<1x3x4x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[ARG1]], %[[ARG2]])
// CHECK-SAME: (tensor<1x3x4x3xf32>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.0023622048182750312,0.0023622048182750312}>>) -> tensor<1x3x4x2xf32>
// CHECK: return %[[CONV]]
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/remove_sharding_custom_call.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: stablehlo-quant-opt %s -tf-stablehlo-remove-sharding-custom-call \
// RUN: -split-input-file | FileCheck %s
// CHECK-LABEL: sharding_custom_call_removed
func.func @sharding_custom_call_removed(%arg0: tensor<3xf32>) -> tensor<3xf32> {
%1 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = ""} : (tensor<3xf32>) -> tensor<3xf32>
return %1 : tensor<3xf32>
}
// CHECK-NOT: custom_call
// -----
// Tests that a custom_call that is not @Sharding is not removed.
// CHECK-LABEL: custom_call_not_removed
func.func @custom_call_not_removed(%arg0: tensor<3xf32>) -> tensor<3xf32> {
%1 = stablehlo.custom_call @NotSharding(%arg0) : (tensor<3xf32>) -> tensor<3xf32>
return %1 : tensor<3xf32>
}
// CHECK: custom_call @NotSharding
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.mlir
+535
View File
@@ -0,0 +1,535 @@
// 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: stablehlo-quant-opt %s -split-input-file \
// RUN: -tf-stablehlo-replace-stablehlo-ops-in-main-function-with-xla-call-module-ops \
// RUN: | FileCheck %s
// Modules with "main" or "serving_default" should properly run this pass and
// convert subgraphs into XLACallModuleOp.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_1
// CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1x3xf32>
// CHECK: return
// CHECK: }
// CHECK: func private @_stablehlo_main_0
// CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<3x64xf32>
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1x64xf32>
// CHECK: return
// CHECK: }
func.func @main(%arg0: tensor<1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x64xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1x3xf32>
%2:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%3 = "tf.XlaCallModule"(%2#0, %0, %1) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "{{.*}}", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64} : (tensor<1x1024xf32>, tensor<1024x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
%4:4 = "tf.CustomAggregator"(%3) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%5 = stablehlo.constant dense<1.000000e+03> : tensor<3x64xf32>
%6 = stablehlo.constant dense<1.000000e+03> : tensor<1x64xf32>
%7:4 = "tf.CustomAggregator"(%4#0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%8 = "tf.XlaCallModule"(%7#0, %5, %6) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1", _stablehlo_version = "{{.*}}", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64} : (tensor<1x3xf32>, tensor<3x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32>
%9:4 = "tf.CustomAggregator"(%6) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x64xf32>) -> (tensor<1x64xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
return %9#0 : tensor<1x64xf32>
}
// CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<{{.*}}>, #tf_type.shape<{{.*}}>], {{.*}}, module = "", platforms = ["CPU", "TPU"], use_shardy_partitioner = false, version = 9 : i64}> {_entry_function = @_stablehlo_main_1
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE_0:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "{{.*}}", _tfl_quant_trait = "fully_quantizable"}
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE_0]])
// CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<{{.*}}>, #tf_type.shape<{{.*}}>], {{.*}}, module = "", platforms = ["CPU", "TPU"], use_shardy_partitioner = false, version = 9 : i64}> {_entry_function = @_stablehlo_main_0
// CHECK: %[[CUSTOM_AGGREGATOR_2:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[CUSTOM_AGGREGATOR_1]]) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE_1:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_2]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1", _stablehlo_version = "{{.*}}", _tfl_quant_trait = "fully_quantizable"}
// CHECK: %[[CUSTOM_AGGREGATOR_3:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE_1:.*]])
// CHECK: return %[[CUSTOM_AGGREGATOR_3]] : tensor<1x64xf32>
// CHECK: }
// CHECK: @composite_dot_general_fn_1
// CHECK-NOT: tf_quant.composite_function
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
%1 = stablehlo.add %0, %arg2 : tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// CHECK: @composite_dot_general_with_relu_fn_1
// CHECK-NOT: tf_quant.composite_function
func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<1x3xf32>, %arg1: tensor<3x64xf32>, %arg2: tensor<1x64xf32>) -> tensor<1x64xf32> {
%0 = stablehlo.constant dense<0.000000e+00> : tensor<1x64xf32>
%1 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf32>, tensor<3x64xf32>) -> tensor<1x64xf32>
%2 = stablehlo.add %1, %arg2 : tensor<1x64xf32>
%3 = stablehlo.maximum %2, %0 : tensor<1x64xf32>
return %3 : tensor<1x64xf32>
}
}
// -----
// Tests that the subgraph in serving_default excluding the tf.Identity is
// converted to a single XlaCallModuleOp.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1654 : i32}, tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_0(%arg0: tensor<i32>, %arg1: tensor<1x1024xf32>)
// CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<0.134728625> : tensor<1x3xf32>
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<-1.280000e+02> : tensor<1x1024xf32>
// CHECK: %[[CONSTANT_2:.*]] = stablehlo.constant dense<0.003921567> : tensor<1x1024xf32>
// CHECK: %[[DIVIDE:.*]] = stablehlo.divide %arg1, %[[CONSTANT_2]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DIVIDE]], %[[CONSTANT_1]]
// CHECK return %[[ADD]]
// CHECK: }
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x1024xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<0.134728625> : tensor<1x3xf32>
%1 = stablehlo.constant dense<-1.280000e+02> : tensor<1x1024xf32>
%2 = stablehlo.constant dense<0.003921567> : tensor<1x1024xf32>
%3 = stablehlo.divide %arg0, %2 : tensor<1x1024xf32>
%4 = stablehlo.add %3, %1 : tensor<1x1024xf32>
%5 = "tf.Identity"(%4) {device = ""} : (tensor<1x1024xf32>) -> tensor<1x1024xf32>
return %5 : tensor<1x1024xf32>
}
// CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]] = "tf.XlaCallModule"(%arg0) <{Sout = [#tf_type.shape<1x1024>], {{.*}}, module = "", platforms = ["CPU", "TPU"], use_shardy_partitioner = false, version = 9 : i64}> {_entry_function = @_stablehlo_main_0, _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _stablehlo_version = "{{.*}}"} : (tensor<1x1024xf32>) -> tensor<1x1024xf32>
// CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP]])
// CHECK: return %[[IDENTITY]]
// CHECK }
}
// -----
// Tests that the first stablehlo.constant is converted to XlaCallModuleOp.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_0
// CHECK: %[[CONSTANT:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: return %[[CONSTANT:.*]]
// CHECK: }
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%1:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%2 = "tf.XlaCallModule"(%1#0, %0) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "{{.*}}", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64} : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
%3:4 = "tf.CustomAggregator"(%2) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
return %3#0 : tensor<1x3xf32>
}
// CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}}, module = "", platforms = ["CPU", "TPU"], use_shardy_partitioner = false, version = 9 : i64}> {_entry_function = @_stablehlo_main_0, _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _stablehlo_version = "{{.*}}"}
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "{{.*}}"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: return %[[CUSTOM_AGGREGATOR_1]]
// CHECK: }
// CHECK: @composite_dot_general_fn_1
// CHECK-NOT: tf_quant.composite_function
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
// Tests to confirm that the StableHLO graph is not replaced if "main" or
// "serving_default" function is not in the module.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK-NOT: func private @_stablehlo_main_
// CHECK-LABEL: @random_name
func.func @random_name(%arg0: tensor<1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%1:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%2 = "tf.XlaCallModule"(%1#0, %0) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64} : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
%3:4 = "tf.CustomAggregator"(%2) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
return %3#0 : tensor<1x3xf32>
}
// CHECK: %[[CONSTANT:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR:.*]], %[[XLA_CALL_MODULE_EXTRACTED_FROM_SUBGRAPH:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: return %[[CUSTOM_AGGREGATOR_1]]
// CHECK: }
// CHECK: @composite_dot_general_fn_1
// CHECK: tf_quant.composite_function
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
// Tests where StableHLO graph in main has a small constant to be duplicated.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_1(%arg0: tensor<i32>) -> tensor<1024x3xf32> attributes {_from_xla_call_module}
// CHECK: %[[CONSTANT1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: return %[[CONSTANT1:.*]]
// CHECK: }
// CHECK: func private @_stablehlo_main_0(%arg0: tensor<i32>
// CHECK-SAME: %[[INPUT1:.*]]: tensor<1024x3xf32>, %[[INPUT2:.*]]: tensor<1024x3xf32>
// CHECK: %[[CONSTANT2:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[INPUT1]], %[[CONSTANT2]] : tensor<1024x3xf32>
// CHECK: %[[MUL:.*]] = stablehlo.multiply %[[INPUT1]], %[[INPUT2]] : tensor<1024x3xf32>
// CHECK: return %[[ADD]], %[[MUL]]
// CHECK: }
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<1024x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1024x3xf32> {tf_saved_model.index_path = ["output1"]}, tensor<1024x3xf32> {tf_saved_model.index_path = ["output2"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%1:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x1024xf32>) -> (tensor<1024x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%2 = "tf.XlaCallModule"(%1#0, %0) {Sout = [#tf_type.shape<1024x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1024x1024xf32>, tensor<1024x3xf32>) -> tensor<1024x3xf32>
%3:4 = "tf.CustomAggregator"(%2) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x3xf32>) -> (tensor<1024x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%4 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%5 = stablehlo.add %3#0, %4 : tensor<1024x3xf32>
%6 = stablehlo.multiply %3#0, %0 : tensor<1024x3xf32>
return %5, %6 : tensor<1024x3xf32>, tensor<1024x3xf32>
}
// CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_1]], %[[SUBGRAPH_1]]) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0"
// CHECK: %[[CUSTOM_AGGREGATOR_2:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: %[[SUBGRAPH_2:.*]]:2 = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_2]], %[[SUBGRAPH_1]]) <{Sout = [#tf_type.shape<1024x3>, #tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @_stablehlo_main_0
// CHECK: return %[[SUBGRAPH_2]]#0, %[[SUBGRAPH_2]]#1
// CHECK: }
// CHECK: @composite_dot_general_fn_1
// CHECK-NOT: tf_quant.composite_function
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
// -----
// Tests where StableHLO graph in main has branches.
// This test makes sure tracing won't stop at op (%1) with multiple uses.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_1(%arg0: tensor<i32>) -> tensor<3x11xf32>
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<3x11xf32>
// CHECK: return %[[CONSTANT_1:.*]]
// CHECK: }
// CHECK: func private @_stablehlo_main_0
// CHECK-SAME: (%arg0: tensor<i32>, %[[INPUT_1:.*]]: tensor<3x11xf32>)
// CHECK-SAME: -> tensor<3x11xf32>
// CHECK: %[[CONSTANT_2:.*]] = stablehlo.constant dense<1.000000e+01> : tensor<3x11xf32>
// CHECK: %[[ADD:.*]] = stablehlo.add %[[INPUT_1]], %[[CONSTANT_2]] : tensor<3x11xf32>
// CHECK: %[[MUL:.*]] = stablehlo.multiply %[[ADD]], %[[CONSTANT_2]] : tensor<3x11xf32>
// CHECK: return %[[MUL]]
// CHECK: }
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<3x3xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<3x11xf32> {tf_saved_model.index_path = ["output1"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<3x11xf32>
// %1 is large enough that it won't be duplicated.
%1 = stablehlo.constant dense<1.000000e+01> : tensor<3x11xf32>
%2:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%3 = "tf.XlaCallModule"(%2#0, %0) {Sout = [#tf_type.shape<3x11>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<3x3xf32>, tensor<3x11xf32>) -> tensor<3x11xf32>
%4:4 = "tf.CustomAggregator"(%3) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<3x11xf32>) -> (tensor<3x11xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%5 = stablehlo.add %4#0, %1 : tensor<3x11xf32>
%6 = stablehlo.multiply %5, %1 : tensor<3x11xf32>
return %6 : tensor<3x11xf32>
}
// CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<3x11>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_1]], %[[SUBGRAPH_1]]) <{Sout = [#tf_type.shape<3x11>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0"
// CHECK: %[[CUSTOM_AGGREGATOR_2:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: %[[SUBGRAPH_2:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_2]]) <{Sout = [#tf_type.shape<3x11>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_0
// CHECK: return %[[SUBGRAPH_2]]
// CHECK: }
// CHECK: @composite_dot_general_fn_1
// CHECK-NOT: tf_quant.composite_function
func.func private @composite_dot_general_fn_1(%arg0: tensor<3x3xf32>, %arg1: tensor<3x11xf32>) -> tensor<3x11xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<3x3xf32>, tensor<3x11xf32>) -> tensor<3x11xf32>
return %0 : tensor<3x11xf32>
}
}
// -----
// Tests where StableHLO graph in main has dead end.
// This test makes sure tracing will include the dead end from the op in the
// same sub graph:
// stablehlo.add and %0 along with its dead end branch are in the same sub
// graph.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func.func private @_stablehlo_main_1(%arg0: tensor<i32>) -> tensor<1024x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<2.000000e+03> : tensor<1024x3xf32>
// CHECK: return %[[CONSTANT_0]]
// CHECK: }
// CHECK: func.func private @_stablehlo_main_0(%arg0: tensor<i32>, %[[ARG_1:.*]]: tensor<1024x3xf32>) -> tensor<1024x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[CONSTANT_2:.*]] = stablehlo.constant dense<5.000000e+01> : tensor<1024x3xf32>
// CHECK: %[[CONSTANT_3:.*]] = stablehlo.constant dense<4.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[REMAINDER:.*]] = stablehlo.remainder %[[CONSTANT_3]], %[[CONSTANT_1]] : tensor<1024x3xf32>
// CHECK: %[[COMPARE:.*]] = stablehlo.compare EQ, %[[REMAINDER]], %[[CONSTANT_2]], NOTYPE : (tensor<1024x3xf32>, tensor<1024x3xf32>) -> tensor<1024x3xi1>
// CHECK: stablehlo.custom_call @shape_assertion(%[[COMPARE]]) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<1024x3xi1>) -> ()
// CHECK: %[[ADD:.*]] = stablehlo.add %[[ARG_1]], %[[CONSTANT_3]]
// CHECK: return %[[ADD]]
// CHECK: }
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<1024x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1024x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<4.000000e+03> : tensor<1024x3xf32>
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%2 = stablehlo.constant dense<5.000000e+01> : tensor<1024x3xf32>
%3 = stablehlo.remainder %0, %1 : tensor<1024x3xf32>
%4 = stablehlo.compare EQ, %3, %2, NOTYPE : (tensor<1024x3xf32>, tensor<1024x3xf32>) -> tensor<1024x3xi1>
stablehlo.custom_call @shape_assertion(%4) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<1024x3xi1>) -> ()
%5 = stablehlo.constant dense<2.000000e+03> : tensor<1024x3xf32>
%6:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x1024xf32>) -> (tensor<1024x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%7 = "tf.XlaCallModule"(%6#0, %5) {Sout = [#tf_type.shape<1024x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1024x1024xf32>, tensor<1024x3xf32>) -> tensor<1024x3xf32>
%8:4 = "tf.CustomAggregator"(%7) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x3xf32>) -> (tensor<1024x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%9 = stablehlo.add %8#0, %0 : tensor<1024x3xf32>
return %9 : tensor<1024x3xf32>
}
// CHECK: %[[SUBGRAPH_0:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[SUBGRAPH_0]]) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_1]]) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @_stablehlo_main_0
// CHECK: return %[[SUBGRAPH_1]] : tensor<1024x3xf32>
// CHECK: }
}
// -----
// Tests where StableHLO graph in main has branch.
// This test makes sure the branch will not be added to subgraph when it reaches
// a tf op:
// stablehlo.add and %0 are not in the same subgraph.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func.func private @_stablehlo_main_2(%arg0: tensor<i32>) -> (tensor<1024x3xf32>, tensor<1024x3xf32>) attributes {_from_xla_call_module} {
// CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<4.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[REMAINDER:.*]] = stablehlo.remainder %[[CONSTANT_0]], %[[CONSTANT_1]] : tensor<1024x3xf32>
// CHECK: return %[[CONSTANT_0]], %[[REMAINDER]]
// CHECK: }
// CHECK: func.func private @_stablehlo_main_1(%arg0: tensor<i32>) -> tensor<1024x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[CONSTANT_2:.*]] = stablehlo.constant dense<2.000000e+03> : tensor<1024x3xf32>
// CHECK: return %[[CONSTANT_2]] : tensor<1024x3xf32>
// CHECK: }
// CHECK: func.func private @_stablehlo_main_0(%arg0: tensor<i32>, %[[ARG_1:.*]]: tensor<1024x3xf32>, %[[ARG_2:.*]]: tensor<1024x3xf32>) -> tensor<1024x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[ADD:.*]] = stablehlo.add %[[ARG_1]], %[[ARG_2]]
// CHECK: return %[[ADD]]
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<1024x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1024x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<4.000000e+03> : tensor<1024x3xf32>
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%2 = stablehlo.remainder %0, %1 : tensor<1024x3xf32>
%3 = "tf.Identity"(%2) {device = ""} : (tensor<1024x3xf32>) -> tensor<1024x3xf32>
%4 = stablehlo.constant dense<2.000000e+03> : tensor<1024x3xf32>
%5:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x1024xf32>) -> (tensor<1024x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%6 = "tf.XlaCallModule"(%5#0, %4) {Sout = [#tf_type.shape<1024x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1024x1024xf32>, tensor<1024x3xf32>) -> tensor<1024x3xf32>
%7:4 = "tf.CustomAggregator"(%6) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x3xf32>) -> (tensor<1024x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%8 = stablehlo.add %7#0, %0 : tensor<1024x3xf32>
return %8 : tensor<1024x3xf32>
}
// CHECK: %[[SUBGRAPH_0:.*]]:2 = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>, #tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_2
// CHECK: %[[IDENTIFY:.*]] = "tf.Identity"(%[[SUBGRAPH_0]]#1) {device = ""} : (tensor<1024x3xf32>) -> tensor<1024x3xf32>
// CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[SUBGRAPH_1]]) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "1.0.0"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: %[[SUBGRAPH_2:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_1]], %[[SUBGRAPH_0]]#0) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @_stablehlo_main_0
// CHECK: return %[[SUBGRAPH_2]] : tensor<1024x3xf32>
// CHECK: }
}
// -----
// Tests where StableHLO graph in main has dead end.
// This test checks tracing will stop if the dead end is too deep (>5):
// stablehlo.add and %0 are not in the same subgraph.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func.func private @_stablehlo_main_1(%arg0: tensor<i32>) -> (tensor<1024x3xf32>, tensor<1024x3xf32>) attributes {_from_xla_call_module} {
// CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<4.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
// CHECK: %[[CONSTANT_2:.*]] = stablehlo.constant dense<5.000000e+01> : tensor<1024x3xf32>
// CHECK: %[[REMAINDER_0:.*]] = stablehlo.remainder %[[CONSTANT_0]], %[[CONSTANT_1]] : tensor<1024x3xf32>
// CHECK: %[[REMAINDER_1:.*]] = stablehlo.remainder %[[REMAINDER_0]], %[[CONSTANT_1]] : tensor<1024x3xf32>
// CHECK: %[[REMAINDER_2:.*]] = stablehlo.remainder %[[REMAINDER_1]], %[[CONSTANT_1]] : tensor<1024x3xf32>
// CHECK: %[[REMAINDER_3:.*]] = stablehlo.remainder %[[REMAINDER_2]], %[[CONSTANT_1]] : tensor<1024x3xf32>
// CHECK: %[[COMPARE:.*]] = stablehlo.compare EQ, %[[REMAINDER_3]], %[[CONSTANT_2]], NOTYPE : (tensor<1024x3xf32>, tensor<1024x3xf32>) -> tensor<1024x3xi1>
// CHECK: stablehlo.custom_call @shape_assertion(%[[COMPARE]]) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<1024x3xi1>) -> ()
// CHECK: %[[CONSTANT_3:.*]] = stablehlo.constant dense<2.000000e+03> : tensor<1024x3xf32>
// CHECK: return %[[CONSTANT_0]], %[[CONSTANT_3]]
// CHECK: }
// CHECK: func.func private @_stablehlo_main_0(%arg0: tensor<i32>, %[[ARG_1:.*]]: tensor<1024x3xf32>, %[[ARG_2:.*]]: tensor<1024x3xf32>) -> tensor<1024x3xf32> attributes {_from_xla_call_module} {
// CHECK: %[[ADD:.*]] = stablehlo.add %[[ARG_1]], %[[ARG_2]]
// CHECK: return %[[ADD]]
// CHECK: }
// CHECK: @serving_default
func.func @serving_default(%arg0: tensor<1024x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1024x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<4.000000e+03> : tensor<1024x3xf32>
%1 = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32>
%2 = stablehlo.constant dense<5.000000e+01> : tensor<1024x3xf32>
%3 = stablehlo.remainder %0, %1 : tensor<1024x3xf32>
%4 = stablehlo.remainder %3, %1 : tensor<1024x3xf32>
%5 = stablehlo.remainder %4, %1 : tensor<1024x3xf32>
%6 = stablehlo.remainder %5, %1 : tensor<1024x3xf32>
%7 = stablehlo.compare EQ, %6, %2, NOTYPE : (tensor<1024x3xf32>, tensor<1024x3xf32>) -> tensor<1024x3xi1>
stablehlo.custom_call @shape_assertion(%7) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<1024x3xi1>) -> ()
%8 = stablehlo.constant dense<2.000000e+03> : tensor<1024x3xf32>
%9:4 = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x1024xf32>) -> (tensor<1024x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%10 = "tf.XlaCallModule"(%9#0, %8) {Sout = [#tf_type.shape<1024x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "{{.*}}", _tfl_quant_trait = "fully_quantizable", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1024x1024xf32>, tensor<1024x3xf32>) -> tensor<1024x3xf32>
%11:4 = "tf.CustomAggregator"(%10) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<1024x3xf32>) -> (tensor<1024x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
%12 = stablehlo.add %11#0, %0 : tensor<1024x3xf32>
return %12 : tensor<1024x3xf32>
}
// CHECK: %[[SUBGRAPH_0:.*]]:2 = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>, #tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[SUBGRAPH_0]]#1) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_version = "{{.*}}"
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])
// CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_1]], %[[SUBGRAPH_0]]#0) <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @_stablehlo_main_0
// CHECK: return %[[SUBGRAPH_1]] : tensor<1024x3xf32>
// CHECK: }
}
// -----
// main function contains PartitionedCall and StatefulPartitionedCall ops which
// is used to preserve aliased functions. This test make sure stablehlo ops in
// each PartitionedCall functions are lifted.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_2
// CHECK: stablehlo.multiply %arg1, %arg2 : tensor<3x3xf32>
// CHECK: return
// CHECK: }
// CHECK: func private @_stablehlo_main_1
// CHECK: stablehlo.add %arg1, %arg2 : tensor<3x3xf32>
// CHECK: return
// CHECK: }
// CHECK: func private @_stablehlo_main_0
// CHECK: stablehlo.constant dense<1.000000e+03> : tensor<3x3xf32>
// CHECK: stablehlo.constant dense<2.000000e+03> : tensor<3x3xf32>
// CHECK: return
// CHECK: }
func.func @main() -> (tensor<3x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<3x3xf32>
%1 = stablehlo.constant dense<2.000000e+03> : tensor<3x3xf32>
%2 = "tf.StatefulPartitionedCall"(%0, %1) <{
config = "", config_proto = "", executor_type = "", f = @some_func
}> {
_collective_manager_ids = [], device = ""
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%3 = "tf.PartitionedCall"(%2, %1) <{
config = "", config_proto = "", executor_type = "", f = @some_other_func
}> {
_collective_manager_ids = [], device = ""
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
return %3 : tensor<3x3xf32>
}
// CHECK: func.func @main
// CHECK: %[[INPUT:.*]]:3 = "tf.XlaCallModule"()
// CHECK-SAME: _entry_function = @_stablehlo_main_0
// CHECK: %[[ADD:.*]] = "tf.StatefulPartitionedCall"(%[[INPUT]]#1, %[[INPUT]]#2)
// CHECK-SAME: f = @some_func
// CHECK: "tf.PartitionedCall"(%[[ADD]], %[[INPUT]]#0)
// CHECK-SAME: f = @some_other_func
// CHECK: return
func.func private @some_func(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> attributes {tf._noinline = true} {
%0 = stablehlo.add %arg0, %arg1 : tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
// CHECK: func.func private @some_func
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @_stablehlo_main_1
// CHECK: return
func.func private @some_other_func(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> attributes {tf._noinline = true} {
%0 = stablehlo.multiply %arg0, %arg1 : tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
// CHECK: func.func private @some_other_func
// CHECK: tf.XlaCallModule
// CHECK-SAME: _entry_function = @_stablehlo_main_2
// CHECK: return
}
// -----
// main function contains PartitionedCall and StatefulPartitionedCall ops that
// have _no_xla_call_module attribute. This test make sure stablehlo ops in
// their functions are NOT lifted.
// CHECK-LABEL: functions_with_no_xla_call_module
module @functions_with_no_xla_call_module attributes {tf_saved_model.semantics} {
// CHECK: func private @_stablehlo_main_0
// CHECK: stablehlo.constant dense<1.000000e+03> : tensor<3x3xf32>
// CHECK: stablehlo.constant dense<2.000000e+03> : tensor<3x3xf32>
func.func @main() -> (tensor<3x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%0 = stablehlo.constant dense<1.000000e+03> : tensor<3x3xf32>
%1 = stablehlo.constant dense<2.000000e+03> : tensor<3x3xf32>
%2 = "tf.StatefulPartitionedCall"(%0, %1) <{
config = "", config_proto = "", executor_type = "", f = @some_func
}> : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
%3 = "tf.PartitionedCall"(%2, %1) <{f = @some_other_func}> : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
return %3 : tensor<3x3xf32>
}
// CHECK: func.func @main
// CHECK: %[[INPUT:.*]]:3 = "tf.XlaCallModule"()
// CHECK-SAME: _entry_function = @_stablehlo_main_0
// CHECK: %[[ADD:.*]] = "tf.StatefulPartitionedCall"(%[[INPUT]]#1, %[[INPUT]]#2)
// CHECK-SAME: f = @some_func
// CHECK: "tf.PartitionedCall"(%[[ADD]], %[[INPUT]]#0)
// CHECK-SAME: f = @some_other_func
// CHECK: return
func.func private @some_func(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> attributes {_no_xla_call_module} {
%0 = stablehlo.add %arg0, %arg1 : tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
// CHECK: func.func private @some_func
// CHECK-NOT: tf.XlaCallModule
func.func private @some_other_func(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> attributes {_no_xla_call_module} {
%0 = stablehlo.multiply %arg0, %arg1 : tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
// CHECK: func.func private @some_other_func
// CHECK-NOT: tf.XlaCallModule
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/restore_function_name.mlir
+66
View File
@@ -0,0 +1,66 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-restore-function-name | FileCheck %s
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1646 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: @serving_default
// CHECK-SAME: %[[ARG0:[^:[:space:]]+]]
// CHECK-SAME: %[[ARG1:[^:[:space:]]+]]
func.func private @serving_default(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> {
%0 = "tf.XlaCallModule"(%arg0, %arg1) {Sout = [#tf_type.shape<1x3>], _entry_function = @main, _original_entry_function = "composite_dot_general_fn_1", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn_1
// CHECK-SAME: _original_entry_function = "composite_dot_general_fn_1"
// CHECK: return %[[CALL]]
}
// CHECK: @composite_dot_general_fn_1
// CHECK-SAME: %[[ARG2:[^:[:space:]]+]]
// CHECK-SAME: %[[ARG3:[^:[:space:]]+]]
func.func private @main(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
// CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK: return %[[DOT]]
}
}
// -----
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1646 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: @serving_default
// CHECK-SAME: %[[ARG0:[^:[:space:]]+]]
// CHECK-SAME: %[[ARG1:[^:[:space:]]+]]
func.func private @serving_default(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> {
%0 = "tf.XlaCallModule"(%arg0, %arg1) {Sout = [#tf_type.shape<1x3>], _entry_function = @main, _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]])
// CHECK-SAME: _entry_function = @main
// CHECK-NOT: _original_entry_function = "composite_dot_general_fn_1"
// CHECK: return %[[CALL]]
}
// CHECK: @main
// CHECK-NOT: @composite_dot_general_fn_1
// CHECK-SAME: %[[ARG2:[^:[:space:]]+]]
// CHECK-SAME: %[[ARG3:[^:[:space:]]+]]
func.func private @main(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
// CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
// CHECK: return %[[DOT]]
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/shape_cstr_legalize_to_hlo.mlir
+124
View File
@@ -0,0 +1,124 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-convert-shape-to-stablehlo-with-constraints --verify-diagnostics | FileCheck %s
// CHECK-LABEL: func.func @shape_cstr_broadcastable
func.func @shape_cstr_broadcastable(%arg0: tensor<2xindex>, %arg1: tensor<2xindex>) {
%0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<2xindex>, tensor<2xindex>
shape.assuming %0 {
}
func.return
// CHECK: %[[DIMS1:.*]] = builtin.unrealized_conversion_cast %arg0 : tensor<2xindex> to tensor<2xi32>
// CHECK-NEXT: %[[DIMS2:.*]] = builtin.unrealized_conversion_cast %arg1 : tensor<2xindex> to tensor<2xi32>
// CHECK-NEXT: %[[ONES:.*]] = stablehlo.constant dense<1> : tensor<2xi32>
// CHECK-NEXT: %[[DIMS1_IS_1:.*]] = stablehlo.compare EQ, %[[DIMS1]], %[[ONES:.*]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[DIMS2_IS_1:.*]] = stablehlo.compare EQ, %[[DIMS2]], %[[ONES:.*]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[EITHER_DIM_IS_1:.*]] = stablehlo.or %[[DIMS1_IS_1]], %[[DIMS2_IS_1]] : tensor<2xi1>
// CHECK-NEXT: %[[DIMS_EQ:.*]] = stablehlo.compare EQ, %[[DIMS1]], %[[DIMS2]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[DIMS_BROADCASTABLE:.*]] = stablehlo.or %[[EITHER_DIM_IS_1]], %[[DIMS_EQ]] : tensor<2xi1>
// CHECK-NEXT: %[[TRUE:.*]] = stablehlo.constant dense<true> : tensor<1xi1>
// CHECK-NEXT: %[[DIM1_BROADCASTABLE:.*]] = stablehlo.slice %[[DIMS_BROADCASTABLE]] [0:1] : (tensor<2xi1>) -> tensor<1xi1>
// CHECK-NEXT: %[[BROADCASTABLE_TEMP:.*]] = stablehlo.and %[[TRUE]], %[[DIM1_BROADCASTABLE]] : tensor<1xi1>
// CHECK-NEXT: %[[DIM2_BROADCASTABLE:.*]] = stablehlo.slice %[[DIMS_BROADCASTABLE]] [1:2] : (tensor<2xi1>) -> tensor<1xi1>
// CHECK-NEXT: %[[ALL_BROADCASTABLE:.*]] = stablehlo.and %[[BROADCASTABLE_TEMP]], %[[DIM2_BROADCASTABLE]] : tensor<1xi1>
// CHECK-NEXT: %[[ALL_BROADCASTABLE_SCALAR:.*]] = stablehlo.reshape %[[ALL_BROADCASTABLE]] : (tensor<1xi1>) -> tensor<i1>
// CHECK-NEXT: stablehlo.custom_call @shape_assertion(%[[ALL_BROADCASTABLE_SCALAR]]) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<i1>) -> ()
// CHECK-NEXT: %[[WITNESS:.*]] = shape.const_witness true
// CHECK-NEXT: shape.assuming %[[WITNESS]] {
// CHECK-NEXT: }
// CHECK-NEXT: return
}
// -----
// CHECK-LABEL: func @shape_cstr_broadcastable_different_dims_1
func.func @shape_cstr_broadcastable_different_dims_1(%arg0: tensor<2xindex>, %arg1: tensor<1xindex>) {
%0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<2xindex>, tensor<1xindex>
shape.assuming %0 {
}
func.return
// CHECK: %[[DIMS1:.*]] = builtin.unrealized_conversion_cast %arg0 : tensor<2xindex> to tensor<2xi32>
// CHECK-NEXT: %[[DIMS2:.*]] = builtin.unrealized_conversion_cast %arg1 : tensor<1xindex> to tensor<1xi32>
// CHECK-NEXT: %[[PAD:.*]] = stablehlo.constant dense<1> : tensor<1xi32>
// CHECK-NEXT: %[[DIMS2_PAD:.*]] = stablehlo.concatenate %[[PAD]], %[[DIMS2]], dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
// CHECK-NEXT: %[[ONES:.*]] = stablehlo.constant dense<1> : tensor<2xi32>
// CHECK-NEXT: %[[DIMS1_IS_1:.*]] = stablehlo.compare EQ, %[[DIMS1]], %[[ONES:.*]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[DIMS2_IS_1:.*]] = stablehlo.compare EQ, %[[DIMS2_PAD]], %[[ONES:.*]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[EITHER_DIM_IS_1:.*]] = stablehlo.or %[[DIMS1_IS_1]], %[[DIMS2_IS_1]] : tensor<2xi1>
// CHECK-NEXT: %[[DIMS_EQ:.*]] = stablehlo.compare EQ, %[[DIMS1]], %[[DIMS2_PAD]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[DIMS_BROADCASTABLE:.*]] = stablehlo.or %[[EITHER_DIM_IS_1]], %[[DIMS_EQ]] : tensor<2xi1>
// CHECK-NEXT: %[[TRUE:.*]] = stablehlo.constant dense<true> : tensor<1xi1>
// CHECK-NEXT: %[[DIM1_BROADCASTABLE:.*]] = stablehlo.slice %[[DIMS_BROADCASTABLE]] [0:1] : (tensor<2xi1>) -> tensor<1xi1>
// CHECK-NEXT: %[[BROADCASTABLE_TEMP:.*]] = stablehlo.and %[[TRUE]], %[[DIM1_BROADCASTABLE]] : tensor<1xi1>
// CHECK-NEXT: %[[DIM2_BROADCASTABLE:.*]] = stablehlo.slice %[[DIMS_BROADCASTABLE]] [1:2] : (tensor<2xi1>) -> tensor<1xi1>
// CHECK-NEXT: %[[ALL_BROADCASTABLE:.*]] = stablehlo.and %[[BROADCASTABLE_TEMP]], %[[DIM2_BROADCASTABLE]] : tensor<1xi1>
// CHECK-NEXT: %[[ALL_BROADCASTABLE_SCALAR:.*]] = stablehlo.reshape %[[ALL_BROADCASTABLE]] : (tensor<1xi1>) -> tensor<i1>
// CHECK-NEXT: stablehlo.custom_call @shape_assertion(%[[ALL_BROADCASTABLE_SCALAR]]) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<i1>) -> ()
// CHECK-NEXT: %[[WITNESS:.*]] = shape.const_witness true
// CHECK-NEXT: shape.assuming %[[WITNESS]] {
// CHECK-NEXT: }
// CHECK-NEXT: return
}
// -----
// CHECK-LABEL: func @shape_cstr_broadcastable_different_dims_2
func.func @shape_cstr_broadcastable_different_dims_2(%arg0: tensor<1xindex>, %arg1: tensor<2xindex>) {
%0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<1xindex>, tensor<2xindex>
shape.assuming %0 {
}
func.return
// CHECK: %[[DIMS1:.*]] = builtin.unrealized_conversion_cast %arg0 : tensor<1xindex> to tensor<1xi32>
// CHECK-NEXT: %[[DIMS2:.*]] = builtin.unrealized_conversion_cast %arg1 : tensor<2xindex> to tensor<2xi32>
// CHECK-NEXT: %[[PAD:.*]] = stablehlo.constant dense<1> : tensor<1xi32>
// CHECK-NEXT: %[[DIMS1_PAD:.*]] = stablehlo.concatenate %[[PAD]], %[[DIMS1]], dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
// CHECK-NEXT: %[[ONES:.*]] = stablehlo.constant dense<1> : tensor<2xi32>
// CHECK-NEXT: %[[DIMS1_IS_1:.*]] = stablehlo.compare EQ, %[[DIMS1_PAD]], %[[ONES:.*]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[DIMS2_IS_1:.*]] = stablehlo.compare EQ, %[[DIMS2]], %[[ONES:.*]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[EITHER_DIM_IS_1:.*]] = stablehlo.or %[[DIMS1_IS_1]], %[[DIMS2_IS_1]] : tensor<2xi1>
// CHECK-NEXT: %[[DIMS_EQ:.*]] = stablehlo.compare EQ, %[[DIMS1_PAD]], %[[DIMS2]] : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1>
// CHECK-NEXT: %[[DIMS_BROADCASTABLE:.*]] = stablehlo.or %[[EITHER_DIM_IS_1]], %[[DIMS_EQ]] : tensor<2xi1>
// CHECK-NEXT: %[[TRUE:.*]] = stablehlo.constant dense<true> : tensor<1xi1>
// CHECK-NEXT: %[[DIM1_BROADCASTABLE:.*]] = stablehlo.slice %[[DIMS_BROADCASTABLE]] [0:1] : (tensor<2xi1>) -> tensor<1xi1>
// CHECK-NEXT: %[[BROADCASTABLE_TEMP:.*]] = stablehlo.and %[[TRUE]], %[[DIM1_BROADCASTABLE]] : tensor<1xi1>
// CHECK-NEXT: %[[DIM2_BROADCASTABLE:.*]] = stablehlo.slice %[[DIMS_BROADCASTABLE]] [1:2] : (tensor<2xi1>) -> tensor<1xi1>
// CHECK-NEXT: %[[ALL_BROADCASTABLE:.*]] = stablehlo.and %[[BROADCASTABLE_TEMP]], %[[DIM2_BROADCASTABLE]] : tensor<1xi1>
// CHECK-NEXT: %[[ALL_BROADCASTABLE_SCALAR:.*]] = stablehlo.reshape %[[ALL_BROADCASTABLE]] : (tensor<1xi1>) -> tensor<i1>
// CHECK-NEXT: stablehlo.custom_call @shape_assertion(%[[ALL_BROADCASTABLE_SCALAR]]) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<i1>) -> ()
// CHECK-NEXT: %[[WITNESS:.*]] = shape.const_witness true
// CHECK-NEXT: shape.assuming %[[WITNESS]] {
// CHECK-NEXT: }
// CHECK-NEXT: return
}
// -----
func.func @shape_cstr_broadcast_too_many_operands(%arg0: tensor<4xindex>, %arg1: tensor<4xindex>, %arg2: tensor<4xindex>) {
// expected-error@+1 {{failed to legalize operation 'shape.cstr_broadcastable' that was explicitly marked illegal}}
%0 = shape.cstr_broadcastable %arg0, %arg1, %arg2 : tensor<4xindex>, tensor<4xindex>, tensor<4xindex>
shape.assuming %0 {
}
func.return
}
// -----
func.func @shape_cstr_broadcastable_input_shape(%arg0: !shape.shape, %arg1: !shape.shape) {
// expected-error@+1 {{failed to legalize operation 'shape.cstr_broadcastable' that was explicitly marked illegal}}
%0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
shape.assuming %0 {
}
func.return
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/unfuse_mhlo_batch_norm.mlir
+44
View File
@@ -0,0 +1,44 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-unfuse-mhlo-batch-norm | FileCheck %s
// CHECK-LABEL: @unfuse_batch_norm
// CHECK-SAME: %[[X:[^:[:space:]]+]]
// CHECK-SAME: %[[SCALE:[^:[:space:]]+]]
// CHECK-SAME: %[[OFFSET:[^:[:space:]]+]]
// CHECK-SAME: %[[MEAN:[^:[:space:]]+]]
// CHECK-SAME: %[[VARIANCE:[^:[:space:]]+]]
func.func @unfuse_batch_norm(
%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: %[[STDDEV:.+]] = mhlo.sqrt %[[VARIANCE_EPS]] : tensor<256xf32>
// CHECK-DAG: %[[STDDEV_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[STDDEV]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[SCALE_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[SCALE]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[OFFSET_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[OFFSET]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK-DAG: %[[MEAN_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[MEAN]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>
// CHECK: %[[X_CENTER:.+]] = mhlo.subtract %[[X]], %[[MEAN_BCAST]] : tensor<4x256xf32>
// CHECK: %[[X_SCALED:.+]] = mhlo.multiply %[[X_CENTER]], %[[SCALE_BCAST]] : tensor<4x256xf32>
// CHECK: %[[X_NORMED:.+]] = mhlo.divide %[[X_SCALED]], %[[STDDEV_BCAST]] : tensor<4x256xf32>
// CHECK: %[[RESULT:.+]] = mhlo.add %[[X_NORMED]], %[[OFFSET_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>
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/unwrap_xla_call_module_op.mlir
+67
View File
@@ -0,0 +1,67 @@
// 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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-unwrap-xla-call-module-op | FileCheck %s
// Tests if XlaCallModule op without quantizable trait that calls function with
// '_from_xla_call_module' trait is unwrapped.
// Tests if XlaCallModule op with quantizable trait is not unwrapped.
// Tests if XlaCallModule op without quantizable trait that calls function
// without '_from_xla_call_module' trait is not unwrapped.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1682 : i32}, tf_saved_model.semantics} {
// CHECK-LABEL: @main_00
// CHECK: %[[ARG0:.*]]: tensor<10x1x1024xf32>
func.func private @main_00(%arg0: tensor<10x1x1024xf32>) -> tensor<6x5xf32> attributes {tf._original_func_name = "main_0"} {
%0 = "tf.Const"() <{value = dense<1.000000e+00> : tensor<10x1024x3xf32>}> : () -> tensor<10x1024x3xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<10x1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
%2 = "tf.XlaCallModule"(%1) <{Sout = [#tf_type.shape<3x10>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @main_0, _stablehlo_version = "1.0.0", _stablehlo_module_attrs = {}, device = ""} : (tensor<10x1x3xf32>) -> tensor<3x10xf32>
%3 = "tf.XlaCallModule"(%2) <{Sout = [#tf_type.shape<6x5>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @main_1, _stablehlo_version = "1.0.0", _stablehlo_module_attrs = {}, device = ""} : (tensor<3x10xf32>) -> tensor<6x5xf32>
return %3 : tensor<6x5xf32>
}
// CHECK: %[[CST:.*]] = "tf.Const"()
// CHECK-NEXT: %[[CALL1:.*]] = "tf.XlaCallModule"(%[[ARG0]], %[[CST]])
// CHECK-SAME: _entry_function = @composite_dot_general_fn_1
// CHECK-SAME: _tfl_quant_trait = "fully_quantizable"
// CHECK-NOT: "tf.XlaCallModule"
// CHECK-NEXT: %[[RESHAPE:.*]] = stablehlo.reshape %[[CALL1]] : (tensor<10x1x3xf32>) -> tensor<3x10xf32>
// CHECK-NEXT: %[[CALL2:.*]] = "tf.XlaCallModule"(%[[RESHAPE]])
// CHECK-SAME: _entry_function = @main_1
// CHECK-NOT: _tfl_quant_trait = "fully_quantizable"
// CHECK-NEXT: return %[[CALL2]]
// CHECK: @composite_dot_general_fn_1
func.func private @composite_dot_general_fn_1(%arg0: tensor<10x1x1024xf32>, %arg1: tensor<10x1024x3xf32>) -> tensor<10x1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
return %0 : tensor<10x1x3xf32>
}
// CHECK: %[[DOT:.*]] = stablehlo.dot_general
// CHECK-NEXT: return %[[DOT]]
// CHECK: @main_0
func.func private @main_0(%arg0: tensor<10x1x3xf32>) -> tensor<3x10xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.reshape %arg0 : (tensor<10x1x3xf32>) -> tensor<3x10xf32>
return %0 : tensor<3x10xf32>
}
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape
// CHECK-NEXT: return %[[RESHAPE]]
// CHECK: @main_1
func.func private @main_1(%arg0: tensor<3x10xf32>) -> tensor<6x5xf32> {
%0 = stablehlo.reshape %arg0 : (tensor<3x10xf32>) -> tensor<6x5xf32>
return %0 : tensor<6x5xf32>
}
// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape
// CHECK-NEXT: return %[[RESHAPE]]
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/xla_call_module_to_call.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: stablehlo-quant-opt %s -split-input-file -tf-stablehlo-xla-call-module-to-call | FileCheck %s
// -----
// Tests composite tf.XlaCallModule is converted to func.call.
module {
// CHECK-LABEL: func.func @main
func.func @main(%arg0: tensor<1x1024xf32>) -> tensor<1x3xf32> {
// CHECK: call @composite_dot_general_fn_1
// CHECK-SAME: (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
// CHECK-NOT: tf.XlaCallModule
%0 = "tf.Const"() <{value = dense<0.5> : tensor<1024x3xf32>}> : () -> tensor<1024x3xf32>
%2 = "tf.XlaCallModule"(%arg0, %0) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %2 : tensor<1x3xf32>
}
// CHECK-LABEL: func.func private @composite_dot_general_fn_1
// CHECK-SAME: -> tensor<1x3xf32>
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
}
tensorflow/compiler/mlir/quantization/stablehlo/tests/pipelines/process_nchw_tensor.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: stablehlo-quant-opt %s -stablehlo-process-nchw-tensor \
// RUN: -split-input-file -verify-diagnostics | FileCheck %s
// Tests that a `convolution(%activation, %weight)` with the activation tensor
// NCHW format is converted to NHWC convolution. Transpose ops are inserted to
// the activation and output to match the function signature. The weight
// constant is transposed.
// CHECK-LABEL: nchw_conv
// CHECK-SAME: %[[ARG:.+]]: tensor<1x8x4x4xf32>
func.func @nchw_conv(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32> {
%0 = stablehlo.constant() {value = dense<7.000000e+00> : tensor<8x8x3x3xf32>} : () -> tensor<8x8x3x3xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x8x4x4xf32>, tensor<8x8x3x3xf32>) -> tensor<1x8x4x4xf32>
return %2 : tensor<1x8x4x4xf32>
}
// CHECK-DAG: %[[CONST:.+]] = stablehlo.constant {{.*}} : tensor<3x3x8x8xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[CONST]]) 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<1x4x4x8xf32>, tensor<3x3x8x8xf32>) -> tensor<1x4x4x8xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[CONV]], dims = [0, 3, 1, 2] : (tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32>
// CHECK: return %[[TRANSPOSE_1]]
// -----
// Tests that a `add(convolution(%activation, %weight), %bias)` with the
// activation tensor of NCHW format is converted to NHWC convolution + add
// operation. Transpose ops are inserted to activations and outputs to match the
// function signature. Constants are also transposed accordingly.
// CHECK-LABEL: nchw_conv_with_bias_add
// CHECK-SAME: %[[ARG:.+]]: tensor<1x2x5x5xf32>
func.func @nchw_conv_with_bias_add(%arg0: tensor<1x2x5x5xf32>) -> tensor<1x4x5x5xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<4x2x3x3xf32>
%1 = stablehlo.constant dense<3.000000e+00> : tensor<1x4x5x5xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x2x5x5xf32>, tensor<4x2x3x3xf32>) -> tensor<1x4x5x5xf32>
%3 = stablehlo.add %2, %1 : tensor<1x4x5x5xf32>
return %3 : tensor<1x4x5x5xf32>
}
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} : tensor<3x3x2x4xf32>
// CHECK-DAG: %[[BIAS_CONST:.+]] = stablehlo.constant {{.*}} : tensor<1x5x5x4xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x2x5x5xf32>) -> tensor<1x5x5x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[WEIGHT_CONST]]) 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<1x5x5x2xf32>, tensor<3x3x2x4xf32>) -> tensor<1x5x5x4xf32>
// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[BIAS_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[ADD]], dims = [0, 3, 1, 2] : (tensor<1x5x5x4xf32>) -> tensor<1x4x5x5xf32>
// CHECK: return %[[TRANSPOSE_1]]
// -----
// Tests that a `add(convolution(%activation, %weight), %bias)` pattern with the
// activation tensor of NCHW format and non-constant bias is converted to NHWC
// convolution, but without the deferred transpose for `stablehlo.add`.
// Transpose ops are inserted to the activation and output of
// `stablehlo.convolution`. The weight constants is transposed.
// CHECK-LABEL: nchw_conv_with_nonconst_bias_add
// CHECK-SAME: %[[ARG_0:.+]]: tensor<1x2x5x5xf32>
// CHECK-SAME: %[[ARG_1:.+]]: tensor<1x4x5x5xf32>
func.func @nchw_conv_with_nonconst_bias_add(%arg0: tensor<1x2x5x5xf32>, %arg1: tensor<1x4x5x5xf32>) -> tensor<1x4x5x5xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<4x2x3x3xf32>
%1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x2x5x5xf32>, tensor<4x2x3x3xf32>) -> tensor<1x4x5x5xf32>
%2 = stablehlo.add %1, %arg1 : tensor<1x4x5x5xf32>
return %2 : tensor<1x4x5x5xf32>
}
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} : tensor<3x3x2x4xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG_0]], dims = [0, 2, 3, 1] : (tensor<1x2x5x5xf32>) -> tensor<1x5x5x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[WEIGHT_CONST]]) 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<1x5x5x2xf32>, tensor<3x3x2x4xf32>) -> tensor<1x5x5x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[CONV]], dims = [0, 3, 1, 2] : (tensor<1x5x5x4xf32>) -> tensor<1x4x5x5xf32>
// CHECK: %[[ADD:.+]] = stablehlo.add %[[TRANSPOSE_1]], %[[ARG_1]] : tensor<1x4x5x5xf32>
// CHECK: return %[[ADD]]
// -----
// Tests that a `reduce_window{max}(add(convolution(%activation, %weight), %bias), %init_value)`
// with the activation tensor of NCHW format is converted to NHWC convolution +
// add + reduce_window (with max) operation. Transpose ops are inserted to
// activation and the final result to match the function signature. Constants
// are also transposed accordingly.
// CHECK-LABEL: nchw_conv_with_bias_add_max_pool
// CHECK-SAME: %[[ARG:.+]]: tensor<1x2x5x5xf32>
func.func @nchw_conv_with_bias_add_max_pool(%arg0: tensor<1x2x5x5xf32>) -> tensor<1x4x2x2xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<4x2x3x3xf32>
%1 = stablehlo.constant dense<3.000000e+00> : tensor<1x4x5x5xf32>
%5 = stablehlo.constant dense<0xFF800000> : tensor<f32> // -inf
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x2x5x5xf32>, tensor<4x2x3x3xf32>) -> tensor<1x4x5x5xf32>
%3 = stablehlo.add %2, %1 : tensor<1x4x5x5xf32>
%4 = "stablehlo.reduce_window"(%3, %5) ({ // max pool
^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
%6 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
stablehlo.return %6 : tensor<f32>
}) {
window_dimensions = array<i64: 1, 1, 2, 2>,
window_strides = array<i64: 1, 1, 2, 2>
} : (tensor<1x4x5x5xf32>, tensor<f32>) -> tensor<1x4x2x2xf32>
return %4 : tensor<1x4x2x2xf32>
}
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} : tensor<3x3x2x4xf32>
// CHECK-DAG: %[[BIAS_CONST:.+]] = stablehlo.constant {{.*}} : tensor<1x5x5x4xf32>
// CHECK-DAG: %[[INIT_VALUE_CONST:.+]] = stablehlo.constant dense<0xFF800000> : tensor<f32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x2x5x5xf32>) -> tensor<1x5x5x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[WEIGHT_CONST]]) 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<1x5x5x2xf32>, tensor<3x3x2x4xf32>) -> tensor<1x5x5x4xf32>
// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[BIAS_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[REDUCE_WINDOW_MAX:.+]] = "stablehlo.reduce_window"(%[[ADD]], %[[INIT_VALUE_CONST:.+]])
// CHECK: <{window_dimensions = array<i64: 1, 2, 2, 1>, window_strides = array<i64: 1, 2, 2, 1>}>
// CHECK: stablehlo.maximum
// CHECK: (tensor<1x5x5x4xf32>, tensor<f32>) -> tensor<1x2x2x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[REDUCE_WINDOW_MAX]], dims = [0, 3, 1, 2] : (tensor<1x2x2x4xf32>) -> tensor<1x4x2x2xf32>
// CHECK: return %[[TRANSPOSE_1]]
// -----
// Tests that a `maximum(add(convolution(%activation, %weight), %bias), %zero)`
// with the activation tensor of NCHW format is converted to NHWC convolution +
// add + maximum operation. Transpose ops are inserted to the activation and the
// final output to match the function signature. Constants are also transpose-
// folded accordingly.
// CHECK-LABEL: nchw_conv_with_bias_add_relu
// CHECK-SAME: %[[ARG:.+]]: tensor<1x2x5x5xf32>
func.func @nchw_conv_with_bias_add_relu(%arg0: tensor<1x2x5x5xf32>) -> tensor<1x4x5x5xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<4x2x3x3xf32>
%5 = stablehlo.constant dense<0.000000e+00> : tensor<1x4x5x5xf32>
%1 = stablehlo.constant dense<3.000000e+00> : tensor<1x4x5x5xf32>
%2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x2x5x5xf32>, tensor<4x2x3x3xf32>) -> tensor<1x4x5x5xf32>
%3 = stablehlo.add %2, %1 : tensor<1x4x5x5xf32>
%4 = stablehlo.maximum %3, %5 : tensor<1x4x5x5xf32>
return %4 : tensor<1x4x5x5xf32>
}
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} : tensor<3x3x2x4xf32>
// CHECK-DAG: %[[ZERO_CONST:.+]] = stablehlo.constant {{.*}} : tensor<1x5x5x4xf32>
// CHECK-DAG: %[[BIAS_CONST:.+]] = stablehlo.constant {{.*}} : tensor<1x5x5x4xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x2x5x5xf32>) -> tensor<1x5x5x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[WEIGHT_CONST]]) 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<1x5x5x2xf32>, tensor<3x3x2x4xf32>) -> tensor<1x5x5x4xf32>
// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[BIAS_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[MAX:.+]] = stablehlo.maximum %[[ADD]], %[[ZERO_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[MAX]], dims = [0, 3, 1, 2] : (tensor<1x5x5x4xf32>) -> tensor<1x4x5x5xf32>
// CHECK: return %[[TRANSPOSE_1]]
// -----
// Tests that a `maximum(add(convolution(%activation, %weight), broadcast(%bias)
// ), %zero)` with the activation tensor of NCHW format is converted to NHWC
// convolution + add + maximum operation. Transpose ops are inserted to the
// first activation, final output, and the bias constant (after the broadcast),
// to match the function signature. Constants are also transpose-folded
// accordingly.
//
// Note that the `transpose` after the `broadcast_in_dim` is not folded by the
// `FoldConstantTransposePass`.
// CHECK-LABEL: nchw_conv_with_broadcasted_bias_add_relu
// CHECK-SAME: %[[ARG:.+]]: tensor<1x2x5x5xf32>
func.func @nchw_conv_with_broadcasted_bias_add_relu(%arg0: tensor<1x2x5x5xf32>) -> tensor<1x4x5x5xf32> {
%0 = stablehlo.constant dense<2.000000e+00> : tensor<4x2x3x3xf32> // weight
%1 = stablehlo.constant dense<3.000000e+00> : tensor<4xf32> // bias
%2 = stablehlo.constant dense<0.000000e+00> : tensor<1x4x5x5xf32> // relu
%3 = stablehlo.broadcast_in_dim %1, dims = [1] : (tensor<4xf32>) -> tensor<1x4x5x5xf32>
%4 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x2x5x5xf32>, tensor<4x2x3x3xf32>) -> tensor<1x4x5x5xf32>
%5 = stablehlo.add %4, %3 : tensor<1x4x5x5xf32>
%6 = stablehlo.maximum %5, %2 : tensor<1x4x5x5xf32>
return %6 : tensor<1x4x5x5xf32>
}
// CHECK-DAG: %[[WEIGHT_CONST:.+]] = stablehlo.constant {{.*}} : tensor<3x3x2x4xf32>
// CHECK-DAG: %[[ZERO_CONST:.+]] = stablehlo.constant {{.*}} : tensor<1x5x5x4xf32>
// CHECK-DAG: %[[BIAS_CONST:.+]] = stablehlo.constant {{.*}} : tensor<4xf32>
// CHECK-DAG: %[[BROADCAST_IN_DIM:.+]] = stablehlo.broadcast_in_dim %[[BIAS_CONST]], dims = [1] : (tensor<4xf32>) -> tensor<1x4x5x5xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x2x5x5xf32>) -> tensor<1x5x5x2xf32>
// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[WEIGHT_CONST]]) 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<1x5x5x2xf32>, tensor<3x3x2x4xf32>) -> tensor<1x5x5x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[BROADCAST_IN_DIM]], dims = [0, 2, 3, 1] : (tensor<1x4x5x5xf32>) -> tensor<1x5x5x4xf32>
// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[TRANSPOSE_1]] : tensor<1x5x5x4xf32>
// CHECK: %[[MAX:.+]] = stablehlo.maximum %[[ADD]], %[[ZERO_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[MAX]], dims = [0, 3, 1, 2] : (tensor<1x5x5x4xf32>) -> tensor<1x4x5x5xf32>
// CHECK: return %[[TRANSPOSE_1]]