[XLA] Fix name_uniquer to avoiding using the name prefixes reserved by the LLVM

x86 backend.

Previously, we allow kernel names such as __llvm_retpoline_, which are reserved
by the LLVM x86 backend. This change modifies the name uniquer to morph names
with prefix __ but not __xla_.

Change the AOT compiler tests to use name prefix __xla_ instead of __.

Add test cases.

PiperOrigin-RevId: 243713430
This commit is contained in:
Bixia Zheng 2019-04-15 16:46:29 -07:00 committed by TensorFlower Gardener
parent 85709f9d72
commit a30dd31e19
3 changed files with 19 additions and 1 deletions

View File

@ -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:

View File

@ -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;
}

View File

@ -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"));