[pull] master from tensorflow:master

tensorflow/compiler/mlir/lite/BUILD

@@ -1390,6 +1390,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:Dialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
@@ -1827,7 +1828,6 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
-        "@eigen_archive//:eigen3",
         "@flatbuffers",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",

tensorflow/compiler/mlir/lite/python/BUILD

@@ -237,7 +237,6 @@ cc_library(
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:ReconcileUnrealizedCasts",
         "@llvm-project//mlir:Support",
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_passes",

tensorflow/compiler/mlir/lite/python/slim_model_importer.cc

@@ -222,6 +222,7 @@ absl::StatusOr<OwningOpRef<ModuleOp>> LoadSlimModel(
 
     mlir::PassManager pm(context);
     pm.addPass(mlir::odml::CreateDropShapeAssertionsPass());
+    pm.addPass(mlir::odml::CreateLegalizeVhloQuantCustomCallsPass());
     pm.addPass(mlir::stablehlo::createVhloLegalizeToStablehloPass());
     if (mlir::failed(pm.run(*module))) {
       return absl::InternalError("Failed to legalize VHLO to StableHLO.");

tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td

@@ -31,9 +31,20 @@ def CreateTFLCastToInt32Op : NativeCodeCall<
 def ConstDefaultResultAccuracyAttr :
   ConstantAttr<MHLO_ResultAccuracyAttr, "::mlir::mhlo::ResultAccuracyMode::DEFAULT">;
 
+def IsBuildableWithArithConstant : Constraint<
+  CPred<"::mlir::arith::ConstantOp::isBuildableWith($0, $1.getType())">
+>;
+
 def : Pat<
   (MHLO_ConstantOp:$output $value),
   (Arith_ConstantOp $value),
+  [(IsBuildableWithArithConstant $value, $output)],
+  [],
+  (addBenefit 10)>;
+
+def : Pat<
+  (MHLO_ConstantOp:$output $value),
+  (TFL_ConstOp $value),
   [(TFL_TensorOf<[AnyType]> $output)]>;
 
 

tensorflow/compiler/mlir/lite/tests/optimize.mlir

@@ -5165,3 +5165,30 @@ func.func @Fuse4DResourceAddIntoDepthwiseConv2D(%arg0: tensor<1x32x32x4xf32>, %a
   // CHECK: %[[DW_CONV:.*]] = "tfl.depthwise_conv_2d"(%arg0, %arg1, %[[NEW_BIAS]])
   // CHECK: return %[[DW_CONV]]
 }
+
+// CHECK-LABEL: @fuse_sum_mul_to_mean
+func.func @fuse_sum_mul_to_mean(%arg0: tensor<2x3x4xf32>) -> tensor<2x1x4xf32> {
+  %cst_axes = arith.constant dense<1> : tensor<1xi32>
+  %cst_factor = arith.constant dense<0.333333333> : tensor<1xf32>
+  %0 = "tfl.sum"(%arg0, %cst_axes) <{keep_dims = true}> : (tensor<2x3x4xf32>, tensor<1xi32>) -> tensor<2x1x4xf32>
+  %1 = "tfl.mul"(%0, %cst_factor) <{fused_activation_function = "NONE"}> : (tensor<2x1x4xf32>, tensor<1xf32>) -> tensor<2x1x4xf32>
+  func.return %1 : tensor<2x1x4xf32>
+
+  // CHECK-NOT: tfl.sum
+  // CHECK-NOT: tfl.mul
+  // CHECK:     "tfl.mean"(%arg0, %{{.*}}) <{keep_dims = true}>
+}
+
+// CHECK-LABEL: @do_not_fuse_sum_mul_with_arbitrary_factor
+func.func @do_not_fuse_sum_mul_with_arbitrary_factor(%arg0: tensor<2x3x4xf32>) -> tensor<2x1x4xf32> {
+  %cst_axes = arith.constant dense<1> : tensor<1xi32>
+  %cst_factor = arith.constant dense<8.000000e-01> : tensor<1xf32>
+  %0 = "tfl.sum"(%arg0, %cst_axes) <{keep_dims = true}> : (tensor<2x3x4xf32>, tensor<1xi32>) -> tensor<2x1x4xf32>
+  %1 = "tfl.mul"(%0, %cst_factor) <{fused_activation_function = "NONE"}> : (tensor<2x1x4xf32>, tensor<1xf32>) -> tensor<2x1x4xf32>
+  func.return %1 : tensor<2x1x4xf32>
+
+  // CHECK:     "tfl.sum"
+  // CHECK:     tfl.mul
+  // CHECK-NOT: "tfl.mean"
+}
+

tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_patterns.td

@@ -58,6 +58,33 @@ multiclass FuseSplatConstIntoSelectOp<Op SelectOp> {
     [(HasRankAtLeast<2> $constant_attr),
      (OperandsBroadcastToOutputType $input1, $input2, $result),
      (HasRankAtMost<5> $constant_value)]>;
+
+  def FuseSplatTflConstLhsInto#SelectOp : Pat<
+    (SelectOp:$result
+      AnyStaticShapeTensor:$input1,
+      (TFL_ConstOp:$constant_value SplatElementsAttr:$constant_attr),
+      AnyStaticShapeTensor:$input2),
+    (TFL_SelectV2Op
+      $input1,
+      (TFL_ConstOp (GetScalarElementsAttrFromSplat $constant_attr)),
+      $input2),
+    // Check if condition or rhs will promote the required broadcasting.
+    [(HasRankAtLeast<2> $constant_attr),
+     (OperandsBroadcastToOutputType $input1, $input2, $result),
+     (HasRankAtMost<5> $constant_value)]>;
+
+  def FuseSplatTflConstRhsInto#SelectOp : Pat<
+    (SelectOp:$result
+      AnyStaticShapeTensor:$input1,
+      AnyStaticShapeTensor:$input2,
+      (TFL_ConstOp:$constant_value SplatElementsAttr:$constant_attr)),
+    (TFL_SelectV2Op
+      $input1, $input2,
+      (TFL_ConstOp (GetScalarElementsAttrFromSplat $constant_attr))),
+    // Check if condition or lhs will promote the required broadcasting.
+    [(HasRankAtLeast<2> $constant_attr),
+     (OperandsBroadcastToOutputType $input1, $input2, $result),
+     (HasRankAtMost<5> $constant_value)]>;
 }
 
 // Pattern for skipping FillOp if it is mainly for broadcasting and the
@@ -83,6 +110,27 @@ multiclass FuseFillOpBroadcastIntoFollowingSelectOp<Op SelectOp> {
   [(OperandsBroadcastToOutputType $input1, $input2, $result),
    (HasRankAtMost<5> $result),
    (IsRankLessThanEqualTo $fill_output, $result)]>;
+
+  def FoldTflFillOpIntoSelectOpRHS#SelectOp : Pat<
+    (SelectOp:$result
+      AnyStaticShapeTensor:$input1,
+      AnyStaticShapeTensor:$input2,
+      (TFL_FillOp:$fill_output $fill_dims,
+        (TFL_ConstOp:$fill_value $val))),
+    (TFL_SelectV2Op $input1, $input2, $fill_value),
+  [(OperandsBroadcastToOutputType $input1, $input2, $result),
+   (HasRankAtMost<5> $result),
+   (IsRankLessThanEqualTo $fill_output, $result)]>;
+
+  def FoldTflFillOpIntoSelectOpLHS#SelectOp : Pat<
+    (SelectOp:$result
+      AnyStaticShapeTensor:$input1,
+      (TFL_FillOp:$fill_output $fill_dims, (TFL_ConstOp:$fill_value $val)),
+      AnyStaticShapeTensor:$input2),
+    (TFL_SelectV2Op $input1, $fill_value, $input2),
+  [(OperandsBroadcastToOutputType $input1, $input2, $result),
+   (HasRankAtMost<5> $result),
+   (IsRankLessThanEqualTo $fill_output, $result)]>;
 }
 
 multiclass FuseBroadcastToIntoSelectOp<Op SelectOp> {

tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc

@@ -21,11 +21,12 @@ limitations under the License.
 #include <algorithm>
 #include <array>
 #include <cassert>
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <functional>
 #include <iterator>
-#include <memory>
 #include <numeric>
 #include <optional>
 #include <utility>
@@ -42,12 +43,14 @@ limitations under the License.
 #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/QuantTypes.h"  // from @llvm-project
+#include "mlir/Dialect/Traits.h"  // from @llvm-project
 #include "mlir/IR/Attributes.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/Location.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Matchers.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
@@ -572,6 +575,45 @@ TypeAttr RescaleQtype(Type input, Attribute factor) {
   return RescaleQuantizedType(input, factor);
 }
 
+// Returns true if the multiplication by factor (1/N) following a sum reduction
+// along axes can be fused into TFL::MeanOp.
+static bool CanFuseSumMulToMean(Value input, Attribute axes, Attribute factor) {
+  RankedTensorType input_type =
+      mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
+  if (!input_type || !input_type.hasStaticShape()) return false;
+
+  auto dense_factor = mlir::dyn_cast_or_null<DenseElementsAttr>(factor);
+  if (!dense_factor || dense_factor.getNumElements() != 1) return false;
+  if (!mlir::isa<FloatType>(dense_factor.getElementType())) return false;
+
+  float factor_val =
+      (*dense_factor.getValues<APFloat>().begin()).convertToFloat();
+  if (factor_val <= 0.0f) return false;
+
+  auto dense_axes = mlir::dyn_cast_or_null<DenseIntElementsAttr>(axes);
+  if (!dense_axes || dense_axes.empty()) return false;
+
+  int64_t rank = input_type.getRank();
+  auto shape = input_type.getShape();
+  int64_t reduction_elements = 1;
+
+  llvm::SmallSet<int64_t, 4> unique_axes;
+  for (const APInt& val : dense_axes.getValues<APInt>()) {
+    int64_t axis = val.getSExtValue();
+    if (axis < 0) axis += rank;
+    if (axis < 0 || axis >= rank) return false;
+    if (unique_axes.insert(axis).second) {
+      reduction_elements *= shape[axis];
+    }
+  }
+
+  if (reduction_elements <= 0) return false;
+
+  float expected_factor = 1.0f / static_cast<float>(reduction_elements);
+  float diff = std::abs(factor_val - expected_factor);
+  return diff <= 1e-4f * expected_factor;
+}
+
 // Returns `true` if reducing `axes` in `input` with `keep_dims=true` results
 // in the specified `shape` and `false` otherwise.
 static bool ShapeMatchesReduceWithKeepAxes(Value input,

tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td

@@ -1076,6 +1076,46 @@ foreach ReduceOp = [TFL_MeanOp, TFL_ReduceMaxOp, TFL_ReduceMinOp,
      (HasOneUse $reduce)]>;
 }
 
+def CanFuseSumMulToMean : Constraint<CPred<
+  "CanFuseSumMulToMean($0, $1, $2)">>;
+
+// Fuse Sum + Mul(1/N) -> Mean
+def FuseSumMulRhsIntoMean : Pat<
+  (TFL_MulOp
+    (TFL_SumOp:$sum $input, (Arith_ConstantOp ElementsAttr:$axes), $keep_dims),
+    (Arith_ConstantOp ElementsAttr:$factor),
+    TFL_AF_None),
+  (TFL_MeanOp $input, (Arith_ConstantOp $axes), $keep_dims),
+  [(CanFuseSumMulToMean $input, $axes, $factor),
+   (HasOneUse $sum)]>;
+
+def FuseSumMulLhsIntoMean : Pat<
+  (TFL_MulOp
+    (Arith_ConstantOp ElementsAttr:$factor),
+    (TFL_SumOp:$sum $input, (Arith_ConstantOp ElementsAttr:$axes), $keep_dims),
+    TFL_AF_None),
+  (TFL_MeanOp $input, (Arith_ConstantOp $axes), $keep_dims),
+  [(CanFuseSumMulToMean $input, $axes, $factor),
+   (HasOneUse $sum)]>;
+
+def FuseTflSumMulRhsIntoMean : Pat<
+  (TFL_MulOp
+    (TFL_SumOp:$sum $input, (TFL_ConstOp ElementsAttr:$axes), $keep_dims),
+    (TFL_ConstOp ElementsAttr:$factor),
+    TFL_AF_None),
+  (TFL_MeanOp $input, (TFL_ConstOp $axes), $keep_dims),
+  [(CanFuseSumMulToMean $input, $axes, $factor),
+   (HasOneUse $sum)]>;
+
+def FuseTflSumMulLhsIntoMean : Pat<
+  (TFL_MulOp
+    (TFL_ConstOp ElementsAttr:$factor),
+    (TFL_SumOp:$sum $input, (TFL_ConstOp ElementsAttr:$axes), $keep_dims),
+    TFL_AF_None),
+  (TFL_MeanOp $input, (TFL_ConstOp $axes), $keep_dims),
+  [(CanFuseSumMulToMean $input, $axes, $factor),
+   (HasOneUse $sum)]>;
+
 
 def IsSame : Constraint<CPred<"$0 == $1">>;
 def HasTwoUse : Constraint<CPred<

tensorflow/core/data/service/dispatcher_state.cc

@@ -203,6 +203,9 @@ void DispatcherState::AcquireIterationClient(
   iteration->num_clients++;
   next_available_iteration_client_id_ =
       std::max(next_available_iteration_client_id_, iteration_client_id + 1);
+  VLOG(3) << "Acquired iteration client for iteration "
+          << iteration->iteration_id
+          << " num_clients: " << iteration->num_clients;
 }
 
 void DispatcherState::ReleaseIterationClient(
@@ -215,6 +218,9 @@ void DispatcherState::ReleaseIterationClient(
   DCHECK_GE(iteration->num_clients, 0);
   iteration->last_client_released_micros =
       release_iteration_client.time_micros();
+  VLOG(3) << "Released iteration client for iteration "
+          << iteration->iteration_id
+          << " num_clients: " << iteration->num_clients;
   iterations_for_client_ids_.erase(iteration_client_id);
 }
 

tensorflow/lite/python/BUILD

@@ -227,7 +227,7 @@ py_library(
         ":util",
         "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:representative_dataset",
-        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/framework:graph_proto_py_proto",
         "//tensorflow/lite/experimental/microfrontend:audio_microfrontend_py",
         "//tensorflow/lite/profiling/proto:model_runtime_info_py",
         "//tensorflow/lite/profiling/proto:profiling_info_py",

tensorflow/lite/python/convert.py

@@ -348,6 +348,38 @@ def convert(
     raise converter_error
 
 
+def convert_mlir_bytecode(
+    conversion_flags: _conversion_flags_pb2.ConverterFlags,
+    model_dir: str,
+    output_file_path: str,
+):
+  """Converts `model_dir` to a TFLite model file directly.
+
+  Args:
+    conversion_flags: Proto describing conversion properties, see
+      `compiler/mlir/lite/converter_flags.proto`.
+    model_dir: Directory containing the MLIR bytecode and weights.
+    output_file_path: Path where the TFLite model should be saved.
+
+  Returns:
+    Status or result of the conversion.
+
+  Raises:
+    ConverterError: When conversion fails.
+  """
+  try:
+    return wrap_converter.wrapped_convert_mlir_bytecode(
+        conversion_flags.SerializeToString(),
+        model_dir,
+        output_file_path,
+    )
+  except Exception as e:
+    converter_error = ConverterError(str(e))
+    for error_data in _metrics_wrapper.retrieve_collected_errors():
+      converter_error.append_error(error_data)
+    raise converter_error from e
+
+
 def build_model_flags(
     change_concat_input_ranges=False,
     allow_nonexistent_arrays=False,
@@ -419,6 +451,7 @@ def build_conversion_flags(
     accumulation_type=None,
     allow_bfloat16=False,
     unfold_large_splat_constant=False,
+    fold_fp16_resource_casts=True,
     supported_backends=None,
     disable_per_channel_quantization=False,
     enable_mlir_dynamic_range_quantizer=False,
@@ -449,6 +482,8 @@ def build_conversion_flags(
     serialize_debug_metadata=False,
     unsafe_fuse_dynamic_shaped_broadcast=False,
     unsafe_single_batch_rank_reduction=False,
+    enable_debug=False,
+    debug_dir=None,
     **_,
 ):
   """Builds protocol buffer describing a conversion of a model.
@@ -519,6 +554,8 @@ def build_conversion_flags(
       inference with the bfloat16 type.
     unfold_large_splat_constant: Whether to unfold large splat constant tensors
       in the flatbuffer model to reduce size.
+    fold_fp16_resource_casts: Whether to fold 16-bit float (fp16/bf16) resource
+      casts.
     supported_backends: List of TFLite backends which needs to check
       compatibility.
     disable_per_channel_quantization: Disable per-channel quantized weights for
@@ -589,6 +626,9 @@ def build_conversion_flags(
       the source model.
     unsafe_single_batch_rank_reduction: When set to true, enable the unsafe
       single batch rank reduction.
+    enable_debug: When set to true, enable debug mode.
+    debug_dir: Directory to save debug output.
+    **_: Additional unused keyword arguments.
 
   Returns:
     conversion_flags: protocol buffer describing the conversion process.
@@ -653,6 +693,8 @@ def build_conversion_flags(
     )
   conversion_flags.allow_bfloat16 = allow_bfloat16
   conversion_flags.unfold_large_splat_constant = unfold_large_splat_constant
+  if hasattr(conversion_flags, "fold_fp16_resource_casts"):
+    conversion_flags.fold_fp16_resource_casts = fold_fp16_resource_casts
   if supported_backends:
     conversion_flags.supported_backends.extend(supported_backends)
   conversion_flags.disable_per_channel_quantization = (
@@ -696,6 +738,11 @@ def build_conversion_flags(
         elide_elementsattrs_if_larger
     )
 
+  if hasattr(conversion_flags, "enable_debug"):
+    conversion_flags.enable_debug = enable_debug
+  if debug_dir is not None and hasattr(conversion_flags, "debug_dir"):
+    conversion_flags.debug_dir = debug_dir
+
   if use_buffer_offset is not None:
     conversion_flags.use_buffer_offset = use_buffer_offset
   if reduce_type_precision is not None: