Switching the presized cuckoo map from using strict mod to Lemire's

David G. Andersen · tensorflower-gardener · commit a47a30018502 · 2016-09-01T12:17:46.000-07:00
uniform range mapping trick. Removes dependency on Eigen's TensorIntDiv, which doesn't work properly on Android, is 10-20% faster on x86, and should be much faster on Cuda if needed. (There are remaining optimizations to force the use of __umulhi on cuda, but this table is designed for CPU, so I don't see a reason to complicate things). OLD: CPU: Intel Haswell with HyperThreading (6 cores) dL1:32KB dL2:256KB dL3:15MB Benchmark Time(ns) CPU(ns) Iterations --------------------------------------------------- BM_CuckooFill/1000 14859 14846 46766 BM_CuckooFill/10M 835154969 834427162 100 BM_CuckooRead/1000 10 10 67484647 BM_CuckooRead/10M 56 56 10000000 NEW: BM_CuckooFill/1000 12385 12374 56240 BM_CuckooFill/10M 696061920 695467681 100 BM_CuckooRead/1000 9 9 7872528 BM_CuckooRead/10M 44 44 15487881 This change will have bad consequences for people who violate the table's requirement that keys be pre-hashed into random-looking uint64's before inserting -- the table will not achieve its full capacity. (It won't return wrong results, and will return an error on insert.) That's a documented requirement, but we'll want to make sure that nobody's misusing it. Updates the TooManyKeys test to be more robust (assertion failure if the table fails to fill during the fill phase), and to pre-hash its keys as it should. Change: 131976392
diff --git a/tensorflow/core/util/presized_cuckoo_map.h b/tensorflow/core/util/presized_cuckoo_map.h
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include <algorithm>
 #include <vector>
-#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/platform/macros.h"
 
