[pull] master from tensorflow:master

ci/official/envs/linux_x86_cuda

@@ -14,7 +14,7 @@
 # ==============================================================================
 source ci/official/envs/linux_x86
 export TF_FORCE_GPU_ALLOW_GROWTH=true
-TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true --local_test_jobs=16"
+TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true"
 TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
 TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow"

ci/official/envs/linux_x86_cuda13_nvcc

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 source ci/official/envs/linux_x86
-TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version --local_test_jobs=16"
+TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version"
 TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda_13_nvcc
 TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow_cuda13"

tensorflow/compiler/mlir/lite/BUILD

@@ -47,6 +47,7 @@ package_group(
         "//third_party/iree/...",
         "//third_party/odml/infra/...",
         "//third_party/odml/litert/...",
+        "//third_party/py/ai_edge_jax/...",
         "//waymo/accelerator/alpine/tools/...",
         "//waymo/ml/compiler/mlir/...",
     ],

tensorflow/compiler/mlir/lite/python/BUILD

@@ -299,12 +299,26 @@ cc_library(
         "flatbuffer_to_mlir.h",
     ],
     deps = [
+        "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
+        "//tensorflow/compiler/mlir/lite:flatbuffer_export",
         "//tensorflow/compiler/mlir/lite:flatbuffer_import",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TranslateLib",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:vhlo_ops",
     ],
 )
 

tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi

