#!/bin/bash

dir=$(cd `dirname $0`;pwd)

is_nv_gpu_host()
{
    gpu_info=$(lspci | grep NVIDIA)
    if [ ${#gpu_info} -gt 0 ];then
        echo "nvidia gpu host"
        return 1
    fi  

    return 0
}

is_gpu_driver_installed()
{
    driver_info=$(lsmod | grep nvidia)
    if [ ${#driver_info} -gt 0 ];then
        return 1
    fi  
    
    return 0
}

nvlink_h800_check()
{
    reset=0

    for cc in {0..2}; do

    echo "check $cc; `date`"
    reset=0

    for i in {0..7}; do
        linksta=$(nvidia-smi nvlink -i $i -s)
        if echo "$linksta" | grep -i "inactive" ; then
            echo "nvlink inactive, reset gpu$i"
            fuser -k /dev/nvidia$i; nvidia-smi -r -i $i
            reset=1
        fi
    done

    output=$(nvidia-smi topo -p2p r)
    echo $output

    for i in {0..7}; do
        for j in {0..7}; do
        status=$(echo "$output" | awk -v i=$((i+2)) -v j=$((j+2)) 'NR==i{print $j}')
        #echo "GPU$i and GPU$j have $status"
        if [ "$status" = "NS" ]; then
            echo "nvlink p2p status NS, reset gpu$i"
            fuser -k /dev/nvidia$i; nvidia-smi -r -i $i
            reset=1
            output=$(nvidia-smi topo -p2p r)
            echo $output
            break
        fi
        done
    done

    if [ "$reset" = "0" ]; then
        echo "all nvlink ok"
        return 0
    fi

    done
}

# consumer gpu: nvidia-smi -pm 1
# server   gpu: nvidia-persistenced --persistence-mode
nvidia_smi_pm_on()
{
    consumer_gpu=$(lspci -d 10de::000300 | grep NVIDIA)
    if [ -n "$consumer_gpu" ];then
        nvidia-smi -pm 1
        echo "set Persistence-M for nvidia gpu driver success"
        return
    fi 

    nvidia-persistenced --persistence-mode
    echo "nvidia-persistenced starts with persistence mode enabled for all devices"
}

is_nv_gpu_host
if [ $? -eq 0 ];then
    echo "not nvidia gpu host. exit"
    exit 0
fi

# h800/h20-96G/h20-141G
gpu_info=$(lspci -d 10de: -Dnn | grep -iE "2324|2329|232c" | awk '{print $1}')
if [ -n "$gpu_info" ];then
    echo  "hopper gpu: config"

    for bdf in $gpu_info; do
        echo "hopper mask gpu ce and set ct value"
        setpci -s $bdf ecap_aer+0x14.w=ffff >/dev/null 2>&1
        setpci -s $bdf cap_exp+0x28.w=6:f >/dev/null 2>&1
    done
fi

is_gpu_driver_installed
if [ $? -eq 0 ];then
    echo "gpu driver not installed. exit"
    exit 0
fi

command -v nvidia-smi >/dev/null 2>&1 && nvidia_smi_pm_on || { echo "no nvidia-smi command"; } 

gpu_info=$(lspci -d 10de:2324)
if [ -n "$gpu_info" ];then
    echo  "h800 gpu: nvlink check"
    # check fabric-manamger service is OK
    if [ "`which systemctl &> /dev/null;echo $?`" == 0 ] && [ "`systemctl is-enabled nvidia-fabricmanager &> /dev/null;echo $?`" == 0 ]; then
        if [ "`systemctl status nvidia-fabricmanager &> /dev/null;echo $?`" == 0 ]; then
            command -v nvidia-smi >/dev/null 2>&1 && nvlink_h800_check
        fi
    fi
fi

gpu_info=$(lspci -d 10de: -Dnn | grep -iE "26ba" | awk '{print $1}')
if [ -n "$gpu_info" ];then
    echo  "L20 gpu: lock clock"
    nvidia-smi -lgc 2520,2520 >/dev/null 2>&1
fi

gpu_info=$(lspci -d 10de: -Dnn | grep -iE "26b5" | awk '{print $1}')
if [ -n "$gpu_info" ];then
    echo  "L40 gpu: lock clock"
    nvidia-smi -lgc 2490,2490 >/dev/null 2>&1
fi

echo "gpu config finish"