diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index fd701ab7166..2c08cb63275 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -163,7 +163,10 @@ def tf_library( header_file = name + ".h" metadata_object_file = name + "_tfcompile_metadata.o" function_object_file = name + "_tfcompile_function.o" - ep = ("__" + native.package_name() + "__" + name).replace("/", "_") + + # The XLA backends morph kernal name prefix __ that is not in the form of + # __xla_. + ep = ("__xla_" + native.package_name() + "__" + name).replace("/", "_") if type(tfcompile_flags) == type(""): flags = tfcompile_flags else: diff --git a/tensorflow/compiler/xla/service/name_uniquer.cc b/tensorflow/compiler/xla/service/name_uniquer.cc index e55b83d17e9..70742b67a28 100644 --- a/tensorflow/compiler/xla/service/name_uniquer.cc +++ b/tensorflow/compiler/xla/service/name_uniquer.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/name_uniquer.h" #include "absl/strings/ascii.h" +#include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/primitive_util.h" @@ -62,6 +63,14 @@ NameUniquer::NameUniquer(const string& separator) { if (primitive_util::IsPrimitiveTypeName(result) && result != "tuple") { result += "_"; } + + if (absl::StartsWith(result, "__") && !absl::StartsWith(result, "__xla_")) { + // Morph name prefix __ that is not __xla_, to avoid using name prefixes + // reserved by the backends, such as __llvm_retpoline_ reserved by the LLVM + // x86 backend. + result[0] = 'a'; + } + return result; } diff --git a/tensorflow/compiler/xla/service/name_uniquer_test.cc b/tensorflow/compiler/xla/service/name_uniquer_test.cc index d0d04147e0c..1007c2aeae8 100644 --- a/tensorflow/compiler/xla/service/name_uniquer_test.cc +++ b/tensorflow/compiler/xla/service/name_uniquer_test.cc @@ -111,6 +111,12 @@ TEST_F(NameUniquerTest, AvoidKeywords) { EXPECT_EQ("s64_", uniquer.GetUniqueName("s64")); EXPECT_EQ("pred_", uniquer.GetUniqueName("pred")); + // Name prefix __xla_ is preserved. + EXPECT_NE(uniquer.GetUniqueName("__xla_").find("__xla_"), std::string::npos); + // Other form of __ prefixes is not preserved to avoid using name prefixes + // reserved by backends. + EXPECT_EQ(uniquer.GetUniqueName("__abx").find("__"), std::string::npos); + // Though a primitive type, "tuple" is not a keyword. EXPECT_EQ("tuple", uniquer.GetUniqueName("tuple"));