[pull] master from tensorflow:master

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 <typename DataType = int32_t>
 static absl::Status FlatBufferIntVectorToArray(
-    int max_size_of_buffer, const flatbuffers::Vector<DataType>* flat_vector,
+    size_t max_size_of_buffer, const flatbuffers::Vector<DataType>* 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<int64_t>(
-          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<int64_t>(
-          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<int64_t>(
-          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<int64_t>(
-          /*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<int64_t>(
-          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,19 +2596,17 @@ absl::Status ParseStablehloGather(const Operator* op,
 
     if (schema_params->start_index_map()) {
       TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray<int64_t>(
-          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();
     }
 
     params->index_vector_dim = schema_params->index_vector_dim();
 
     if (schema_params->slice_sizes()) {
       TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray<int64_t>(
-          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();
     }
 

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<int64_t>(update_window_dims_count, 1));
+    const auto inserted_window_dims = builder_.CreateVector(
+        std::vector<int64_t>(inserted_window_dims_count, 1));
+    const auto scatter_dims =
+        builder_.CreateVector(std::vector<int64_t>(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<int64_t> 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<int64_t>({1}),
+          /*start_index_map=*/builder_.CreateVector<int64_t>({1}),
+          /*index_vector_dim=*/0,
+          /*slice_sizes=*/builder_.CreateVector<int64_t>({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

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",

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::func::FuncOp>(
       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::func::FuncOp>(
       mlir::mhlo::createLegalizeEinsumToDotGeneralPass());
@@ -128,7 +139,13 @@ void AddHloToTfLiteLegalizationPasses(mlir::OpPassManager& pass_manager) {
       mlir::odml::CreateUniformQuantizedStableHloToTflPass());
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       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::func::FuncOp>(
+      mlir::odml::createLegalizeCompositeToCustomOpPass());
 }
 
 void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager,
@@ -148,19 +165,23 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager,
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::TFL::CreateOptimizePass());
 
-  // Quantization
-  pass_manager.addNestedPass<mlir::func::FuncOp>(
-      mlir::TFL::CreatePrepareQuantizePass(pass_config.quant_specs));
-  pass_manager.addNestedPass<mlir::func::FuncOp>(
-      mlir::TFL::CreateQuantizePass(pass_config.quant_specs));
-  pass_manager.addNestedPass<mlir::func::FuncOp>(
-      mlir::TFL::CreatePostQuantizePass(/*emit_quant_adaptor_ops=*/true));
-
   if (!pass_config.unfold_batch_matmul) {
     pass_manager.addNestedPass<mlir::func::FuncOp>(
         mlir::TFL::CreateOptimizeBatchMatmulPass());
+    pass_manager.addNestedPass<mlir::func::FuncOp>(
+        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::func::FuncOp>(
+      mlir::TFL::CreatePostQuantizePass(/*emit_quant_adaptor_ops=*/true));
+  pass_manager.addNestedPass<mlir::func::FuncOp>(
+      mlir::createCanonicalizerPass());
+
   // Some optimizations need to happen on the quantized graph.
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::TFL::CreateOptimizePass());
@@ -176,6 +197,8 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager,
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::createCanonicalizerPass());
   pass_manager.addNestedPass<mlir::func::FuncOp>(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<mlir::TFL::TFLDialect, mlir::mhlo::MhloDialect,
-                  mlir::stablehlo::StablehloDialect, mlir::vhlo::VhloDialect>();
+  registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::mhlo::MhloDialect,
+                  mlir::stablehlo::StablehloDialect, mlir::vhlo::VhloDialect,
+                  mlir::quant::QuantDialect,
+                  mlir::quantfork::QuantizationForkDialect>();
   context->appendDialectRegistry(registry);
   context->loadAllAvailableDialects();