notebooks/integer_hash_timing.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "integer_hash_timing.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPHAxMVPPUxZlwgWQbZyvFt"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "3Ba3IMGz32R0"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import time\n",
        "import sys"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class HashModP64(tf.Module):\n",
        "    \"\"\"Universal hash for integers via affine-mod-prime transformation.\"\"\"\n",
        "    def __init__(\n",
        "            self,\n",
        "            num_hashes,\n",
        "            prime = 2**31 - 1,\n",
        "            seed = None\n",
        "            ):\n",
        "        \"\"\"Hash function constructor.\n",
        "\n",
        "        Arguments:\n",
        "            num_hashes: Integer. Number of hash values to compute for each\n",
        "                input.\n",
        "            prime: Integer. Large prime number used to compute the modulus.\n",
        "            seed: Integer. Seed for repeatable hashing across class instances.\n",
        "        \"\"\"\n",
        "        self._num_hashes = num_hashes\n",
        "        self._prime = prime\n",
        "        self._seed = seed\n",
        "        initializer_a = tf.random_uniform_initializer(\n",
        "            minval=1, maxval=self._prime - 1, seed=self._seed)\n",
        "        initializer_b = tf.random_uniform_initializer(\n",
        "            minval=0, maxval=self._prime - 1, seed=self._seed)\n",
        "        self._a = tf.Variable(\n",
        "            initializer_a(shape=[self._num_hashes], dtype=tf.int64))\n",
        "        self._b = tf.Variable(\n",
        "            initializer_b(shape=[self._num_hashes], dtype=tf.int64))\n",
        "\n",
        "    @tf.function(\n",
        "        input_signature=[tf.TensorSpec(shape=None, dtype=tf.int64)])\n",
        "    def hash(self, x):\n",
        "        \"\"\"Calculates hashes for input tensor x.\n",
        "\n",
        "        Arguments:\n",
        "            x: tf.Tensor of rank R, type tf.int64 and arbitrary shape\n",
        "            [N0, N1, ... NR].\n",
        "\n",
        "        Returns:\n",
        "            A rank R+1 tf.Tensor of shape [N0, N1, ... NR, num_hashes] that\n",
        "            contains tf.int64 hash codes.\n",
        "        \"\"\"\n",
        "        affine = tf.tensordot(x, self._a, axes=0) + self._b\n",
        "        return tf.math.floormod(affine, self._prime)\n",
        "\n",
        "class HashBitMix64(tf.Module):\n",
        "    \"\"\"Universal hash for integers via mixing with bitwise operations.\"\"\"\n",
        "    def __init__(\n",
        "            self, \n",
        "            num_hashes,\n",
        "            seed = None\n",
        "            ):\n",
        "        \"\"\"Hash function constructor.\n",
        "\n",
        "        Arguments:\n",
        "            num_hashes: Integer. Number of hash values to compute for each\n",
        "                input.\n",
        "            seed: Integer. Seed for repeatable hashing across class instances.\n",
        "        \"\"\"\n",
        "        if num_hashes < 1:\n",
        "            raise ValueError(\"num_hashes must be >= 1 but is %s\" % num_hashes)\n",
        "\n",
        "        self._num_hashes = num_hashes\n",
        "        self._seed = seed\n",
        "        self._c1 = tf.convert_to_tensor(0xbf58476d1ce4e5b9, dtype=tf.uint64)\n",
        "        self._c2 = tf.convert_to_tensor(0x94d049bb133111eb, dtype=tf.uint64)\n",
        "        initializer_xor = tf.random_uniform_initializer(\n",
        "            minval=1, maxval=2**31 - 1, seed=self._seed)\n",
        "        self._xor_constants = tf.Variable(\n",
        "            initializer_xor(shape=[self._num_hashes], dtype=tf.int64))\n",
        "\n",
        "    @tf.function(\n",
        "        input_signature=[tf.TensorSpec(shape=None, dtype=tf.int64)])\n",
        "    def hash(self, x):\n",
        "        \"\"\"Calculates hashes for input tensor x.\n",
        "\n",
        "        Arguments:\n",
        "            x: tf.Tensor of rank R, type tf.int64 and arbitrary shape\n",
        "            [N0, N1, ... NR].\n",
        "\n",
        "        Returns:\n",
        "            A rank R+1 tf.Tensor of shape [N0, N1, ... NR, num_hashes] that\n",
        "            contains tf.int64 hash codes.\n",
        "        \"\"\"\n",
        "        x = tf.stack(\n",
        "            [tf.bitwise.bitwise_xor(x, self._xor_constants[i])\n",
        "            for i in range(self._num_hashes)], axis=-1)\n",
        "        # Use uint64 mixer (constants found via annealing by David Stafford).\n",
        "        x = tf.bitcast(x, tf.uint64)  # Does not copy data.\n",
        "        x = tf.bitwise.bitwise_xor(x, tf.bitwise.right_shift(x, 30))\n",
        "        x = x * self._c1\n",
        "        x = tf.bitwise.bitwise_xor(x, tf.bitwise.right_shift(x, 27))\n",
        "        x = x * self._c2\n",
        "        x = tf.bitwise.bitwise_xor(x, tf.bitwise.right_shift(x, 31))\n",
        "        return tf.bitcast(x, tf.int64)\n"
      ],
      "metadata": {
        "id": "GXWMEwW436sw"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class HashModP32(tf.Module):\n",
        "    \"\"\"Universal hash for integers via affine-mod-prime transformation.\"\"\"\n",
        "    def __init__(\n",
        "            self,\n",
        "            num_hashes,\n",
        "            prime = 2**31 - 1,\n",
        "            seed = None\n",
        "            ):\n",
        "        \"\"\"Hash function constructor.\n",
        "\n",
        "        Arguments:\n",
        "            num_hashes: Integer. Number of hash values to compute for each\n",
        "                input.\n",
        "            prime: Integer. Large prime number used to compute the modulus.\n",
        "            seed: Integer. Seed for repeatable hashing across class instances.\n",
        "        \"\"\"\n",
        "        self._num_hashes = num_hashes\n",
        "        self._prime = prime\n",
        "        self._seed = seed\n",
        "        initializer_a = tf.random_uniform_initializer(\n",
        "            minval=1, maxval=self._prime - 1, seed=self._seed)\n",
        "        initializer_b = tf.random_uniform_initializer(\n",
        "            minval=0, maxval=self._prime - 1, seed=self._seed)\n",
        "        self._a = tf.Variable(\n",
        "            initializer_a(shape=[self._num_hashes], dtype=tf.int32))\n",
        "        self._b = tf.Variable(\n",
        "            initializer_b(shape=[self._num_hashes], dtype=tf.int32))\n",
        "\n",
        "    @tf.function(\n",
        "        input_signature=[tf.TensorSpec(shape=None, dtype=tf.int32)])\n",
        "    def hash(self, x):\n",
        "        \"\"\"Calculates hashes for input tensor x.\n",
        "\n",
        "        Arguments:\n",
        "            x: tf.Tensor of rank R, type tf.int32 and arbitrary shape\n",
        "            [N0, N1, ... NR].\n",
        "\n",
        "        Returns:\n",
        "            A rank R+1 tf.Tensor of shape [N0, N1, ... NR, num_hashes] that\n",
        "            contains tf.int64 hash codes.\n",
        "        \"\"\"\n",
        "        affine = tf.tensordot(x, self._a, axes=0) + self._b\n",
        "        return tf.math.floormod(affine, self._prime)\n",
        "\n",
        "class HashBitMix32(tf.Module):\n",
        "    \"\"\"Universal hash for integers via mixing with bitwise operations.\"\"\"\n",
        "    def __init__(\n",
        "            self, \n",
        "            num_hashes,\n",
        "            seed = None\n",
        "            ):\n",
        "        \"\"\"Hash function constructor.\n",
        "\n",
        "        Arguments:\n",
        "            num_hashes: Integer. Number of hash values to compute for each\n",
        "                input.\n",
        "            seed: Integer. Seed for repeatable hashing across class instances.\n",
        "        \"\"\"\n",
        "        if num_hashes < 1:\n",
        "            raise ValueError(\"num_hashes must be >= 1 but is %s\" % num_hashes)\n",
        "\n",
        "        self._num_hashes = num_hashes\n",
        "        self._seed = seed\n",
        "        self._c1 = tf.convert_to_tensor(0x85ebca6b, dtype=tf.uint32)\n",
        "        self._c2 = tf.convert_to_tensor(0xc2b2ae35, dtype=tf.uint32)\n",
        "        initializer_xor = tf.random_uniform_initializer(\n",
        "            minval=1, maxval=2**31 - 1, seed=self._seed)\n",
        "        self._xor_constants = tf.Variable(\n",
        "            initializer_xor(shape=[self._num_hashes], dtype=tf.int32))\n",
        "\n",
        "    @tf.function(\n",
        "        input_signature=[tf.TensorSpec(shape=None, dtype=tf.int32)])\n",
        "    def hash(self, x):\n",
        "        \"\"\"Calculates hashes for input tensor x.\n",
        "\n",
        "        Arguments:\n",
        "            x: tf.Tensor of rank R, type tf.int32 and arbitrary shape\n",
        "            [N0, N1, ... NR].\n",
        "\n",
        "        Returns:\n",
        "            A rank R+1 tf.Tensor of shape [N0, N1, ... NR, num_hashes] that\n",
        "            contains tf.int64 hash codes.\n",
        "        \"\"\"\n",
        "        x = tf.stack(\n",
        "            [tf.bitwise.bitwise_xor(x, self._xor_constants[i])\n",
        "            for i in range(self._num_hashes)], axis=-1)\n",
        "        # Use uint64 mixer (constants found via annealing by David Stafford).\n",
        "        x = tf.bitcast(x, tf.uint32)  # Does not copy data.\n",
        "        x = tf.bitwise.bitwise_xor(x, tf.bitwise.right_shift(x, 16))\n",
        "        x = x * self._c1\n",
        "        x = tf.bitwise.bitwise_xor(x, tf.bitwise.right_shift(x, 13))\n",
        "        x = x * self._c2\n",
        "        x = tf.bitwise.bitwise_xor(x, tf.bitwise.right_shift(x, 16))\n",
        "        return tf.bitcast(x, tf.int32)"
      ],
      "metadata": {
        "id": "dZN0f2O5DoWz"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def run_benchmark64(device,\n",
        "                  num_hashes=100,\n",
        "                  batch_size=512,\n",
        "                  num_tokens=1000,\n",
        "                  experiment_reps=500):\n",
        "    with tf.device(device):\n",
        "        input_shape = [batch_size, num_tokens]\n",
        "        initializer = tf.random_uniform_initializer(\n",
        "            minval=int(-1e6), maxval=int(1e6))\n",
        "        x = [tf.Variable(initializer(shape=input_shape, dtype=tf.int64))\n",
        "             for _ in range(experiment_reps)]\n",
        "        h1 = HashBitMix64(num_hashes)\n",
        "        h2 = HashModP64(num_hashes)\n",
        "        h1_times = []\n",
        "        for xi in x:\n",
        "            t0 = time.time()\n",
        "            y = h1.hash(xi)\n",
        "            t1 = time.time()\n",
        "            h1_times.append(t1-t0)\n",
        "        h2_times = []\n",
        "        for xi in x:\n",
        "            t0 = time.time()\n",
        "            y = h2.hash(xi)\n",
        "            t1 = time.time()\n",
        "            h2_times.append(t1-t0)\n",
        "    return (h1_times, h2_times)\n",
        "\n",
        "def run_benchmark32(device,\n",
        "                  num_hashes=100,\n",
        "                  batch_size=512,\n",
        "                  num_tokens=1000,\n",
        "                  experiment_reps=500):\n",
        "    with tf.device(device):\n",
        "        input_shape = [batch_size, num_tokens]\n",
        "        initializer = tf.random_uniform_initializer(\n",
        "            minval=int(-1e6), maxval=int(1e6))\n",
        "        x = [tf.Variable(initializer(shape=input_shape, dtype=tf.int32))\n",
        "             for _ in range(experiment_reps)]\n",
        "        h1 = HashBitMix32(num_hashes)\n",
        "        h2 = HashModP32(num_hashes)\n",
        "        h1_times = []\n",
        "        for xi in x:\n",
        "            t0 = time.time()\n",
        "            y = h1.hash(xi)\n",
        "            t1 = time.time()\n",
        "            h1_times.append(t1-t0)\n",
        "        h2_times = []\n",
        "        for xi in x:\n",
        "            t0 = time.time()\n",
        "            y = h2.hash(xi)\n",
        "            t1 = time.time()\n",
        "            h2_times.append(t1-t0)\n",
        "    return (h1_times, h2_times)"
      ],
      "metadata": {
        "id": "cCk8xFyI3-EA"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "h1_times_cpu_64, h2_times_cpu_64 = run_benchmark64('/cpu:0')\n",
        "print(\"On CPU (64-bit)\")\n",
        "print(\"HashBitMix: \",np.mean(h1_times_cpu_64), np.std(h1_times_cpu_64))\n",
        "print(\"HashModP: \",np.mean(h2_times_cpu_64), np.std(h2_times_cpu_64))\n",
        "sys.stdout.flush()\n",
        "\n",
        "h1_times_gpu_64, h2_times_gpu_64 = run_benchmark64('/device:GPU:0')\n",
        "print(\"On GPU (64-bit)\")\n",
        "print(\"HashBitMix: \",np.mean(h1_times_gpu_64), np.std(h1_times_gpu_64))\n",
        "print(\"HashModP: \",np.mean(h2_times_gpu_64), np.std(h2_times_gpu_64))\n",
        "sys.stdout.flush()\n",
        "\n",
        "h1_times_cpu_32, h2_times_cpu_32 = run_benchmark32('/cpu:0')\n",
        "print(\"On CPU (32-bit)\")\n",
        "print(\"HashBitMix: \",np.mean(h1_times_cpu_32), np.std(h1_times_cpu_32))\n",
        "print(\"HashModP: \",np.mean(h2_times_cpu_32), np.std(h2_times_cpu_32))\n",
        "sys.stdout.flush()\n",
        "\n",
        "h1_times_gpu_32, h2_times_gpu_32 = run_benchmark32('/device:GPU:0')\n",
        "print(\"On GPU (32-bit)\")\n",
        "print(\"HashBitMix: \",np.mean(h1_times_gpu_32), np.std(h1_times_gpu_32))\n",
        "print(\"HashModP: \",np.mean(h2_times_gpu_32), np.std(h2_times_gpu_32))\n",
        "sys.stdout.flush()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UwSGA77N4Byx",
        "outputId": "e51a3971-a212-49d5-e4a2-e9b86e674252"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "On CPU (64-bit)\n",
            "HashBitMix:  1.2728994584083557 0.0606345520431441\n",
            "HashModP:  0.822510507106781 0.019799892577717637\n",
            "On GPU (64-bit)\n",
            "HashBitMix:  0.09299386167526245 0.023977337576366914\n",
            "HashModP:  0.8131655917167664 0.026527320797214995\n",
            "On CPU (32-bit)\n",
            "HashBitMix:  0.8894901666641235 0.03794307909790895\n",
            "HashModP:  0.3682930154800415 0.010422571491485655\n",
            "On GPU (32-bit)\n",
            "HashBitMix:  0.567615369796753 0.04233825338727441\n",
            "HashModP:  0.35730077362060547 0.00836473357131244\n"
          ]
        }
      ]
    }
  ]
}