@@ -17,6 +17,7 @@ def Convert(model_flags_proto_txt_raw: object, converter_flags_proto_txt_raw: ob
 def ConvertMlirBytecode(converter_flags_proto_txt_raw: object, model_dir_txt_raw: object, output_file_path_raw: object) -> object: ...
 def ExperimentalMlirQuantizeModel(input_contents_txt_raw: object, disable_per_channel: bool = ..., fully_quantize: bool = ..., inference_type: int = ..., input_data_type: int = ..., output_data_type: int = ..., enable_numeric_verify: bool = ..., enable_whole_model_verify: bool = ..., op_blocklist: object = ..., node_blocklist: object = ..., enable_variable_quantization: bool = ..., disable_per_channel_for_dense_layers: bool = ..., debug_options_proto_txt_raw: object = ...) -> object: ...
 def ExperimentalMlirSparsifyModel(input_contents_txt_raw: object) -> object: ...
-def FlatBufferToMlir(arg0: str, arg1: bool) -> str: ...
+def FlatBufferToMlir(model: str, input_is_filepath: bool = ..., bytecode: bool = ..., cl_options: list[str] = ...) -> object: ...
+def MlirToFlatBuffer(mlir: str, input_is_filepath: bool = ..., emit_builtin_tflite_ops: bool = ..., emit_select_tf_ops: bool = ..., emit_custom_ops: bool = ..., emit_stablehlo_ops: bool = ...) -> bytes: ...
 def RegisterCustomOpdefs(custom_opdefs_txt_raw: object) -> object: ...
 def RetrieveCollectedErrors() -> list: ...

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

@@ -16,10 +16,10 @@ limitations under the License.
 
 #include <Python.h>
 
-#include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
@@ -517,8 +517,20 @@ std::vector<std::string> RetrieveCollectedErrors() {
 }
 
 std::string FlatBufferFileToMlir(const std::string& model,
-                                 bool input_is_filepath) {
-  return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath);
+                                 bool input_is_filepath, bool bytecode,
+                                 const std::vector<std::string>& cl_options) {
+  return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath, bytecode,
+                                            cl_options);
+}
+
+std::string MlirToFlatBufferFile(const std::string& mlir,
+                                 bool input_is_filepath,
+                                 bool emit_builtin_tflite_ops,
+                                 bool emit_select_tf_ops, bool emit_custom_ops,
+                                 bool emit_stablehlo_ops) {
+  return ::tensorflow::MlirToFlatBufferFile(
+      mlir, input_is_filepath, emit_builtin_tflite_ops, emit_select_tf_ops,
+      emit_custom_ops, emit_stablehlo_ops);
 }
 
 PyObject* ConvertMlirBytecode(PyObject* converter_flags_proto_txt_raw,

tensorflow/compiler/mlir/lite/python/converter_python_api.h

@@ -64,8 +64,17 @@ PyObject* RegisterCustomOpdefs(PyObject* list);
 std::vector<std::string> RetrieveCollectedErrors();
 
 // Returns MLIR string dump of the given Flatbuffer model.
-std::string FlatBufferFileToMlir(const std::string& model,
-                                 bool input_is_filepath);
+std::string FlatBufferFileToMlir(
+    const std::string& model, bool input_is_filepath, bool bytecode = false,
+    const std::vector<std::string>& cl_options = {});
+
+// Converts MLIR (text or bytecode) to a TFLite Flatbuffer.
+std::string MlirToFlatBufferFile(const std::string& mlir,
+                                 bool input_is_filepath,
+                                 bool emit_builtin_tflite_ops = true,
+                                 bool emit_select_tf_ops = false,
+                                 bool emit_custom_ops = true,
+                                 bool emit_stablehlo_ops = false);
 
 // Convert slim model to TfLite flatbuffer streamed directly to a file.
 PyObject* ConvertMlirBytecode(PyObject* converter_flags_proto_txt_raw,

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

@@ -17,6 +17,7 @@ limitations under the License.
 #include <vector>
 
 #include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
 #include "tensorflow/compiler/mlir/lite/python/converter_python_api.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
 #include "tensorflow/python/lib/core/pybind11_lib.h"
@@ -119,12 +120,39 @@ PYBIND11_MODULE(_pywrap_converter_api, m, py::mod_gil_not_used()) {
     )pbdoc");
   m.def(
       "FlatBufferToMlir",
-      [](const std::string& model, bool input_is_filepath) {
-        return tflite::FlatBufferFileToMlir(model, input_is_filepath);
+      [](const std::string& model, bool input_is_filepath, bool bytecode,
+         const std::vector<std::string>& cl_options) {
+        std::string res = tflite::FlatBufferFileToMlir(model, input_is_filepath,
+                                                       bytecode, cl_options);
+        if (bytecode) {
+          return py::object(py::bytes(res));
+        } else {
+          return py::object(py::str(res));
+        }
       },
+      py::arg("model"), py::arg("input_is_filepath") = false,
+      py::arg("bytecode") = false,
+      py::arg("cl_options") = std::vector<std::string>(),
       R"pbdoc(
       Returns MLIR dump of the given TFLite model.
     )pbdoc");
+  m.def(
+      "MlirToFlatBuffer",
+      [](const std::string& mlir, bool input_is_filepath,
+         bool emit_builtin_tflite_ops, bool emit_select_tf_ops,
+         bool emit_custom_ops, bool emit_stablehlo_ops) {
+        std::string res = tflite::MlirToFlatBufferFile(
+            mlir, input_is_filepath, emit_builtin_tflite_ops,
+            emit_select_tf_ops, emit_custom_ops, emit_stablehlo_ops);
+        return py::bytes(res);
+      },
+      py::arg("mlir"), py::arg("input_is_filepath") = false,
+      py::arg("emit_builtin_tflite_ops") = true,
+      py::arg("emit_select_tf_ops") = false, py::arg("emit_custom_ops") = true,
+      py::arg("emit_stablehlo_ops") = false,
+      R"pbdoc(
+      Converts MLIR (text or bytecode) into a TFLite FlatBuffer binary tensor.
+    )pbdoc");
   m.def(
       "ConvertMlirBytecode",
       [](py::object converter_flags_proto_txt_raw, py::object model_dir_txt_raw,

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

@@ -18,21 +18,40 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/attributes.h"
+#include "absl/base/const_init.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
+#include "mlir/Bytecode/BytecodeWriter.h"  // from @llvm-project
+#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/IR/AsmState.h"  // from @llvm-project
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
 #include "mlir/IR/Location.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
 #include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Tools/mlir-translate/Translation.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h"
 #include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
 
 namespace tensorflow {
 namespace {
@@ -53,8 +72,33 @@ static mlir::OwningOpRef<mlir::ModuleOp> FlatBufferFileToMlirTranslation(
 }  // namespace
 
 std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer,
-                                 bool input_is_filepath) {
-  // referred logic from mlir::mlirTranslateMain().
+                                 bool input_is_filepath, bool bytecode,
+                                 const std::vector<std::string>& cl_options) {
+  ABSL_CONST_INIT static absl::Mutex cl_mutex(absl::kConstInit);
+  absl::MutexLock lock(&cl_mutex);
+
+  // Reset options from any previous invocation.
+  llvm::cl::ResetAllOptionOccurrences();
+  mlir::registerAsmPrinterCLOptions();
+
+  if (!cl_options.empty()) {
+    std::vector<const char*> argv;
+    argv.reserve(cl_options.size() + 1);
+    argv.push_back("flatbuffer_to_mlir");
+    for (const auto& opt : cl_options) {
+      argv.push_back(opt.c_str());
+    }
+    std::string cl_errors;
+    llvm::raw_string_ostream cl_err_stream(cl_errors);
+    if (!llvm::cl::ParseCommandLineOptions(
+            argv.size(), argv.data(), "flatbuffer_to_mlir", &cl_err_stream)) {
+      cl_err_stream.flush();
+      if (!cl_errors.empty()) {
+        llvm::errs() << "Failed to parse MLIR options: " << cl_errors << "\n";
+      }
+      return "";
+    }
+  }
 
   std::string errorMessage;
   std::unique_ptr<llvm::MemoryBuffer> input;
@@ -78,17 +122,102 @@ std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer,
   llvm::SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc());
 
+  std::string diagnostic_str;
+  llvm::raw_string_ostream diag_os(diagnostic_str);
+  mlir::SourceMgrDiagnosticHandler diag_handler(sourceMgr, &context, diag_os);
+
   mlir::OwningOpRef<mlir::ModuleOp> module =
       FlatBufferFileToMlirTranslation(&sourceMgr, &context);
-  if (!module || failed(verify(*module))) return "";
+  if (!module || failed(verify(*module))) {
+    diag_os.flush();
+    if (!diagnostic_str.empty()) {
+      llvm::errs() << diagnostic_str << "\n";
+    }
+    return "";
+  }
 
   std::string mlir_output;
   llvm::raw_string_ostream output_stream(mlir_output);
-  // Dump MLIR with eliding large elements.
-  module->print(
-      output_stream,
-      mlir::OpPrintingFlags().useLocalScope().elideLargeElementsAttrs());
+  if (bytecode) {
+    if (mlir::failed(mlir::writeBytecodeToFile(*module, output_stream))) {
+      llvm::errs() << "Failed to write MLIR bytecode.\n";
+      return "";
+    }
+  } else {
+    mlir::OpPrintingFlags flags;
+    module->print(output_stream, flags);
+  }
+  output_stream.flush();
   return mlir_output;
 }
 
+std::string MlirToFlatBufferFile(const std::string& mlir_file_or_buffer,
+                                 bool input_is_filepath,
+                                 bool emit_builtin_tflite_ops,
+                                 bool emit_select_tf_ops, bool emit_custom_ops,
+                                 bool emit_stablehlo_ops) {
+  std::string errorMessage;
+  std::unique_ptr<llvm::MemoryBuffer> input;
+  if (input_is_filepath) {
+    input = mlir::openInputFile(mlir_file_or_buffer, &errorMessage);
+    if (!input) {
+      llvm::errs() << errorMessage << "\n";
+      return "";
+    }
+  } else {
+    input =
+        llvm::MemoryBuffer::getMemBuffer(mlir_file_or_buffer, "mlir", false);
+    if (!input) {
+      llvm::errs() << "Can't get llvm::MemoryBuffer\n";
+      return "";
+    }
+  }
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::quant::QuantDialect,
+                  mlir::quantfork::QuantizationForkDialect,
+                  mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
+                  mlir::func::FuncDialect, mlir::vhlo::VhloDialect,
+                  mlir::stablehlo::StablehloDialect>();
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+  context.printOpOnDiagnostic(true);
+
+  llvm::SourceMgr sourceMgr;
+  sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc());
+
+  std::string diagnostic_str;
+  llvm::raw_string_ostream diag_os(diagnostic_str);
+  mlir::SourceMgrDiagnosticHandler diag_handler(sourceMgr, &context, diag_os);
+
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      mlir::parseSourceFile<mlir::ModuleOp>(sourceMgr, &context);
+  if (!module || failed(verify(*module))) {
+    diag_os.flush();
+    if (!diagnostic_str.empty()) {
+      llvm::errs() << diagnostic_str << "\n";
+    } else {
+      llvm::errs() << "Failed to parse MLIR source.\n";
+    }
+    return "";
+  }
+
+  std::string serialized_flatbuffer;
+  tensorflow::OpOrArgLocNameMapper op_or_arg_name_mapper;
+  tflite::FlatbufferExportOptions options;
+  options.converter_flags.set_force_select_tf_ops(!emit_builtin_tflite_ops);
+  options.converter_flags.set_enable_select_tf_ops(emit_select_tf_ops);
+  options.converter_flags.set_allow_custom_ops(emit_custom_ops);
+  options.converter_flags.set_use_buffer_offset(true);
+  options.op_or_arg_name_mapper = &op_or_arg_name_mapper;
+
+  if (!tflite::MlirToFlatBufferTranslateFunction(
+          *module, options, &serialized_flatbuffer, emit_stablehlo_ops)) {
+    llvm::errs() << "MlirToFlatBufferTranslateFunction failed.\n";
+    return "";
+  }
+  return serialized_flatbuffer;
+}
+
 }  // namespace tensorflow

tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h

@@ -17,14 +17,24 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_
 
 #include <string>
+#include <vector>
 
 namespace tensorflow {
 
 // Translates the given FlatBuffer filename or buffer into MLIR and returns
-// translated MLIR as string.
-std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer,
-                                 bool input_is_filepath);
-
+// translated MLIR as string or bytecode.
+std::string FlatBufferFileToMlir(
+    const std::string& model_file_or_buffer, bool input_is_filepath,
+    bool bytecode = false, const std::vector<std::string>& cl_options = {});
+
+// Translates the given MLIR filename or buffer into a TFLite FlatBuffer
+// binary string.
+std::string MlirToFlatBufferFile(const std::string& mlir_file_or_buffer,
+                                 bool input_is_filepath,
+                                 bool emit_builtin_tflite_ops = true,
+                                 bool emit_select_tf_ops = false,
+                                 bool emit_custom_ops = true,
+                                 bool emit_stablehlo_ops = false);
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_