chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/python/BUILD

@@ -0,0 +1,111 @@
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+bool_flag(
+    name = "disable_mlir",
+    build_setting_default = False,
+)
+
+config_setting(
+    name = "disable_mlir_config",
+    flag_values = {":disable_mlir": "True"},
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "mlir",
+    srcs = ["mlir.cc"],
+    hdrs = ["mlir.h"],
+    deps = [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "//tensorflow/cc/saved_model:bundle_v2",
+        "//tensorflow/cc/saved_model:loader",
+        "//tensorflow/compiler/mlir/tensorflow/translate/tools:parsers",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:bridge_passes",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:RegisterAllPasses",  # buildcleaner: keep
+        "//tensorflow/core/common_runtime:graph_constructor",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:register",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:tfe_context_internal",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_import_options",
+        "@xla//xla/mlir_hlo:all_passes",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:import_model",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:import_utils",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tensorflow:mlprogram_util",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_passes",
+        "//tensorflow/compiler/mlir/tensorflow:translate_lib",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_tf",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib_proto_parsing",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tflite_portable_logging",
+        "//tensorflow/core/common_runtime/eager:context",
+        # (yongtang) The graph_optimization_pass_registration needs to be part
+        # of a shared object that will be loaded whenever `import tensorflow`
+        # is run. The natural place is libtensorflow_framework.so.
+        # While adding graph_optimization_pass_registration to
+        # libtensorflow_framework.so is possible with some modification in
+        # dependency, many tests will fail due to multiple copies of LLVM.
+        # See https://github.com/tensorflow/tensorflow/pull/39231 for details.
+        # Alternatively, we place graph_optimization_pass_registration here
+        # because:
+        # - tensorflow/python/_pywrap_mlir.so already depends on LLVM anyway
+        # - tensorflow/python/_pywrap_mlir.so always loaded as part of python
+        #   binding
+        # TODO: It might be still preferrable to place graph_optimization_pass
+        # as part of the libtensorflow_framework.so, as it is the central
+        # place for core related components.
+        "//tensorflow/compiler/mlir/tensorflow/transforms:graph_optimization_pass_registration",  # buildcleaner: keep
+    ],
+    alwayslink = 1,
+)
+
+tf_cc_test(
+    name = "mlir_test",
+    srcs = ["mlir_test.cc"],
+    deps = [
+        ":mlir",
+        "//tensorflow/c:safe_ptr",
+        "//tensorflow/c:tf_status_headers",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+filegroup(
+    name = "pywrap_mlir_hdrs",
+    srcs = [
+        "mlir.h",
+    ],
+    visibility = [
+        "//tensorflow/python:__pkg__",
+    ],
+)

tensorflow/compiler/mlir/python/mlir.cc

@@ -0,0 +1,425 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/python/mlir.h"
+
+#include <memory>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "llvm/Support/LogicalResult.h"
+#include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Bytecode/BytecodeWriter.h"  // from @llvm-project
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/IR/AsmState.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "stablehlo/dialect/Register.h"  // from @stablehlo
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/tfe_context_internal.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/cc/saved_model/bundle_v2.h"
+#include "tensorflow/cc/saved_model/loader.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_import_options.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/tools/parsers.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/import_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/mlprogram_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h"
+#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/common_runtime/function_body.h"
+#include "tensorflow/core/common_runtime/function_def_utils.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+namespace {
+// All the passes we will make available to Python by default.
+// TODO(tf): this should be sharded instead of being monolithic like that.
+static void RegisterPasses() {
+  static bool unique_registration = [] {
+    mlir::registerAllPasses();
+    mlir::registerTensorFlowPasses();
+    mlir::TFDevice::registerTensorFlowDevicePasses();
+    mlir::mhlo::registerAllMhloPasses();
+    // These are in compiler/mlir/xla and not part of the above MHLO
+    // passes.
+    mlir::mhlo::registerTfXlaPasses();
+    mlir::mhlo::registerLegalizeTFPass();
+    mlir::quant::stablehlo::registerBridgePasses();
+    mlir::tf_saved_model::registerTensorFlowSavedModelPasses();
+    tensorflow::RegisterMlProgramPasses();
+    return true;
+  }();
+  (void)unique_registration;
+}
+
+// Runs pass pipeline `pass_pipeline` on `module` if `pass_pipeline` is not
+// empty.
+std::string RunPassPipelineOnModule(mlir::ModuleOp module,
+                                    const std::string& pass_pipeline,
+                                    bool show_debug_info, TF_Status* status) {
+  RegisterPasses();
+  if (!pass_pipeline.empty()) {
+    mlir::PassManager pm(module.getContext());
+    std::string error;
+    llvm::raw_string_ostream error_stream(error);
+    if (failed(mlir::parsePassPipeline(pass_pipeline, pm, error_stream))) {
+      TF_SetStatus(status, TF_INVALID_ARGUMENT,
+                   ("Invalid pass_pipeline: " + error_stream.str()).c_str());
+      return "// error";
+    }
+
+    mlir::StatusScopedDiagnosticHandler statusHandler(module.getContext());
+    if (failed(pm.run(module))) {
+      tsl::Set_TF_Status_from_Status(status, statusHandler.ConsumeStatus());
+      return "// error";
+    }
+  }
+  return MlirModuleToString(module, show_debug_info);
+}
+
+}  // anonymous namespace
+
+static std::string ImportGraphDefImpl(const std::string& proto,
+                                      const std::string& pass_pipeline,
+                                      bool show_debug_info,
+                                      GraphDebugInfo& debug_info,
+                                      GraphImportConfig& specs,
+                                      TF_Status* status) {
+  GraphDef graphdef;
+  auto s = tensorflow::LoadProtoFromBuffer(proto, &graphdef);
+  if (!s.ok()) {
+    tsl::Set_TF_Status_from_Status(status, s);
+    return "// error";
+  }
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+  mlir::MLIRContext context(registry);
+  GraphConstructorOptions options;
+  Graph graph(OpRegistry::Global());
+  absl::Status graph_status = ConvertGraphDefToGraph(options, graphdef, &graph);
+  auto module = tensorflow::tf2xla::v2::ConvertGraphToTfExecutor(
+      graph, debug_info, graph.flib_def(), specs, &context);
+  if (!module.ok() || !graph_status.ok()) {
+    tsl::Set_TF_Status_from_Status(status, module.status());
+    return "// error";
+  }
+
+  return RunPassPipelineOnModule(module->get(), pass_pipeline, show_debug_info,
+                                 status);
+}
+
+std::string ImportFunction(const std::string& functiondef_proto,
+                           const std::string& pass_pipeline,
+                           bool show_debug_info, TFE_Context* tfe_context,
+                           TF_Status* status) {
+  FunctionDef functiondef;
+  auto s = tensorflow::LoadProtoFromBuffer(functiondef_proto, &functiondef);
+  if (!s.ok()) {
+    tsl::Set_TF_Status_from_Status(status, s);
+    return "// error";
+  }
+
+  const std::string& function_name = functiondef.signature().name();
+  EagerContext* cpp_context = ContextFromInterface(unwrap(tfe_context));
+  FunctionLibraryDefinition& flib_def = *cpp_context->FuncLibDef();
+  const tensorflow::FunctionDef* fdef = flib_def.Find(function_name);
+  if (fdef == nullptr) {
+    s = absl::NotFoundError(
+        absl::StrCat("Cannot find function ", function_name));
+    tsl::Set_TF_Status_from_Status(status, s);
+    return "// error";
+  }
+
+  std::unique_ptr<tensorflow::FunctionBody> fbody;
+  s = FunctionDefToBodyHelper(*fdef, tensorflow::AttrSlice(), &flib_def,
+                              &fbody);
+  if (!s.ok()) {
+    tsl::Set_TF_Status_from_Status(status, s);
+    return "// error";
+  }
+
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+  mlir::MLIRContext context(registry);
+
+  tensorflow::GraphImportConfig specs;
+  specs.graph_func_name = fbody->record->fdef().signature().name();
+  specs.enable_shape_inference = false;
+  specs.graph_as_function = true;
+  for (const auto* control_ret_node : fbody->control_ret_nodes)
+    specs.control_outputs.push_back(control_ret_node->name());
+  auto module = tensorflow::tf2xla::v2::ConvertGraphToTfExecutor(
+      *fbody->graph, {}, flib_def, specs, &context);
+  if (!module.ok()) {
+    tsl::Set_TF_Status_from_Status(status, module.status());
+    return "// error";
+  }
+
+  return RunPassPipelineOnModule(module->get(), pass_pipeline, show_debug_info,
+                                 status);
+}
+
+std::string ImportGraphDef(const std::string& proto,
+                           const std::string& pass_pipeline,
+                           bool show_debug_info, TF_Status* status) {
+  GraphDebugInfo debug_info;
+  GraphImportConfig specs;
+  return ImportGraphDefImpl(proto, pass_pipeline, show_debug_info, debug_info,
+                            specs, status);
+}
+
+std::string ImportGraphDef(const std::string& proto,
+                           const std::string& pass_pipeline,
+                           bool show_debug_info, absl::string_view input_names,
+                           absl::string_view input_data_types,
+                           absl::string_view input_data_shapes,
+                           absl::string_view output_names, TF_Status* status) {
+  GraphDebugInfo debug_info;
+  GraphImportConfig specs;
+  auto s = ParseInputArrayInfo(input_names, input_data_types, input_data_shapes,
+                               &specs.inputs);
+  if (!s.ok()) {
+    tsl::Set_TF_Status_from_Status(status, s);
+    return "// error";
+  }
+  if (!output_names.empty()) {
+    specs.outputs = absl::StrSplit(output_names, ',');
+  }
+  return ImportGraphDefImpl(proto, pass_pipeline, show_debug_info, debug_info,
+                            specs, status);
+}
+
+std::string ExperimentalConvertSavedModelToMlir(
+    const std::string& saved_model_path, const std::string& exported_names_str,
+    bool show_debug_info, TF_Status* status) {
+  // Load the saved model into a SavedModelV2Bundle.
+
+  tensorflow::SavedModelV2Bundle bundle;
+  auto load_status =
+      tensorflow::SavedModelV2Bundle::Load(saved_model_path, &bundle);
+  if (!load_status.ok()) {
+    tsl::Set_TF_Status_from_Status(status, load_status);
+    return "// error";
+  }
+
+  // Convert the SavedModelV2Bundle to an MLIR module.
+
+  std::vector<std::string> exported_names =
+      absl::StrSplit(exported_names_str, ',', absl::SkipEmpty());
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+  mlir::MLIRContext context(registry);
+  auto module_or = ConvertSavedModelToMlir(
+      &bundle, &context, absl::Span<std::string>(exported_names));
+  if (!module_or.status().ok()) {
+    tsl::Set_TF_Status_from_Status(status, module_or.status());
+    return "// error";
+  }
+
+  return MlirModuleToString(*std::move(module_or).value(), show_debug_info);
+}
+
+std::string ExperimentalConvertSavedModelV1ToMlirLite(
+    const std::string& saved_model_path, const std::string& exported_names_str,
+    const std::string& tags, bool upgrade_legacy, bool show_debug_info,
+    TF_Status* status) {
+  std::unordered_set<std::string> tag_set =
+      absl::StrSplit(tags, ',', absl::SkipEmpty());
+
+  std::vector<std::string> exported_names =
+      absl::StrSplit(exported_names_str, ',', absl::SkipEmpty());
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+  mlir::MLIRContext context(registry);
+
+  tensorflow::MLIRImportOptions import_options;
+  import_options.upgrade_legacy = upgrade_legacy;
+  auto module_or = SavedModelSignatureDefsToMlirImportLite(
+      saved_model_path, tag_set, absl::Span<std::string>(exported_names),
+      &context, import_options);
+  if (!module_or.status().ok()) {
+    tsl::Set_TF_Status_from_Status(status, module_or.status());
+    return "// error";
+  }
+
+  return MlirModuleToString(*module_or.value(), show_debug_info);
+}
+
+std::string ExperimentalConvertSavedModelV1ToMlir(
+    const std::string& saved_model_path, const std::string& exported_names_str,
+    const std::string& tags, bool lift_variables,
+    bool include_variables_in_initializers, bool upgrade_legacy,
+    bool show_debug_info, TF_Status* status) {
+  // Load the saved model into a SavedModelBundle.
+
+  std::unordered_set<std::string> tag_set =
+      absl::StrSplit(tags, ',', absl::SkipEmpty());
+
+  tensorflow::SavedModelBundle bundle;
+  auto load_status =
+      tensorflow::LoadSavedModel({}, {}, saved_model_path, tag_set, &bundle);
+  if (!load_status.ok()) {
+    tsl::Set_TF_Status_from_Status(status, load_status);
+    return "// error";
+  }
+
+  // Convert the SavedModelBundle to an MLIR module.
+  std::vector<std::string> exported_names =
+      absl::StrSplit(exported_names_str, ',', absl::SkipEmpty());
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+  mlir::MLIRContext context(registry);
+  tensorflow::MLIRImportOptions import_options;
+  import_options.upgrade_legacy = upgrade_legacy;
+  import_options.lift_variables = lift_variables;
+  import_options.include_variables_in_initializers =
+      include_variables_in_initializers;
+  auto module_or =
+      ConvertSavedModelV1ToMlir(bundle, absl::Span<std::string>(exported_names),
+                                &context, import_options);
+  if (!module_or.status().ok()) {
+    tsl::Set_TF_Status_from_Status(status, module_or.status());
+    return "// error";
+  }
+
+  // Run the tf standard pipeline by default and then, run passes that lift
+  // variables if the flag is set on the module.
+  mlir::OwningOpRef<mlir::ModuleOp> module = std::move(module_or).value();
+  mlir::PassManager pm(&context);
+  std::string error;
+  llvm::raw_string_ostream error_stream(error);
+
+  mlir::TF::StandardPipelineOptions tf_options;
+  mlir::TF::CreateTFStandardPipeline(pm, tf_options);
+
+  mlir::StatusScopedDiagnosticHandler diagnostic_handler(&context);
+  if (failed(pm.run(*module))) {
+    tsl::Set_TF_Status_from_Status(status, diagnostic_handler.ConsumeStatus());
+    return "// error";
+  }
+  return MlirModuleToString(*module, show_debug_info);
+}
+
+std::string ExperimentalRunPassPipeline(const std::string& mlir_txt,
+                                        const std::string& pass_pipeline,
+                                        bool show_debug_info,
+                                        TF_Status* status) {
+  RegisterPasses();
+  mlir::DialectRegistry registry;
+  mlir::RegisterAllTensorFlowDialects(registry);
+  mlir::stablehlo::registerAllDialects(registry);
+  registry.insert<mlir::shape::ShapeDialect>();
+  mlir::MLIRContext context(registry);
+  mlir::OwningOpRef<mlir::ModuleOp> module;
+  {
+    mlir::StatusScopedDiagnosticHandler diagnostic_handler(&context);
+    module = mlir::parseSourceString<mlir::ModuleOp>(mlir_txt, &context);
+    if (!module) {
+      tsl::Set_TF_Status_from_Status(status,
+                                     diagnostic_handler.ConsumeStatus());
+      return "// error";
+    }
+  }
+
+  // Run the pass_pipeline on the module.
+  mlir::PassManager pm(&context);
+  std::string error;
+  llvm::raw_string_ostream error_stream(error);
+  if (failed(mlir::parsePassPipeline(pass_pipeline, pm, error_stream))) {
+    TF_SetStatus(status, TF_INVALID_ARGUMENT,
+                 ("Invalid pass_pipeline: " + error_stream.str()).c_str());
+    return "// error";
+  }
+
+  mlir::StatusScopedDiagnosticHandler diagnostic_handler(&context);
+  if (failed(pm.run(*module))) {
+    tsl::Set_TF_Status_from_Status(status, diagnostic_handler.ConsumeStatus());
+    return "// error";
+  }
+  return MlirModuleToString(*module, show_debug_info);
+}
+
+void ExperimentalWriteBytecode(const std::string& filename,
+                               const std::string& mlir_txt, TF_Status* status) {
+  mlir::DialectRegistry registry;
+  mlir::RegisterAllTensorFlowDialects(registry);
+  mlir::stablehlo::registerAllDialects(registry);
+  registry.insert<mlir::shape::ShapeDialect>();
+  mlir::MLIRContext context(registry);
+  mlir::OwningOpRef<mlir::ModuleOp> module;
+  mlir::StatusScopedDiagnosticHandler diagnostic_handler(&context);
+  {
+    module = mlir::parseSourceString<mlir::ModuleOp>(mlir_txt, &context);
+    if (!module) {
+      tsl::Set_TF_Status_from_Status(status,
+                                     diagnostic_handler.ConsumeStatus());
+      return;
+    }
+  }
+  mlir::FallbackAsmResourceMap fallback_resource_map;
+  mlir::BytecodeWriterConfig writer_config(fallback_resource_map);
+  // TODO(jpienaar): Make this an option to the call.
+  writer_config.setDesiredBytecodeVersion(1);
+  std::string error;
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      mlir::openOutputFile(filename, &error);
+  if (!error.empty()) {
+    TF_SetStatus(status, TF_INVALID_ARGUMENT,
+                 ("Unable to create output file " + error).c_str());
+    return;
+  }
+  outputFile->keep();
+  if (failed(mlir::writeBytecodeToFile(*module, outputFile->os(),
+                                       writer_config))) {
+    tsl::Set_TF_Status_from_Status(status, diagnostic_handler.ConsumeStatus());
+  }
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/python/mlir.h

@@ -0,0 +1,114 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+// Functions for getting information about kernels registered in the binary.
+// Migrated from previous SWIG file (mlir.i) authored by aminim@.
+#ifndef TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_H_
+#define TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/tf_status.h"
+
+namespace tensorflow {
+
+// Simple wrapper to support tf.mlir.experimental.convert_graph_def.
+// Load a GraphDef (binary or textual proto format), convert to MLIR, and
+// (optionally) optimize the module before returning it as a string.
+// This is an early experimental API, ideally we should return a wrapper object
+// around a Python binding to the MLIR module.
+std::string ImportGraphDef(const std::string &proto,
+                           const std::string &pass_pipeline,
+                           bool show_debug_info, TF_Status *status);
+
+// Simple wrapper to support tf.mlir.experimental.convert_function.
+// Load FunctionDef (binary or textual proto format), convert to MLIR, and
+// (optionally) optimize the module before returning it as a string.
+// This is an early experimental API, ideally we should return a wrapper object
+// around a Python binding to the MLIR module.
+std::string ImportFunction(const std::string &functiondef_proto,
+                           const std::string &pass_pipeline,
+                           bool show_debug_info, TFE_Context *context,
+                           TF_Status *status);
+
+// This wrapper passes the graph_def taking names of input nodes, the shapes and
+// types of its inputs and the output nodes as parameters to MLIR.
+std::string ImportGraphDef(const std::string &proto,
+                           const std::string &pass_pipeline,
+                           bool show_debug_info, absl::string_view(input_names),
+                           absl::string_view(input_data_types),
+                           absl::string_view(input_data_shapes),
+                           absl::string_view(output_names), TF_Status *status);
+
+// Load a SavedModel and return a textual MLIR string corresponding to it.
+//
+// Args:
+//   saved_model_path: File path from which to load the SavedModel.
+//   exported_names_str: Comma-separated list of names to export.
+//                       Empty means "export all".
+//
+// Returns:
+//   A string of textual MLIR representing the raw imported SavedModel.
+std::string ExperimentalConvertSavedModelToMlir(
+    const std::string &saved_model_path, const std::string &exported_names_str,
+    bool show_debug_info, TF_Status *status);
+
+// Load a SavedModel V1 and return a textual MLIR string corresponding to it
+// without any MLIR graph transformation.
+//
+// Args:
+//   saved_model_path: File path from which to load the SavedModel.
+//   tags: Tags to identify MetaGraphDef that need to be loaded.
+//   upgrade_legacy: Boolean flag that indicates whether to upgrade legacy
+//                   graphs
+//
+// Returns:
+//   A string of textual MLIR representing the raw imported SavedModel.
+std::string ExperimentalConvertSavedModelV1ToMlirLite(
+    const std::string &saved_model_path, const std::string &exported_names_str,
+    const std::string &tags, bool upgrade_legacy, bool show_debug_info,
+    TF_Status *status);
+
+// Load a SavedModel V1 and return a textual MLIR string corresponding to it.
+//
+// Args:
+//   saved_model_path: File path from which to load the SavedModel.
+//   tags: Tags to identify MetaGraphDef that need to be loaded.
+//   lift_variables: Boolean flag that indicates whether to hoist variables
+//                   after loading the SavedModel.
+//
+// Returns:
+//   A string of textual MLIR representing the raw imported SavedModel.
+std::string ExperimentalConvertSavedModelV1ToMlir(
+    const std::string &saved_model_path, const std::string &exported_names_str,
+    const std::string &tags, bool lift_variables,
+    bool include_variables_in_initializers, bool upgrade_legacy,
+    bool show_debug_info, TF_Status *status);
+
+std::string ExperimentalRunPassPipeline(const std::string &mlir_txt,
+                                        const std::string &pass_pipeline,
+                                        bool show_debug_info,
+                                        TF_Status *status);
+
+// Writes the input textual MLIR as bytecode to output file.
+void ExperimentalWriteBytecode(const std::string &filename,
+                               const std::string &mlir_txt, TF_Status *status);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_H_

tensorflow/compiler/mlir/python/mlir_test.cc

@@ -0,0 +1,88 @@
+/* Copyright 2024 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/python/mlir.h"
+
+#include <string>
+
+#include <gtest/gtest.h>
+#include "absl/strings/match.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/c/safe_ptr.h"
+#include "tensorflow/c/tf_status.h"
+
+namespace tensorflow {
+
+namespace {
+
+class MlirTest : public ::testing::Test {};
+
+TEST_F(MlirTest, ImportGraphDef) {
+  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
+  std::string input_graphdef = R"pb(
+    node {
+      name: "Const"
+      op: "Const"
+      attr {
+        key: "dtype"
+        value { type: DT_INT32 }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT32
+            tensor_shape { dim { size: 1 } }
+            int_val: 1
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_1"
+      op: "Const"
+      attr {
+        key: "dtype"
+        value { type: DT_INT32 }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT32
+            tensor_shape { dim { size: 1 } }
+            int_val: 2
+          }
+        }
+      }
+    }
+    node {
+      name: "Add"
+      op: "Add"
+      input: "Const"
+      input: "Const_1"
+      attr {
+        key: "T"
+        value { type: DT_INT32 }
+      }
+    }
+  )pb";
+
+  std::string result = ImportGraphDef(input_graphdef, /*pass_pipeline=*/"",
+                                      /*show_debug_info=*/false, status.get());
+
+  EXPECT_EQ(TF_GetCode(status.get()), TF_OK);
+  EXPECT_TRUE(absl::StrContains(result, "tf.Const"));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/mlir/python/mlir_wrapper/BUILD

@@ -0,0 +1,54 @@
+load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_python_pybind_extension(
+    name = "mlir_wrapper",
+    srcs = [
+        "attrs.cc",
+        "basic_classes.cc",
+        "builders.cc",
+        "mlir_wrapper.cc",
+        "mlir_wrapper.h",
+        "ops.cc",
+        "types.cc",
+    ],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "mlir_wrapper.pyi",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@llvm-project//llvm:FileCheckLib",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@pybind11",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "filecheck_wrapper",
+    srcs = ["filecheck_wrapper.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "filecheck_wrapper.pyi",
+    ],
+    starlark_only = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@llvm-project//llvm:FileCheckLib",
+        "@llvm-project//llvm:Support",
+        "@pybind11",
+    ],
+)

tensorflow/compiler/mlir/python/mlir_wrapper/attrs.cc

@@ -0,0 +1,26 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+void init_attrs(py::module& m) {
+  py::class_<mlir::IntegerAttr, mlir::Attribute>(m, "IntegerAttr")
+      .def("get",
+           py::overload_cast<mlir::Type, int64_t>(&mlir::IntegerAttr::get));
+}

tensorflow/compiler/mlir/python/mlir_wrapper/basic_classes.cc

@@ -0,0 +1,53 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include "llvm/FileCheck/FileCheck.h"
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+void init_basic_classes(py::module& m) {
+  py::class_<mlir::MLIRContext>(m, "MLIRContext").def(py::init<>());
+
+  py::class_<mlir::Location>(m, "Location");
+
+  py::class_<mlir::UnknownLoc>(m, "UnknownLoc")
+      .def("get", [](mlir::MLIRContext* context) -> mlir::Location {
+        return mlir::UnknownLoc::get(context);
+      });
+
+  py::class_<mlir::Region>(m, "Region")
+      .def("back", &mlir::Region::back, py::return_value_policy::reference)
+      .def("front", &mlir::Region::front, py::return_value_policy::reference)
+      .def("add_block", [](mlir::Region& r) { r.push_back(new mlir::Block); })
+      .def("push_back", &mlir::Region::push_back)
+      .def("size", [](mlir::Region& r) { return r.getBlocks().size(); })
+      .def("front", &mlir::Region::front, py::return_value_policy::reference);
+  py::class_<mlir::Block::iterator>(m, "Block_Iterator");
+  py::class_<mlir::Block>(m, "Block")
+      .def("new", ([]() { return new mlir::Block; }),
+           py::return_value_policy::reference)
+      .def("end", &mlir::Block::end)
+      .def("addArgument", [](mlir::Block& block, mlir::Type type) {
+        return block.addArgument(type, block.getParent()->getLoc());
+      });
+
+  py::class_<mlir::Value>(m, "Value").def("getType", &mlir::Value::getType);
+  py::class_<mlir::OpResult, mlir::Value>(m, "OpResult");
+  py::class_<mlir::BlockArgument, mlir::Value>(m, "BlockArgument");
+}

tensorflow/compiler/mlir/python/mlir_wrapper/builders.cc

@@ -0,0 +1,55 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+
+#include <vector>
+
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+void init_builders(py::module& m) {
+  py::class_<mlir::Builder>(m, "Builder")
+      .def(py::init<mlir::MLIRContext*>())
+      .def("getFunctionType",
+           [](mlir::Builder& b, std::vector<mlir::Type> inputs,
+              std::vector<mlir::Type> outputs) {
+             return b.getFunctionType(llvm::ArrayRef<mlir::Type>(inputs),
+                                      llvm::ArrayRef<mlir::Type>(outputs));
+           });
+  py::class_<mlir::OpBuilder>(m, "OpBuilder")
+      .def(py::init<mlir::MLIRContext*>())
+      .def(py::init<mlir::Region&>())
+      .def(py::init<mlir::Operation*>())
+      .def(py::init<mlir::Block*, mlir::Block::iterator>())
+      .def("getUnknownLoc", &mlir::OpBuilder::getUnknownLoc)
+      .def("setInsertionPoint",
+           py::overload_cast<mlir::Block*, mlir::Block::iterator>(
+               &mlir::OpBuilder::setInsertionPoint))
+      .def("saveInsertionPoint", &mlir::OpBuilder::saveInsertionPoint)
+      .def("restoreInsertionPoint", &mlir::OpBuilder::restoreInsertionPoint)
+      .def(
+          "create",
+          [](mlir::OpBuilder& opb, mlir::OperationState& state) {
+            return opb.create(state);
+          },
+          py::return_value_policy::reference)
+      .def("getContext", &mlir::OpBuilder::getContext,
+           py::return_value_policy::reference);
+
+  py::class_<mlir::OpBuilder::InsertPoint>(m, "OpBuilder_InsertionPoint")
+      .def("getBlock", &mlir::OpBuilder::InsertPoint::getBlock);
+}

tensorflow/compiler/mlir/python/mlir_wrapper/filecheck_wrapper.cc

@@ -0,0 +1,37 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include <string>
+
+#include "llvm/FileCheck/FileCheck.h"
+#include "llvm/Support/SourceMgr.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+PYBIND11_MODULE(filecheck_wrapper, m) {
+  m.def("check", [](std::string input, std::string check) {
+    llvm::FileCheckRequest fcr;
+    llvm::FileCheck fc(fcr);
+    llvm::SourceMgr SM = llvm::SourceMgr();
+    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),
+                          llvm::SMLoc());
+    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(check),
+                          llvm::SMLoc());
+    fc.readCheckFile(SM, llvm::StringRef(check));
+    return fc.checkInput(SM, llvm::StringRef(input));
+  });
+}

tensorflow/compiler/mlir/python/mlir_wrapper/filecheck_wrapper.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+def check(arg0: str, arg1: str) -> bool: ...

tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.cc

@@ -0,0 +1,65 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+
+#include <string>
+
+#include "llvm/Support/SourceMgr.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+PYBIND11_MODULE(mlir_wrapper, m) {
+  m.def("preloadTensorFlowDialects", [](mlir::MLIRContext &context) {
+    mlir::DialectRegistry registry;
+    mlir::RegisterAllTensorFlowDialects(registry);
+    context.appendDialectRegistry(registry);
+    context.loadAllAvailableDialects();
+  });
+
+  m.def("verify", [](std::string input) {
+    llvm::SourceMgr SM = llvm::SourceMgr();
+    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),
+                          llvm::SMLoc());
+    mlir::DialectRegistry registry;
+    mlir::RegisterAllTensorFlowDialects(registry);
+    mlir::MLIRContext ctx(registry);
+    ctx.loadAllAvailableDialects();
+    auto module = mlir::parseSourceFile<mlir::ModuleOp>(SM, &ctx);
+    if (!module) {
+      return false;
+    }
+    if (failed(mlir::verify(*module))) {
+      module->emitError("Invalid MLIR module: failed verification.");
+      return false;
+    }
+    return true;
+  });
+
+  init_basic_classes(m);
+  init_types(m);
+  init_builders(m);
+  init_ops(m);
+  init_attrs(m);
+}

tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h

@@ -0,0 +1,30 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_WRAPPER_MLIR_WRAPPER_H_
+#define TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_WRAPPER_MLIR_WRAPPER_H_
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+
+namespace py = pybind11;
+
+void init_basic_classes(py::module& m);
+void init_types(py::module& m);
+void init_builders(py::module& m);
+void init_ops(py::module& m);
+void init_attrs(py::module& m);
+
+#endif  // TENSORFLOW_COMPILER_MLIR_PYTHON_MLIR_WRAPPER_MLIR_WRAPPER_H_

tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.pyi

@@ -0,0 +1,193 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+from typing import overload
+
+class Attribute:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Block:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def addArgument(self, *args, **kwargs): ...
+    def end(self) -> Block_Iterator: ...
+    def new(self) -> Block: ...
+
+class BlockArgument(Value):
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Block_Iterator:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Builder:
+    def __init__(self, arg0: MLIRContext) -> None: ...
+    def getFunctionType(self, arg0: list[Type], arg1: list[Type]) -> FunctionType: ...
+
+class FloatType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getBF16(self) -> FloatType: ...
+    def getF16(self) -> FloatType: ...
+    def getF32(self) -> FloatType: ...
+    def getF64(self) -> FloatType: ...
+
+class FuncOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: str, arg1: FunctionType) -> FuncOp: ...
+    def getArguments(self) -> list[BlockArgument]: ...
+    def getBody(self) -> Region: ...
+    def getName(self) -> str: ...
+    def getType(self) -> FunctionType: ...
+
+class FunctionType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getResults(self) -> list[Type]: ...
+
+class IntegerAttr(Attribute):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self, arg0: int) -> IntegerAttr: ...
+
+class IntegerType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self, arg0: int) -> IntegerType: ...
+
+class Location:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class MLIRContext:
+    def __init__(self) -> None: ...
+
+class ModuleOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self) -> ModuleOp: ...
+    def dump(self) -> None: ...
+    def getAsStr(self) -> str: ...
+    def push_back(self, arg0) -> None: ...
+
+class OpBuilder:
+    @overload
+    def __init__(self, arg0: MLIRContext) -> None: ...
+    @overload
+    def __init__(self, arg0: Region) -> None: ...
+    @overload
+    def __init__(self, arg0) -> None: ...
+    @overload
+    def __init__(self, arg0: Block, arg1: Block_Iterator) -> None: ...
+    def create(self, *args, **kwargs): ...
+    def getContext(self) -> MLIRContext: ...
+    def getUnknownLoc(self) -> Location: ...
+    def restoreInsertionPoint(self, arg0) -> None: ...
+    def saveInsertionPoint(self, *args, **kwargs): ...
+    def setInsertionPoint(self, arg0: Block, arg1: Block_Iterator) -> None: ...
+
+class OpBuilder_InsertionPoint:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getBlock(self) -> Block: ...
+
+class OpResult(Value):
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Operation:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def dump(self) -> None: ...
+    def getNumResults(self) -> int: ...
+    def getRegion(self, arg0: int) -> Region: ...
+    def getResult(self, arg0: int) -> OpResult: ...
+
+class OperationState:
+    def __init__(self, arg0: Location, arg1: str) -> None: ...
+    def addOperands(self, arg0: list[Value]) -> None: ...
+    def addRegion(self) -> Region: ...
+    def addTypes(self, arg0: list[Type]) -> None: ...
+
+class RankedTensorType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self, arg0: Type) -> RankedTensorType: ...
+
+class Region:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def add_block(self) -> None: ...
+    def back(self, *args, **kwargs): ...
+    def front(self, *args, **kwargs): ...
+    def push_back(self, arg0) -> None: ...
+    def size(self) -> int: ...
+
+class ReturnOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: list[Value]) -> Operation: ...
+
+class Tf_AddV2Op:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_AnyOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value, arg3: bool) -> Operation: ...
+
+class Tf_ConstOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1) -> Operation: ...
+
+class Tf_EqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_GreaterEqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_GreaterOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_LegacyCallOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: list[Type], arg2: list[Value], arg3: str) -> Operation: ...
+
+class Tf_LessEqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_LessOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_NegOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value) -> Operation: ...
+
+class Tf_NotEqualOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Tf_SubOp:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def create(self, arg0: Location, arg1: Value, arg2: Value) -> Operation: ...
+
+class Type:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class UnknownLoc:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self) -> Location: ...
+
+class UnrankedTensorType(Type):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get(self) -> UnrankedTensorType: ...
+
+class Value:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def getType(self, *args, **kwargs): ...
+
+def preloadTensorFlowDialects(arg0) -> None: ...
+def verify(arg0: str) -> bool: ...

