#!/bin/bash
#
# Copyright (c) 2024, Barcelona Supercomputing Center (BSC)
# SPDX-License-Identifier: MIT
#
# Based on a script provided by Xavier Martorell
# Rewritten by Rodrigo Arias Mallo <rodrigo.arias@bsc.es>

# Stop on first error
set -e

function is_module_loaded() # {{{
{
  lsmod | grep -wq "$1"
} # }}}
function check_environment() # {{{
{
  if ! command -v dma-ctl &> /dev/null; then
    echo "error: dma-ctl not found in PATH" >&2
    exit 1
  fi

  if ! command -v vivado &> /dev/null; then
    echo "error: vivado not found in PATH" >&2
    exit 1
  fi
} # }}}
function create_qdma_queue() # {{{
{
  if [ ! -d "$pcidir" ]; then
    echo "missing pci directory: $pcidir" >&2
    exit 1
  fi

  if [ ! -d "$pcidir/qdma" ]; then
    echo "missing qdma directory: $pcidir/qdma" >&2
    exit 1
  fi

  if [ ! -f "$pcidir/qdma/qmax" ]; then
    echo "missing qmax file: $pcidir/qdma/qmax" >&2
    exit 1
  fi

  if [ ! -r "$pcidir/qdma/qmax" ]; then
    echo "cannot read qmax file: $pcidir/qdma/qmax" >&2
    exit 1
  fi

  # There should be two queues
  local qmax=$(cat "$pcidir/qdma/qmax")
  if [ "$qmax" != 2 ]; then
    if [ -w "$pcidir/qdma/qmax" ]; then
      echo 2 | dd of="$pcidir/qdma/qmax"
    else
      echo 2 | sudo dd of="$pcidir/qdma/qmax"
    fi
  fi

  # Create the two queues if they don't exist
  if [ ! -c "/dev/${qdmadev}-MM-1" ]; then
    dma-ctl "${qdmadev}" q add mode mm idx 1 dir bi
    dma-ctl "${qdmadev}" q start idx 1 dir bi
  fi
  if [ ! -c "/dev/${qdmadev}-MM-0" ]; then
    dma-ctl "${qdmadev}" q add mode mm idx 0 dir bi
    dma-ctl "${qdmadev}" q start idx 0 dir bi
  fi

  # Wait for udev to process the new devices
  udevadm settle

  # Ensure we have write access. On some clusters this is automatically done
  # by udev rules, on others we are expect to use sudo.
  for f in /dev/${qdmadev}-MM-{0,1} ${pcidir}/resource{0,0_wc,2,2_wc}; do
    test -w "$f" || sudo chmod go+rw "$f"
  done

  sleep 2
} # }}}
function do_cpu_reset() # {{{
{
  if [ "$model" == "hun" ]; then
    # UartBootEn (bit2) + system reset (bit0)
    dma-ctl "${qdmadev}" reg write bar 2 0x0 0x0
    sleep 0.2
    dma-ctl "${qdmadev}" reg write bar 2 0x0 0x1
  elif [ "$model" == "ox" ]; then
    dma-ctl "${qdmadev}" reg write bar 2 0x0 0x0
    sleep 0.2
  fi
} # }}}
function do_cpu_release() # {{{
{
  if [ "$model" == "hun" ]; then
    # Release Ariane's reset
    dma-ctl "${qdmadev}" reg write bar 2 0x0 0x3
  elif [ "$model" == "ox" ]; then
    dma-ctl "${qdmadev}" reg write bar 2 0x0 0x1
    sleep 1
  fi
} # }}}
function copy_by_dma() # {{{
{
  ifile="$1"
  address="$2"

  ofile="/dev/${qdmadev}-MM-1"
  #bs=$((8*1024*1024)) # 8 MiB
  bs=$((1*1024*1024)) # 1 MiB

  total_size=$(stat --format "%s" "$ifile")

  nblocks=$(( ($total_size + $bs - 1) / $bs ))
  skip=0

  # Using just a single dd command doesn't seem to work. My hypothesis is that
  # the driver doesn't allow multiple writes without setting the fseek position.
  # So we keep copying blocks until the end of the file, one at a time.

  while [ "$skip" -lt "$nblocks" ]; do
    dst=$(($address + $skip * $bs))
    dd if="$ifile" skip=$skip count=1 bs=$bs of="$ofile" seek=$dst oflag=seek_bytes status=none
    let skip=$skip+1
  done
  #dma-to-device -d "$ofile" -s "$total_size" -a "$address" -f "$ifile"

} # }}}
function load_file_in_memory() # {{{
{
  file="$1"
  address_hex="$2"
  address=$(($address_hex))

  total_size=$(stat --format "%s" "$file")
  md5sum=$(md5sum "$file" | cut -d' ' -f1)

  # Previous tests...
  #strace -f dma-to-device -d /dev/${qdmadev}-MM-1 -a "$address" -s $((8*1024*1024)) -f "$file"
  #strace -f dd if="$file" bs=16M seek="${address}" oflag=seek_bytes of=/dev/${qdmadev}-MM-1 status=progress conv=sync
  #strace -f fpgakit/fpgadd -i "$file" -a "$address" -d /dev/${qdmadev}-MM-1 -c 1024 -s 1024
  #ID=08 ./load_image.sh "$file" "$address"

  # Now dd seems to work fine, but I will leave this as fallback:
  copy_by_dma "$file" "$address"

  #dd if="$file" bs=8M seek="${address}" oflag=seek_bytes of=/dev/${qdmadev}-MM-1 status=none

  printf "loaded '%s' at 0x%x with size %d and md5 %s\n" "$file" "$address" "$total_size" "$md5sum" >&2
} # }}}
function do_boot_only() # {{{
{
  do_cpu_reset

  ./load_image.sh ${OSBI} $((0x80000000)) &&

  rm -f ${OSBI} &&

  sleep 2 &&

  # #Release Ariane's reset
  dma-ctl "${qdmadev}" reg write bar 2 0x0 0x3 &&

  sleep 10 &&

  echo mount  -o nolock -o rw -o retrans=10 10.0.2.2:/media/sda2/scratch/xavim/point /root &&

  echo mount  -o nolock -o rw -o retrans=10 192.168.0.16:/media/sda2/scratch/xavim/point /root &&

  if [ ! -c "/dev/${qdmadev}-MM-0" ] ; then
     /home/tools/drivers/`/bin/hostname`/dma_ip_drivers-onic-gamma/create-queue-qdma.sh -2
  fi &&

  sleep 2 #&&

  # uncomment to enable eth-over-pcie
  #sudo ifconfig onic${onicid}s0f0 10.0.2.2 netmask 255.255.255.0 mtu 9000 up
  sudo ifconfig onic${onicid}s0f0 10.0.2.2 netmask 255.255.255.0 up
} # }}}
function do_reload_fs() # {{{
{
  do_cpu_reset

  #~xavim/LAGARTO_LINUX-4.1/./load_image.sh \
  #  /home/xavim/ARIANE_LINUX-3.0/recovery/fedora-fs-dx-java-cucu-0.108.raw.recovered \
  #  $((0x13ff00000)) &&

  load_file_in_memory /home/xavim/ARIANE_LINUX-3.0/recovery/fedora-fs-dx-java-cucu-0.108.raw.recovered $((0x13ff00000))

  sleep 1

  # Load OpenSBI + kernel
  load_file_in_memory ${OSBI} $((0x80000000))

  rm -f ${OSBI}

  sleep 2

  do_cpu_release
  create_qdma_queue

# uncomment to enable eth-over-pcie
#     sudo ifconfig onic${onicid}s0f0 10.0.2.2 netmask 255.255.255.0 mtu 9000 up
} # }}}
function upload_bitstream_file()  # {{{
{
  if [ -z "$jtagserial" ]; then
    >&2 echo "JTAG serial required"
    usage
  fi

  if [ -z "$bitstream" ]; then
    >&2 echo "bitstream file required"
    usage
  fi

  script=$(mktemp vivado-XXXXXXXXXX.tcl)
  cat > "$script" <<EOF
open_hw_manager
connect_hw_server -url localhost:3121
current_hw_target "localhost:3121/xilinx_tcf/Xilinx/${jtagserial}A"
open_hw_target
set dev [lindex [get_hw_devices] 0]
current_hw_device \$dev
set_property PROGRAM.FILE ${bitstream} \$dev
program_hw_devices \$dev
exit
EOF

  vivado -nolog -nojournal -mode batch -source "$script"
  rm "$script"

  killall hw_server
} # }}}
function unload_modules() # {{{
{
  drvlist="$1"

  # Unload modules
  for mod in $drvlist; do
    if is_module_loaded "$mod"; then
      sudo rmmod $mod
    fi
  done

} # }}}
function remove_pci_devices() # {{{
{
  for slot in $(lspci -mm -d 10ee: | awk '{printf "0000:%s\n",$1}'); do
    devdir="/sys/bus/pci/devices/$slot"
    if [ -d $devdir ]; then
      if [ -w "$devdir/remove" ]; then
        echo 1 | dd "of=$devdir/remove"
      else
        echo 1 | sudo dd "of=$devdir/remove"
      fi
    fi
  done

} # }}}
function rescan_pci_devices() # {{{
{
  if [ -w /sys/bus/pci/rescan ]; then
    echo 1 | dd of=/sys/bus/pci/rescan
  else
    echo 1 | sudo dd of=/sys/bus/pci/rescan
  fi
} # }}}
function load_qdma_modules() # {{{
{
  drv="$DMA_IP_DRIVERS/QDMA/linux-kernel/bin/qdma-pf.ko"
  hw_buffers="$((0x4FFE0000))"

  if [ ! -r "$drv" ]; then
    echo "error: missing $"
    exit 1
  fi

  sudo insmod "$drv" "hw_buffers=$hw_buffers"
  sleep 4
} # }}}
function select_pcidev()  # {{{
{
  if [ -z "$pcidev" ]; then
    >&2 echo -e "error: missing PCI device (hint: lspci -d 10ee:902f)"
    usage
  fi

  # Ensure it is ok
  local matches=$(lspci -s "$pcidev")
  if [ -z "$matches" ]; then
    >&2 echo "no match for PCI device '$pcidev'"
    exit 1
  fi

  local n="$(echo "$matches" | wc -l)"
  if [ "$n" -gt 1 ]; then
    >&2 echo "multiple matches for PCI device '$pcidev'"
    exit 1
  fi

  # Fill the PCI device with the domain
  local fulldev=$(lspci -s "$pcidev" -D | cut -d' ' -f1)
  pcidir="/sys/bus/pci/devices/$fulldev"

  if [ ! -d "$pcidir" ]; then
    >&2 echo "cannot find PCI dir: $pcidir"
    exit 1
  fi

  # Set the PCI device to the full device
  pcidev="$fulldev"

  # Find slot
  slot=$(lspci -s "$pcidev" -vm | grep PhySlot | cut -f2)

  if [ -z "$slot" ]; then
    >&2 echo "cannot find physical slot for PCI '$pcidev'"
    exit 1
  fi

  local devid=$(echo "$pcidev" | cut -d: -f2- | tr -d ':.')
  qdmadev="qdma${devid}"

} # }}}
function preload_hook() #{{{
{
  case "$hostname" in
    cucu)
      unload_modules "xocl xclmgmt qdma_pf xdma" # qdma_vf not removable
      remove_pci_devices
      ;;
    fpgan*)
      ;;
    *)
      echo "hostname $hostname not known"
      exit 1
      ;;
  esac
} #}}}
function postload_hook() #{{{
{
  rescan_pci_devices

  case "$hostname" in
    cucu)
      unload_modules "qdma_pf xdma" # qdma_vf not removable
      remove_pci_devices
      load_qdma_modules
      rescan_pci_devices
      ;;
    fpgan*)
      ;;
    *)
      echo "hostname $hostname not known"
      exit 1
      ;;
  esac

  create_qdma_queue
} #}}}
function load_bitstream()  # {{{
{
  preload_hook
  upload_bitstream_file
  postload_hook
} # }}}

