Dynamic Lookup Table

[chrisflesher]

Hello,

I was interested in creating an unbounded voxel map. I think this requires some data structure that can convert a voxel location (key) to data storage location (value). I found this discussion which talks about using searchsorted and cuckoo hash table to do this.

I saw that BCOO sparse gather is working recently but not sure that I can use this since keys could contain negative integers?

I took a crack at implementing a data structure that can might be able to scale to larger data sizes. Is something like this a good idea or am I totally barking up the wrong tree? Is there a library that does this kind of stuff already?

import functools
import math
import typing

import jax
import jax.numpy as jnp
import numpy as onp


class PartitionTable(typing.NamedTuple):
    """Hash lookup table that translates keys to a unique data storage index.

    There are 3 basic operations:
    1. Initialization: Memory is over-allocated to make hash collisions statistically
       unlikely (using a Chernoff bound). If a hash collision occurs the table should
       be regenerated using a different hash prime numbers (not implemented yet).
    2. Insert: Multiple keys are inserted into the table by selecting a bucket
       using a hash function and adding keys so they are left-justified.
    3. Query: Return a unique data storage index for the given key.

    Attributes:
        values: Key storage (num_buckets, num_slots)
        counts: Number of slots filled (num_buckets)

    Reference:
        Bender, Michael A., et al. "Modern Hashing Made Simple." 2024 Symposium
        on Simplicity in Algorithms (SOSA). Society for Industrial and Applied
        Mathematics, 2024.
    """

    keys: jax.Array
    counts: jax.Array

    @classmethod
    def init(cls, capacity: int, c: float = 0.):
        """Initialize a partition table.

        Args:
            capacity: Maximum number of keys that can be inserted
            c: Fudge factor that allocates more memory to reduce hash collision
        """
        ndim = 3
        num_buckets = int(math.ceil(capacity / math.log(capacity)**3))
        num_slots = int(math.ceil(math.log(capacity)**3 + c*math.log(capacity)**2))
        key_dtype = onp.int32
        return cls(
            keys=jnp.full((num_buckets, num_slots, ndim), _invalid_key(key_dtype), dtype=key_dtype),
            counts=jnp.zeros(num_buckets, dtype=jnp.uint32))

    @property
    def capacity(self) -> int:
        """Maximum number of keys that can be inserted."""
        return self.num_buckets * self.num_slots

    @property
    def num_buckets(self) -> int:
        """Number of buckets."""
        return self.keys.shape[0]

    @property
    def num_slots(self) -> int:
        """Number of slots per bucket."""
        return self.keys.shape[1]


@functools.partial(jax.jit, donate_argnums=(0,))
def insert(table: PartitionTable, keys: jax.Array) -> PartitionTable:
    """Add new keys to a partition table."""
    unique_keys = jnp.unique(keys, axis=0, size=keys.shape[0], fill_value=_invalid_key(table.keys.dtype))
    return jax.vmap(_insert, in_axes=(None, None, 0))(table, unique_keys, jnp.arange(table.num_buckets))


@jax.jit
def query(table: PartitionTable, key: jax.Array) -> PartitionTable:
    """Return data storage index or -1 if the key doesn't exist."""
    key_exists, bucket_index, slot_index = _query(table, key)
    return jnp.where(key_exists, bucket_index * table.num_slots + slot_index, -1)


def _hash_function(key: jax.Array, mod: int) -> jax.Array:
    primes = onp.array([1, 2654435761, 805459861], dtype=onp.uint32)
    result = jnp.bitwise_xor(key[0] * primes[0], key[1] * primes[1])
    result = jnp.bitwise_xor(result, key[2] * primes[2])
    return jnp.mod(result, mod)


def _insert(table: PartitionTable, keys: jax.Array, bucket_index: jax.Array) -> PartitionTable:
    keys_exist, bucket_indices, slot_indices = jax.vmap(_query, in_axes=(None, 0))(table, keys)
    insert_ok = jnp.where(bucket_indices == bucket_index, ~keys_exist, False)
    insert_indices = jnp.where(insert_ok, table.counts[bucket_index] + jnp.cumsum(insert_ok) - 1, -1)
    insert_keys = jnp.where(insert_ok[:, jnp.newaxis], keys, table.keys[bucket_index, insert_indices])
    return PartitionTable(
        counts=table.counts[bucket_index] + jnp.sum(insert_ok),
        keys=table.keys[bucket_index].at[insert_indices].set(insert_keys))


def _invalid_key(key_dtype: jnp.dtype) -> jax.Array:
    return jnp.iinfo(key_dtype).min


def _query(table: PartitionTable, key: jax.Array) -> jax.Array:
    """Return if the key is in the table, also where it is stored."""
    bucket_index = _hash_function(key, table.num_buckets)
    bucket_query = jnp.all(table.keys[bucket_index] == key, axis=1)
    slot_index = jnp.argmax(bucket_query)
    key_exists = jnp.logical_and(bucket_query[slot_index], ~jnp.array_equal(key, _invalid_key(table.keys.dtype)))
    return key_exists, bucket_index, slot_index


if __name__ == '__main__':
    table = PartitionTable.init(capacity=10000)
    table = insert(table, keys=jnp.array([[0, 0, 0], [0, 0, 1], [0, 0, 0]], jnp.int32))
    print(table.keys[:, :2])

[adam-hartshorne]

You could call PyTorch code for this sort of operating using torch2jax.

https://github.com/rdyro/torch2jax