Add bitmap implementation to tensorflow/core/lib/core

Change: 133839671

4 files changed, 307 insertions, 0 deletions

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index b677cab512..e6efcffe70 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -139,6 +139,7 @@ cc_library(
     hdrs = [
         # TODO(josh11b): Make many of these internal.
         "lib/core/arena.h",
+        "lib/core/bitmap.h",
         "lib/core/bits.h",
         "lib/core/casts.h",
         "lib/core/coding.h",
@@ -1361,6 +1362,7 @@ tf_cc_tests(
     srcs = [
         "lib/core/arena_test.cc",
         "lib/core/bit_cast_test.cc",
+        "lib/core/bitmap_test.cc",
         "lib/core/blocking_counter_test.cc",
         "lib/core/coding_test.cc",
         "lib/core/notification_test.cc",
diff --git a/tensorflow/core/lib/core/bitmap.cc b/tensorflow/core/lib/core/bitmap.cc
new file mode 100644
index 0000000000..110d31044a
--- /dev/null
+++ b/tensorflow/core/lib/core/bitmap.cc
@@ -0,0 +1,71 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/lib/core/bitmap.h"
+
+#include <string.h>
+#include <strings.h>
+
+namespace tensorflow {
+namespace core {
+
+void Bitmap::Reset(size_t n) {
+  const size_t num_words = NumWords(n);
+  if (num_words != NumWords(nbits_)) {
+    // Reallocate.
+    Word* w = new Word[num_words];
+    delete[] word_;
+    word_ = w;
+  }
+  memset(word_, 0, sizeof(word_[0]) * num_words);
+  nbits_ = n;
+}
+
+// Return 1+index of the first set bit in w; return 0 if w == 0.
+static size_t FindFirstSet(uint64 w) { return ffs(w); }
+
+size_t Bitmap::FirstUnset(size_t start) const {
+  if (start >= nbits_) {
+    return nbits_;
+  }
+
+  // Mask to or-into first word to account for bits to skip in that word.
+  size_t mask = (1ull << (start % kBits)) - 1;
+  const size_t nwords = NumWords(nbits_);
+  for (size_t i = start / kBits; i < nwords; i++) {
+    Word word = word_[i] | mask;
+    mask = 0;  // Only ignore bits in the first word we process.
+    size_t r = FindFirstSet(~word);
+    if (r) {
+      size_t result = i * kBits + (r - 1);
+      if (result > nbits_) result = nbits_;
+      return result;
+    }
+  }
+
+  return nbits_;
+}
+
+string Bitmap::ToString() const {
+  string result;
+  result.resize(bits());
+  for (int i = 0; i < nbits_; i++) {
+    result[i] = get(i) ? '1' : '0';
+  }
+  return result;
+}
+
+}  // namespace core
+}  // namespace tensorflow
diff --git a/tensorflow/core/lib/core/bitmap.h b/tensorflow/core/lib/core/bitmap.h
new file mode 100644
index 0000000000..b30479fa1b
--- /dev/null
+++ b/tensorflow/core/lib/core/bitmap.h
@@ -0,0 +1,106 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_TENSORFLOW_CORE_LIB_CORE_BITMAP_H_
+#define THIRD_PARTY_TENSORFLOW_CORE_LIB_CORE_BITMAP_H_
+
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace core {
+
+class Bitmap {
+ public:
+  // Create a bitmap that holds 0 bits.
+  Bitmap();
+
+  // Create a bitmap that holds n bits, all initially zero.
+  explicit Bitmap(size_t n);
+
+  ~Bitmap();
+
+  Bitmap(const Bitmap&) = delete;
+  Bitmap& operator=(const Bitmap&) = delete;
+
+  // Return the number of bits that the bitmap can hold.
+  size_t bits() const;
+
+  // Replace contents of *this with a bitmap of n bits, all set to zero.
+  void Reset(size_t n);
+
+  // Return the contents of the ith bit.
+  // REQUIRES: i < bits()
+  bool get(size_t i) const;
+
+  // Set the contents of the ith bit to true.
+  // REQUIRES: i < bits()
+  void set(size_t i);
+
+  // Set the contents of the ith bit to false.
+  // REQUIRES: i < bits()
+  void clear(size_t i);
+
+  // Return the smallest i such that i >= start and !get(i).
+  // Returns bits() if no such i exists.
+  size_t FirstUnset(size_t start) const;
+
+  // Returns the bitmap as an ascii string of '0' and '1' characters, bits()
+  // characters in length.
+  string ToString() const;
+
+ private:
+  typedef uint32 Word;
+  static const size_t kBits = 32;
+
+  // Return the number of words needed to store n bits.
+  static size_t NumWords(size_t n) { return (n + kBits - 1) / kBits; }
+
+  // Return the mask to use for the ith bit in a word.
+  static Word Mask(size_t i) { return 1ull << i; }
+
+  size_t nbits_;  // Length of bitmap in bits.
+  Word* word_;
+};
+
+// Implementation details follow.  Clients should ignore.
+
+inline Bitmap::Bitmap() : nbits_(0), word_(nullptr) {}
+
+inline Bitmap::Bitmap(size_t n) : Bitmap() { Reset(n); }
+
+inline Bitmap::~Bitmap() { delete[] word_; }
+
+inline size_t Bitmap::bits() const { return nbits_; }
+
+inline bool Bitmap::get(size_t i) const {
+  DCHECK_LT(i, nbits_);
+  return word_[i / kBits] & Mask(i % kBits);
+}
+
+inline void Bitmap::set(size_t i) {
+  DCHECK_LT(i, nbits_);
+  word_[i / kBits] |= Mask(i % kBits);
+}
+
+inline void Bitmap::clear(size_t i) {
+  DCHECK_LT(i, nbits_);
+  word_[i / kBits] &= ~Mask(i % kBits);
+}
+
+}  // namespace core
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_TENSORFLOW_CORE_LIB_CORE_BITMAP_H_
diff --git a/tensorflow/core/lib/core/bitmap_test.cc b/tensorflow/core/lib/core/bitmap_test.cc
new file mode 100644
index 0000000000..5046c5e2c3
--- /dev/null
+++ b/tensorflow/core/lib/core/bitmap_test.cc
@@ -0,0 +1,128 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/lib/core/bitmap.h"
+
+#include "tensorflow/core/lib/random/simple_philox.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace core {
+namespace {
+
+// Return next size to test after n.
+size_t NextSize(size_t n) { return n + ((n < 75) ? 1 : 25); }
+
+static void MakeRandomBitmap(random::SimplePhilox* rnd, Bitmap* bitmap) {
+  size_t n = rnd->Uniform(200);
+  bitmap->Reset(n);
+  for (size_t i = 0; i < n; i++) {
+    if (rnd->OneIn(2)) bitmap->set(i);
+  }
+}
+
+TEST(BitmapTest, Basic) {
+  for (size_t n = 0; n < 200; n = NextSize(n)) {
+    Bitmap bits(n);
+    for (size_t i = 0; i < n; i++) {
+      EXPECT_FALSE(bits.get(i)) << n << " " << i << " " << bits.ToString();
+      bits.set(i);
+      EXPECT_TRUE(bits.get(i)) << n << " " << i << " " << bits.ToString();
+      bits.clear(i);
+      EXPECT_FALSE(bits.get(i)) << n << " " << i << " " << bits.ToString();
+    }
+  }
+}
+
+TEST(BitmapTest, ToString) {
+  Bitmap bits(10);
+  bits.set(1);
+  bits.set(3);
+  EXPECT_EQ(bits.ToString(), "0101000000");
+}
+
+TEST(BitmapTest, FirstUnset) {
+  for (size_t n = 0; n < 200; n = NextSize(n)) {
+    for (size_t p = 0; p <= 100; p++) {
+      for (size_t q = 0; q <= 100; q++) {
+        // Generate a bitmap of length n with long runs of ones.
+        Bitmap bitmap(n);
+        // Set first p bits to 1.
+        int one_count = 0;
+        size_t i = 0;
+        while (i < p && i < n) {
+          one_count++;
+          bitmap.set(i);
+          i++;
+        }
+        // Fill rest with a pattern of 0 followed by q 1s.
+        while (i < n) {
+          i++;
+          for (int j = 0; j < q && i < n; j++, i++) {
+            one_count++;
+            bitmap.set(i);
+          }
+        }
+
+        // Now use FirstUnset to iterate over unset bits and verify
+        // that all encountered bits are clear.
+        int seen = 0;
+        size_t pos = 0;
+        while (true) {
+          pos = bitmap.FirstUnset(pos);
+          if (pos == n) break;
+          ASSERT_FALSE(bitmap.get(pos)) << pos << " " << bitmap.ToString();
+          seen++;
+          pos++;
+        }
+        EXPECT_EQ(seen, n - one_count) << " " << bitmap.ToString();
+      }
+    }
+  }
+}
+
+TEST(BitmapTest, FirstUnsetRandom) {
+  random::PhiloxRandom philox(301, 17);
+  random::SimplePhilox rnd(&philox);
+  for (int iter = 0; iter < 10000; iter++) {
+    Bitmap bitmap;
+    MakeRandomBitmap(&rnd, &bitmap);
+
+    // Count number of unset bits in bitmap.
+    size_t zero_bits = 0;
+    for (size_t i = 0; i < bitmap.bits(); i++) {
+      if (!bitmap.get(i)) zero_bits++;
+    }
+
+    // Now use FirstUnset to iterate over unset bits and verify
+    // that all encountered bits are clear.
+    int seen = 0;
+    size_t pos = 0;
+    while (true) {
+      pos = bitmap.FirstUnset(pos);
+      if (pos == bitmap.bits()) break;
+      ASSERT_FALSE(bitmap.get(pos)) << pos << " " << bitmap.ToString();
+      seen++;
+      pos++;
+    }
+
+    EXPECT_EQ(seen, zero_bits) << " " << bitmap.ToString();
+  }
+}
+
+}  // namespace
+}  // namespace core
+}  // namespace tensorflow