bitstream=
bootloader=
kernel=
initrd=
rootfs=
bootrom=
resetcpu=
verbose=
pcidev=
model=ox
# Internal
slot=
pcidir=
qdmadev=

bootloader_addr="${FPGACTL_BOOTLOADER_ADDR:-0x80000000}"
kernel_addr="${FPGACTL_KERNEL_ADDR:-0x84000000}"
initrd_addr="${FPGACTL_INITRD_ADDR:-0x8c300000}"
rootfs_addr="${FPGACTL_ROOTFS_ADDR:-0x140000000}"
bootrom_addr="${FPGACTL_BOOTROM_ADDR:-0x60000100}"

hostname="${hostname:-$(hostname)}"
echo "hostname=$hostname"

function usage()
{
  echo "" >&2
  echo "Usage: $0 [-p pcidev] [-v] [-w bitstream] [-j serial] [-b bootloader] [-k kernel] [-i initrd] [-R bootroom] " >&2
  echo "" >&2
  echo "First writes the bitstream if given. Then loads the rest of files" >&2
  echo "into memory and restarts the CPU." >&2
  echo "" >&2
  echo "Options" >&2
  echo "  -p pcidev      Select PCI device (same format as lspci -s)." >&2
  echo "                 Read from \$FPGACTL_PCIDEV if not given." >&2
  echo "  -w bitstream   Write the bitstream file to the FPGA" >&2
  echo "  -j serial      JTAG serial (can be found by lsusb -v)" >&2
  echo "                 Read from \$FPGACTL_SERIAL if not given." >&2
  echo "  -b bootloader  Load the bootloader file in $bootloader_addr" >&2
  echo "  -k kernel      Load the kernel file in $kernel_addr" >&2
  echo "  -i initrd      Load the initrd file in $initrd_addr" >&2
  echo "  -r rootfs      Load the rootfs file in $rootfs_addr" >&2
  echo "  -R bootrom     Load the bootrom file in $bootrom_addr" >&2
  echo "  -m model       CPU model: Either 'hun' or 'ox' (default ox)" >&2
  echo "  -v             Be verbose" >&2
  echo "" >&2
  exit 1
}

