[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:
parent
85709f9d72
commit
a30dd31e19
@ -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:
|
||||
|
@ -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;
|
||||
}
|
||||
|
||||
|
@ -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"));
|
||||
|
||||
|
Loading…
x
Reference in New Issue
Block a user