#!/bin/bash
set -eu

# Hat tip to https://gist.github.com/koallen/32709a244d77a2c0f8e17ed79a4092ed

MLNX_SRIOV_DEVICES=""
[ -e /etc/default/mlnx_sriov ] && . /etc/default/mlnx_sriov

# params
#   - device sys path (e.g. /sys/class/infiniband/mlx5_0)
#   - device id (eg, 0 from mlx5_0)
#   - vf number (eg, 1)
configure_vf () {
  local -r device="$1"
  local -r devid="$2"
  local -r vf="$3"
  printf "   %s\n"  "Configuring virtual function $vf"
  # enable the virtual function
  local -r policy_file=${device}/device/sriov/$vf/policy
  if [[ -e ${policy_file} ]]; then
    echo Follow > ${device}/device/sriov/$vf/policy

    # assign GUID to virtual card and port
    (( first_part=vf/100 )) || :
    (( second_part=vf - first_part * 100 ))  || :
    local -r ip_last_seg=$(hostname -i | cut -d. -f4)
    (( ip_last_seg_first=ip_last_seg / 100 )) || :
    (( ip_last_seg_second=ip_last_seg - ip_last_seg_first*100 )) || :
    local -r guid_prefix="ca:$(printf "%02d" $devid):fe:$(printf "%02d" $first_part):$(printf "%02d" $second_part):$(printf "%02d" $ip_last_seg_first):$(printf "%02d" $ip_last_seg_second)"
    echo "$guid_prefix:90" > ${device}/device/sriov/$vf/node
    echo "$guid_prefix:91" > ${device}/device/sriov/$vf/port
  fi
  # reload driver to make the change effective
  pcie_addr="$(readlink -f ${device}/device/virtfn${vf} | awk -F/ '{print $NF}')"
  echo $pcie_addr > /sys/bus/pci/drivers/mlx5_core/unbind
  echo $pcie_addr > /sys/bus/pci/drivers/mlx5_core/bind
  return 0;
}

# params
#   - device sys path (e.g. /sys/class/infiniband/mlx5_0)
#   - number of virtual functions (e.g. 10)
configure_dev () {
  local -r device="$1"
  local -r num_of_vfs="$2"
  local -r devid=$(echo ${device} | cut -d_ -f2)
  if [[ ! -e ${device}/device/sriov_vf_device ]]; then
    printf "   Warning: mlx5_${devid} is no PF; skipping"
    return
  fi
  local -r dev_maxvfs="$(< ${device}/device/sriov_totalvfs)"
  if (( num_of_vfs > dev_maxvfs )); then
    printf "   %s\n" "Warning: more VFs (${num_of_vfs}) requested than available (${dev_maxvfs}) on mlx5_${devid}; skipping"
    return
  fi
  echo 0 > ${device}/device/sriov_drivers_autoprobe
  local -r dev_numvfs="$(< ${device}/device/sriov_numvfs)"
  if  (( dev_numvfs < num_of_vfs )); then
    echo $num_of_vfs > ${device}/device/sriov_numvfs
  fi
  for (( vf = 0; vf < num_of_vfs; ++vf)); do
    configure_vf "${device}" "${devid}" "${vf}"
  done
}

usage () {
  cat <<-USAGE
		Usage:  $0 [device num_vfs]+
		
		Either specify pairs of device and virtual functions or nothing for all devices
		USAGE
  exit 1;
}

process ()  {
  # if specific devices are provided, only those will be configured
  # otherwise, all devices supporting SR-IOV will be configured
  if (( $# == 0 )); then
    printf "%s\n" "Configuring SR-IOV for all supported devices"
    for dev in /sys/class/infiniband/*; do
      local -r totalvfs_path="/sys/class/infiniband/$dev/device/sriov_totalvfs"
      [[ -e "$totalvfs_path" ]] || continue
      local -r totalvfs="$(< ${totalvfs_path})"
      if (( $totalvfs > 0 )); then
        printf " %s\n" "Configuring for $dev ($totalvfs)"
        configure_dev $dev $totalvfs
      fi
    done
  elif ! (( $# % 2 )); then
    printf "%s\n" "Configuring SR-IOV for specified devices"
    while (( $# )); do
      printf " %s" "Configuring for $1"
      dev=/sys/class/infiniband/$1
      num_of_vfs=$2
      if [[ ! -e $dev/device/sriov_vf_device ]]; then
        printf ": %s\n" "(no SR-IOV, skipping)"
      else
        printf "%s\n" "."
        configure_dev $dev $num_of_vfs
      fi
      shift 2
    done
  else
    usage
  fi
}


if (( $# > 0 )); then
  case "$1" in
  -h|--help) usage;;
  esac
fi

(( $# == 0 )) && set -- ${MLNX_SRIOV_DEVICES} || :
process "$@"