[Discuss] Reducing allocations in HnswUtil::markRooted

[viswanathk]

Description

I was going through the nightly benchmark runs and came across this https://blunders.io/jfr-demo/indexing-1kb-vectors-2024.11.17.18.04.47/allocations-drill-down

Almost 7% of the allocations seem to be due to usage of ArrayDeque and autoboxing that's required to build the stack in this code:

private static Component markRooted(
      HnswGraph hnswGraph,
      int level,
      FixedBitSet connectedNodes,
      FixedBitSet notFullyConnected,
      int maxConn,
      int entryPoint)
      throws IOException {
    // Start at entry point and search all nodes on this level
    // System.out.println("markRooted level=" + level + " entryPoint=" + entryPoint);
    Deque<Integer> stack = new ArrayDeque<>();
    stack.push(entryPoint);
    int count = 0;
    while (!stack.isEmpty()) {
      int node = stack.pop();
      if (connectedNodes.get(node)) {
        continue;
      }
      count++;
      connectedNodes.set(node);
      hnswGraph.seek(level, node);
      int friendOrd;
      int friendCount = 0;
      while ((friendOrd = hnswGraph.nextNeighbor()) != NO_MORE_DOCS) {
        ++friendCount;
        stack.push(friendOrd);
      }
      if (friendCount < maxConn && notFullyConnected != null) {
        notFullyConnected.set(node);
      }
    }
    return new Component(entryPoint, count);
  }

Is this something that's worth improving - seems like a low hanging fix? (I can contribute)

If it's worth improving, then:

1. If there's a theoretical max depth for the stack, then the stack can be implemented with primitives avoiding the autoboxing cost, and reducing the allocation cost.
2. In case there's no reasonable max depth, then starting off with a higher capacity for ArrayDeque should reduce some allocations. This code in ArrayDeque::grow looks quite wasteful in terms of allocations.

    private void grow(int needed) {
        int oldCapacity = this.elements.length;
        int jump = oldCapacity < 64 ? oldCapacity + 2 : oldCapacity >> 1;
        int newCapacity;
        if (jump < needed || (newCapacity = oldCapacity + jump) - 2147483639 > 0) {
            newCapacity = this.newCapacity(needed, jump);
        }

        Object[] es = this.elements = Arrays.copyOf(this.elements, newCapacity);
        if (this.tail < this.head || this.tail == this.head && es[this.head] != null) {
            int newSpace = newCapacity - oldCapacity;
            System.arraycopy(es, this.head, es, this.head + newSpace, oldCapacity - this.head);
            int i = this.head;

            for(int to = this.head += newSpace; i < to; ++i) {
                es[i] = null;
            }
        }

    }

Thoughts?

[viswanathk]

cc: @msokolov

[vigyasharma]

Interesting. The "theoretical" max depth of this stack would be the size of this graph. I suppose the stack does get large, which would explain a high no. of Array::grow calls? Would be interesting to see the impact of just setting initial capacity to half of the size of graph.

[msokolov]

I do think it's worth improving. Another way could be to measure empirically the stack depth - maybe it scales in a predictable way with total number of vectors? And then we can use that prediction as an initial size. We could also think of using a primitive array that grows (implement our own primitive ArrayDeque).