ozIMMU - DGEMM on Int8 Tensor Core

This library intercepts function calls for cuBLAS DGEMM functions and executes ozIMMU instead

Build

git clone https://github.com/enp1s0/ozIMMU --recursive
cd ozIMMU
mkdir build
cd build
cmake ..
make -j4

Usage

1. Set an environmental variable to hijack the function calls

export LD_PRELOAD=/path/to/ozIMMU/build/libozimmu.so

2. Set an environmental variable to choose the compute mode

export OZIMMU_COMPUTE_MODE=fp64_int8_9

The supported compute modes are here.

3. Execute the application

Supported compute mode

Mode	Tensor Core type	Num splits
dgemm	--	--	Disable hijacking
sgemm	--	--	Use SGEMM internally
fp64_int8_3	Int8 TC	3
fp64_int8_4	Int8 TC	4
fp64_int8_5	Int8 TC	5
fp64_int8_6	Int8 TC	6
fp64_int8_7	Int8 TC	7
fp64_int8_8	Int8 TC	8
fp64_int8_9	Int8 TC	9
fp64_int8_10	Int8 TC	10
fp64_int8_11	Int8 TC	11
fp64_int8_12	Int8 TC	12
fp64_int8_13	Int8 TC	13
fp64_int8_14	Int8 TC	14
fp64_int8_15	Int8 TC	15
fp64_int8_16	Int8 TC	16
fp64_int8_17	Int8 TC	17
fp64_int8_18	Int8 TC	18
fp64_int8_auto	Int8 TC	AUTO	fp64_int8_3..18 / dgemm

Optional environmental variables

# Show info log
export OZIMMU_INFO=1

# Show error and warning log
export OZIMMU_ERROR=1

# Show CULiP ( https://github.com/enp1s0/CULiP ) log
export OZIMMU_ENABLE_CULIP_PROFILING=1

# Choose malloc mode
export OZIMMU_MALLOC_ASYNC=1

# Set AUTO mode mantissa loss threshold
export OZIMMU_AUTO_AVG_MANTISSA_LOSS_THRESHOLD=1.5

# Set ozIMMU intercept threshold.
export OZIMMU_INTERCEPT_THRESHOLD_M=128
export OZIMMU_INTERCEPT_THRESHOLD_N=128
export OZIMMU_INTERCEPT_THRESHOLD_K=128
# The ozIMMU gemm function is executed if `m`, `n`, and `k` are larger or equal to `OZIMMU_INTERCEPT_THRESHOLD_M`, `N`, and `K`.
# Otherwise, the original cuBLAS function is executed.

Citation

@article{ootomo2024dgemm,
    author = {Hiroyuki Ootomo and Katsuhisa Ozaki and Rio Yokota},
    title = {DGEMM on integer matrix multiplication unit},
    journal = {The International Journal of High Performance Computing Applications},
    year = {2024},
    doi = {10.1177/10943420241239588},
    URL = {https://doi.org/10.1177/10943420241239588}
}

License

MIT