#!/bin/bash

## Usage #############################
# source ipex-llm-init
# Example:
# source ipex-llm-init
######################################

function enable_iomp {
    ENABLE_IOMP=1
}

function disable_iomp {
    ENABLE_IOMP=0
}

function enable_jemalloc {
    ENABLE_JEMALLOC=1
    ENABLE_TCMALLOC=0
}

function disable_jemalloc {
    ENABLE_JEMALLOC=0
}

function enable_tcmalloc {
    ENABLE_TCMALLOC=1
    ENABLE_JEMALLOC=0
}

function disable_tcmalloc {
    ENABLE_TCMALLOC=0
}

function enable_gpu {
    ENABLE_GPU=1
}

function disable_gpu {
    ENABLE_GPU=0
    unset_gpu_envs
}

function unset_gpu_envs {
    unset USE_XETLA
    unset ENABLE_SDP_FUSION
    unset SYCL_CACHE_PERSISTENT
    unset BIGDL_LLM_XMX_DISABLED
    unset SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS
}

function display-var {
    echo "+++++ Env Variables +++++"
    echo "Internal:"
    echo "    ENABLE_IOMP     = ${ENABLE_IOMP}"
    echo "    ENABLE_GPU      = ${ENABLE_GPU}"
    echo "    ENABLE_JEMALLOC = ${ENABLE_JEMALLOC}"
    echo "    ENABLE_TCMALLOC = ${ENABLE_TCMALLOC}"
    echo "    LIB_DIR    = ${LIB_DIR}"
    echo "    BIN_DIR    = ${BIN_DIR}"
    echo "    LLM_DIR    = ${LLM_DIR}"
    echo ""
    echo "Exported:"
    echo "    LD_PRELOAD             = ${LD_PRELOAD}"
    echo "    OMP_NUM_THREADS        = ${OMP_NUM_THREADS}"
    echo "    MALLOC_CONF            = ${MALLOC_CONF}"
    echo "    USE_XETLA              = ${USE_XETLA}"
    echo "    ENABLE_SDP_FUSION      = ${ENABLE_SDP_FUSION}"
    echo "    SYCL_CACHE_PERSISTENT  = ${SYCL_CACHE_PERSISTENT}"
    echo "    BIGDL_LLM_XMX_DISABLED = ${BIGDL_LLM_XMX_DISABLED}"
    echo "    SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS = ${SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS}"
    echo "+++++++++++++++++++++++++"
}

function display-help {
    echo "Usage: source ipex-llm-init [--option]"
    echo ""
    echo "ipex-llm-init is a tool to automatically configure and run the subcommand under"
    echo "environment variables for accelerating IPEX-LLM."
    echo ""
    echo "Optional options:"
    echo "    -h, --help                Display this help message and exit."
    echo "    -o, --gomp                Disable intel-openmp and use default openmp (i.e. gomp)"
    echo "    -j, --jemalloc            Use jemalloc as allocator"
    echo "    -t, --tcmalloc            Use tcmalloc as allocator"
    echo "    -c, --disable-allocator   Use the system default allocator"
    echo "    -g, --gpu                 Enable OneAPI and other settings for GPU support"
    echo "    -d, --debug               Print all internal and exported variables (for debug)"
    echo "    --device <device_type>    Specify the device type (Max, Flex, Arc, iGPU)"
}

function display-error {
    echo "Invalid Option: -$1" 1>&2
    echo ""
    display-help
}

enable_iomp
disable_gpu
disable_jemalloc
disable_tcmalloc
LD_PRELOAD=""

OPTIND=1
DEVICE=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help)
            display-help
            return 0
            ;;
        -o|--gomp)
            disable_iomp
            shift
            ;;
        -j|--jemalloc)
            enable_jemalloc
            shift
            ;;
        -t|--tcmalloc)
            enable_tcmalloc
            shift
            ;;
        -c|--disable-allocator)
            disable_jemalloc
            disable_tcmalloc
            shift
            ;;
        -g|--gpu)
            enable_gpu
            shift
            ;;
        -d|--debug)
            display-var
            return 0
            ;;
        --device)
            DEVICE="$2"
            shift 2
            ;;
        *)
            display-error "$1"
            return 1
            ;;
    esac
