#!/bin/bash
# SkillFishOS — live Compute Unit (WGP) manager for the AMD BC-250.
# Routes GPU Compute Units at runtime by writing the CC and SPI WGP masks via umr.
# RLC_PG_ALWAYS_ON is deliberately left alone: see apply_rows.
# 1 WGP = 2 CU; 5 WGP per shader array x 4 arrays (2 SE x 2 SH) = 40 CU.
# No WGP is forced on. Any pair can be switched off, down to one per row.
# Our own implementation. The facts it rests on, the register names and fields
# and the rule that 1 WGP is 2 CU, come from the AMD headers (MIT licence) and
# from gfx_v10_0.c in the Linux kernel.
#
# There used to be a third "fact" here: that the driver locks WGP0-2 on and
# they can never be disabled. It is not true, and it cost a user with a
# defective CU in that range any way of switching it off. Measured on the dev
# board on 16/09/2026 by writing the masks through umr directly:
#
#   a row set to 0x1b keeps WGP2 off with WGP3 on, through a GPU load and
#   fifteen seconds, with nothing putting the bit back, and with
#   RLC_PG_ALWAYS_ON_WGP_MASK still at 0x1f
#
#   0x18 leaves only WGP3 and WGP4, so WGP0, WGP1 and WGP2 go off together
#
#   the silicon honours it rather than just reading it back: vkpeak fp32 gives
#   10156 GFLOPS at 40 CU, 8145 at 32, 4091 at 16, linear within a fifth of a
#   percent, and returns to 10158 when everything is switched back on
#
# See issue #77. Reported by tom lima.
set -u
UMR=/usr/local/bin/umr

# SENZA umr QUESTO STRUMENTO NON PUO' FARE NIENTE, E DEVE DIRLO.
#
# Prima non lo diceva, ed e' il guasto peggiore che si possa avere: ogni
# chiamata a umr qui sotto ha 2>/dev/null, quindi se il programma manca la
# lettura torna vuota, il conteggio fa zero e lo script stampava "0/40"
# USCENDO CON SUCCESSO. Il servizio risultava "active (exited)", il HUD
# mostrava 0/40, e l'utente restava a 24 CU convinto di averne 40 — cioe' senza
# la caratteristica principale della distribuzione, e senza un solo messaggio
# che glielo dicesse.
#
# Non e' un caso di scuola: umr NON esiste nei repository Debian (apt-cache
# policy umr non trova niente). Sulla BC-250 di sviluppo sta in /usr/local/bin
# e non appartiene a nessun pacchetto, quindi ci finisce nelle immagini fatte
# con eggs, che clonano quella macchina. Chi invece installa i nostri .deb da
# apt su un altro sistema, o costruisce la ISO con live-build, non ce l'ha.
if [ ! -x "$UMR" ]; then
    echo "skillfish-cu: manca $UMR, senza il quale non si possono instradare le CU." >&2
    echo "              umr non e' nei repository Debian: va installato a parte." >&2
    echo "              Le CU restano quelle imposte dal driver (24 su 40)." >&2
    exit 3
fi
INST=0
ASIC=cyan_skillfish.gfx1013
CC=mmCC_GC_SHADER_ARRAY_CONFIG
SPI=mmSPI_PG_ENABLE_STATIC_WGP_MASK
ROWS=(0:0 0:1 1:0 1:1)
MIN_WGP=0x01          # a row with no WGP has nothing to run on: keep one
FULL=0x1f            # all 5 WGP
STOCK=0x07           # WGP0-2: the mapping the driver comes up with on its own
RUN=/run/skillfish
STATE=$RUN/cu_active
BOOTCONF=/etc/skillfish/cu-boot.conf

popcount2(){ local m=$1 c=0 i; for i in 0 1 2 3 4; do (( (m>>i)&1 )) && c=$((c+2)); done; echo "$c"; }

