Fix FunctionDefsEqual to handle multiple attributes.

The fix is to change StringPiece to string in our temporary maps, since the
backing memory for the StringPiece was actually dead by the time we used the
temporary maps.

Note that FunctionDef::attr() is a proto2::Map, which has a value_type that
behaves similarly to std::pair.  Here's a re-written snippet of the original
GetSetAttrs code, which makes the bug more apparent:

  std::map<StringPiece, AttrValue> GetSetAttrs(const FunctionDef& fdef) {
    std::map<StringPiece, AttrValue> attrs;
	  for (std::pair<string, AttrValue> pair : fdef.attr()) {
		  attrs[pair.first] = pair.second;
	  }
		return attrs;
	}

The problem is that the `pair` that captures each entry of fdef.attr() creates a
copy of the underlying string, and we're converting that into a StringPiece.
The underlying memory for the string is dead on each iteration of the loop.

I've also added a test that fails with the old implementation, and passes with
the new implementation.

PiperOrigin-RevId: 161265130

2 files changed, 24 insertions, 8 deletions

diff --git a/tensorflow/core/framework/function.cc b/tensorflow/core/framework/function.cc
index c6a20bf3ce..1774f74ca8 100644
--- a/tensorflow/core/framework/function.cc
+++ b/tensorflow/core/framework/function.cc
@@ -729,11 +729,11 @@ namespace {
 // Returns the name -> attr mapping of fdef's attrs that have a value set. In
 // Python, it's possible to access unset attrs, which returns a default value
 // and adds an unset attr to the map.
-std::map<StringPiece, AttrValue> GetSetAttrs(const FunctionDef& fdef) {
-  std::map<StringPiece, AttrValue> set_attrs;
-  for (auto iter : fdef.attr()) {
-    if (iter.second.value_case() != AttrValue::VALUE_NOT_SET) {
-      set_attrs[iter.first] = iter.second;
+std::map<string, AttrValue> GetSetAttrs(const FunctionDef& fdef) {
+  std::map<string, AttrValue> set_attrs;
+  for (auto pair : fdef.attr()) {
+    if (pair.second.value_case() != AttrValue::VALUE_NOT_SET) {
+      set_attrs[pair.first] = pair.second;
     }
   }
   return set_attrs;
@@ -753,8 +753,8 @@ bool FunctionDefsEqual(const FunctionDef& f1, const FunctionDef& f2) {
   f2.signature().SerializeToString(&sig2);
   if (sig1 != sig2) return false;
 
-  std::map<StringPiece, AttrValue> f1_attrs = GetSetAttrs(f1);
-  std::map<StringPiece, AttrValue> f2_attrs = GetSetAttrs(f2);
+  std::map<string, AttrValue> f1_attrs = GetSetAttrs(f1);
+  std::map<string, AttrValue> f2_attrs = GetSetAttrs(f2);
   if (f1_attrs.size() != f2_attrs.size()) return false;
   for (auto iter1 : f1_attrs) {
     auto iter2 = f2_attrs.find(iter1.first);
diff --git a/tensorflow/core/framework/function_test.cc b/tensorflow/core/framework/function_test.cc
index 1173384a1e..140dbd8932 100644
--- a/tensorflow/core/framework/function_test.cc
+++ b/tensorflow/core/framework/function_test.cc
@@ -1156,7 +1156,7 @@ TEST(FunctionLibraryDefinitionTest, GetAttr_Gradient) {
 // TODO(skyewm): this could be more thorough
 TEST(FunctionDefsEqualTest, TestFunctionDefsEqual) {
   // Equal functions
-  FunctionDef fdef1 = test::function::XTimesTwo();
+  const FunctionDef fdef1 = test::function::XTimesTwo();
   FunctionDef fdef2 = test::function::XTimesTwo();
   EXPECT_TRUE(FunctionDefsEqual(fdef1, fdef2));
 
@@ -1183,6 +1183,22 @@ TEST(FunctionDefsEqualTest, TestFunctionDefsEqual) {
   fdef2 = test::function::XTimesTwo();
   (*fdef2.mutable_ret())["y"] = "y:z:1";  // originally is "y:z:0"
   EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+
+  // Different attributes
+  fdef2 = test::function::XTimesTwo();
+  SetAttrValue(&fdef2, "ExtraAttr", true);
+  EXPECT_FALSE(FunctionDefsEqual(fdef1, fdef2));
+
+  // Multiple equivalent attributes; the two functions should be equal.
+  fdef2 = test::function::XTimesTwo();
+  FunctionDef fdef3 = test::function::XTimesTwo();
+  SetAttrValue(&fdef2, "Foo", true);
+  SetAttrValue(&fdef3, "Foo", true);
+  SetAttrValue(&fdef2, "Bar", 123);
+  SetAttrValue(&fdef3, "Bar", 123);
+  SetAttrValue(&fdef2, "Baz", "abc");
+  SetAttrValue(&fdef3, "Baz", "abc");
+  EXPECT_TRUE(FunctionDefsEqual(fdef2, fdef3));
 }
 
 }  // end namespace