tensorflow/compiler/mlir/python/mlir_wrapper/ops.cc

@@ -0,0 +1,196 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+void init_ops(py::module& m) {
+  py::class_<mlir::Operation, std::unique_ptr<mlir::Operation, py::nodelete>>(
+      m, "Operation")
+      .def("getRegion", &mlir::Operation::getRegion,
+           py::return_value_policy::reference)
+      .def("getResult", &mlir::Operation::getResult)
+      .def("dump", &mlir::Operation::dump)
+      .def("getNumResults", &mlir::Operation::getNumResults);
+
+  py::class_<mlir::OperationState>(m, "OperationState")
+      .def(py::init([](mlir::Location loc, std::string name) {
+        return mlir::OperationState(loc, llvm::StringRef(name));
+      }))
+      .def("addTypes",
+           [](mlir::OperationState& state, std::vector<mlir::Type> tys) {
+             state.addTypes(mlir::ArrayRef<mlir::Type>(tys));
+           })
+      .def("addOperands",
+           [](mlir::OperationState& os, std::vector<mlir::Value> ops) {
+             os.addOperands(mlir::ArrayRef<mlir::Value>(ops));
+           })
+      .def("addRegion", py::overload_cast<>(&mlir::OperationState::addRegion),
+           py::return_value_policy::reference);
+
+  py::class_<mlir::ModuleOp>(m, "ModuleOp")
+      .def("create",
+           [](mlir::Location loc) { return mlir::ModuleOp::create(loc); })
+      .def("push_back",
+           [](mlir::ModuleOp& m, mlir::func::FuncOp f) { m.push_back(f); })
+      .def("dump", &mlir::ModuleOp::dump)
+      .def("getAsStr", [](mlir::ModuleOp& m) {
+        std::string str;
+        llvm::raw_string_ostream os(str);
+        m.print(os);
+        return os.str();
+      });
+
+  py::class_<mlir::func::FuncOp>(m, "FuncOp")
+      .def("create",
+           [](mlir::Location location, std::string name,
+              mlir::FunctionType type) {
+             auto func = mlir::func::FuncOp::create(location, name, type);
+             func.addEntryBlock();
+             return func;
+           })
+      .def(
+          "getBody",
+          [](mlir::func::FuncOp& f) -> mlir::Region& { return f.getBody(); },
+          py::return_value_policy::reference)
+      .def("getArguments",
+           [](mlir::func::FuncOp& f) { return f.getArguments().vec(); })
+      .def("getName", [](mlir::func::FuncOp& f) { return f.getName().str(); })
+      .def("getType", &mlir::func::FuncOp::getFunctionType);
+
+  py::class_<mlir::func::ReturnOp>(m, "ReturnOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              std::vector<mlir::Value> values) -> mlir::Operation* {
+             return mlir::func::ReturnOp::create(
+                        opb, loc, mlir::ArrayRef<mlir::Value>(values))
+                 .getOperation();
+           });
+
+  // mlir::TF::AddOp
+  py::class_<mlir::TF::AddV2Op>(m, "Tf_AddV2Op")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::AddV2Op::create(opb, loc, x, y).getOperation();
+           });
+
+  py::class_<mlir::TF::AnyOp>(m, "Tf_AnyOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value input,
+              mlir::Value reduction_indices,
+              bool keep_dims = false) -> mlir::Operation* {
+             return mlir::TF::AnyOp::create(opb, loc, opb.getI1Type(), input,
+                                            reduction_indices, keep_dims)
+                 .getOperation();
+           });
+
+  // mlir::TF::ConstOp
+  py::class_<mlir::TF::ConstOp>(m, "Tf_ConstOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              mlir::Attribute value) -> mlir::Operation* {
+             return mlir::TF::ConstOp::create(opb, loc, value).getOperation();
+           });
+
+  // mlir::TF::EqualOp
+  py::class_<mlir::TF::EqualOp>(m, "Tf_EqualOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::EqualOp::create(opb, loc, x, y,
+                                              opb.getBoolAttr(true))
+                 .getOperation();
+           });
+
+  // mlir::TF::GreaterEqualOp
+  py::class_<mlir::TF::GreaterEqualOp>(m, "Tf_GreaterEqualOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::GreaterEqualOp::create(opb, loc, x, y)
+                 .getOperation();
+           });
+
+  // mlir::TF::GreaterOp
+  py::class_<mlir::TF::GreaterOp>(m, "Tf_GreaterOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::GreaterOp::create(opb, loc, x, y).getOperation();
+           });
+
+  // mlir::TF::LegacyCallOp
+  py::class_<mlir::TF::LegacyCallOp>(m, "Tf_LegacyCallOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              std::vector<mlir::Type> output, std::vector<mlir::Value> args,
+              std::string f) -> mlir::Operation* {
+             return mlir::TF::LegacyCallOp::create(
+                        opb, loc, mlir::ArrayRef<mlir::Type>(output),
+                        mlir::ArrayRef<mlir::Value>(args),
+                        /*args_attrs=*/nullptr,
+                        /*res_attrs=*/nullptr, mlir::StringRef(f))
+                 .getOperation();
+           });
+
+  // mlir::TF::LessEqualOp
+  py::class_<mlir::TF::LessEqualOp>(m, "Tf_LessEqualOp")
+      .def(
+          "create",
+          [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+             mlir::Value y) -> mlir::Operation* {
+            return mlir::TF::LessEqualOp::create(opb, loc, x, y).getOperation();
+          });
+
+  // mlir::TF::LessOp
+  py::class_<mlir::TF::LessOp>(m, "Tf_LessOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::LessOp::create(opb, loc, x, y).getOperation();
+           });
+
+  // mlir::TF::NegOp
+  py::class_<mlir::TF::NegOp>(m, "Tf_NegOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc,
+              mlir::Value x) -> mlir::Operation* {
+             return mlir::TF::NegOp::create(opb, loc, x).getOperation();
+           });
+
+  py::class_<mlir::TF::NotEqualOp>(m, "Tf_NotEqualOp")
+      .def("create", [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+                        mlir::Value y) {
+        return mlir::TF::NotEqualOp::create(
+                   opb, loc, x, y, mlir::BoolAttr::get(opb.getContext(), true))
+            .getOperation();
+      });
+
+  // mlir::TF::SubOp
+  py::class_<mlir::TF::SubOp>(m, "Tf_SubOp")
+      .def("create",
+           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value x,
+              mlir::Value y) -> mlir::Operation* {
+             return mlir::TF::SubOp::create(opb, loc, x, y).getOperation();
+           });
+}

