Responded to more suggestions

Chillee · Chillee · commit 99674d15a9f6 · 2019-05-28T19:18:26.000-07:00
diff --git a/c10/util/sparse_bitset.h b/c10/util/sparse_bitset.h
@@ -154,12 +154,12 @@
      Copy &= ~0UL << BitPos;
 
      if (Copy != 0)
-       return WordPos * BITWORD_SIZE + __builtin_ctzl(Copy);
+       return WordPos * BITWORD_SIZE + llvm::countTrailingZeros(Copy);
 
      // Check subsequent words.
      for (unsigned i = WordPos+1; i < BITWORDS_PER_ELEMENT; ++i)
        if (Bits[i] != 0)
-         return i * BITWORD_SIZE + __builtin_ctzl(Bits[i]);
+         return i * BITWORD_SIZE + llvm::countTrailingZeros(Bits[i]);
      return -1;
    }
 
diff --git a/torch/csrc/jit/passes/utils/memory_dag.cpp b/torch/csrc/jit/passes/utils/memory_dag.cpp
@@ -8,17 +8,11 @@
 namespace torch {
 namespace jit {
 namespace {
-ska::flat_hash_map<const Element*, unsigned> comprMap;
 std::vector<const Element*> decomprMap;
 } // namespace
-
-unsigned Element::toIndex(const Element* x) {
-  if (comprMap.count(x)) {
-    return comprMap[x];
-  }
-  comprMap[x] = comprMap.size();
-  decomprMap.push_back(x);
-  return comprMap[x];
+unsigned Element::indexCount = 0;
+Element::Element(const Value* value_) : value(value_), index(indexCount++) {
+  decomprMap.push_back(this);
 }
 
 const Element* Element::toElement(unsigned x) {
@@ -54,7 +48,7 @@ void collectAllContainedMemoryLocations(
     const Element* elem,
     MemoryLocations& cont) {
   // we have already recursed on this element
-  unsigned compIdx = Element::toIndex(elem);
+  unsigned compIdx = elem->index;
   if (cont.test(compIdx)) {
     return;
   }
@@ -101,18 +95,17 @@ bool MemoryDAG::mayContainAlias(
 
 // Make `v` point at `to`.
 void MemoryDAG::makePointerTo(Element* from, Element* to) {
-  from->pointsTo.set(Element::toIndex(to));
-  to->pointedFrom.set(Element::toIndex(from));
+  from->pointsTo.set(to->index);
+  to->pointedFrom.set(from->index);
 }
 
 void MemoryDAG::addToContainedElements(Element* elem, Element* container) {
-  container->contained_elements.set(Element::toIndex(elem));
+  container->contained_elements.set(elem->index);
 }
 
 // Give `v` a fresh alias (i.e. it does not point to any value)
 Element* MemoryDAG::makeFreshValue(const Value* v) {
-  auto el = torch::make_unique<Element>();
-  el->value = v;
+  auto el = torch::make_unique<Element>(v);
 
   auto rawPtr = el.get();
   elements_.emplace(rawPtr, std::move(el));
@@ -135,8 +128,8 @@ const MemoryLocations& Element::getMemoryLocations() const {
 // traversing in the direction `dir`.`fn` will be run on each element.
 void Element::bfs(BfsDirection dir, MemoryLocations& res) const {
   std::queue<unsigned> queue;
-  std::unordered_set<int> seen;
-  queue.push(Element::toIndex(this));
+  ska::flat_hash_set<int> seen;
+  queue.push(this->index);
   while (!queue.empty()) {
     const auto el = queue.front();
     queue.pop();
diff --git a/torch/csrc/jit/passes/utils/memory_dag.h b/torch/csrc/jit/passes/utils/memory_dag.h
@@ -126,6 +126,9 @@ struct Element {
   MemoryLocations pointedFrom;
 
   MemoryLocations contained_elements;
+  static unsigned indexCount;
+  signed index;
+  Element(const Value* value_);
 
   // Return the unique memory locations that `Element` might represent.
   TORCH_API const MemoryLocations& getMemoryLocations() const;
@@ -138,8 +141,7 @@ struct Element {
   // traversing in the direction `dir`.`fn` will be run on each element.
   void bfs(BfsDirection dir, MemoryLocations& res) const;
 
-  // Converts to and from the compressed index representation
-  static unsigned toIndex(const Element* x);
+  // Converts from the compressed index representation
   static const Element* toElement(unsigned x);
 };
 } // namespace jit