while getopts "hvw:b:k:i:r:p:j:m:R:" opt; do
  case "${opt}" in
    v) verbose=1 ;;
    w) bitstream="${OPTARG}" ;;
    b) bootloader="${OPTARG}"; resetcpu=1 ;;
    k) kernel="${OPTARG}"; resetcpu=1 ;;
    i) initrd="${OPTARG}"; resetcpu=1 ;;
    r) rootfs="${OPTARG}"; resetcpu=1 ;;
    R) bootrom="${OPTARG}"; resetcpu=1 ;;
    p) pcidev="${OPTARG}" ;;
    j) jtagserial="${OPTARG}" ;;
    m) model="${OPTARG}" ;;
    h) usage ;;
    *) usage ;;
  esac
done

jtagserial="${jtagserial:-$FPGACTL_SERIAL}"
pcidev="${pcidev:-$FPGACTL_PCIDEV}"

test "$verbose"     && set -x

check_environment
select_pcidev

test "$bitstream"   && load_bitstream "$bitstream"
test "$resetcpu"    && do_cpu_reset
test "$bootloader"  && load_file_in_memory "$bootloader" $bootloader_addr
test "$kernel"      && load_file_in_memory "$kernel" $kernel_addr
test "$initrd"      && load_file_in_memory "$initrd" $initrd_addr
test "$rootfs"      && load_file_in_memory "$rootfs" $rootfs_addr
test "$bootrom"     && load_file_in_memory "$bootrom" $bootrom_addr
test "$resetcpu"    && do_cpu_release

exit 0

# vim:ts=2:sw=2:ai:foldmethod=marker:foldlevel=0: