#!/system/bin/sh
# On-device single PREFILL pass via qnn-net-run against the A16W8 trunk context binary.
#
# Unlike the decode graph this has NO KV in or out: one forward over a fixed SEQ=128 window
# produces the full [1,128,1536] hidden. So verifying it is one net-run per prompt, not a loop.
#
# Layout on device (BASE=/data/local/tmp/gemma):
#   $BASE/bin/qnn-net-run, $BASE/lib/*.so, $BASE/dsp/*.so
#   $BASE/artifacts/gemma4_trunk_a16w8_v79.bin
#   $BASE/tstep/<per-prompt input .raw files>
#   $BASE/tout/                      (net-run output dir -> hidden.raw)
set -e
BASE=/data/local/tmp/gemma
export LD_LIBRARY_PATH=$BASE/lib:/system/lib64:/vendor/lib64
export ADSP_LIBRARY_PATH="$BASE/dsp;/vendor/dsp/cdsp;/vendor/lib/rfsa/cdsp;/vendor/dsp"
export LD_PRELOAD=/system/lib64/libbinder.so   # Android 16 linker fix

STEP=$BASE/tstep
OUT=$BASE/tout
rm -rf "$OUT"; mkdir -p "$OUT"

LINE="inputs_embeds:=$STEP/inputs_embeds.raw"
LINE="$LINE per_layer_inputs:=$STEP/per_layer_inputs.raw"
LINE="$LINE position_ids:=$STEP/position_ids.raw"
LINE="$LINE full_mask:=$STEP/full_mask.raw"
LINE="$LINE sliding_mask:=$STEP/sliding_mask.raw"
echo "$LINE" > "$STEP/in.txt"

"$BASE/bin/qnn-net-run" \
  --backend "$BASE/lib/libQnnHtp.so" \
  --retrieve_context "$BASE/artifacts/gemma4_trunk_a16w8_v79.bin" \
  --input_list "$STEP/in.txt" \
  --output_dir "$OUT" \
  --use_native_input_files \
  --use_native_output_files \
  --log_level error

echo "TRUNK_OK"