@@ -44,6 +43,32 @@ namespace tensorflow {
 // a good cuckoo path with less data movement (see
 // https://www.cs.cmu.edu/~dga/papers/cuckoo-eurosys14.pdf )
 
+namespace presized_cuckoo_map {
+// Utility function to compute (x * y) >> 64, or "multiply high".
+// On x86-64, this is a single instruction, but not all platforms
+// support the __uint128_t type, so we provide a generic
+// implementation as well.
+inline uint64 multiply_high_u64(uint64 x, uint64 y) {
+#if defined(__SIZEOF_INT128__)
+  return (uint64)(((__uint128_t)x * (__uint128_t)y) >> 64);
+#else
+  // For platforms without int128 support, do it the long way.
+  uint64 x_lo = x & 0xffffffff;
+  uint64 x_hi = x >> 32;
+  uint64 buckets_lo = y & 0xffffffff;
+  uint64 buckets_hi = y >> 32;
+  uint64 prod_hi = x_hi * buckets_hi;
+  uint64 prod_lo = x_lo * buckets_lo;
+  uint64 prod_mid1 = x_hi * buckets_lo;
+  uint64 prod_mid2 = x_lo * buckets_hi;
+  uint64 carry =
+      ((prod_mid1 & 0xffffffff) + (prod_mid2 & 0xffffffff) + (prod_lo >> 32)) >>
+      32;
+  return prod_hi + (prod_mid1 >> 32) + (prod_mid2 >> 32) + carry;
+#endif
+}
+}
+
 template <class value>
 class PresizedCuckooMap {
  public:
@@ -67,18 +92,14 @@ class PresizedCuckooMap {
     }
     buckets_.clear();
     buckets_.resize(num_buckets_, empty_bucket);
-#if !defined(__GCUDACC__) && !defined(__GCUDACC_HOST__) && \
-    !defined(IS_MOBILE_PLATFORM)
-    buckets_divisor_ = Eigen::internal::TensorIntDivisor<uint64>(num_buckets_);
-#endif
   }
 
   // Returns false if k is already in table or if the table
   // is full; true otherwise.
   bool InsertUnique(const key_type k, const value& v) {
     uint64 tk = key_transform(k);
-    uint64 b1 = fast_mod_by_buckets(tk);
-    uint64 b2 = fast_mod_by_buckets(h2(tk));
+    uint64 b1 = fast_map_to_buckets(tk);
+    uint64 b2 = fast_map_to_buckets(h2(tk));
 
     // Merged find and duplicate checking.
     uint64 target_bucket = 0;
@@ -107,8 +128,8 @@ class PresizedCuckooMap {
   // Returns true if found.  Sets *out = value.
   bool Find(const key_type k, value* out) const {
     uint64 tk = key_transform(k);
-    return FindInBucket(k, fast_mod_by_buckets(tk), out) ||
-           FindInBucket(k, fast_mod_by_buckets(h2(tk)), out);
+    return FindInBucket(k, fast_map_to_buckets(tk), out) ||
+           FindInBucket(k, fast_map_to_buckets(h2(tk)), out);
   }
 
  private:
@@ -180,9 +201,9 @@ class PresizedCuckooMap {
       return e;
     }
 
-    bool empty() { return head_ == tail_; }
+    bool empty() const { return head_ == tail_; }
 
-    bool full() { return ((tail_ + 1) % kMaxQueueSize) == head_; }
+    bool full() const { return ((tail_ + 1) % kMaxQueueSize) == head_; }
 
     void reset() { head_ = tail_ = 0; }
 
@@ -210,13 +231,13 @@ class PresizedCuckooMap {
     return m * ((h >> 32) | (h << 32));
   }
 
-  // alt_bucket identifies the "other" bucket for key k, whether
+  // alt_bucket identifies the "other" bucket for key k, where
   // other is "the one that isn't bucket b"
   inline uint64 alt_bucket(key_type k, uint64 b) const {
-    if (fast_mod_by_buckets(k) != b) {
-      return fast_mod_by_buckets(k);
+    if (fast_map_to_buckets(k) != b) {
+      return fast_map_to_buckets(k);
     }
-    return fast_mod_by_buckets(h2(k));
+    return fast_map_to_buckets(h2(k));
   }
 
   inline void InsertInternal(key_type k, const value& v, uint64 b, int slot) {
@@ -306,22 +327,17 @@ class PresizedCuckooMap {
     return false;
   }
 
-  inline uint64 fast_mod_by_buckets(uint64 x) const {
-// Omitting the optimized bucket mod for CUDA platforms
-// until Eigen supports 2^63 divisors on GPU.
-#if !defined(__GCUDACC__) && !defined(__GCUDACC_HOST__) && \
-    !defined(IS_MOBILE_PLATFORM)
-    x &= ~(1ULL << 63);  // Fast div can only handle 2^63-1
-    return x - num_buckets_ * (x / buckets_divisor_);
-#else
-    return x % num_buckets_;
-#endif
+  inline uint64 fast_map_to_buckets(uint64 x) const {
+    // Map x (uniform in 2^64) to the range [0, num_buckets_ -1]
+    // using Lemire's alternative to modulo reduction:
+    // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
+    // Instead of x % N, use (x * N) >> 64.
+    return presized_cuckoo_map::multiply_high_u64(x, num_buckets_);
   }
 
   // Set upon initialization: num_entries / kLoadFactor / kSlotsPerBucket.
   uint64 num_buckets_;
   std::vector<Bucket> buckets_;
-  Eigen::internal::TensorIntDivisor<uint64> buckets_divisor_;  // for fast mod
 
   std::unique_ptr<CuckooPathQueue> cpq_;
   CuckooPathEntry visited_[kVisitedListSize];
diff --git a/tensorflow/core/util/presized_cuckoo_map_test.cc b/tensorflow/core/util/presized_cuckoo_map_test.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/util/presized_cuckoo_map.h"
+#include <array>
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/fingerprint.h"
 #include "tensorflow/core/platform/test.h"
@@ -22,6 +23,25 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+TEST(PresizedCuckooMapTest, MultiplyHigh) {
+  struct Testcase {
+    uint64 x;
+    uint64 y;
+    uint64 result;
+  };
+  std::array<Testcase, 7> testcases{
+      {{0, 0, 0},
+       {0xffffffff, 0xffffffff, 0},
+       {0x2, 0xf000000000000000, 1},
+       {0x3, 0xf000000000000000, 2},
+       {0x3, 0xf000000000000001, 2},
+       {0x3, 0xffffffffffffffff, 2},
+       {0xffffffffffffffff, 0xffffffffffffffff, 0xfffffffffffffffe}}};
+  for (auto &tc : testcases) {
+    EXPECT_EQ(tc.result, presized_cuckoo_map::multiply_high_u64(tc.x, tc.y));
+  }
+}
+
 TEST(PresizedCuckooMapTest, Basic) {
   PresizedCuckooMap<int> pscm(1000);
   EXPECT_TRUE(pscm.InsertUnique(1, 2));
@@ -34,24 +54,31 @@ TEST(PresizedCuckooMapTest, TooManyItems) {
   static constexpr int kTableSize = 1000;
   PresizedCuckooMap<int> pscm(kTableSize);
   for (uint64 i = 0; i < kTableSize; i++) {
-    EXPECT_TRUE(pscm.InsertUnique(i, i));
+    uint64 key =
+        Fingerprint64(string(reinterpret_cast<char *>(&i), sizeof(int64)));
+    ASSERT_TRUE(pscm.InsertUnique(key, i));
   }
   // Try to over-fill the table.  A few of these
   // inserts will succeed, but should start failing.
   uint64 failed_at = 0;
   for (uint64 i = kTableSize; i < (2 * kTableSize); i++) {
-    if (!pscm.InsertUnique(i, i)) {
+    uint64 key =
+        Fingerprint64(string(reinterpret_cast<char *>(&i), sizeof(int64)));
+    if (!pscm.InsertUnique(key, i)) {
       failed_at = i;
       break;
     }
   }
   // Requirement 1:  Table must return failure when it's full.
   EXPECT_NE(failed_at, 0);
+
   // Requirement 2:  Table must preserve all items inserted prior
   // to the failure.
   for (uint64 i = 0; i < failed_at; i++) {
     int out;
-    EXPECT_TRUE(pscm.Find(i, &out));
+    uint64 key =
+        Fingerprint64(string(reinterpret_cast<char *>(&i), sizeof(int64)));
+    EXPECT_TRUE(pscm.Find(key, &out));
     EXPECT_EQ(out, i);
   }
 }
