From 29779f823c581300422be4274799e52640bbf571 Mon Sep 17 00:00:00 2001 From: Christian Aurich Date: Wed, 26 Aug 2026 20:12:14 -0300 Subject: [PATCH 1/9] Fix AttributeError in tnp around, heaviside and average `_scalar` was fixed to call `math_ops.cast` because `Tensor` only has an `astype` method after `enable_numpy_methods_on_tensor()` has run. Five more `astype` calls, in four other functions, still fail without that opt-in, which is what a plain `import tensorflow` gives you: >>> tnp.around([1.5, 2.5]) AttributeError: EagerTensor object has no attribute 'astype'. `around` casts back to the argument's dtype on every path, so it fails on any input, taking `tnp.round` and `Tensor.__round__` with it. `heaviside` and `average` fail on integer input, `_with_index_update_helper` on any. `astype` is installed as an alias of `math_ops.cast`, so calling `cast` directly is the same operation without the dependency on the opt-in. Neither module calls `astype` after this change. The regression tests go in the module added for the `_scalar` fix, plus an `np_array_ops` counterpart. Both deliberately skip the opt-in, which is what makes these paths observable in a test. --- tensorflow/python/ops/numpy_ops/BUILD | 12 ++ .../python/ops/numpy_ops/np_array_ops.py | 11 +- .../np_array_ops_no_numpy_methods_test.py | 111 ++++++++++++++++++ .../python/ops/numpy_ops/np_math_ops.py | 8 +- .../np_math_ops_no_numpy_methods_test.py | 31 +++++ 5 files changed, 168 insertions(+), 5 deletions(-) create mode 100644 tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py diff --git a/tensorflow/python/ops/numpy_ops/BUILD b/tensorflow/python/ops/numpy_ops/BUILD index 01fe3fe3a74d04..a9f146815d06bc 100644 --- a/tensorflow/python/ops/numpy_ops/BUILD +++ b/tensorflow/python/ops/numpy_ops/BUILD @@ -266,6 +266,18 @@ cuda_py_strict_test( ], ) +cuda_py_strict_test( + name = "np_array_ops_no_numpy_methods_test", + srcs = ["np_array_ops_no_numpy_methods_test.py"], + deps = [ + ":np_array_ops", + "//tensorflow/python/framework:constant_op", + "//tensorflow/python/framework:ops", + "//tensorflow/python/platform:client_testlib", + "//third_party/py/numpy", + ], +) + cuda_py_strict_test( name = "np_math_ops_no_numpy_methods_test", srcs = ["np_math_ops_no_numpy_methods_test.py"], diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 42caa5e7d514fd..85d12f133b612e 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -853,12 +853,15 @@ def around(a, decimals=0): # pylint: disable=missing-docstring # Use float as the working dtype when a.dtype is exact (e.g. integer), # because `decimals` can be negative. float_dtype = np_utils.result_type(float) - a = a.astype(float_dtype) + # `Tensor.astype` only exists once `enable_numpy_methods_on_tensor()` + # has been called, and is an alias of `math_ops.cast`; calling `cast` + # directly keeps `around` working without that opt-in. + a = math_ops.cast(a, float_dtype) factor = math_ops.cast(factor, float_dtype) a = math_ops.multiply(a, factor) a = math_ops.round(a) a = math_ops.divide(a, factor) - return a.astype(dtype) + return math_ops.cast(a, dtype) setattr(np_arrays.ndarray, '__round__', around) @@ -2177,7 +2180,9 @@ def _with_index_update_helper(update_method, a, slice_spec, updates): a_dtype = a.dtype a, updates = _promote_dtype_binary(a, updates) result_t = _slice_helper(a, slice_spec, update_method, updates) - return result_t.astype(a_dtype) + # See the note in `around`: `astype` depends on an opt-in that this + # module-level helper does not require of its callers. + return math_ops.cast(result_t, a_dtype) setattr(np_arrays.ndarray, '_numpy_style_getitem', _getitem) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py new file mode 100644 index 00000000000000..0a0c69f0f1f91b --- /dev/null +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py @@ -0,0 +1,111 @@ +# 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. +# ============================================================================== +"""Tests for tf numpy array methods that must not depend on Tensor methods. + +`np_math_ops.enable_numpy_methods_on_tensor()` adds numpy methods such as +`astype` to `Tensor`. The other numpy_ops test modules call it from their +`__main__`, so a code path that only works once those methods are installed +still passes there. This module deliberately does not call it, which is the +configuration a user gets from a plain `import tensorflow`. +""" + +import numpy as np + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.ops.numpy_ops import np_array_ops +from tensorflow.python.platform import test + + +class ArrayWithoutNumpyMethodsOnTensorTest(test.TestCase): + + def testAroundAcceptsFloatInputs(self): + # `around` casts back to the argument's dtype on every path, so it is + # broken for float arguments too, not only the promoted integer ones. + arg = np.array([1.234, 5.678, -2.345], dtype=np.float32) + for decimals in (0, 1, 2): + actual = np_array_ops.around(arg, decimals) + np.testing.assert_allclose( + np.asarray(actual), + np.around(arg, decimals), + rtol=1e-6, + atol=1e-6, + err_msg='around({}, {})'.format(arg, decimals), + ) + + def testAroundAcceptsIntegerInputs(self): + # An integer argument takes the promoting branch, which computes in a + # float dtype and casts back to the integer dtype at the end. + for arg in ( + np.array([11, 25, 37], dtype=np.int32), + np.array([11, 25, 37], dtype=np.int64), + ): + for decimals in (0, 1): + actual = np_array_ops.around(arg, decimals) + np.testing.assert_array_equal( + np.asarray(actual), + np.around(arg, decimals), + err_msg='around({}, {})'.format(arg, decimals), + ) + + def testAroundPreservesDtype(self): + for dtype in (np.int32, np.int64, np.float32, np.float64): + arg = np.array([1, 2, 3], dtype=dtype) + self.assertEqual(np_array_ops.around(arg).dtype, dtype) + + def testRoundAcceptsFloatInputs(self): + arg = np.array([1.234, 5.678], dtype=np.float32) + np.testing.assert_allclose( + np.asarray(np_array_ops.round(arg, 1)), + np.round(arg, 1), + rtol=1e-6, + atol=1e-6, + ) + + def testBuiltinRoundOnTensor(self): + # `around` is installed as `Tensor.__round__` at import time, so the + # builtin `round()` reaches the same code path. + tensor = constant_op.constant([1.234, 5.678], dtype='float32') + np.testing.assert_allclose( + np.asarray(round(tensor, 1)), + np.around(np.array([1.234, 5.678], dtype=np.float32), 1), + rtol=1e-6, + atol=1e-6, + ) + + def testIndexUpdateHelpersPreserveDtype(self): + # The `_with_index_*` helpers are attached to `Tensor` at import time, + # rather than by `enable_numpy_methods_on_tensor()`, so they must not + # require the opt-in either. They are attached as `functools.partial` + # objects, which are not descriptors, so the tensor is passed explicitly + # instead of being bound as `self`. + tensor = constant_op.constant([1, 2, 3, 4], dtype='int32') + updates = constant_op.constant([9, 9], dtype='int32') + cases = [ + ('update', tensor._with_index_update, np.array([9, 9, 3, 4])), + ('add', tensor._with_index_add, np.array([10, 11, 3, 4])), + ] + for name, helper, expected in cases: + actual = helper(tensor, slice(0, 2), updates) + self.assertEqual(actual.dtype, tensor.dtype, msg=name) + np.testing.assert_array_equal( + np.asarray(actual), expected, err_msg=name + ) + + +if __name__ == '__main__': + ops.enable_eager_execution() + # Intentionally not calling `np_math_ops.enable_numpy_methods_on_tensor()`. + test.main() diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index b31cc630874afb..18e885f5dbdc0d 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -443,7 +443,9 @@ def f(x1, x2): y = _bin_op(f, x1, x2) if not np.issubdtype(y.dtype.as_numpy_dtype, np.inexact): - y = y.astype(np_utils.result_type(float)) + # See the note in `_scalar`: `astype` is unavailable without the + # `enable_numpy_methods_on_tensor()` opt-in. + y = math_ops.cast(y, np_utils.result_type(float)) return y @@ -1557,7 +1559,9 @@ def average(a, axis=None, weights=None, returned=False): # pylint: disable=miss default_float_type = np_utils.result_type(float) if weights is None: # Treat all weights as 1 if not np.issubdtype(a.dtype.as_numpy_dtype, np.inexact): - a = a.astype(np_utils.result_type(a.dtype, default_float_type)) + # See the note in `_scalar`: `astype` is unavailable without the + # `enable_numpy_methods_on_tensor()` opt-in. + a = math_ops.cast(a, np_utils.result_type(a.dtype, default_float_type)) avg = math_ops.reduce_mean(a, axis=axis) if returned: if axis is None: diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py index 73833e41036f04..18c50d31fbdbce 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py @@ -95,6 +95,37 @@ def testUnaryOpsAcceptIntegerScalars(self): err_msg='{}(2)'.format(name), ) + def testHeavisideAcceptsIntegerInputs(self): + # `heaviside` promotes an exact result to a float dtype, which must not + # depend on `enable_numpy_methods_on_tensor()`. + x1 = np.array([-1, 0, 1], dtype=np.int32) + for x2 in (np.int32(1), np.array([5, 5, 5], dtype=np.int32)): + actual = np_math_ops.heaviside(x1, x2) + np.testing.assert_allclose( + np.asarray(actual), + np.heaviside(x1, x2), + rtol=1e-6, + atol=1e-6, + err_msg='heaviside({}, {})'.format(x1, x2), + ) + + def testAverageAcceptsIntegerInputs(self): + # The unweighted branch of `average` promotes an integer argument to a + # float dtype before reducing. + for arg in ( + [1, 2, 3], + np.array([1, 2, 3], dtype=np.int32), + np.array([1, 2, 3], dtype=np.int64), + ): + actual = np_math_ops.average(arg) + np.testing.assert_allclose( + np.asarray(actual), + np.average(np.asarray(arg)), + rtol=1e-6, + atol=1e-6, + err_msg='average({})'.format(arg), + ) + def testFloatInputsAreUnchanged(self): arg = np.array([1.5, 2.5, 3.5], dtype=np.float32) for name, tf_fun, np_fun in _PROMOTING_UNARY_OPS: From 108033bb9b2866774839d50309286396748089c6 Mon Sep 17 00:00:00 2001 From: Christian Aurich Date: Thu, 27 Aug 2026 10:36:43 -0300 Subject: [PATCH 2/9] Call _with_index_update_helper directly in its regression test The `_with_index_*` methods are attached to `Tensor` as `functools.partial` objects, and `functools.partial` became a descriptor in Python 3.13. On 3.13 and later it binds the tensor as the helper's first argument; before 3.13 it does not, so the tensor has to be passed explicitly. Reaching the helper through the method therefore needs opposite call sites on either side of that version boundary, and TensorFlow builds against both. Call `_with_index_update_helper` directly instead. That is the function this change fixes, and it takes the same arguments on every supported version. --- .../np_array_ops_no_numpy_methods_test.py | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py index 0a0c69f0f1f91b..803699f0503dbb 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py @@ -85,23 +85,26 @@ def testBuiltinRoundOnTensor(self): atol=1e-6, ) - def testIndexUpdateHelpersPreserveDtype(self): - # The `_with_index_*` helpers are attached to `Tensor` at import time, - # rather than by `enable_numpy_methods_on_tensor()`, so they must not - # require the opt-in either. They are attached as `functools.partial` - # objects, which are not descriptors, so the tensor is passed explicitly - # instead of being bound as `self`. + def testIndexUpdateHelperPreservesDtype(self): + # `_with_index_update_helper` backs the `_with_index_*` methods, which are + # attached to `Tensor` at import time rather than by + # `enable_numpy_methods_on_tensor()`, so it must not require the opt-in + # either. It is called directly here because the methods are attached as + # `functools.partial` objects, and whether those bind the tensor as their + # first argument differs between Python versions. tensor = constant_op.constant([1, 2, 3, 4], dtype='int32') updates = constant_op.constant([9, 9], dtype='int32') cases = [ - ('update', tensor._with_index_update, np.array([9, 9, 3, 4])), - ('add', tensor._with_index_add, np.array([10, 11, 3, 4])), + (np_array_ops._UpdateMethod.UPDATE, np.array([9, 9, 3, 4])), + (np_array_ops._UpdateMethod.ADD, np.array([10, 11, 3, 4])), ] - for name, helper, expected in cases: - actual = helper(tensor, slice(0, 2), updates) - self.assertEqual(actual.dtype, tensor.dtype, msg=name) + for update_method, expected in cases: + actual = np_array_ops._with_index_update_helper( + update_method, tensor, slice(0, 2), updates + ) + self.assertEqual(actual.dtype, tensor.dtype, msg=str(update_method)) np.testing.assert_array_equal( - np.asarray(actual), expected, err_msg=name + np.asarray(actual), expected, err_msg=str(update_method) ) From c9071ec29e013864cc166b788e0502fb586ab6f3 Mon Sep 17 00:00:00 2001 From: Bhatu Date: Mon, 31 Aug 2026 17:58:27 -0700 Subject: [PATCH 3/9] Fix use-after-free in ConstraintPropagator::PropagateFusionBoundary caused by map reallocation. In ConstraintPropagator::PropagateFusionBoundary, copying ConstraintState objects by value before mutating states_ prevents use-after-free reads when absl::flat_hash_map rehashes during insertion. PiperOrigin-RevId: 974180232 --- third_party/xla/xla/tests/constraint_propagator.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index 8ce31e5b30cff4..15399c6d887e8f 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -1001,11 +1001,11 @@ absl::Status ConstraintPropagator::PropagateFusionBoundary( const HloInstruction* fused_root = fusion_instruction->fused_expression_root(); if (fused_root != nullptr) { - // Backward: outer constraint on fusion result flows into inner root. ConstraintState fusion_state = states_[fusion_instruction]; + ConstraintState root_state = states_[fused_root]; + // Backward: outer constraint on fusion result flows into inner root. states_[fused_root].Merge(fusion_state); // Forward: internal constraint computed on root flows out to fusion result. - ConstraintState root_state = states_[fused_root]; states_[fusion_instruction].Merge(root_state); } @@ -1016,13 +1016,13 @@ absl::Status ConstraintPropagator::PropagateFusionBoundary( if (fused_param == nullptr) { continue; } + ConstraintState operand_state = states_[operand]; + ConstraintState param_state = states_[fused_param]; // Backward: constraints accumulated on the internal parameter flow out // to the caller operand. - ConstraintState param_state = states_[fused_param]; states_[operand].Merge(param_state); // Forward: constraints established on the caller operand flow into the // internal parameter. - ConstraintState operand_state = states_[operand]; states_[fused_param].Merge(operand_state); } From da04a19a8f9c33cf4bc54bcac9de9bcab273961b Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 31 Aug 2026 17:59:09 -0700 Subject: [PATCH 4/9] Add vr200 backend to XLA test build rules PiperOrigin-RevId: 974180531 --- third_party/xla/build_tools/lint/tags.py | 3 +++ third_party/xla/xla/tests/backend_defs.bzl | 3 +++ 2 files changed, 6 insertions(+) diff --git a/third_party/xla/build_tools/lint/tags.py b/third_party/xla/build_tools/lint/tags.py index ffb5c2bbd74707..a909af93153f8e 100644 --- a/third_party/xla/build_tools/lint/tags.py +++ b/third_party/xla/build_tools/lint/tags.py @@ -88,6 +88,7 @@ "requires-gpu-sm90-only": "Requires exactly sm90.", "requires-gpu-sm100-only": "Requires exactly sm100.", "requires-gpu-sm103-only": "Requires exactly sm103.", + "requires-gpu-sm107-only": "Requires exactly sm107.", "requires-gpu-sm120-only": "Requires exactly sm120.", "full": ( "Test requires a full GPU, not a partitioned one. No effect in OSS." @@ -116,6 +117,7 @@ "xla_gb200": "Runs on a gb200.", "xla_gb300": "Runs on a gb300.", "xla_rtx6000pro": "Runs on an rtx6000pro.", + "xla_vr200": "Runs on a vr200.", "xla_device_p100": "Runs on a p100.", "xla_device_v100": "Runs on a v100.", "xla_device_a100": "Runs on an a100.", @@ -124,6 +126,7 @@ "xla_device_gb200": "Runs on a gb200.", "xla_device_gb300": "Runs on a gb300.", "xla_device_rtx6000pro": "Runs on an rtx6000pro.", + "xla_device_vr200": "Runs on a vr200.", # Below tags are consumed by `xla_test`. "test_migrated_to_hlo_runner_pjrt": ( "Adds the appropriate `xla/tests:pjrt_$BACKEND_client_registry` to the" diff --git a/third_party/xla/xla/tests/backend_defs.bzl b/third_party/xla/xla/tests/backend_defs.bzl index 8b2eb3df05510b..e6ee22a9bf6226 100644 --- a/third_party/xla/xla/tests/backend_defs.bzl +++ b/third_party/xla/xla/tests/backend_defs.bzl @@ -34,6 +34,7 @@ NVIDIA_GPU_BACKENDS = [ "b200", "gb200", "gb300", + "vr200", ] + if_google([], ["rtx6000pro"]) # The generic "gpu" backend includes the actual backends in this list. @@ -44,6 +45,7 @@ NVIDIA_GPU_DEFAULT_BACKENDS = [ "b200", "gb200", "gb300", + "vr200", ] + if_google([], ["rtx6000pro"]) AMD_GPU_DEFAULT_BACKENDS = ["amdgpu_any"] @@ -93,6 +95,7 @@ def prepare_nvidia_gpu_backend_data(backends, disabled_backends, backend_tags, b "b200": (10, 0), "gb200": (10, 0), "gb300": (10, 3), + "vr200": (10, 7), "rtx6000pro": (12, 0), } for gpu_backend in NVIDIA_GPU_BACKENDS: From 8e227aa0e44cc213b1a55ca8dd704b467b8b6782 Mon Sep 17 00:00:00 2001 From: Majid Dadashi Date: Mon, 31 Aug 2026 18:36:07 -0700 Subject: [PATCH 5/9] Enable QSV and QDQ quantization pipeline in StableHLO-to-TFLite converter - Enable PropagateQParamsPass, BiasQuantizerPass, and FuseQDQPass in the StableHLO pipeline. - Inline private functions and run symbol DCE prior to lowering quant annotations. - Add CleanupOptimizationBarrierPass and ReconcileUnrealizedCastsPass to the optimization pipeline. - Update PostQuantizePass and FuseQDQPass to support resource constants and per-axis quantization. - Add test coverage for post-quantization optimization passes in optimize-after-quantization.mlir. PiperOrigin-RevId: 974194477 --- tensorflow/compiler/mlir/lite/python/BUILD | 4 + .../lite/python/stablehlo_tfl_pipeline.cc | 47 ++- .../quantization_lib/quantization_utils.cc | 44 ++- .../quantization_lib/quantization_utils.h | 20 +- .../lite/quantization/ir/QuantizeUtils.cc | 347 ++++++++++++++++-- .../stablehlo/transforms/stablehlo_passes.h | 2 + .../tests/optimize-after-quantization.mlir | 29 +- .../lower_quant_annotations_helper.cc | 84 ----- .../lower_quant_annotations_helper.h | 167 ++++++++- .../lower_quant_annotations_pass.cc | 289 ++++++++++++++- .../transforms/optimize_batch_matmul_pass.cc | 16 +- .../compiler/mlir/lite/transforms/passes.h | 1 + .../mlir/lite/transforms/post_quantize.cc | 317 +++++++++++++--- .../transforms/quantization/fuse_qdq_pass.cc | 79 +++- .../lite/transforms/reduce_type_precision.cc | 4 +- 15 files changed, 1227 insertions(+), 223 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD index 27db8de43e7089..44b2ef7416e28e 100644 --- a/tensorflow/compiler/mlir/lite/python/BUILD +++ b/tensorflow/compiler/mlir/lite/python/BUILD @@ -91,7 +91,9 @@ cc_library( "//tensorflow/compiler/mlir/lite:tensorflow_lite", "//tensorflow/compiler/mlir/lite/core:macros", "//tensorflow/compiler/mlir/lite/debug", + "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps", "//tensorflow/compiler/mlir/lite/stablehlo:drop_shape_assertions", + "//tensorflow/compiler/mlir/lite/stablehlo:legalize_vhlo_quant_custom_calls", "//tensorflow/compiler/mlir/lite/stablehlo:prepare_hlo", "//tensorflow/compiler/mlir/lite/stablehlo:tfl_legalize_hlo", "//tensorflow/compiler/mlir/lite/stablehlo:unfold_splat_constant_pass", @@ -103,6 +105,8 @@ cc_library( "@llvm-project//mlir:FuncExtensions", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", + "@llvm-project//mlir:QuantOps", + "@llvm-project//mlir:ReconcileUnrealizedCasts", "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", "@stablehlo//:stablehlo_ops", diff --git a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc index a48f944c2de7c9..59504630a1db05 100644 --- a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc +++ b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc @@ -23,8 +23,10 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "llvm/Support/raw_ostream.h" #include "mlir/Bytecode/BytecodeWriter.h" // from @llvm-project +#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h" // from @llvm-project #include "mlir/Dialect/Func/Extensions/InlinerExtension.h" // from @llvm-project #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/IR/OperationSupport.h" // from @llvm-project #include "mlir/Pass/PassInstrumentation.h" // from @llvm-project @@ -40,6 +42,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/debug/debug.h" #include "tensorflow/compiler/mlir/lite/flatbuffer_export.h" #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" +#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h" #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h" #include "tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.h" #include "tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.h" @@ -56,7 +59,7 @@ namespace mlir::TFL { void AddSkipToTflitePasses(mlir::OpPassManager& pass_manager) { pass_manager.addNestedPass( mlir::odml::CreateLegalizeChloToTflPass()); - pass_manager.addPass(mlir::odml::CreateCompositeLoweringPass()); + pass_manager.addPass(mlir::createInlinerPass()); pass_manager.addPass(mlir::TFL::CreateLowerQuantAnnotationsPass()); pass_manager.addPass(mlir::createSymbolDCEPass()); } @@ -65,6 +68,10 @@ void AddHloOptimizationPasses(mlir::OpPassManager& pass_manager) { // Drop shape assertion custom calls before VHLO legalization pass_manager.addPass(mlir::odml::CreateDropShapeAssertionsPass()); + // Legalize VHLO quant custom calls to StableHLO custom calls before VHLO + // legalization + pass_manager.addPass(mlir::odml::CreateLegalizeVhloQuantCustomCallsPass()); + // VHLO -> StableHLO pass_manager.addPass(mlir::stablehlo::createVhloLegalizeToStablehloPass()); @@ -108,6 +115,10 @@ void AddHloOptimizationPasses(mlir::OpPassManager& pass_manager) { // StableHLO -> MHLO bridge pass_manager.addPass(mlir::mhlo::createStablehloLegalizeToHloPass()); + // Composite lowering when IR is in MHLO + pass_manager.addPass(mlir::odml::CreateCompositeLoweringPass()); + pass_manager.addPass(mlir::createSymbolDCEPass()); + // MHLO algebraic optimizations pass_manager.addNestedPass( mlir::mhlo::createLegalizeEinsumToDotGeneralPass()); @@ -128,7 +139,13 @@ void AddHloToTfLiteLegalizationPasses(mlir::OpPassManager& pass_manager) { mlir::odml::CreateUniformQuantizedStableHloToTflPass()); pass_manager.addNestedPass( mlir::odml::CreatePrepareHloPass()); + pass_manager.addPass(mlir::odml::CreateUnfoldSplatConstantPass()); pass_manager.addPass(mlir::odml::CreateLegalizeHloToTfLitePass()); + + // Legalize remaining MHLO ops to StableHLO + pass_manager.addPass(mlir::mhlo::createHloLegalizeToStablehloPass()); + pass_manager.addNestedPass( + mlir::odml::createLegalizeCompositeToCustomOpPass()); } void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager, @@ -148,19 +165,23 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager, pass_manager.addNestedPass( mlir::TFL::CreateOptimizePass()); - // Quantization - pass_manager.addNestedPass( - mlir::TFL::CreatePrepareQuantizePass(pass_config.quant_specs)); - pass_manager.addNestedPass( - mlir::TFL::CreateQuantizePass(pass_config.quant_specs)); - pass_manager.addNestedPass( - mlir::TFL::CreatePostQuantizePass(/*emit_quant_adaptor_ops=*/true)); - if (!pass_config.unfold_batch_matmul) { pass_manager.addNestedPass( mlir::TFL::CreateOptimizeBatchMatmulPass()); + pass_manager.addNestedPass( + mlir::TFL::CreateOptimizePass()); } + // Quantization + pass_manager.addPass(mlir::TFL::CreatePropagateQParamsPass()); + pass_manager.addPass(mlir::TFL::CreateBiasQuantizerPass()); + pass_manager.addPass(mlir::TFL::CreateFuseQDQPass()); + + pass_manager.addNestedPass( + mlir::TFL::CreatePostQuantizePass(/*emit_quant_adaptor_ops=*/true)); + pass_manager.addNestedPass( + mlir::createCanonicalizerPass()); + // Some optimizations need to happen on the quantized graph. pass_manager.addNestedPass( mlir::TFL::CreateOptimizePass()); @@ -176,6 +197,8 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager, pass_manager.addNestedPass( mlir::createCanonicalizerPass()); pass_manager.addNestedPass(mlir::createCSEPass()); + pass_manager.addPass(mlir::TFL::CreateCleanupOptimizationBarrierPass()); + pass_manager.addPass(mlir::createReconcileUnrealizedCastsPass()); } absl::Status ConvertStableHloToTFLite( @@ -185,8 +208,10 @@ absl::Status ConvertStableHloToTFLite( mlir::MLIRContext* context = module->getContext(); mlir::DialectRegistry registry; mlir::func::registerInlinerExtension(registry); - registry.insert(); + registry.insert(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); diff --git a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc index fa86a8486913a0..a132320460fdce 100644 --- a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc +++ b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc @@ -334,18 +334,27 @@ TypeAttr RescaleQuantizedType(const Type input, const Attribute factor) { const auto element_type = quant::QuantizedType::getQuantizedElementType(input); if (!element_type) return {}; - if (auto qtype = dyn_cast(element_type)) { + if (const auto qtype = + dyn_cast(element_type)) { const ArrayRef scales = qtype.getScales(); - // Broadcasting hasn't been implemented yet. - if (static_cast(scales.size()) != factor_values.getNumElements()) + const int64_t num_factors = factor_values.getNumElements(); + if (static_cast(scales.size()) != num_factors && num_factors != 1) return {}; SmallVector new_scales; new_scales.reserve(scales.size()); - auto scales_iter = scales.begin(); - for (const auto& f : factor_values) { - new_scales.push_back(*scales_iter * - std::fabs(FloatAttr::getValueAsDouble(f))); - ++scales_iter; + if (num_factors == 1) { + const double factor_val = + std::fabs(FloatAttr::getValueAsDouble(*factor_values.begin())); + for (const double scale : scales) { + new_scales.push_back(scale * factor_val); + } + } else { + auto scales_iter = scales.begin(); + for (const auto& f : factor_values) { + new_scales.push_back(*scales_iter * + std::fabs(FloatAttr::getValueAsDouble(f))); + ++scales_iter; + } } // We are assuming symmetric quantization. auto new_ele_type = quant::UniformQuantizedPerAxisType::get( @@ -356,8 +365,21 @@ TypeAttr RescaleQuantizedType(const Type input, const Attribute factor) { quant::QuantizedType::castToExpressedType(input))) { return TypeAttr::get(new_type); } + } else if (const auto qtype = + dyn_cast(element_type)) { + if (factor_values.getNumElements() != 1) return {}; + const double factor_val = + std::fabs(FloatAttr::getValueAsDouble(*factor_values.begin())); + const double new_scale = qtype.getScale() * factor_val; + auto new_ele_type = quant::UniformQuantizedType::get( + qtype.getFlags(), qtype.getStorageType(), qtype.getExpressedType(), + new_scale, qtype.getZeroPoint(), qtype.getStorageTypeMin(), + qtype.getStorageTypeMax()); + if (const auto new_type = new_ele_type.castFromExpressedType( + quant::QuantizedType::castToExpressedType(input))) { + return TypeAttr::get(new_type); + } } - // Currently, we only support per-axis quantized type. return {}; } @@ -511,8 +533,8 @@ Type GetUniformQuantizedPerAxisTypeForWeight( const int dim_size = shape[quant_dim]; const int slice_size = - std::accumulate(std::next(shape.begin(), quant_dim + 1), shape.end(), 1, - std::multiplies()); + std::accumulate(std::next(shape.begin(), quant_dim + 1), shape.end(), + int64_t{1}, std::multiplies()); SmallVector mins(dim_size, std::numeric_limits::max()); SmallVector maxs(dim_size, std::numeric_limits::min()); const auto fp = dyn_cast(attr); diff --git a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h index e039138f7c8a90..52f234d1a7ab81 100644 --- a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h +++ b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h @@ -43,6 +43,7 @@ limitations under the License. #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project // IWYU pragma: keep #include "mlir/IR/IRMapping.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project @@ -51,6 +52,7 @@ limitations under the License. #include "mlir/IR/Operation.h" // from @llvm-project #include "mlir/IR/OperationSupport.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project @@ -662,10 +664,26 @@ class QuantizationPattern : public RewritePattern { for (int i = 0, e = quantized_op->getNumOperands(); i < e; ++i) { auto def = quantized_op->getOperand(i).getDefiningOp(); if (auto q = llvm::dyn_cast_or_null(def)) { - DenseFPElementsAttr attr; + ElementsAttr attr; if (!matchPattern(q.getOperand(), m_Constant(&attr))) { continue; } + if (auto resourceAttr = + mlir::dyn_cast(attr)) { + if (AsmResourceBlob* blob = + resourceAttr.getRawHandle().getBlob()) { + ArrayRef ptr = blob->getData(); + if (DenseElementsAttr::isValidRawBuffer(resourceAttr.getType(), + ptr)) { + attr = DenseElementsAttr::getFromRawBuffer( + resourceAttr.getType(), ptr); + } + } + } + if (!mlir::isa( + mlir::getElementTypeOrSelf(attr.getType()))) { + continue; + } auto cst = arith::ConstantOp::create(rewriter, quantized_op->getLoc(), attr); quantizing_op->setOperand(i, cst.getResult()); diff --git a/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc b/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc index af0d21594ae957..7bc5887455258f 100644 --- a/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc +++ b/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc @@ -15,15 +15,31 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "llvm/ADT/STLExtras.h" #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinDialect.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h" -using namespace mlir; -using namespace mlir::quantfork; +namespace mlir { +namespace quantfork { /// Converts a possible primitive, real expressed value attribute to a /// corresponding storage attribute (typically FloatAttr -> IntegerAttr). @@ -32,10 +48,10 @@ using namespace mlir::quantfork; /// Returns a converter Attribute or nullptr if conversion is not possible. static Attribute convertPrimitiveValueAttr( Attribute origRealValue, quant::QuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter, - Type &outConvertedType) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter, + Type& outConvertedType) { if (mlir::isa(origRealValue)) { - FloatAttr floatAttr = mlir::cast(origRealValue); + const FloatAttr floatAttr = mlir::cast(origRealValue); outConvertedType = quantizedElementType.getStorageType(); return IntegerAttr::get(quantizedElementType.getStorageType(), converter.quantizeFloatToInt(floatAttr.getValue())); @@ -50,10 +66,10 @@ static Attribute convertPrimitiveValueAttr( static DenseElementsAttr convertDenseFPElementsAttr( DenseFPElementsAttr realFPElementsAttr, quant::QuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter) { return realFPElementsAttr.mapValues( quantizedElementType.getStorageType(), - [&converter](const APFloat &realVal) { + [&converter](const APFloat& realVal) { return converter.quantizeFloatToInt(realVal); }); } @@ -64,12 +80,12 @@ static DenseElementsAttr convertDenseFPElementsAttr( static SparseElementsAttr convertSparseElementsAttr( SparseElementsAttr realSparseAttr, quant::QuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter) { DenseElementsAttr realDenseAttr = realSparseAttr.getValues(); if (!mlir::isa(realDenseAttr)) { return nullptr; } - DenseElementsAttr quantDenseAttr = + const DenseElementsAttr quantDenseAttr = convertDenseFPElementsAttr(mlir::cast(realDenseAttr), quantizedElementType, converter); if (!quantDenseAttr) { @@ -78,7 +94,7 @@ static SparseElementsAttr convertSparseElementsAttr( // Cast from an expressed-type-based type to storage-type-based type, // preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>). - ShapedType newSparseType = mlir::dyn_cast_or_null( + const ShapedType newSparseType = mlir::dyn_cast_or_null( quantizedElementType.castExpressedToStorageType( realSparseAttr.getType())); if (!newSparseType) { @@ -88,17 +104,294 @@ static SparseElementsAttr convertSparseElementsAttr( quantDenseAttr); } +static Attribute quantizeResourceAttrPerAxisLegacy( + DenseResourceElementsAttr resourceAttr, + quant::UniformQuantizedPerAxisType quantizedElementType, + Type& outConvertedType) { + const ShapedType type = resourceAttr.getType(); + const int32_t quantDim = quantizedElementType.getQuantizedDimension(); + const uint32_t storageBitWidth = + quantizedElementType.getStorageTypeIntegralWidth(); + const bool isSigned = quantizedElementType.isSigned(); + const ArrayRef scales = quantizedElementType.getScales(); + const ArrayRef zeroPoints = quantizedElementType.getZeroPoints(); + const double clampMin = + static_cast(quantizedElementType.getStorageTypeMin()); + const double clampMax = + static_cast(quantizedElementType.getStorageTypeMax()); + + const std::string newKey = + (llvm::Twine(resourceAttr.getRawHandle().getKey()) + "_quant_axis_" + + llvm::Twine(quantDim) + "_w_" + llvm::Twine(storageBitWidth)) + .str(); + + const Type storageElemType = + IntegerType::get(resourceAttr.getContext(), storageBitWidth, + isSigned ? IntegerType::Signed : IntegerType::Signless); + const auto resType = RankedTensorType::get(type.getShape(), storageElemType); + + auto& manager = DenseResourceElementsHandle::getManagerInterface( + resourceAttr.getContext()); + if (const auto* entry = manager.getBlobManager().lookup(newKey)) { + if (entry->getBlob()) { + auto* dialect = + resourceAttr.getContext()->getLoadedDialect(); + const DenseResourceElementsHandle handle( + const_cast(entry), dialect); + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, handle); + } + } + + const AsmResourceBlob* blob = resourceAttr.getRawHandle().getBlob(); + if (!blob && resourceAttr.getRawHandle().getResource()) { + blob = resourceAttr.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return nullptr; + + const size_t numElements = type.getNumElements(); + const size_t elemByteSize = std::max(1, storageBitWidth / 8); + const size_t outByteSize = (storageBitWidth == 4) ? (numElements + 1) / 2 + : (storageBitWidth == 2) + ? (numElements + 3) / 4 + : numElements * elemByteSize; + + auto rawOutputBlob = mlir::HeapAsmResourceBlob::allocate( + outByteSize, /*align=*/64, /*dataIsMutable=*/true); + + const std::size_t dimSize = type.getDimSize(quantDim); + if (dimSize != scales.size()) { + return nullptr; + } + SmallVector converters; + converters.reserve(dimSize); + for (int i = 0, e = dimSize; i != e; ++i) { + converters.emplace_back(scales[i], zeroPoints[i], APFloat(clampMin), + APFloat(clampMax), storageBitWidth, isSigned); + } + + const auto shape = type.getShape(); + const int64_t chunkSize = + std::accumulate(std::next(shape.begin(), quantDim + 1), shape.end(), + int64_t{1}, std::multiplies()); + + const ArrayRef rawFloat( + reinterpret_cast(blob->getData().data()), numElements); + char* outData = const_cast(rawOutputBlob.getDataAs().data()); + + if (storageBitWidth == 8) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + outInt8[elemIdx] = static_cast(q.getSExtValue()); + } + } else if (storageBitWidth == 4) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + const int8_t val = static_cast(q.getSExtValue()) & 0x0F; + if (elemIdx % 2 == 0) { + outInt8[elemIdx / 2] |= val; + } else { + outInt8[elemIdx / 2] |= (val << 4); + } + } + } else if (storageBitWidth == 2) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + const int8_t val = static_cast(q.getSExtValue()) & 0x03; + const int shift = (elemIdx % 4) * 2; + outInt8[elemIdx / 4] |= (val << shift); + } + } else { + const MutableArrayRef outInt16(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + outInt16[elemIdx] = static_cast(q.getSExtValue()); + } + } + + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, newKey, + std::move(rawOutputBlob)); +} + +static Attribute quantizeResourceAttrPerAxisFast( + DenseResourceElementsAttr resourceAttr, + quant::UniformQuantizedPerAxisType quantizedElementType, + Type& outConvertedType) { + const ShapedType type = resourceAttr.getType(); + const int32_t quantDim = quantizedElementType.getQuantizedDimension(); + const uint32_t storageBitWidth = + quantizedElementType.getStorageTypeIntegralWidth(); + const bool isSigned = quantizedElementType.isSigned(); + const ArrayRef scales = quantizedElementType.getScales(); + const ArrayRef zeroPoints = quantizedElementType.getZeroPoints(); + const double clampMin = + static_cast(quantizedElementType.getStorageTypeMin()); + const double clampMax = + static_cast(quantizedElementType.getStorageTypeMax()); + + const std::string newKey = + (llvm::Twine(resourceAttr.getRawHandle().getKey()) + "_quant_axis_" + + llvm::Twine(quantDim) + "_w_" + llvm::Twine(storageBitWidth)) + .str(); + + const Type storageElemType = + IntegerType::get(resourceAttr.getContext(), storageBitWidth, + isSigned ? IntegerType::Signed : IntegerType::Signless); + const auto resType = RankedTensorType::get(type.getShape(), storageElemType); + + auto& manager = DenseResourceElementsHandle::getManagerInterface( + resourceAttr.getContext()); + if (const auto* entry = manager.getBlobManager().lookup(newKey)) { + if (entry->getBlob()) { + auto* dialect = + resourceAttr.getContext()->getLoadedDialect(); + const DenseResourceElementsHandle handle( + const_cast(entry), dialect); + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, handle); + } + } + + const AsmResourceBlob* blob = resourceAttr.getRawHandle().getBlob(); + if (!blob && resourceAttr.getRawHandle().getResource()) { + blob = resourceAttr.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return nullptr; + + const size_t numElements = type.getNumElements(); + const size_t elemByteSize = std::max(1, storageBitWidth / 8); + const size_t outByteSize = (storageBitWidth == 4) ? (numElements + 1) / 2 + : (storageBitWidth == 2) + ? (numElements + 3) / 4 + : numElements * elemByteSize; + + auto rawOutputBlob = mlir::HeapAsmResourceBlob::allocate( + outByteSize, /*align=*/64, /*dataIsMutable=*/true); + + const std::size_t dimSize = type.getDimSize(quantDim); + if (dimSize != scales.size()) { + return nullptr; + } + + std::vector invScales(dimSize); + std::vector zps(dimSize); + for (size_t i = 0; i < dimSize; ++i) { + invScales[i] = 1.0f / static_cast(scales[i]); + zps[i] = static_cast(zeroPoints[i]); + } + const float fMin = static_cast(clampMin); + const float fMax = static_cast(clampMax); + + const auto shape = type.getShape(); + const int64_t chunkSize = + std::accumulate(std::next(shape.begin(), quantDim + 1), shape.end(), + int64_t{1}, std::multiplies()); + + const ArrayRef rawFloat( + reinterpret_cast(blob->getData().data()), numElements); + char* outData = const_cast(rawOutputBlob.getDataAs().data()); + + if (storageBitWidth == 8) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + outInt8[elemIdx] = static_cast(val); + } + } else if (storageBitWidth == 4) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + const int8_t nibble = static_cast(val) & 0x0F; + if (elemIdx % 2 == 0) { + outInt8[elemIdx / 2] |= nibble; + } else { + outInt8[elemIdx / 2] |= (nibble << 4); + } + } + } else if (storageBitWidth == 2) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + const int8_t twoBit = static_cast(val) & 0x03; + const int shift = (elemIdx % 4) * 2; + outInt8[elemIdx / 4] |= (twoBit << shift); + } + } else { + const MutableArrayRef outInt16(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + outInt16[elemIdx] = static_cast(val); + } + } + + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, newKey, + std::move(rawOutputBlob)); +} + +static Attribute quantizeResourceAttrPerAxis( + DenseResourceElementsAttr resourceAttr, + quant::UniformQuantizedPerAxisType quantizedElementType, + Type& outConvertedType, bool useLegacySlowQuantize = false) { + if (useLegacySlowQuantize) { + return quantizeResourceAttrPerAxisLegacy(resourceAttr, quantizedElementType, + outConvertedType); + } + return quantizeResourceAttrPerAxisFast(resourceAttr, quantizedElementType, + outConvertedType); +} + /// Converts a real expressed Attribute to a corresponding Attribute containing /// quantized storage values assuming the given uniform quantizedElementType and /// converter. -Attribute mlir::quantfork::quantizeAttrUniform( +Attribute quantizeAttrUniform( Attribute realValue, quant::UniformQuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter, - Type &outConvertedType) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter, + Type& outConvertedType) { // Fork to handle different variants of constants supported. if (mlir::isa(realValue)) { // Dense tensor or vector constant. - auto converted = + const auto converted = convertDenseFPElementsAttr(mlir::cast(realValue), quantizedElementType, converter); outConvertedType = converted.getType(); @@ -106,7 +399,7 @@ Attribute mlir::quantfork::quantizeAttrUniform( } if (mlir::isa(realValue)) { // Sparse tensor or vector constant. - auto converted = + const auto converted = convertSparseElementsAttr(mlir::cast(realValue), quantizedElementType, converter); outConvertedType = converted.getType(); @@ -122,22 +415,27 @@ Attribute mlir::quantfork::quantizeAttrUniform( /// quantizedElementType.getStorageType(). /// Returns nullptr if the conversion is not supported. /// On success, stores the converted type in outConvertedType. -Attribute mlir::quantfork::quantizeAttr( - Attribute realValue, quant::QuantizedType quantizedElementType, - Type &outConvertedType) { - if (auto uniformQuantized = +Attribute quantizeAttr(Attribute realValue, + quant::QuantizedType quantizedElementType, + Type& outConvertedType) { + if (const auto uniformQuantized = mlir::dyn_cast(quantizedElementType)) { - mlir::quant::ir::UniformQuantizedValueConverter converter(uniformQuantized); + const mlir::quant::ir::UniformQuantizedValueConverter converter( + uniformQuantized); return quantizeAttrUniform(realValue, uniformQuantized, converter, outConvertedType); } - if (auto uniformQuantizedPerAxis = + if (const auto uniformQuantizedPerAxis = mlir::dyn_cast( quantizedElementType)) { + if (const auto resourceAttr = + mlir::dyn_cast(realValue)) { + return quantizeResourceAttrPerAxis(resourceAttr, uniformQuantizedPerAxis, + outConvertedType); + } mlir::quant::ir::UniformQuantizedPerAxisValueConverter converter( uniformQuantizedPerAxis); - auto converted = converter.convert(realValue); - // TODO: why we need this outConvertedType? remove it? + const auto converted = converter.convert(realValue); if (converted) { outConvertedType = converted.getType(); } @@ -145,3 +443,6 @@ Attribute mlir::quantfork::quantizeAttr( } return nullptr; } + +} // namespace quantfork +} // namespace mlir diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h index 7ec56df1e24504..ade2a0a45e8562 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h @@ -20,6 +20,8 @@ limitations under the License. #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project namespace mlir::odml { diff --git a/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir b/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir index 97c120796d2c5d..0e7b2ec499daf3 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir @@ -12,21 +12,40 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================== -// RUN: litert-opt %s -tfl-prepare-quantize -canonicalize -tfl-quantize -canonicalize -tfl-optimize -canonicalize | FileCheck %s +// RUN: litert-opt %s -tfl-optimize -tfl-propagate-qparams -tfl-bias-quantizer -tfl-fuse-qdq -canonicalize | FileCheck %s // CHECK-LABEL: fuseMulIntoPerTensorConv2dWithQDQs func.func @fuseMulIntoPerTensorConv2dWithQDQs(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x8x7x3xf32> { - %cst = arith.constant dense<1.5> : tensor<3xf32> + %cst = arith.constant dense<1.5> : tensor %cst_0 = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32> %w = arith.constant dense<2.0> : tensor<3x3x3x3xf32> %q = "tfl.quantize"(%w) {qtype = tensor<3x3x3x3x!quant.uniform>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3x!quant.uniform> %dq = "tfl.dequantize"(%q) : (tensor<3x3x3x3x!quant.uniform>) -> tensor<3x3x3x3xf32> %0 = "tfl.conv_2d"(%arg0, %dq, %cst_0) {dilation_h_factor = 2 : i32, dilation_w_factor = 3 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 4 : i32, stride_w = 5 : i32} : (tensor<256x32x32x3xf32>, tensor<3x3x3x3xf32>, tensor<3xf32>) -> tensor<256x8x7x3xf32> - %1 = "tfl.mul"(%0, %cst) {fused_activation_function = "NONE"} : (tensor<256x8x7x3xf32>, tensor<3xf32>) -> tensor<256x8x7x3xf32> + %1 = "tfl.mul"(%0, %cst) {fused_activation_function = "NONE"} : (tensor<256x8x7x3xf32>, tensor) -> tensor<256x8x7x3xf32> func.return %1 : tensor<256x8x7x3xf32> - // CHECK: %[[weight:.*]] = arith.constant dense<3.000000e+00> : tensor<3x3x3x3xf32> - // CHECK: %[[bias:.*]] = arith.constant dense<[1.500000e+00, 3.000000e+00, 4.500000e+00]> + // CHECK: %[[bias:.*]] = arith.constant dense<[1.500000e+00, 3.000000e+00, 4.500000e+00]> : tensor<3xf32> + // CHECK: %[[qweight:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<3x3x3x3x!quant.uniform>, value = dense<21> : tensor<3x3x3x3xi8>}> + // CHECK: %[[weight:.*]] = "tfl.dequantize"(%[[qweight]]) + // CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %[[weight]], %[[bias]]) + // CHECK: return %[[conv]] : tensor<256x8x7x3xf32> +} + +// CHECK-LABEL: fuseMulIntoPerAxisConv2dWithQDQsBroadcastScalar +func.func @fuseMulIntoPerAxisConv2dWithQDQsBroadcastScalar(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x8x7x3xf32> { + %cst = arith.constant dense<1.5> : tensor + %cst_0 = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32> + %w = arith.constant dense<2.0> : tensor<3x3x3x3xf32> + %q = "tfl.quantize"(%w) {qtype = tensor<3x3x3x3x!quant.uniform>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3x!quant.uniform> + %dq = "tfl.dequantize"(%q) : (tensor<3x3x3x3x!quant.uniform>) -> tensor<3x3x3x3xf32> + %0 = "tfl.conv_2d"(%arg0, %dq, %cst_0) {dilation_h_factor = 2 : i32, dilation_w_factor = 3 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 4 : i32, stride_w = 5 : i32} : (tensor<256x32x32x3xf32>, tensor<3x3x3x3xf32>, tensor<3xf32>) -> tensor<256x8x7x3xf32> + %1 = "tfl.mul"(%0, %cst) {fused_activation_function = "NONE"} : (tensor<256x8x7x3xf32>, tensor) -> tensor<256x8x7x3xf32> + func.return %1 : tensor<256x8x7x3xf32> + + // CHECK: %[[bias:.*]] = arith.constant dense<[1.500000e+00, 3.000000e+00, 4.500000e+00]> : tensor<3xf32> + // CHECK: %[[qweight:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<3x3x3x3x!quant.uniform>, value = dense<20> : tensor<3x3x3x3xi8>}> + // CHECK: %[[weight:.*]] = "tfl.dequantize"(%[[qweight]]) // CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %[[weight]], %[[bias]]) // CHECK: return %[[conv]] : tensor<256x8x7x3xf32> } diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc index 03b7e46f2d44ea..fd6149ad9d54fb 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc @@ -17,9 +17,7 @@ limitations under the License. #include #include -#include -#include "llvm/Support/Casting.h" #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project @@ -30,91 +28,9 @@ limitations under the License. #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project -#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo namespace mlir::TFL { -LogicalResult FillCompositeParams(stablehlo::CompositeOp op, - SmallVector& scales, - SmallVector& zero_points, - int& num_bits, bool& is_signed, - bool& is_narrow_range) { - auto scale_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("scale")); - if (scale_attr == nullptr) { - return failure(); - } - for (auto float_attr : scale_attr.getValues()) { - scales.push_back(float_attr.getValue().convertToDouble()); - } - - auto zero_point_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("zero_point")); - if (zero_point_attr == nullptr) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back(0); - } - } else if (zero_point_attr.isSplat()) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back( - zero_point_attr.getSplatValue().getInt()); - } - } else { - for (IntegerAttr zp : zero_point_attr.getValues()) { - zero_points.push_back(zp.getInt()); - } - } - - auto dtype_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("dtype")); - if (dtype_attr == nullptr) { - return failure(); - } - std::string dtype = dtype_attr.getValue().str(); - if (dtype == "i2") { - num_bits = 2; - is_signed = true; - } else if (dtype == "i4") { - num_bits = 4; - is_signed = true; - } else if (dtype == "ui4") { - num_bits = 4; - is_signed = false; - } else if (dtype == "i8") { - num_bits = 8; - is_signed = true; - } else if (dtype == "i16") { - num_bits = 16; - is_signed = true; - } else { - return failure(); - } - auto narrow_range_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("narrow_range")); - if (narrow_range_attr == nullptr) { - return failure(); - } - is_narrow_range = narrow_range_attr.getValue(); - - return success(); -} - -bool IsDrqFakeQuant(stablehlo::CompositeOp op) { - if (op.getName() != "quant.fake_quant") { - return false; - } - SmallVector scales; - SmallVector zero_points; - int num_bits; - bool is_signed; - bool is_narrow_range; - if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, - is_narrow_range))) { - return false; - } - return scales.empty() && zero_points.empty(); -} - LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range, bool is_signed, MLIRContext* ctx, Type& storage_type, int64_t& qmin, diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h index 35d2f206ad8395..a7ebd380742df7 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h @@ -18,7 +18,9 @@ limitations under the License. #include +#include "llvm/Support/Casting.h" #include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project @@ -28,13 +30,172 @@ limitations under the License. namespace mlir::TFL { -LogicalResult FillCompositeParams(stablehlo::CompositeOp op, +template +LogicalResult FillCompositeParams(CompositeOpType op, SmallVector& scales, SmallVector& zero_points, int& num_bits, bool& is_signed, - bool& is_narrow_range); + bool& is_narrow_range) { + auto scale_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("scale")); + if (scale_attr == nullptr) { + return failure(); + } + for (auto float_attr : scale_attr.template getValues()) { + scales.push_back(float_attr.getValue().convertToDouble()); + } -bool IsDrqFakeQuant(stablehlo::CompositeOp op); + auto zero_point_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("zero_point")); + if (zero_point_attr == nullptr) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back(0); + } + } else if (zero_point_attr.isSplat()) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back( + zero_point_attr.template getSplatValue().getInt()); + } + } else { + for (IntegerAttr zp : zero_point_attr.template getValues()) { + zero_points.push_back(zp.getInt()); + } + } + + auto dtype_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("dtype")); + if (dtype_attr == nullptr) { + return failure(); + } + auto dtype = dtype_attr.getValue(); + + if (dtype == "i2") { + num_bits = 2; + is_signed = true; + } else if (dtype == "i4") { + num_bits = 4; + is_signed = true; + } else if (dtype == "ui4") { + num_bits = 4; + is_signed = false; + } else if (dtype == "i8") { + num_bits = 8; + is_signed = true; + } else if (dtype == "i16") { + num_bits = 16; + is_signed = true; + } else { + return failure(); + } + auto narrow_range_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("narrow_range")); + if (narrow_range_attr == nullptr) { + return failure(); + } + is_narrow_range = narrow_range_attr.getValue(); + + return success(); +} + +template +bool IsDrqFakeQuant(CompositeOpType op) { + if (op.getName() != "quant.fake_quant") { + return false; + } + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, + is_narrow_range))) { + return false; + } + return scales.empty() && zero_points.empty(); +} + +template <> +inline LogicalResult FillCompositeParams( + stablehlo::CustomCallOp op, SmallVector& scales, + SmallVector& zero_points, int& num_bits, bool& is_signed, + bool& is_narrow_range) { + auto scale_attr = + llvm::dyn_cast_or_null(op->getAttr("scale")); + if (scale_attr == nullptr) { + return failure(); + } + for (auto float_attr : scale_attr.getValues()) { + scales.push_back(float_attr.getValue().convertToDouble()); + } + + auto zero_point_attr = + llvm::dyn_cast_or_null(op->getAttr("zero_point")); + if (zero_point_attr == nullptr) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back(0); + } + } else if (zero_point_attr.isSplat()) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back( + zero_point_attr.getSplatValue().getInt()); + } + } else { + for (IntegerAttr zp : zero_point_attr.getValues()) { + zero_points.push_back(zp.getInt()); + } + } + + auto dtype_attr = llvm::dyn_cast_or_null(op->getAttr("dtype")); + if (dtype_attr == nullptr) { + return failure(); + } + auto dtype = dtype_attr.getValue(); + + if (dtype == "i2") { + num_bits = 2; + is_signed = true; + } else if (dtype == "i4") { + num_bits = 4; + is_signed = true; + } else if (dtype == "ui4") { + num_bits = 4; + is_signed = false; + } else if (dtype == "i8") { + num_bits = 8; + is_signed = true; + } else if (dtype == "i16") { + num_bits = 16; + is_signed = true; + } else { + return failure(); + } + auto narrow_range_attr = + llvm::dyn_cast_or_null(op->getAttr("narrow_range")); + if (narrow_range_attr == nullptr) { + return failure(); + } + is_narrow_range = narrow_range_attr.getValue(); + + return success(); +} + +template <> +inline bool IsDrqFakeQuant( + stablehlo::CustomCallOp op) { + if (op.getCallTargetName() != "quant.fake_quant") { + return false; + } + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, + is_narrow_range))) { + return false; + } + return scales.empty() && zero_points.empty(); +} LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range, bool is_signed, MLIRContext* ctx, diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc index 272ce6c866436f..9ecf7dca971e65 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc @@ -20,9 +20,8 @@ limitations under the License. #include #include "llvm/Support/Casting.h" -#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project -#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project +#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project // IWYU pragma: keep #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project @@ -41,7 +40,6 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h" #include "tensorflow/compiler/mlir/lite/transforms/passes.h" // IWYU pragma: keep #include "tensorflow/compiler/mlir/lite/utils/utils.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" namespace mlir { @@ -235,15 +233,16 @@ class RewriteDequantizeCompositeOp Value operand = composite_op.getOperand(num_operands - 1); mlir::Operation* producer_op = operand.getDefiningOp(); - // Check if the producer is an arith.constant - if (auto const_op = - llvm::dyn_cast_or_null(producer_op)) { - // We found a constant (Int4/Int8). + // Check if the producer is a constant (arith.constant or + // stablehlo.constant) + if (producer_op && producer_op->hasAttr("value")) { + // We found a constant (Int4/Int8) with ElementsAttr (DenseElementsAttr + // or DenseResourceElementsAttr). // Instead of casting or hacking the constant, we create a valid // TFL::QConstOp. This op natively maps "Integer Data" -> "Quantized // Type". - auto value_attr = llvm::dyn_cast(const_op.getValue()); + auto value_attr = producer_op->getAttrOfType("value"); if (!value_attr) { return failure(); // Should not happen for tensor constants } @@ -251,7 +250,7 @@ class RewriteDequantizeCompositeOp // Create tfl.qconst // Arguments: Type (Result), TypeAttr (qtype), ElementsAttr (value) auto qconst_op = rewriter.create( - const_op.getLoc(), + producer_op->getLoc(), qtensor_type, // The Result Type (!quant.uniform...) TypeAttr::get(qtensor_type), // The "qtype" attribute value_attr // The reuse of the i4/i8 data @@ -260,7 +259,7 @@ class RewriteDequantizeCompositeOp // Use the output of this new constant tfl_quantize_input = qconst_op.getResult(); - // Note: We leave the old arith.constant alone. + // Note: We leave the old constant alone. // If it has no other uses, the cleanup pass (DCE) will remove it // automatically. @@ -439,9 +438,277 @@ void LowerQuantAnnotationsPass::runOnOperation() { signalPassFailure(); } + class RewriteQuantizeCustomCallOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, + PatternRewriter& rewriter) const final { + if (op.getCallTargetName() != "quant.quantize") { + return failure(); + } + + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, + is_signed, is_narrow_range))) { + return op.emitError( + "quantize custom call does not contain the required attributes."); + } + + ShapedType input_shaped_type = + cast(op.getOperand(0).getType()); + Type input_element_type = input_shaped_type.getElementType(); + + Type quantized_element_type; + if (scales.size() == 1) { + quantized_element_type = GetPerTensorQuantizedTensorType( + rewriter, scales[0], zero_points[0], + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } else { + int32_t quantized_dimension; + if (auto quantized_dimension_attr = llvm::dyn_cast_or_null( + op->getAttr("quantization_dimension"))) { + quantized_dimension = + quantized_dimension_attr.getValue().getSExtValue(); + } else { + return op.emitError( + "quantization_dimension attribute is missing from the custom " + "call."); + } + quantized_element_type = GetPerAxisQuantizedTensorType( + rewriter, scales, zero_points, quantized_dimension, + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } + + RankedTensorType output_type = RankedTensorType::get( + input_shaped_type.getShape(), quantized_element_type); + TFL::QuantizeOp tfl_quantize_op = + TFL::QuantizeOp::create(rewriter, op.getLoc(), output_type, + /*input=*/op.getOperand(0), + /*qtype=*/TypeAttr::get(output_type)); + + rewriter.replaceOp(op, tfl_quantize_op.getOutput()); + return success(); + } + }; + + class RewriteDequantizeCustomCallOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::CustomCallOp custom_call_op, + PatternRewriter& rewriter) const final { + if (custom_call_op.getCallTargetName() != "quant.dequantize") { + return failure(); + } + Type output_type = custom_call_op.getType(0); + + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + + if (failed(FillCompositeParams(custom_call_op, scales, zero_points, + num_bits, is_signed, is_narrow_range))) { + return failure(); + } + + ShapedType output_shaped_type = cast(output_type); + Type output_element_type = output_shaped_type.getElementType(); + + Type quantized_element_type; + if (scales.size() == 1) { + quantized_element_type = GetPerTensorQuantizedTensorType( + rewriter, scales[0], zero_points[0], + /*expressed_type=*/output_element_type, num_bits, + custom_call_op->getLoc(), is_narrow_range, is_signed); + } else { + int32_t quantized_dimension; + if (auto quantized_dimension_attr = llvm::dyn_cast_or_null( + custom_call_op->getAttr("quantization_dimension"))) { + quantized_dimension = + quantized_dimension_attr.getValue().getSExtValue(); + } else { + return failure(); + } + quantized_element_type = GetPerAxisQuantizedTensorType( + rewriter, scales, zero_points, quantized_dimension, + /*expressed_type=*/output_element_type, num_bits, + custom_call_op->getLoc(), is_narrow_range, is_signed); + } + + ShapedType input_shaped_type = + cast(custom_call_op.getOperand(0).getType()); + RankedTensorType qtensor_type = RankedTensorType::get( + input_shaped_type.getShape(), quantized_element_type); + + auto custom_call_operand = custom_call_op.getOperand(0); + + Value tfl_quantize_input; + if (mlir::dyn_cast(custom_call_operand)) { + // Find the function enclosing this custom call op. + func::FuncOp func_op = GetEnclosingFunction(custom_call_op); + if (func_op == nullptr) { + return failure(); + } + + // Find the operand index of the input of the custom call op. + int arg_idx = -1; + for (int i = 0; i < func_op.getNumArguments(); ++i) { + if (func_op.getBody().front().getArgument(i) == custom_call_operand) { + arg_idx = i; + break; + } + } + if (arg_idx == -1) { + return failure(); + } + + // create a new set of operand types for the function with the type of + // the operand that feeds the custom call op changed. + SmallVector new_func_input_types; + auto func_input_types = func_op.getFunctionType().getInputs(); + for (int i = 0; i < func_input_types.size(); ++i) { + if (i != arg_idx) { + new_func_input_types.push_back(func_input_types[i]); + } else { + new_func_input_types.push_back(qtensor_type); + } + } + + auto new_func_type = + mlir::FunctionType::get(func_op.getContext(), new_func_input_types, + func_op.getFunctionType().getResults()); + + rewriter.startOpModification(func_op); + // Update the function type. + func_op.setType(new_func_type); + + // Update the block argument type. + func_op.getBody().front().getArgument(arg_idx).setType(qtensor_type); + rewriter.finalizeOpModification(func_op); + + tfl_quantize_input = func_op.getBody().front().getArgument(arg_idx); + } else { + // Get the producer of the input to dequantize + int num_operands = custom_call_op.getNumOperands(); + Value operand = custom_call_op.getOperand(num_operands - 1); + mlir::Operation* producer_op = operand.getDefiningOp(); + + if (producer_op && producer_op->hasAttr("value")) { + auto value_attr = producer_op->getAttrOfType("value"); + if (!value_attr) { + return failure(); + } + + auto qconst_op = rewriter.create( + producer_op->getLoc(), qtensor_type, TypeAttr::get(qtensor_type), + value_attr); + + tfl_quantize_input = qconst_op.getResult(); + } else if (producer_op) { + rewriter.startOpModification(producer_op); + for (OpResult result : producer_op->getResults()) { + if (result == custom_call_op.getOperand(num_operands - 1)) { + result.setType(qtensor_type); + break; + } + } + rewriter.finalizeOpModification(producer_op); + + tfl_quantize_input = operand; + } else { + return failure(); + } + } + + TFL::DequantizeOp tfl_dequantize_op = TFL::DequantizeOp::create( + rewriter, custom_call_op.getLoc(), output_type, + /*input=*/tfl_quantize_input); + rewriter.replaceOp(custom_call_op, tfl_dequantize_op.getOutput()); + return success(); + } + }; + + class RewriteFakeQuantCustomCallOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + public: + LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, + PatternRewriter& rewriter) const final { + if (op.getCallTargetName() != "quant.fake_quant" || IsDrqFakeQuant(op)) { + return failure(); + } + + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, + is_signed, is_narrow_range))) { + return op.emitError( + "fake quant custom call does not contain the required attributes."); + } + + ShapedType input_shaped_type = + cast(op.getOperand(0).getType()); + Type input_element_type = input_shaped_type.getElementType(); + + Type quantized_element_type; + if (scales.size() == 1) { + quantized_element_type = GetPerTensorQuantizedTensorType( + rewriter, scales[0], zero_points[0], + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } else { + int32_t quantized_dimension; + if (auto quantized_dimension_attr = llvm::dyn_cast_or_null( + op->getAttr("quantization_dimension"))) { + quantized_dimension = + quantized_dimension_attr.getValue().getSExtValue(); + } else { + return op.emitError( + "quantization_dimension attribute is missing from the custom " + "call."); + } + quantized_element_type = GetPerAxisQuantizedTensorType( + rewriter, scales, zero_points, quantized_dimension, + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } + + RankedTensorType output_type = RankedTensorType::get( + input_shaped_type.getShape(), quantized_element_type); + TFL::QuantizeOp tfl_quantize_op = + TFL::QuantizeOp::create(rewriter, op.getLoc(), output_type, + /*input=*/op.getOperand(0), + /*qtype=*/TypeAttr::get(output_type)); + + TFL::DequantizeOp tfl_dequantize_op = + TFL::DequantizeOp::create(rewriter, op.getLoc(), input_shaped_type, + tfl_quantize_op.getOutput()); + + rewriter.replaceOp(op, tfl_dequantize_op.getOutput()); + return success(); + } + }; + RewritePatternSet patterns(&ctx); patterns.add(&ctx); + RewriteFakeQuantCompositeOp, RewriteQuantizeCustomCallOp, + RewriteDequantizeCustomCallOp, RewriteFakeQuantCustomCallOp>( + &ctx); if (failed( applyPatternsGreedily(module, std::move(patterns), greedy_config))) { diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index 668493eca931e7..a3f758306afacc 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -91,17 +91,21 @@ struct ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs rhs = reshape.getInput(); } - DenseElementsAttr dense_constant; - if (matchPattern(rhs, m_Constant(&dense_constant))) { - constant = dense_constant; + ElementsAttr elements_constant; + if (matchPattern(rhs, m_Constant(&elements_constant))) { + constant = elements_constant; } else if (auto dq = rhs.getDefiningOp()) { Value q_input = dq.getInput(); - if (auto q = q_input.getDefiningOp()) { - if (matchPattern(q.getInput(), m_Constant(&dense_constant))) { - constant = dense_constant; + if (matchPattern(q_input, m_Constant(&elements_constant))) { + constant = elements_constant; + } else if (auto q = q_input.getDefiningOp()) { + if (matchPattern(q.getInput(), m_Constant(&elements_constant))) { + constant = elements_constant; } } else if (auto pseudo_q = q_input.getDefiningOp()) { constant = pseudo_q.getValue(); + } else if (auto const_op = q_input.getDefiningOp()) { + constant = const_op.getValue(); } } diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h index 84b7476d569fed..c9e71b5b75d6c5 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.h +++ b/tensorflow/compiler/mlir/lite/transforms/passes.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include "absl/container/flat_hash_set.h" #include "mlir/Pass/Pass.h" // from @llvm-project diff --git a/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc b/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc index 853d6801cc8a75..3859bb49924cea 100644 --- a/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc +++ b/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc @@ -15,30 +15,45 @@ limitations under the License. // This transformation pass applies some clean up steps after quantization. +#include +#include +#include #include +#include +#include #include #include #include +#include #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" #include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/OpDefinition.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project #include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" #include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h" #include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h" #include "tensorflow/compiler/mlir/lite/transforms/passes.h" +#include "tensorflow/compiler/mlir/lite/utils/utils.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h" //===----------------------------------------------------------------------===// @@ -360,6 +375,27 @@ struct RemoveVolatileOps : public OpRewritePattern { } }; +static bool MatchElementsAttr(Value val, ElementsAttr& attr) { + if (matchPattern(val, m_Constant(&attr))) { + return true; + } + if (Operation* op = val.getDefiningOp()) { + if (auto const_op = llvm::dyn_cast(op)) { + attr = const_op.getValue(); + return true; + } + if (auto qconst_op = llvm::dyn_cast(op)) { + attr = qconst_op.getValue(); + return true; + } + if (auto arith_const = llvm::dyn_cast(op)) { + attr = mlir::dyn_cast(arith_const.getValue()); + return attr != nullptr; + } + } + return false; +} + // Fold the constant quantized Transpose ops. struct FoldTransposeOp : public OpRewritePattern { explicit FoldTransposeOp(MLIRContext* context) @@ -394,16 +430,49 @@ struct FoldTransposeOp : public OpRewritePattern { } } + void ComputePermutationRaw(ArrayRef perm, + ArrayRef output_shape, + const char* raw_input, int element_byte_size, + int output_axis, char*& raw_output, + SmallVectorImpl& current_input_index, + ArrayRef input_shape) const { + const int num_dimensions = output_shape.size(); + assert(output_axis < num_dimensions); + const int input_axis = perm[output_axis]; + for (int i = 0; i < output_shape[output_axis]; ++i) { + current_input_index[input_axis] = i; + const bool is_last_axis = output_axis == num_dimensions - 1; + if (is_last_axis) { + uint64_t input_flat_index = 0; + uint64_t stride = 1; + for (int d = num_dimensions - 1; d >= 0; --d) { + input_flat_index += current_input_index[d] * stride; + stride *= input_shape[d]; + } + memcpy(raw_output, raw_input + input_flat_index * element_byte_size, + element_byte_size); + raw_output += element_byte_size; + } else { + ComputePermutationRaw(perm, output_shape, raw_input, element_byte_size, + output_axis + 1, raw_output, current_input_index, + input_shape); + } + } + } + LogicalResult matchAndRewrite(TransposeOp op, PatternRewriter& rewriter) const override { Operation* def_op = op.getInput().getDefiningOp(); auto qconst_op = llvm::dyn_cast_or_null(def_op); if (qconst_op == nullptr) return failure(); - DenseIntElementsAttr perm_tensor; - if (!matchPattern(op.getPerm(), m_Constant(&perm_tensor))) return failure(); + ElementsAttr perm_attr; + if (!MatchElementsAttr(op.getPerm(), perm_attr)) return failure(); + auto int_perm_attr = mlir::dyn_cast(perm_attr); + if (!int_perm_attr) return failure(); - auto output_element_type = getElementTypeOrSelf(op.getOutput().getType()); + auto result_type = mlir::cast(op.getOutput().getType()); + auto output_element_type = result_type.getElementType(); if (!mlir::isa(output_element_type) && !mlir::isa(output_element_type)) { return failure(); @@ -411,48 +480,164 @@ struct FoldTransposeOp : public OpRewritePattern { ElementsAttr input_tensor = qconst_op.getValue(); - assert(perm_tensor.getType().getRank() == 1); const int num_dimensions = input_tensor.getShapedType().getRank(); - assert(perm_tensor.getType().getNumElements() == num_dimensions); - ArrayRef input_shape = input_tensor.getShapedType().getShape(); - auto output_type = mlir::cast(op.getOutput().getType()); SmallVector perm; + for (const APInt& it : int_perm_attr.getValues()) { + perm.push_back(it.getSExtValue()); + } + if (perm.size() != num_dimensions) return failure(); + SmallVector output_shape; for (int i = 0; i < num_dimensions; ++i) { - perm.push_back(perm_tensor.getValues()[i].getInt()); output_shape.push_back(input_shape[perm[i]]); + assert(!result_type.hasStaticShape() || + result_type.getShape()[i] == output_shape[i]); + } + + if (auto dense_input = mlir::dyn_cast(input_tensor)) { + std::vector new_values; + new_values.reserve(input_tensor.getShapedType().getNumElements()); + std::vector input_indices(num_dimensions); + ComputePermutation(dense_input, perm, output_shape, num_dimensions, + /*output_axis=*/0, &input_indices, &new_values); + RankedTensorType values_type; + if (mlir::isa(output_element_type)) { + values_type = RankedTensorType::get( + output_shape, + mlir::cast(output_element_type) + .getStorageType()); + } else { + values_type = RankedTensorType::get( + output_shape, + mlir::cast(output_element_type) + .getStorageType()); + } - // Check that the derived output shape matches the static shape. - assert(!output_type.hasStaticShape() || - output_type.getShape()[i] == output_shape[i]); + rewriter.replaceOpWithNewOp( + op, TypeAttr::get(result_type), + DenseIntElementsAttr::get(values_type, new_values)); + return success(); } - std::vector new_values; - new_values.reserve(input_tensor.getShapedType().getNumElements()); - std::vector input_indices(num_dimensions); - ComputePermutation(input_tensor, perm, output_shape, num_dimensions, - /*output_axis=*/0, &input_indices, &new_values); - auto result_type = - RankedTensorType::get(output_shape, output_type.getElementType()); - RankedTensorType values_type; - if (mlir::isa(output_element_type)) { - values_type = RankedTensorType::get( - output_shape, - mlir::cast(output_type.getElementType()) - .getStorageType()); - } else { - values_type = RankedTensorType::get( - output_shape, mlir::cast( - output_type.getElementType()) - .getStorageType()); + if (auto dense_res = + mlir::dyn_cast(input_tensor)) { + AsmResourceBlob* blob = dense_res.getRawHandle().getBlob(); + if (!blob && dense_res.getRawHandle().getResource()) { + blob = dense_res.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return failure(); + + uint32_t storage_bit_width = 8; + if (auto u = mlir::dyn_cast( + output_element_type)) { + storage_bit_width = u.getStorageTypeIntegralWidth(); + } else if (auto p = mlir::dyn_cast( + output_element_type)) { + storage_bit_width = p.getStorageTypeIntegralWidth(); + } + + size_t num_elements = result_type.getNumElements(); + if (storage_bit_width == 4) { + size_t out_byte_size = (num_elements + 1) / 2; + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + out_byte_size, /*align=*/64, /*dataIsMutable=*/true); + char* raw_output = + const_cast(raw_output_blob.getDataAs().data()); + std::memset(raw_output, 0, out_byte_size); + + const uint8_t* raw_input_u8 = + reinterpret_cast(blob->getData().data()); + std::vector unpacked_input(num_elements); + for (size_t i = 0; i < num_elements; ++i) { + uint8_t byte = raw_input_u8[i / 2]; + uint8_t nibble = (i % 2 == 0) ? (byte & 0x0F) : ((byte >> 4) & 0x0F); + unpacked_input[i] = static_cast(nibble); + } + + std::vector unpacked_output(num_elements); + char* unpacked_output_ptr = unpacked_output.data(); + SmallVector current_input_index(num_dimensions, 0); + ComputePermutationRaw(perm, output_shape, unpacked_input.data(), + /*element_byte_size=*/1, + /*output_axis=*/0, unpacked_output_ptr, + current_input_index, input_shape); + + for (size_t i = 0; i < num_elements; ++i) { + uint8_t nibble = static_cast(unpacked_output[i]) & 0x0F; + if (i % 2 == 0) { + raw_output[i / 2] |= nibble; + } else { + raw_output[i / 2] |= (nibble << 4); + } + } + + DenseResourceElementsAttr new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_res.getRawHandle().getKey(), + std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); + } else if (storage_bit_width == 2) { + size_t out_byte_size = (num_elements + 3) / 4; + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + out_byte_size, /*align=*/64, /*dataIsMutable=*/true); + char* raw_output = + const_cast(raw_output_blob.getDataAs().data()); + std::memset(raw_output, 0, out_byte_size); + + const uint8_t* raw_input_u8 = + reinterpret_cast(blob->getData().data()); + std::vector unpacked_input(num_elements); + for (size_t i = 0; i < num_elements; ++i) { + uint8_t byte = raw_input_u8[i / 4]; + uint8_t val = (byte >> ((i % 4) * 2)) & 0x03; + unpacked_input[i] = static_cast(val); + } + + std::vector unpacked_output(num_elements); + char* unpacked_output_ptr = unpacked_output.data(); + SmallVector current_input_index(num_dimensions, 0); + ComputePermutationRaw(perm, output_shape, unpacked_input.data(), + /*element_byte_size=*/1, + /*output_axis=*/0, unpacked_output_ptr, + current_input_index, input_shape); + + for (size_t i = 0; i < num_elements; ++i) { + uint8_t val = static_cast(unpacked_output[i]) & 0x03; + raw_output[i / 4] |= (val << ((i % 4) * 2)); + } + + DenseResourceElementsAttr new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_res.getRawHandle().getKey(), + std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); + } + + const int element_byte_size = std::max(1, storage_bit_width / 8); + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + blob->getData().size(), /*align=*/64, /*dataIsMutable=*/true); + char* raw_output = + const_cast(raw_output_blob.getDataAs().data()); + const char* raw_input = blob->getData().data(); + + SmallVector current_input_index(num_dimensions, 0); + ComputePermutationRaw(perm, output_shape, raw_input, element_byte_size, + /*output_axis=*/0, raw_output, current_input_index, + input_shape); + + DenseResourceElementsAttr new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_res.getRawHandle().getKey(), + std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); } - rewriter.replaceOpWithNewOp( - op, TypeAttr::get(result_type), - DenseIntElementsAttr::get(values_type, new_values)); - return success(); + return failure(); } }; @@ -471,10 +656,6 @@ struct FoldReshapeOp : public OpRewritePattern { return rewriter.notifyMatchFailure(op, "input is not a QConstOp."); } - auto dense_elements = - mlir::dyn_cast_or_null(qconst_op.getValue()); - if (dense_elements == nullptr) return failure(); - auto output_element_type = getElementTypeOrSelf(op.getType()); if (!mlir::isa(output_element_type)) { return rewriter.notifyMatchFailure(op, "output type is not quantized."); @@ -488,9 +669,10 @@ struct FoldReshapeOp : public OpRewritePattern { // If the result type isn't static, tries to derive the result type from // the #2 operand. if (!result_type.hasStaticShape()) { - DenseIntElementsAttr shape_elements; - if (!matchPattern(op.getShape(), m_Constant(&shape_elements))) - return failure(); + ElementsAttr shape_attr; + if (!MatchElementsAttr(op.getShape(), shape_attr)) return failure(); + auto shape_elements = mlir::dyn_cast(shape_attr); + if (!shape_elements) return failure(); SmallVector shape_data; for (const APInt& it : shape_elements.getValues()) { @@ -499,24 +681,55 @@ struct FoldReshapeOp : public OpRewritePattern { result_type = RankedTensorType::get(shape_data, input_type.getElementType()); } + RankedTensorType values_type; - if (mlir::isa(output_element_type)) { + if (auto uniform_qtype = + mlir::dyn_cast(output_element_type)) { + values_type = RankedTensorType::get(result_type.getShape(), + uniform_qtype.getStorageType()); + } else { values_type = RankedTensorType::get( result_type.getShape(), - mlir::cast(result_type.getElementType()) + mlir::cast(output_element_type) .getStorageType()); - } else { - values_type = - RankedTensorType::get(result_type.getShape(), - mlir::cast( - result_type.getElementType()) - .getStorageType()); } - DenseElementsAttr reshaped_elements = dense_elements.reshape(values_type); - rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), - reshaped_elements); - return success(); + ElementsAttr value_attr = qconst_op.getValue(); + if (auto dense_elements = mlir::dyn_cast(value_attr)) { + DenseElementsAttr reshaped_elements = dense_elements.reshape(values_type); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + reshaped_elements); + return success(); + } + + if (auto dense_resource_elements = + mlir::dyn_cast(value_attr)) { + AsmResourceBlob* blob = dense_resource_elements.getRawHandle().getBlob(); + if (!blob && dense_resource_elements.getRawHandle().getResource()) { + blob = dense_resource_elements.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return failure(); + + DenseResourceElementsAttr new_res_attr; + if (qconst_op.getOutput().hasOneUse()) { + new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_resource_elements.getRawHandle().getKey(), + std::move(*blob)); + } else { + auto new_blob = mlir::HeapAsmResourceBlob::allocate( + blob->getData().size(), /*align=*/64, true); + memcpy(const_cast(new_blob.getData().data()), + blob->getData().data(), blob->getData().size()); + new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_resource_elements.getRawHandle().getKey(), + std::move(new_blob)); + } + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); + } + + return failure(); } }; diff --git a/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc b/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc index 79f832ad31ab50..a293d9ff9f6a7d 100644 --- a/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc @@ -78,6 +78,14 @@ LogicalResult IsDrqTensor(mlir::Value value, mlir::Value& fq_input) { return success(); } } + if (auto custom_call_op = llvm::dyn_cast_or_null( + value.getDefiningOp())) { + if (IsDrqFakeQuant(custom_call_op)) { + int num_operands = custom_call_op.getNumOperands(); + fq_input = custom_call_op.getOperand(num_operands - 1); + return success(); + } + } return failure(); } @@ -368,6 +376,21 @@ class RemoveUnusedFQ : public OpRewritePattern { } }; +class RemoveUnusedCustomCallFQ + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, + PatternRewriter& rewriter) const final { + if (IsDrqFakeQuant(op) && op->getUses().empty()) { + rewriter.eraseOp(op); + return success(); + } + return rewriter.notifyMatchFailure( + op, "is not a drq fake quant custom call with no uses."); + } +}; + // Pushes a drq fake quant op forward through a pad op. // This is to allow DRQ FQ to be fused into the DRQ op. // drq_fake_quant(input) -> pad -> output @@ -404,19 +427,45 @@ class PushForwardDrqFQ : public OpRewritePattern { TFL::PadOp::create(rewriter, pad_op.getLoc(), pad_op.getType(), float_input, pad_op.getPadding()); - // Create a new drq fake quant op. - // Operands are the same, except for the last one. - SmallVector new_drq_operands; - for (mlir::Value operand : drq_fq_op.getOperands().drop_back()) { - new_drq_operands.push_back(operand); - } - new_drq_operands.push_back(new_pad_op.getResult()); + Operation* new_drq_fq_op = rewriter.clone(*drq_fq_op.getOperation()); + new_drq_fq_op->setOperand(new_drq_fq_op->getNumOperands() - 1, + new_pad_op.getResult()); + + rewriter.replaceOp(pad_op, new_drq_fq_op->getResult(0)); + return success(); + } +}; + +// Fixes keep_num_dims option of FC if output dims is different from input dims +// though keep_num_dims is true. It happens when FC's input has changed after +// quantization, e.g. by IsDrqTensor(). +// Sets keep_num_dims to false if that's the case. Otherwise, it's not +// compatible with GPU. See CheckGpuDelegateCompatibility() in +// third_party/tensorflow/lite/tools/versioning/gpu_compatibility.cc. +// Note that if FC is followed by Reshape, the keep_num_dims will be set to true +// with a correct shape later by EnableFullyConnectedKeepNumDimsBeforeReshape() +// in optimize pass. +struct FixFullyConnectedKeepNumDims + : public OpRewritePattern { + explicit FixFullyConnectedKeepNumDims(MLIRContext* context) + : OpRewritePattern(context, /*benefit=*/0) {} + + LogicalResult matchAndRewrite(TFL::FullyConnectedOp fc, + PatternRewriter& rewriter) const override { + if (!fc.getKeepNumDims()) return failure(); - auto new_drq_fq_op = stablehlo::CompositeOp::create( - rewriter, drq_fq_op.getLoc(), pad_op.getType(), new_drq_operands, - drq_fq_op->getAttrs()); + auto input_ty = + mlir::dyn_cast_or_null(fc.getInput().getType()); + auto fc_ty = mlir::dyn_cast_or_null(fc.getType(0)); + if (!input_ty || !fc_ty) return failure(); + + auto input_shape = input_ty.getShape(); + auto fc_shape = fc_ty.getShape(); + if (input_shape.size() == fc_shape.size()) { + return failure(); + } - rewriter.replaceOp(pad_op, new_drq_fq_op.getResult(0)); + rewriter.modifyOpInPlace(fc, [&]() { fc.setKeepNumDims(false); }); return success(); } }; @@ -484,7 +533,7 @@ class QuantizeConstPattern : public OpRewritePattern { : OpRewritePattern(context) {} LogicalResult matchAndRewrite(QuantizeOp op, PatternRewriter& rewriter) const override { - DenseFPElementsAttr attr; + ElementsAttr attr; if (matchPattern(op.getInput(), m_Constant(&attr))) { auto qtype = op.getQtypeAttr(); Attribute quantized_attr = mlir::TFL::Quantize(attr, qtype.getValue()); @@ -524,8 +573,10 @@ void FuseQDQPass::runOnOperation() { mlir::ModuleOp module = getOperation(); RewritePatternSet patterns(ctx); - patterns.add(ctx); + patterns + .add(ctx); // Configure the greedy pattern rewrite driver. GreedyRewriteConfig greedy_config; diff --git a/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc b/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc index 865ae58a04bfde..aa51e8f92b9e1f 100644 --- a/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc +++ b/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc @@ -80,7 +80,7 @@ class CheckRangeAndConvertI8ToI4 : public OpRewritePattern { mlir::RankedTensorType::get(const_type.getShape(), builder.getI4Type()); auto newAttr = DenseElementsAttr::getFromRawBuffer( shaped_type, mlir::cast(op.getValue()).getRawData()); - rewriter.replaceOpWithNewOp(op, newAttr); + rewriter.replaceOpWithNewOp(op, newAttr); return success(); } @@ -108,7 +108,7 @@ class SanitizeGatherOpOutputToI4 : public OpRewritePattern { Builder builder(op.getContext()); auto new_gather_op = TFL::GatherOp::create(rewriter, op.getLoc(), - /*result=*/ + /*resultTypes=*/ mlir::cast(op.getResult().getType()) .clone(builder.getI4Type()), /*operands=*/op.getOperands(), op->getAttrs()); From fbf1d24f7cdf0b37f6b708237444ca02edf7961c Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Mon, 31 Aug 2026 19:45:36 -0700 Subject: [PATCH 6/9] [XLA:CPU] Use 10x longer timeout for CPU collective calls in sanitizer builds. Under sanitizers (ASan, TSan, MSan, HWASan), multi-threaded CPU execution and JIT compilation incur significant overhead, which can cause collective call rendezvous to exceed the default 40-second timeout. Scale CPU collective call timeouts by 10x in sanitizer builds to prevent flakiness. PiperOrigin-RevId: 974216046 --- third_party/xla/xla/BUILD | 1 + third_party/xla/xla/debug_options_flags.cc | 17 ++++++++++++++--- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 0325cc3a0d4984..fc24f6ecb62a80 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1311,6 +1311,7 @@ cc_library( "//xla/tsl/util:command_line_flags", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", + "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 66c478dd3879da..73534629847b27 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/call_once.h" +#include "absl/base/config.h" // IWYU pragma: keep #include "absl/base/no_destructor.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -564,9 +565,19 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_native_emitter_tune_unroll_factor_for_loops(false); opts.set_xla_gpu_experimental_use_ragged_dot_fusion(false); - opts.set_xla_cpu_collective_call_warn_stuck_seconds(20); - opts.set_xla_cpu_collective_call_terminate_timeout_seconds(40); - opts.set_xla_cpu_collective_timeout_seconds(30 * 60); +#if defined(ABSL_HAVE_ADDRESS_SANITIZER) || \ + defined(ABSL_HAVE_HWADDRESS_SANITIZER) || \ + defined(ABSL_HAVE_MEMORY_SANITIZER) || defined(ABSL_HAVE_THREAD_SANITIZER) + constexpr int kSanitizerMultiplier = 10; +#else + constexpr int kSanitizerMultiplier = 1; +#endif + + opts.set_xla_cpu_collective_call_warn_stuck_seconds(20 * + kSanitizerMultiplier); + opts.set_xla_cpu_collective_call_terminate_timeout_seconds( + 40 * kSanitizerMultiplier); + opts.set_xla_cpu_collective_timeout_seconds(30 * 60 * kSanitizerMultiplier); opts.set_xla_keep_shardings_after_spmd(false); opts.set_xla_enable_hlo_sharding_v3(false); From eb01ac25a8e3229729a599b12dc4e656e5c42721 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 22:57:56 -0700 Subject: [PATCH 7/9] Automated Code Change PiperOrigin-RevId: 974283733 --- tensorflow/core/kernels/gather_nd_op.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/gather_nd_op.h b/tensorflow/core/kernels/gather_nd_op.h index 352289b4c554d0..1295b5e8c5b7f0 100644 --- a/tensorflow/core/kernels/gather_nd_op.h +++ b/tensorflow/core/kernels/gather_nd_op.h @@ -156,7 +156,7 @@ absl::Status DoGatherNd( using CPUDevice = Eigen::ThreadPoolDevice; const bool check_bad_indices = - ((std::is_same::value && + ((std::is_same_v && bad_indices_policy == BadIndicesPolicy::kDefault) || bad_indices_policy == BadIndicesPolicy::kError); if (check_bad_indices && bad_i >= 0) { From 0b7857da8bfda1b262aa52957846bff34794fb83 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 23:00:02 -0700 Subject: [PATCH 8/9] Automated Code Change PiperOrigin-RevId: 974284429 --- .../xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc | 1 + third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc b/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc index 7ebcc63fd453b5..27bfad7ff755d7 100644 --- a/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include #include #include diff --git a/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc b/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc index 23f399566ea852..0ad31a3f932563 100644 --- a/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include From e45f2985f35113d457e5707a4af2c9f66413b0ae Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Mon, 31 Aug 2026 23:08:13 -0700 Subject: [PATCH 9/9] Add range check to a few StableHLO ops PiperOrigin-RevId: 974287919 --- .../lite/core/api/flatbuffer_conversions.cc | 26 ++--- .../core/api/flatbuffer_conversions_test.cc | 100 ++++++++++++++++++ .../lite/core/api/flatbuffer_conversions.cc | 26 ++--- .../core/api/flatbuffer_conversions_test.cc | 93 ++++++++++++++++ 4 files changed, 215 insertions(+), 30 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc index aace2080a0481f..2c4b62b39ec179 100644 --- a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc @@ -406,12 +406,12 @@ void CheckParsePointerParams(const Operator* op, TFLITE_DCHECK(builtin_data != nullptr); } -// Copies the contents from the flatbuffer int vector `flatbuffer` into the +// Copies the contents from the flatbuffer int vector `flat_vector` into the // int array `buffer`. `flat_vector` and `buffer` represent the same // configuration operation for a given operation. template static absl::Status FlatBufferIntVectorToArray( - int max_size_of_buffer, const flatbuffers::Vector* flat_vector, + size_t max_size_of_buffer, const flatbuffers::Vector* flat_vector, DataType* buffer, const char* op_name) { if (!flat_vector) { auto error_message = absl::StrFormat( @@ -2498,7 +2498,7 @@ absl::Status ParseStablehloScatter(const Operator* op, if (schema_params->update_window_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->update_window_dims()->size() * sizeof(int64_t), + sizeof(params->update_window_dims), schema_params->update_window_dims(), params->update_window_dims, "stablehlo_scatter")); params->num_update_window_dims = @@ -2507,7 +2507,7 @@ absl::Status ParseStablehloScatter(const Operator* op, if (schema_params->inserted_window_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->inserted_window_dims()->size() * sizeof(int64_t), + sizeof(params->inserted_window_dims), schema_params->inserted_window_dims(), params->inserted_window_dims, "stablehlo_scatter")); params->num_inserted_window_dims = @@ -2516,8 +2516,7 @@ absl::Status ParseStablehloScatter(const Operator* op, if (schema_params->scatter_dims_to_operand_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->scatter_dims_to_operand_dims()->size() * - sizeof(int64_t), + sizeof(params->scatter_dims_to_operand_dims), schema_params->scatter_dims_to_operand_dims(), params->scatter_dims_to_operand_dims, "stablehlo_scatter")); params->num_scatter_dims_to_operand_dims = @@ -2579,8 +2578,7 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params != nullptr) { if (schema_params->offset_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - /*max_size_of_buffer=*/schema_params->offset_dims()->size() * - sizeof(int64_t), + /*max_size_of_buffer=*/sizeof(params->offset_dims), /*flat_vector=*/schema_params->offset_dims(), /*buffer=*/params->offset_dims, /*op_name=*/"stablehlo_gather")); @@ -2589,7 +2587,7 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params->collapsed_slice_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->collapsed_slice_dims()->size() * sizeof(int64_t), + sizeof(params->collapsed_slice_dims), schema_params->collapsed_slice_dims(), params->collapsed_slice_dims, "stablehlo_gather")); params->num_collapsed_slice_dims = @@ -2598,9 +2596,8 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params->start_index_map()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->start_index_map()->size() * sizeof(int64_t), - schema_params->start_index_map(), params->start_index_map, - "stablehlo_gather")); + sizeof(params->start_index_map), schema_params->start_index_map(), + params->start_index_map, "stablehlo_gather")); params->num_start_index_map = schema_params->start_index_map()->size(); } @@ -2608,9 +2605,8 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params->slice_sizes()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->slice_sizes()->size() * sizeof(int64_t), - schema_params->slice_sizes(), params->slice_sizes, - "stablehlo_gather")); + sizeof(params->slice_sizes), schema_params->slice_sizes(), + params->slice_sizes, "stablehlo_gather")); params->num_slice_sizes = schema_params->slice_sizes()->size(); } diff --git a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc index 7bfd7fa15df945..cd6f17154219d2 100644 --- a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc +++ b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc @@ -49,18 +49,24 @@ using tflite::BuiltinOperator_CUSTOM; using tflite::BuiltinOperator_FULLY_CONNECTED; using tflite::BuiltinOperator_RESHAPE; using tflite::BuiltinOperator_SQUEEZE; +using tflite::BuiltinOperator_STABLEHLO_GATHER; using tflite::BuiltinOperator_STABLEHLO_PAD; using tflite::BuiltinOperator_STABLEHLO_REDUCE_WINDOW; +using tflite::BuiltinOperator_STABLEHLO_SCATTER; +using tflite::BuiltinOptions2_StablehloGatherOptions; using tflite::BuiltinOptions2_StablehloPadOptions; using tflite::BuiltinOptions2_StablehloReduceWindowOptions; +using tflite::BuiltinOptions2_StablehloScatterOptions; using tflite::BuiltinOptions_Conv2DOptions; using tflite::BuiltinOptions_FullyConnectedOptions; using tflite::BuiltinOptions_NONE; using tflite::BuiltinOptions_ReshapeOptions; using tflite::CreateReshapeOptions; using tflite::CreateSqueezeOptions; +using tflite::CreateStablehloGatherOptions; using tflite::CreateStablehloPadOptions; using tflite::CreateStablehloReduceWindowOptions; +using tflite::CreateStablehloScatterOptions; using tflite::FullyConnectedOptionsWeightsFormat; using tflite::Padding_SAME; using tflite::TensorType_BFLOAT16; @@ -871,5 +877,99 @@ TEST_F(StablehloPadFlatbufferConversionsTest, DeathTests) { ""); } +class StablehloScatterFlatbufferConversionsTest + : public FlatbufferConversionsTest { + protected: + static constexpr int kMaxDims = + TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT; + + const Operator* BuildScatterOperator(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + const auto update_window_dims = builder_.CreateVector( + std::vector(update_window_dims_count, 1)); + const auto inserted_window_dims = builder_.CreateVector( + std::vector(inserted_window_dims_count, 1)); + const auto scatter_dims = + builder_.CreateVector(std::vector(scatter_dims_count, 1)); + return BuildTestOperator( + BuiltinOptions2_StablehloScatterOptions, + CreateStablehloScatterOptions( + builder_, /*indices_are_sorted=*/true, update_window_dims, + inserted_window_dims, scatter_dims, + /*index_vector_dim=*/0, /*unique_indices=*/false, + /*update_computation_subgraph_index=*/0) + .Union()); + } + + void ExpectTooManyDimensions(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + TfLiteStablehloScatterParams* output_data = nullptr; + const auto status = ParseOpData( + BuildScatterOperator(update_window_dims_count, + inserted_window_dims_count, scatter_dims_count), + BuiltinOperator_STABLEHLO_SCATTER, &mock_allocator_, + (void**)&output_data); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(status.message(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_scatter'.")); + } +}; + +TEST_F(StablehloScatterFlatbufferConversionsTest, AcceptsMaximumDimensions) { + TfLiteStablehloScatterParams* output_data = nullptr; + const auto status = + ParseOpData(BuildScatterOperator(kMaxDims, kMaxDims, kMaxDims), + BuiltinOperator_STABLEHLO_SCATTER, &mock_allocator_, + (void**)&output_data); + EXPECT_TRUE(status.ok()); + ASSERT_NE(output_data, nullptr); + EXPECT_EQ(output_data->num_update_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_inserted_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_scatter_dims_to_operand_dims, kMaxDims); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyUpdateWindowDimensions) { + ExpectTooManyDimensions(kMaxDims + 1, 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyInsertedWindowDimensions) { + ExpectTooManyDimensions(1, kMaxDims + 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyScatterDimensions) { + ExpectTooManyDimensions(1, 1, kMaxDims + 1); +} + +TEST_F(FlatbufferConversionsTest, + ParseStablehloGatherRejectsTooManyDimensions) { + std::vector too_many_dims( + TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT + 1, 1); + const Operator* op = BuildTestOperator( + BuiltinOptions2_StablehloGatherOptions, + CreateStablehloGatherOptions( + builder_, /*offset_dims=*/builder_.CreateVector(too_many_dims), + /*collapsed_slice_dims=*/builder_.CreateVector({1}), + /*start_index_map=*/builder_.CreateVector({1}), + /*index_vector_dim=*/0, + /*slice_sizes=*/builder_.CreateVector({1}), + /*indices_are_sorted=*/true) + .Union()); + void* output_data = nullptr; + const auto status = ParseOpData(op, BuiltinOperator_STABLEHLO_GATHER, + &mock_allocator_, &output_data); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(status.message(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_gather'.")); +} + } // namespace flatbuffer_conversions } // namespace tflite_file diff --git a/tensorflow/lite/core/api/flatbuffer_conversions.cc b/tensorflow/lite/core/api/flatbuffer_conversions.cc index d47767c79b0c3b..f0d93ff9d2bfc6 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions.cc @@ -77,12 +77,12 @@ void CheckParsePointerParams(const Operator* op, ErrorReporter* error_reporter, TFLITE_DCHECK(builtin_data != nullptr); } -// Copies the contents from the flatbuffer int vector `flatbuffer` into the +// Copies the contents from the flatbuffer int vector `flat_vector` into the // int array `buffer`. `flat_vector` and `buffer` represent the same // configuration operation for a given operation. template static TfLiteStatus FlatBufferIntVectorToArray( - int max_size_of_buffer, const flatbuffers::Vector* flat_vector, + size_t max_size_of_buffer, const flatbuffers::Vector* flat_vector, DataType* buffer, ErrorReporter* error_reporter, const char* op_name) { if (!flat_vector) { TF_LITE_REPORT_ERROR(error_reporter, @@ -2248,7 +2248,7 @@ TfLiteStatus ParseStablehloScatter(const Operator* op, if (schema_params->update_window_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->update_window_dims()->size() * sizeof(int64_t), + sizeof(params->update_window_dims), schema_params->update_window_dims(), params->update_window_dims, error_reporter, "stablehlo_scatter")); params->num_update_window_dims = @@ -2257,7 +2257,7 @@ TfLiteStatus ParseStablehloScatter(const Operator* op, if (schema_params->inserted_window_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->inserted_window_dims()->size() * sizeof(int64_t), + sizeof(params->inserted_window_dims), schema_params->inserted_window_dims(), params->inserted_window_dims, error_reporter, "stablehlo_scatter")); params->num_inserted_window_dims = @@ -2266,8 +2266,7 @@ TfLiteStatus ParseStablehloScatter(const Operator* op, if (schema_params->scatter_dims_to_operand_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->scatter_dims_to_operand_dims()->size() * - sizeof(int64_t), + sizeof(params->scatter_dims_to_operand_dims), schema_params->scatter_dims_to_operand_dims(), params->scatter_dims_to_operand_dims, error_reporter, "stablehlo_scatter")); @@ -2332,8 +2331,7 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params != nullptr) { if (schema_params->offset_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - /*max_size_of_buffer=*/schema_params->offset_dims()->size() * - sizeof(int64_t), + /*max_size_of_buffer=*/sizeof(params->offset_dims), /*flat_vector=*/schema_params->offset_dims(), /*buffer=*/params->offset_dims, /*error_reporter=*/error_reporter, /*op_name=*/"stablehlo_gather")); @@ -2342,7 +2340,7 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params->collapsed_slice_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->collapsed_slice_dims()->size() * sizeof(int64_t), + sizeof(params->collapsed_slice_dims), schema_params->collapsed_slice_dims(), params->collapsed_slice_dims, error_reporter, "stablehlo_gather")); params->num_collapsed_slice_dims = @@ -2351,9 +2349,8 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params->start_index_map()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->start_index_map()->size() * sizeof(int64_t), - schema_params->start_index_map(), params->start_index_map, - error_reporter, "stablehlo_gather")); + sizeof(params->start_index_map), schema_params->start_index_map(), + params->start_index_map, error_reporter, "stablehlo_gather")); params->num_start_index_map = schema_params->start_index_map()->size(); } @@ -2361,9 +2358,8 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params->slice_sizes()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->slice_sizes()->size() * sizeof(int64_t), - schema_params->slice_sizes(), params->slice_sizes, error_reporter, - "stablehlo_gather")); + sizeof(params->slice_sizes), schema_params->slice_sizes(), + params->slice_sizes, error_reporter, "stablehlo_gather")); params->num_slice_sizes = schema_params->slice_sizes()->size(); } diff --git a/tensorflow/lite/core/api/flatbuffer_conversions_test.cc b/tensorflow/lite/core/api/flatbuffer_conversions_test.cc index 0903795850a417..093022deeb3798 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions_test.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions_test.cc @@ -943,4 +943,97 @@ TEST_F(StablehloPadFlatbufferConversionsTest, DeathTests) { ""); } +class StablehloScatterFlatbufferConversionsTest + : public FlatbufferConversionsTest { + protected: + static constexpr int kMaxDims = + TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT; + + const Operator* BuildScatterOperator(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + const auto update_window_dims = builder_.CreateVector( + std::vector(update_window_dims_count, 1)); + const auto inserted_window_dims = builder_.CreateVector( + std::vector(inserted_window_dims_count, 1)); + const auto scatter_dims = + builder_.CreateVector(std::vector(scatter_dims_count, 1)); + return BuildTestOperator( + BuiltinOptions2_StablehloScatterOptions, + CreateStablehloScatterOptions( + builder_, /*indices_are_sorted=*/true, update_window_dims, + inserted_window_dims, scatter_dims, + /*index_vector_dim=*/0, /*unique_indices=*/false, + /*update_computation_subgraph_index=*/0) + .Union()); + } + + void ExpectTooManyDimensions(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + TfLiteStablehloScatterParams* output_data = nullptr; + EXPECT_EQ(ParseOpData(BuildScatterOperator(update_window_dims_count, + inserted_window_dims_count, + scatter_dims_count), + BuiltinOperator_STABLEHLO_SCATTER, &mock_reporter_, + &mock_allocator_, (void**)&output_data), + kTfLiteError); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(mock_reporter_.GetString(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_scatter'.")); + } +}; + +TEST_F(StablehloScatterFlatbufferConversionsTest, AcceptsMaximumDimensions) { + TfLiteStablehloScatterParams* output_data = nullptr; + EXPECT_EQ(ParseOpData(BuildScatterOperator(kMaxDims, kMaxDims, kMaxDims), + BuiltinOperator_STABLEHLO_SCATTER, &mock_reporter_, + &mock_allocator_, (void**)&output_data), + kTfLiteOk); + ASSERT_NE(output_data, nullptr); + EXPECT_EQ(output_data->num_update_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_inserted_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_scatter_dims_to_operand_dims, kMaxDims); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyUpdateWindowDimensions) { + ExpectTooManyDimensions(kMaxDims + 1, 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyInsertedWindowDimensions) { + ExpectTooManyDimensions(1, kMaxDims + 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyScatterDimensions) { + ExpectTooManyDimensions(1, 1, kMaxDims + 1); +} + +TEST_F(FlatbufferConversionsTest, + ParseStablehloGatherRejectsTooManyDimensions) { + std::vector too_many_dims( + TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT + 1, 1); + const Operator* op = BuildTestOperator( + BuiltinOptions2_StablehloGatherOptions, + CreateStablehloGatherOptions( + builder_, /*offset_dims=*/builder_.CreateVector(too_many_dims), + /*collapsed_slice_dims=*/builder_.CreateVector({1}), + /*start_index_map=*/builder_.CreateVector({1}), + /*index_vector_dim=*/0, + /*slice_sizes=*/builder_.CreateVector({1}), + /*indices_are_sorted=*/true) + .Union()); + void* output_data = nullptr; + EXPECT_EQ(ParseOpData(op, BuiltinOperator_STABLEHLO_GATHER, &mock_reporter_, + &mock_allocator_, &output_data), + kTfLiteError); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(mock_reporter_.GetString(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_gather'.")); +} + } // namespace tflite