diff --git a/docs/worktrees.md b/docs/worktrees.md new file mode 100644 index 00000000..38b4c895 --- /dev/null +++ b/docs/worktrees.md @@ -0,0 +1,112 @@ +# Parallel worktrees + +Quick setup for working on multiple branches at the same time (e.g. one for +feature work, one for code review, one for a bug fix) without thrashing a +single checkout. Container-dedicated: paths assume the +`ghcr.io/psal-postech/torchsim-ci` layout. + +## Why a script + +Three things have to line up for parallel worktrees to actually work in this +repo: + +1. **Worktree-scoped env vars.** `PyTorchSimFrontend/extension_config.py` + anchors output / log / config paths on `TORCHSIM_DIR`. Without an override + every worktree dumps into the same `outputs/` and `togsim_results/`. +2. **`PYTHONPATH` override.** `pip install -e PyTorchSimDevice` writes a + single editable record into conda's site-packages that points at one + worktree. The override forces `import torch_openreg` to resolve to the + active worktree's code and `.so` first. +3. **Branch tracking.** `git worktree add` from a remote ref sets upstream to + that ref, so `git push` would target `develop`. The script unsets it so + first `git push -u origin ` creates the right remote branch. + +The script bundles all three. + +## Create a worktree + +```bash +scripts/setup_worktree.sh [base-ref] +``` + +`` becomes the branch suffix and the dir suffix: + +| Command | Worktree dir | Branch | +|---|---|---| +| `setup_worktree.sh feature` | `/workspace/PyTorchSim-feature` | `feature/scratch` | +| `setup_worktree.sh review` | `/workspace/PyTorchSim-review` | `refactor/scratch` (rename after) | +| `setup_worktree.sh bugfix/issue-198` | `/workspace/PyTorchSim-bugfix` | `bugfix/issue-198` | +| `setup_worktree.sh feature origin/master` | `/workspace/PyTorchSim-feature` | `feature/scratch` (off master) | + +Default base is `origin/develop` (per `CONTRIBUTING.md`). + +## Activate + +```bash +cd /workspace/PyTorchSim-bugfix +source .envrc +# → Activated worktree: /workspace/PyTorchSim-bugfix +# → prompt: [torchsim:PyTorchSim-bugfix] ... +``` + +`.envrc` is local to each worktree and not committed. Re-source it whenever +you open a new shell in that worktree. + +## Build once per worktree + +The compiled `_C.cpython-*.so` lives under `PyTorchSimDevice/torch_openreg/` +and is not shared across worktrees. After activation: + +```bash +(cd PyTorchSimDevice && python setup.py build_ext --inplace) +``` + +Use `build_ext --inplace` instead of `pip install -e` so the editable +record in `/opt/conda/lib/python3.11/site-packages` keeps pointing at +whichever worktree it already pointed at — `PYTHONPATH` from `.envrc` does +the per-worktree routing. (Running `pip install -e` again rewrites that +record and will pin "default" Python to the new worktree.) + +## What the env looks like + +Worktree-scoped (auto-set by `.envrc`): + +| Var | Value | +|---|---| +| `TORCHSIM_DIR` | `$PWD` of the worktree | +| `TORCHSIM_DUMP_PATH` | `$PWD/outputs` | +| `TORCHSIM_LOG_PATH` | `$PWD/togsim_results` | +| `TOGSIM_CONFIG` | `$PWD/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml` | +| `PYTHONPATH` | `$PWD/PyTorchSimDevice:$PWD:$PYTHONPATH` | + +Shared (container-dedicated, set the same in every `.envrc`): + +| Var | Value | +|---|---| +| `GEM5_PATH` | `/gem5/release/gem5.opt` | +| `TORCHSIM_LLVM_PATH` | `/riscv-llvm/bin` | +| `RISCV` | `/workspace/riscv` | + +## Cleanup + +```bash +git worktree remove /workspace/PyTorchSim-feature +git branch -D feature/scratch # if you do not want to keep the branch +``` + +`git worktree list` shows the current set. + +## Gotchas + +- **Do not commit `.envrc`.** It is per-worktree state. Add to your + personal global gitignore if needed. +- **Editable install conflict.** If you run `pip install -e PyTorchSimDevice` + in worktree A, then again in worktree B, Python's default `import + torch_openreg` flips to B. With `PYTHONPATH` from `.envrc` this still + resolves correctly in either worktree's shell, but a shell with no + `.envrc` sourced will see whichever was installed last. +- **TOGSim FIFO files.** `/tmp/togsim_fifo_` is keyed on PID, not on + worktree — concurrent runs from different worktrees do not collide. +- **`_C.cpython-*.so` rebuild on PyTorch update.** If you `pip install` + a different torch version in the conda env, every worktree's `.so` is + stale; rebuild each. diff --git a/scripts/setup_worktree.sh b/scripts/setup_worktree.sh new file mode 100755 index 00000000..3c4c5f95 --- /dev/null +++ b/scripts/setup_worktree.sh @@ -0,0 +1,99 @@ +#!/usr/bin/env bash +# Create a sibling git worktree for parallel work and wire up per-worktree env. +# +# Container-dedicated: assumes the ghcr.io/psal-postech/torchsim-ci container +# layout. Shared binaries (gem5, LLVM, riscv toolchain) live at known paths +# and are NOT duplicated per worktree. +# +# Usage: +# scripts/setup_worktree.sh [base-ref] +# +# Examples: +# scripts/setup_worktree.sh feature # off origin/develop, branch feature/scratch +# scripts/setup_worktree.sh bugfix/issue-198 # off origin/develop, branch bugfix/issue-198 +# scripts/setup_worktree.sh review origin/master # off origin/master, branch review/scratch +# +# Result: +# /workspace/PyTorchSim- new worktree +# /workspace/PyTorchSim-<...>/.envrc per-worktree env (source it) +# +# After creation: +# cd /workspace/PyTorchSim-<...> +# source .envrc +# (cd PyTorchSimDevice && python setup.py build_ext --inplace) # build the .so once +set -euo pipefail + +if [[ $# -lt 1 || $# -gt 2 ]]; then + sed -n '3,20p' "$0" + exit 1 +fi + +PURPOSE="$1" +BASE_REF="${2:-origin/develop}" + +# Branch name: use the purpose as-is if it already has a slash, else append /scratch. +if [[ "$PURPOSE" == */* ]]; then + BRANCH="$PURPOSE" + SUFFIX="${PURPOSE%%/*}" # before first slash, used in dir name +else + BRANCH="${PURPOSE}/scratch" + SUFFIX="$PURPOSE" +fi + +REPO_ROOT="$(git -C "$(dirname "${BASH_SOURCE[0]}")/.." rev-parse --show-toplevel)" +PARENT_DIR="$(dirname "$REPO_ROOT")" +WT_DIR="${PARENT_DIR}/$(basename "$REPO_ROOT")-${SUFFIX}" + +if [[ -e "$WT_DIR" ]]; then + echo "error: $WT_DIR already exists" >&2 + exit 1 +fi + +# Make sure base-ref is up to date if it's a remote ref. +if [[ "$BASE_REF" == origin/* ]]; then + git -C "$REPO_ROOT" fetch origin "${BASE_REF#origin/}" --depth=1 +fi + +git -C "$REPO_ROOT" worktree add "$WT_DIR" -b "$BRANCH" "$BASE_REF" + +# Default `worktree add` from a remote ref sets upstream to that remote ref, +# which means `git push` would target develop/master. Unset so the new branch +# pushes to its own name on first `git push -u origin `. +git -C "$WT_DIR" branch --unset-upstream || true + +# Per-worktree env. Container-dedicated paths for shared binaries. +cat > "$WT_DIR/.envrc" <<'ENVRC' +#!/usr/bin/env bash +# Source this from the worktree root: source .envrc +_self="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)" + +# Worktree-scoped: override defaults from PyTorchSimFrontend/extension_config.py +export TORCHSIM_DIR="$_self" +export TORCHSIM_DUMP_PATH="$_self/outputs" +export TORCHSIM_LOG_PATH="$_self/togsim_results" +export TOGSIM_CONFIG="$_self/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.yml" + +# Make `import torch_openreg` resolve to THIS worktree's .so first, +# overriding the conda-wide editable install that points at the main worktree. +export PYTHONPATH="$_self/PyTorchSimDevice:$_self:${PYTHONPATH:-}" + +# Container-dedicated shared binaries. +export GEM5_PATH="/gem5/release/gem5.opt" +export TORCHSIM_LLVM_PATH="/riscv-llvm/bin" +export RISCV="/workspace/riscv" + +# Prompt hint so you do not lose track of which worktree this shell is on. +export PS1="[torchsim:$(basename "$_self")] ${PS1:-\\w\\$ }" + +unset _self +echo "Activated worktree: $TORCHSIM_DIR" +ENVRC + +echo +echo "Created worktree: $WT_DIR" +echo "Branch: $BRANCH (base: $BASE_REF)" +echo +echo "Next:" +echo " cd $WT_DIR" +echo " source .envrc" +echo " (cd PyTorchSimDevice && python setup.py build_ext --inplace) # build the .so once"