read_row(){ # se sh -> decimal mask (low 5 bits)
  local v
  v=$($UMR -i $INST -r "cyan_skillfish.*.$SPI" -b "$1" "$2" 0xffffffff 2>/dev/null | grep -o '0x[0-9a-fA-F]\{8\}' | head -1)
  echo $(( ${v:-0} & 0x1f ))
}

read_total(){ local t=0 r se sh; for r in "${ROWS[@]}"; do se=${r%:*}; sh=${r#*:}; t=$((t + $(popcount2 $(read_row "$se" "$sh")))); done; echo "$t"; }

cmd_get(){ # JSON: per-row masks + active CU total
  local r se sh m parts=()
  for r in "${ROWS[@]}"; do se=${r%:*}; sh=${r#*:}; m=$(read_row "$se" "$sh"); parts+=("\"$se.$sh\":$m"); done
  # "floor" stays in the JSON because the Control Center reads it, and it is
  # now what it always should have been: nothing is forced on.
  printf '{"active_cu":%s,"max_cu":40,"floor":0,"min_wgp":%d,"rows":{%s}}\n' "$(read_total)" $((MIN_WGP)) "$(IFS=,; echo "${parts[*]}")"
}

write_row(){ # se sh mask(decimal)
  local m=$(( $3 & 0x1f ))
  (( m == 0 )) && m=$MIN_WGP
  # A write that fails must say so. These calls used to go to /dev/null with
  # nobody looking, so a routing that never happened still reported success.
  # Spotted by tom lima in PR #80.
  $UMR -i $INST -w "$ASIC.$CC" 0x0 -b "$1" "$2" 0xffffffff >/dev/null 2>&1 || return 1
  $UMR -i $INST -w "$ASIC.$SPI" "$(printf '0x%x' $m)" -b "$1" "$2" 0xffffffff >/dev/null 2>&1 || return 1
}

apply_rows(){ # m00 m01 m10 m11 (decimal/hex)
  # Four masks, each one a real one, checked before anything reaches a
  # register. Also from PR #80.
  [ "$#" -eq 4 ] || return 2
  local arg
  for arg in "$@"; do
    [[ "$arg" =~ ^(0x[0-9a-fA-F]{1,2}|[0-9]{1,2})$ ]] || return 2
    (( arg >= 0 && arg <= 31 )) || return 2
  done
  local r se sh idx=1
  $UMR -i $INST -w "$ASIC.$CC" 0x0 >/dev/null 2>&1 || return 1
  for r in "${ROWS[@]}"; do se=${r%:*}; sh=${r#*:}
    local m=$(( ${!idx} & 0x1f )); (( m == 0 )) && m=$MIN_WGP
    write_row "$se" "$sh" "$m" || return 1
    idx=$((idx+1))
  done
  # ⚠️ RLC_PG_ALWAYS_ON_WGP_MASK IS DELIBERATELY NOT WRITTEN HERE.
  #
  # It used to be, with the union of the rows. We then changed kernels to stop
  # the driver writing it: userpatches-350 carries Ariel's CU unlock, the one
  # whose subject line ends "no RLC_PG write", because writing it while the RLC
  # firmware is running makes the RLC wait forever for a ready ACK from WGPs
  # that VBIOS harvested off, and the board hangs on the first ROCm/HSA launch.
  # Doing it from here is that same write, at a worse moment. The always-on
  # mask is redundant anyway: our ppfeaturemask already keeps global power
  # gating off, and the kernel disables GFXOFF for gfx1013 outright.
  #
  # This is also why the union is gone. It existed only to feed this write, and
  # computing it as a numeric maximum instead of an OR was a real bug that tom
  # lima found in PR #80. The fix was right; the code it fixed is no longer
  # needed.
  update_state
}

apply_uniform(){ local m=$1; apply_rows "$m" "$m" "$m" "$m"; }

update_state(){
  mkdir -p "$RUN" 2>/dev/null
  cmd_get > "$STATE.json" 2>/dev/null
  echo "$(read_total)/40" > "$STATE" 2>/dev/null
  chmod 644 "$STATE" "$STATE.json" 2>/dev/null
}

need_root(){ [ "$(id -u)" = 0 ] || { echo "skillfish-cu: needs root" >&2; exit 1; }; }

righe_salvate(){ # the saved mapping, or nothing
  [ -r "$BOOTCONF" ] || return 0
  sed -n 's/^rows=//p' "$BOOTCONF" | head -1
}

cmd_boot_get(){
  local rows; rows=$(righe_salvate)
  if [ -n "$rows" ] && [ "$(echo $rows | wc -w)" = 4 ]; then
    printf '{"keep":true,"rows":[%s]}\n' "$(echo $rows | tr ' ' ',')"
  else
    printf '{"keep":false,"rows":[]}\n'
  fi
}

cmd_boot_save(){ # m00 m01 m10 m11
  mkdir -p "$(dirname "$BOOTCONF")" 2>/dev/null
  printf '# scritto da skillfish-cu: la mappatura da rimettere all%s avvio\nrows=%d %d %d %d\n' \
    "'" $(( $1 & 0x1f )) $(( $2 & 0x1f )) $(( $3 & 0x1f )) $(( $4 & 0x1f )) > "$BOOTCONF"
  chmod 644 "$BOOTCONF" 2>/dev/null
  echo "tenuta per il prossimo avvio: $1 $2 $3 $4"
}

cmd_boot(){ # what the unit runs
  # ⚠️ All 40 CU stay the default. A saved mapping is used only when somebody
  # asked for one, and only when it is four fields: a truncated or hand-edited
  # file must not leave the board with fewer CU than its owner expects.
  local rows; rows=$(righe_salvate)
  if [ -n "$rows" ] && [ "$(echo $rows | wc -w)" = 4 ]; then
    apply_rows $rows && riporta
  else
    apply_uniform $FULL
    riporta 40
  fi
}

# Secondo presidio: umr puo' esserci ed essere comunque inutile — driver diverso,
# scheda che non e' una BC-250, registri non accessibili. Anche in quel caso
# prima si usciva con successo dichiarando un numero sbagliato. Qui il numero
# ottenuto si confronta con quello chiesto, e se non torna lo si dice.
riporta(){ # riporta <cu_attese>
    local ottenute; ottenute=$(read_total)
    echo "$ottenute/40"
    if [ "$ottenute" -lt 1 ]; then
        echo "skillfish-cu: nessuna CU leggibile: umr non sta parlando con la GPU." >&2
        exit 4
    fi
    if [ -n "${1:-}" ] && [ "$ottenute" -lt "$1" ]; then
        echo "skillfish-cu: chieste $1 CU, ottenute $ottenute." >&2
        exit 5
    fi
    return 0
}

case "${1:-}" in
  get)              cmd_get ;;
  max|enable-all)   need_root; apply_uniform $FULL  && riporta 40 ;;
  stock|disable-extra) need_root; apply_uniform $STOCK && riporta 24 ;;
  boot)             need_root; cmd_boot ;;
  boot-save)        need_root; cmd_boot_save "${2:?}" "${3:?}" "${4:?}" "${5:?}" ;;
  boot-clear)       need_root; rm -f "$BOOTCONF"; echo "all'avvio tornano tutte accese" ;;
  boot-get)         cmd_boot_get ;;
  set)              need_root; apply_uniform "${2:?mask}" && riporta ;;   # 0x01..0x1f, all rows
  set-rows)         need_root; apply_rows "${2:?}" "${3:?}" "${4:?}" "${5:?}" && riporta ;;
  state)            need_root; update_state; cat "$STATE" ;;
  *) echo "usage: skillfish-cu {get|max|stock|set <0x01..0x1f>|set-rows m00 m01 m10 m11|state|boot|boot-save m00 m01 m10 m11|boot-clear|boot-get}"; exit 1 ;;
esac