done

# Find ipex-llm-init dir
if [ ! -z $BASH_SOURCE ]; then
    # using bash
    if [ "$BASH_SOURCE" = "$0" ]; then
        echo "Error: Incorrect usage: ipex-llm-init must be sourced."
        exit 1
    fi
    BIN_DIR="$(dirname $BASH_SOURCE)"
else
    # using zsh
    if [ "$zsh_eval_context" = "toplevel" ]; then
        echo "Error: Incorrect usage: ipex-llm-init must be sourced."
        exit 1
    fi
    BIN_DIR="$(dirname ${(%):-%N})"
fi

LIB_DIR=$(dirname ${BIN_DIR})/lib

if [ "${ENABLE_IOMP}" -eq 1 ]; then
    file="${LIB_DIR}/libiomp5.so"
    if [ -f ${file} ]; then
        echo "found intel-openmp in ${file}"
        LD_PRELOAD=$(echo ${LD_PRELOAD} ${file})
        export OMP_NUM_THREADS=`cat /proc/cpuinfo | grep "cpu cores" | uniq | awk '{print $4}'`
    fi
else
    unset OMP_NUM_THREADS
fi

if [ "${ENABLE_GPU}" -eq 1 ]; then
    file="/opt/intel/oneapi/setvars.sh"
    if [ -f "${file}" ]; then
        echo "found oneapi in ${file}"
        # set +xv
        source "${file}" --force
        # set -xv
        case "${DEVICE}" in
            Max|MAX)
                conda install -c conda-forge -y gperftools=2.10
                export LD_PRELOAD=$(echo ${LD_PRELOAD} ${CONDA_PREFIX}/lib/libtcmalloc.so)
                export SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1
                export SYCL_CACHE_PERSISTENT=1
                export ENABLE_SDP_FUSION=1
                ;;
            Flex|FLEX|Arc|ARC)
                export USE_XETLA=OFF
                export SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1
                export SYCL_CACHE_PERSISTENT=1
                ;;
            iGPU|IGPU|MTL)
                export SYCL_CACHE_PERSISTENT=1
                export BIGDL_LLM_XMX_DISABLED=1
                ;;
            *)
                echo "Error: Invalid device type specified for GPU."
                echo ""
                display-help
                return 1
                ;;
        esac
    else
        echo "Error: ${file} not found"
        return 1
    fi
else
    unset_gpu_envs
fi


LLM_DIR=$(dirname $(python3 -c "import ipex_llm; print(ipex_llm.__file__)"))

if [ "${ENABLE_JEMALLOC}" -eq 1 ]; then
    file="${LLM_DIR}/libs/libjemalloc.so"
    if [ -f ${file} ]; then
        echo "found jemalloc in ${file}"
        LD_PRELOAD=$(echo ${LD_PRELOAD} ${file})
        export MALLOC_CONF="oversize_threshold:1,background_thread:false,metadata_thp:always,dirty_decay_ms:-1,muzzy_decay_ms:-1"
    fi
else
    unset MALLOC_CONF
fi

if [ "${ENABLE_TCMALLOC}" -eq 1 ]; then
    if [ "${DEVICE}" = "Arc" ] || [ "${DEVICE}" = "ARC" ]; then
        echo "Warning: We do not recommend enabling tcmalloc on ARC, as this will cause segmentation fault"
    fi
    file="${LLM_DIR}/libs/libtcmalloc.so"
    if [ -f ${file} ]; then
        echo "found tcmalloc in ${file}"
        LD_PRELOAD=$(echo ${LD_PRELOAD} ${file})
    fi
fi

export LD_PRELOAD=${LD_PRELOAD}

display-var
echo "Complete."