wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD
+60
View File
@@ -0,0 +1,60 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
)
cc_library(
name = "graph_to_tf_executor_registration",
srcs = [
"graph_to_tf_executor_registration.cc",
],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow/translate/tools:file_tf_mlir_translate",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/compiler/mlir/tools:translate_cl_options",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
"@tsl//tsl/platform:protobuf",
"@xla//xla/client:client_library",
"@xla//xla/client:compile_only_client",
"@xla//xla/service/cpu:cpu_compiler",
"@xla//xla/service/cpu:cpu_transfer_manager",
"@xla//xla/stream_executor:platform_manager",
"@xla//xla/stream_executor/host:host_platform",
"@xla//xla/stream_executor/host:host_platform_id",
"@xla//xla/tsl/platform:status",
],
alwayslink = 1,
)
tf_cc_test(
name = "graph_to_tf_executor_registration_test",
size = "small",
srcs = ["graph_to_tf_executor_registration_test.cc"],
deps = [
":graph_to_tf_executor_registration",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TranslateLib",
],
)
tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc
+193
View File
@@ -0,0 +1,193 @@
/* 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 <memory>
#include <utility>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/tools/file_tf_mlir_translate.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include "tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/stream_executor/host/host_platform_id.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
using tsl::Status;
static constexpr char kMlirToGraphCompilationCheckName[] =
"mlir-to-graph-compilation-check";
// Use CPU arbitrarily in order to check that a graph compiles at all
static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT";
static Status CompileGraph(tensorflow::Graph* graph,
xla::CompileOnlyClient* client) {
if (!graph || !client) {
return Status(absl::StatusCode::kInvalidArgument,
"Invalid graph or client");
}
tensorflow::FunctionDefLibrary flib;
auto flib_def = std::make_unique<tensorflow::FunctionLibraryDefinition>(
tensorflow::OpRegistry::Global(), flib);
tensorflow::XlaCompiler::Options options;
options.device_type = tensorflow::DeviceType(kArbitraryDeviceName);
options.client = client;
options.flib_def = flib_def.get();
tensorflow::XlaCompiler compiler(options);
std::unique_ptr<tensorflow::Graph> graph_copy =
std::make_unique<tensorflow::Graph>(tensorflow::OpRegistry::Global());
tensorflow::CopyGraph(*graph, graph_copy.get());
tensorflow::XlaCompiler::CompileOptions compile_options;
tensorflow::XlaCompiler::CompilationResult result;
return compiler.CompileGraph(compile_options,
kMlirToGraphCompilationCheckName,
std::move(graph_copy), {}, &result);
}
static mlir::OwningOpRef<mlir::ModuleOp> GraphdefToMlirTranslateFunction(
llvm::StringRef input, mlir::MLIRContext* context) {
tensorflow::GraphdefToMlirOptions options{
debug_info_file, xla_compile_device_type,
prune_unused_nodes, convert_legacy_fed_inputs,
graph_as_function, upgrade_legacy,
enable_shape_inference, unconditionally_use_set_output_shapes,
enable_soft_placement, set_original_tf_func_name};
auto module_or = tensorflow::GraphdefToMlirTranslateFunction(
input, input_arrays, input_dtypes, input_shapes, output_arrays,
control_output_arrays, options, context);
if (!module_or.status().ok()) return nullptr;
return std::move(module_or).value();
}
static mlir::TranslateToMLIRRegistration GraphdefToMlirTranslate(
"graphdef-to-mlir", "graphdef-to-mlir", GraphdefToMlirTranslateFunction);
static mlir::LogicalResult MlirToGraphTranslateFunction(
mlir::ModuleOp module, llvm::raw_ostream& output) {
if (!module) return mlir::failure();
tensorflow::GraphExportConfig confs;
confs.export_entry_func_to_flib = export_entry_func_to_flib;
confs.export_original_tf_func_name = export_original_tf_func_name;
std::unique_ptr<tensorflow::FunctionLibraryDefinition> flib_def;
auto graph =
std::make_unique<tensorflow::Graph>(tensorflow::OpRegistry::Global());
absl::flat_hash_set<tensorflow::Node*> control_ret_nodes;
auto status = tensorflow::tf2xla::v2::ConvertTfExecutorToGraph(
module, confs, &graph, flib_def.get(), &control_ret_nodes);
if (!status.ok()) {
LOG(ERROR) << "Export to Graph failed: " << status;
return mlir::failure();
}
// Use Host platform, which should always exist, to make sure graphs compile.
auto platform = stream_executor::PlatformManager::PlatformWithId(
stream_executor::host::kHostPlatformId);
if (!platform.ok()) {
return mlir::failure();
}
auto client =
xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform.value());
tensorflow::XlaOpRegistry::RegisterCompilationKernels();
// Verify that the resulting graph can compile.
if (client.ok() && !CompileGraph(graph.get(), client.value()).ok()) {
return mlir::failure();
}
auto graphdef = std::make_unique<tensorflow::GraphDef>();
// Print the graph to the output after going through GraphDef conversion.
// The DumpGraphToFile would do this anyway so just skip straight to it.
graph->ToGraphDef(graphdef.get());
output << tsl::LegacyUnredactedDebugString(*graphdef);
return mlir::success();
}
static mlir::TranslateFromMLIRRegistration mlir_to_graph_translate(
/*name=*/"mlir-to-graph", /*description=*/"convert mlir to graph",
MlirToGraphTranslateFunction, [](mlir::DialectRegistry& registry) {
mlir::RegisterAllTensorFlowDialects(registry);
});
static llvm::LogicalResult MlirToGraphdefTranslateFunction(
mlir::ModuleOp module, llvm::raw_ostream& output) {
if (!module) return mlir::failure();
tensorflow::GraphExportConfig confs;
confs.export_entry_func_to_flib = export_entry_func_to_flib;
confs.export_original_tf_func_name = export_original_tf_func_name;
tensorflow::FunctionLibraryDefinition flib_def(
tensorflow::OpRegistry::Global(), tensorflow::FunctionDefLibrary());
auto graph =
std::make_unique<tensorflow::Graph>(tensorflow::OpRegistry::Global());
absl::flat_hash_set<tensorflow::Node*> control_ret_nodes;
auto status = tensorflow::tf2xla::v2::ConvertTfExecutorToGraph(
module, confs, &graph, &flib_def, &control_ret_nodes);
if (!status.ok()) {
LOG(ERROR) << "Export to Graph failed: " << status;
return mlir::failure();
}
tensorflow::GraphDef graphdef;
graph->ToGraphDef(&graphdef);
output << tsl::LegacyUnredactedDebugString(graphdef);
return mlir::success();
}
static mlir::TranslateFromMLIRRegistration mlir_to_graphdef_translate(
"mlir-to-graphdef", "mlir-to-graphdef", MlirToGraphdefTranslateFunction,
[](mlir::DialectRegistry& registry) {
mlir::RegisterAllTensorFlowDialects(registry);
});
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow
tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration_test.cc
+98
View File
@@ -0,0 +1,98 @@
/* 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 <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/match.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace testing {
using mlir::LogicalResult;
using mlir::StringRef;
using mlir::Translation;
using mlir::TranslationParser;
class MlirTranslationTest : public ::testing::Test {
private:
static constexpr char kMlirToGraphFlag[] = "-mlir-to-graph";
public:
MlirTranslationTest() : translation_(RegisterTranslation()) {
// Create fake command line args so that the parser gets chosen.
std::vector<const char*> argv = {""};
argv.push_back(kMlirToGraphFlag);
llvm::cl::ParseCommandLineOptions(argv.size(), &argv[0],
"TF MLIR translation test\n");
}
LogicalResult Translate(StringRef source, std::string& sink) {
auto source_manager = std::make_shared<llvm::SourceMgr>();
auto source_buffer = llvm::MemoryBuffer::getMemBuffer(source);
source_manager->AddNewSourceBuffer(std::move(source_buffer), llvm::SMLoc());
mlir::MLIRContext context;
llvm::raw_string_ostream os(sink);
return (**translation_)(source_manager, os, &context);
}
private:
llvm::cl::opt<const Translation*, false, TranslationParser>*
RegisterTranslation() {
// Can only register once per process.
static const auto requested_translation =
new llvm::cl::opt<const Translation*, false, TranslationParser>(
llvm::cl::desc("Translation to perform"));
return requested_translation;
}
llvm::cl::opt<const Translation*, false, TranslationParser>* translation_;
};
TEST_F(MlirTranslationTest, TranslatesMlirToGraph) {
static constexpr char kMlirSource[] = R"(
func.func @main() -> (tensor<1x2xf16>, tensor<2xf16>) {
%graph:2 = tf_executor.graph {
%0:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<1.0> : tensor<1x2xf16>} : () -> tensor<1x2xf16> loc("const1")
%1:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<[1.0, 2.0]> : tensor<2xf16>} : () -> tensor<2xf16> loc("const2")
tf_executor.fetch %0#0, %1#0 : tensor<1x2xf16>, tensor<2xf16>
}
func.return %graph#0, %graph#1 : tensor<1x2xf16>, tensor<2xf16>
})";
std::string result;
auto status = Translate(kMlirSource, result);
ASSERT_TRUE(status.succeeded());
EXPECT_TRUE(absl::StrContains(result, "node {"));
}
} // namespace testing
} // namespace v2
} // namespace tf2xla
} // namespace tensorflow