From a29691ddf2420ce7d42565014e2e96393b23750f Mon Sep 17 00:00:00 2001
From: Ayush Dubey <ayushd@google.com>
Date: Thu, 27 Jun 2019 16:53:28 -0700
Subject: [PATCH] Do not change scope_id in AllocatorAttributes::Merge when it
 matches other.

Before this change, if attr1 and attr2 had the same scope_id value,
attr1.Merge(attr2) would result in attr1.scope_id being set to 0.  This change
makes Merge do nothing to scope_id if it is the same, which seems like the
natural behavior of "merging" the same values.

This change also introduces a CHECK to ensure that, when merging, both scope
ids are not non-zero.

PiperOrigin-RevId: 255506545
---
 tensorflow/core/framework/allocator.h       | 10 +++--
 tensorflow/core/framework/allocator_test.cc | 46 +++++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/framework/allocator.h b/tensorflow/core/framework/allocator.h
index 861e855c3d5..73196570f75 100644
--- a/tensorflow/core/framework/allocator.h
+++ b/tensorflow/core/framework/allocator.h
@@ -288,9 +288,13 @@ struct AllocatorAttributes {
   bool gpu_compatible() const { return value & (0x1 << 2); }
   void Merge(AllocatorAttributes other) {
     value |= other.value;
-    scope_id = (scope_id > 0 && other.scope_id == 0)
-                   ? scope_id
-                   : ((scope_id == 0) ? other.scope_id : 0);
+    if (scope_id != other.scope_id) {
+      CHECK(scope_id == 0 || other.scope_id == 0)
+          << "At least one scope_id should be zero to merge "
+             "AllocatorAttributes but found this.scope_id="
+          << scope_id << " and other.scope_id=" << other.scope_id;
+      scope_id = scope_id == 0 ? other.scope_id : scope_id;
+    }
   }
   // Returns true if the fields set in *this is a subset of or equal to
   // those set in other.
diff --git a/tensorflow/core/framework/allocator_test.cc b/tensorflow/core/framework/allocator_test.cc
index 3d03b2da1d3..3caab02eeba 100644
--- a/tensorflow/core/framework/allocator_test.cc
+++ b/tensorflow/core/framework/allocator_test.cc
@@ -86,6 +86,52 @@ TEST(AllocatorAttributesTest, IsEqualOrLessRestrictiveThan) {
   EXPECT_FALSE(a.IsEqualOrLessRestrictiveThan(b));
 }
 
+TEST(AllocatorAttributesTest, Merge) {
+  AllocatorAttributes a, b;
+
+  // Merging nic_compatible=True and nic_compatible=False results in
+  // nic_compatible=True.
+  EXPECT_EQ(a.value, 0);
+  EXPECT_EQ(b.value, 0);
+  EXPECT_FALSE(a.nic_compatible());
+  EXPECT_FALSE(b.nic_compatible());
+  b.set_nic_compatible(true);
+  a.Merge(b);
+  EXPECT_TRUE(a.nic_compatible());
+  EXPECT_TRUE(b.nic_compatible());
+
+  // a.Merge(b) does not change b.
+  EXPECT_EQ(a.scope_id, 0);
+  EXPECT_EQ(b.scope_id, 0);
+  a.scope_id = 1;
+  a.Merge(b);
+  EXPECT_EQ(a.scope_id, 1);
+  EXPECT_EQ(b.scope_id, 0);
+
+  // If a.scope_id=1 and b.scope_id=0, then b.Merge(a) results in b.scope_id=1.
+  a.scope_id = 1;
+  b.scope_id = 0;
+  b.Merge(a);
+  EXPECT_EQ(a.scope_id, 1);
+  EXPECT_EQ(b.scope_id, 1);
+
+  // If a.scope_id and b.scope_id are same, then merge leaves them unchanged.
+  a.scope_id = 2;
+  b.scope_id = 2;
+  a.Merge(b);
+  EXPECT_EQ(a.scope_id, 2);
+  EXPECT_EQ(b.scope_id, 2);
+}
+
+TEST(AllocatorAttributesDeathTest, MergeDifferentScopeIds) {
+  AllocatorAttributes a, b;
+  // If a.scope_id and b.scope_id are both positive but different, then
+  // a.Merge(b) should cause a CHECK failure.
+  a.scope_id = 3;
+  b.scope_id = 4;
+  EXPECT_DEATH({ a.Merge(b); }, "");
+}
+
 TEST(CPUAllocatorTest, Simple) {
   EnableCPUAllocatorStats(true);
   Allocator* a = cpu_allocator();