tensorflow/compiler/mlir/python/mlir_wrapper/types.cc

@@ -0,0 +1,62 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include <cstdint>
+#include <vector>
+
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/python/mlir_wrapper/mlir_wrapper.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+
+void init_types(py::module& m) {
+  // Type
+  py::class_<mlir::Type> Type(m, "Type");
+
+  // Type Sub-classes
+  py::class_<mlir::FunctionType, mlir::Type>(m, "FunctionType")
+      .def("getResults",
+           [](mlir::FunctionType& ft) { return ft.getResults().vec(); });
+
+  py::class_<mlir::FloatType, mlir::Type>(m, "FloatType")
+      .def("getBF16",
+           [](mlir::MLIRContext* context) -> mlir::FloatType {
+             return mlir::BFloat16Type::get(context);
+           })
+      .def("getF16",
+           [](mlir::MLIRContext* context) -> mlir::FloatType {
+             return mlir::Float16Type::get(context);
+           })
+      .def("getF32",
+           [](mlir::MLIRContext* context) -> mlir::FloatType {
+             return mlir::Float32Type::get(context);
+           })
+      .def("getF64", [](mlir::MLIRContext* context) -> mlir::FloatType {
+        return mlir::Float64Type::get(context);
+      });
+
+  py::class_<mlir::IntegerType, mlir::Type>(m, "IntegerType")
+      .def("get", [](mlir::MLIRContext* context, unsigned width) {
+        return mlir::IntegerType::get(context, width,
+                                      mlir::IntegerType::Signless);
+      });
+
+  py::class_<mlir::UnrankedTensorType, mlir::Type>(m, "UnrankedTensorType")
+      .def("get", &mlir::UnrankedTensorType::get);
+
+  py::class_<mlir::RankedTensorType, mlir::Type>(m, "RankedTensorType")
+      .def("get", [](std::vector<int64_t> shape, mlir::Type ty) {
+        return mlir::RankedTensorType::get(mlir::ArrayRef<int64_t>(shape), ty);
+      });
+}