Bit magic to find multiplicative inverse of integers

[paul-tqh-nguyen]

This line has a TODO comment requesting to conjure up some bit magic to calculate the multiplicative inverse of integers.

I might've found an approach with good results. It might be worth implementing.

It's demonstrated in this script (click to expand).

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import time
import os

from mlir_graphblas import MlirJitEngine
from mlir_graphblas.tests.jit_engine_test_utils import *

from contextlib import contextmanager


@contextmanager
def timer():
    elapsed_time_container = []
    start_time = time.time()
    yield elapsed_time_container
    end_time = time.time()
    elapsed_time = end_time - start_time
    elapsed_time_container.append(elapsed_time)
    return


@contextmanager
def temp_plt_figure(*args, **kwargs):
    figure = plt.figure(*args, **kwargs)
    yield figure
    plt.close(figure)
    return


engine = MlirJitEngine()

mlir_text = """

#trait_add = {
 indexing_maps = [
   affine_map<(i, j) -> (i, j)>,
   affine_map<(i, j) -> (i, j)>
 ],
 iterator_types = ["parallel", "parallel"]
}

func @fast_func(%matrix: tensor<?x?xi32>) -> tensor<?x?xi32> {
  %c0_i32 = constant 0 : i32
  %c1_i32 = constant 1 : i32
  %c31_i32 = constant 31 : i32
  %answer = linalg.generic #trait_add
    ins(%matrix: tensor<?x?xi32>)
    outs(%matrix: tensor<?x?xi32>) {
      ^bb(%a: i32, %s: i32):
        %mask = shift_right_signed %a, %c31_i32 : i32
        %mask_plus_a = addi %mask, %a : i32
        %abs_a = xor %mask_plus_a, %mask : i32
        %abs_a_is_one_i1 = cmpi "eq", %abs_a, %c1_i32 : i32
        %abs_a_is_one_i32 = sexti %abs_a_is_one_i1 : i1 to i32
        %result = and %a, %abs_a_is_one_i32 : i32
        linalg.yield %result : i32
  } -> tensor<?x?xi32>
  return %answer : tensor<?x?xi32>
}

func @slow_func(%matrix: tensor<?x?xi32>) -> tensor<?x?xi32> {
  %c1_i32 = constant 1 : i32
  %answer = linalg.generic #trait_add
    ins(%matrix: tensor<?x?xi32>)
    outs(%matrix: tensor<?x?xi32>) {
      ^bb(%a: i32, %s: i32):
        %result = divi_signed %c1_i32, %a : i32
        linalg.yield %result : i32
  } -> tensor<?x?xi32>
  return %answer : tensor<?x?xi32>
}

"""

engine.add(mlir_text, GRAPHBLAS_PASSES)

CSV_LOCATION = "./chart_data.csv"

df = (
    pd.read_csv(CSV_LOCATION, index_col="element_counts")
    if os.path.isfile(CSV_LOCATION)
    else pd.DataFrame(
        {
            "element_counts": [],
            "bit_magic": [],
            "normal": [],
        }
    ).set_index("element_counts")
)

for i in range(1, 20):
    dim_size = i * 1_000
    num_elements = dim_size ** 2

    if num_elements not in df.index:

        dense_input_tensor = np.arange(num_elements).reshape([dim_size, dim_size])
        dense_input_tensor -= int(dim_size ** 2 / 2)
        dense_input_tensor = dense_input_tensor.astype(np.int32)
        dense_input_tensor[dense_input_tensor == 0] = 1

        expected_ans = (1 / dense_input_tensor).astype(np.int32)

        num_runs = 10
        
        with timer() as fast_ans_time_container:
            for _ in range(num_runs):
                fast_ans = engine.fast_func(dense_input_tensor)

        with timer() as slow_ans_time_container:
            for _ in range(num_runs):
                slow_ans = engine.slow_func(dense_input_tensor)

        assert np.all(expected_ans == fast_ans)
        assert np.all(expected_ans == slow_ans)

        (fast_time,) = fast_ans_time_container
        (slow_time,) = slow_ans_time_container

        fast_time /= num_runs
        slow_time /= num_runs

        row = pd.Series(
            {"bit_magic": fast_time, "normal": slow_time}, name=num_elements
        )
        df = df.append(row)
        
        df.reset_index().to_csv(CSV_LOCATION, index=False)

with temp_plt_figure(figsize=(20.0, 10.0)) as figure:
    plot = figure.add_subplot(111)
    element_counts = df.index
    fast_times = df.bit_magic
    slow_times = df.normal
    fast_to_slow_ratios = fast_times / slow_times
    plot.plot(
        element_counts,
        fast_times,
        label="Bit Magic Optimization",
        color="green",
        marker="o",
    )
    plot.plot(
        element_counts,
        slow_times,
        label="Normal Integer Division",
        color="red",
        marker="o",
    )
    plot.plot(
        element_counts,
        fast_to_slow_ratios,
        label="Ratio of Bit Magic Running Time to Normal Division Running Time",
        color="black",
        marker="o",
    )
    plot.set_xlim(0, element_counts[-1])
    plot.set_ylim(
        0, max(1.0, fast_times.max(), slow_times.max(), fast_to_slow_ratios.max())
    )
    plot.set_xlabel("Number of Dense Matrix Elements")
    plt.legend(loc="best")
    figure.savefig("./chart.png")

The contents of the generated CSV can be seen here (click to expand).

(mlirgraphblas) pnguyen@demouser-DGX-Station:/home/pnguyen/code/mlir-graphblas$ cat chart_data.csv 
element_counts,bit_magic,normal
1000000,0.0031572580337524414,0.004949522018432617
4000000,0.012293648719787598,0.01949498653411865
9000000,0.028962373733520508,0.04381330013275146
16000000,0.05139765739440918,0.07783379554748535
25000000,0.08090529441833497,0.1214491605758667
36000000,0.11553986072540283,0.1751626253128052
49000000,0.15684835910797118,0.23777360916137696
64000000,0.20537080764770507,0.31049931049346924
81000000,0.2590143918991089,0.39394965171813967
100000000,0.3208931922912598,0.48620307445526123
121000000,0.38777225017547606,0.58694486618042
144000000,0.46244988441467283,0.7004130363464356
169000000,0.5414215087890625,0.8211428880691528
196000000,0.6268703460693359,0.9502758502960205
225000000,0.7183868169784546,1.089606475830078
256000000,0.8166777849197387,1.2439750909805298
289000000,0.9624243259429932,1.4112451314926147
324000000,1.0389467239379884,1.5837181329727172
361000000,1.1860792875289916,1.7596060037612915
400000000,1.2864532947540284,1.9521419525146484
(mlirgraphblas) pnguyen@demouser-DGX-Station:/home/pnguyen/code/mlir-graphblas$ 

Here's a chart showing the performance of the bit magic vs using signed integer division:

It looks like we can save 1/3 of the time.