#!/bin/bash
# Make `tile_kernels.torch` importable on a host without TileLang.
#
# Why this is needed: `tile_kernels.torch` holds pure-torch reference implementations of nearly
# every kernel the model forward calls (quantisation, SwiGLU, MoE routing, RoPE, MHC, Engram),
# and none of them reference TileLang. But the package's `__init__` re-exported every compiled
# kernel eagerly, so `import tile_kernels.torch` dragged in the whole TileLang stack and failed
# where that stack cannot be built -- which is exactly the case on a non-NVIDIA card.
#
# The three edits, applied to the installed package:
#   1. top-level `__init__` resolves its re-exports on first access (PEP 562)
#   2. `quant/__init__` does the same, because importing `quant.common` runs it first
#   3. `quant/common.py` stops requiring TileLang for its host-side half (config dataclasses,
#      shape helpers) while keeping the `@T.macro` fragments for the kernels
#
# Nothing changes on a host that has TileLang: `scripts/verify_tile_kernels_patch.py` checks
# that the surface is identical.
#
# Usage, inside the target environment:
#     bash scripts/enable_torch_kernel_refs.sh            # patch in place
#     bash scripts/enable_torch_kernel_refs.sh --revert   # restore the originals
#     PYTHON=/path/to/python bash scripts/enable_torch_kernel_refs.sh
#
# This patches a third-party package, which is why it is a script rather than a checkout-local
# shim: the intent is for these changes to go upstream to FlagOS. Until then, re-running this
# after a reinstall is the whole maintenance burden, and --revert makes it reversible.

set -euo pipefail

PYTHON="${PYTHON:-python3}"
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

TARGET="$("$PYTHON" - <<'PY'
import pathlib

import tile_kernels

print(pathlib.Path(tile_kernels.__file__).parent)
PY
)"

say() { printf '\n\033[1m== %s\033[0m\n' "$*"; }

say "tile_kernels at $TARGET"

if [ "${1:-}" = "--revert" ]; then
    "$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/lazify_package_init.py" \
        "$TARGET/quant/__init__.py" --revert
    "$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_quant_common.py" "$TARGET" --revert
    "$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_tile_kernels.py" "$TARGET" --revert
    say "reverted"
    exit 0
fi

say "1/3 lazy top-level __init__"
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_tile_kernels.py" "$TARGET"

say "2/3 lazy quant/__init__"
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/lazify_package_init.py" "$TARGET/quant/__init__.py"

say "3/3 quant/common.py without a hard TileLang import"
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_quant_common.py" "$TARGET"

say "verify: unchanged where TileLang is installed"
"$PYTHON" "$REPO_ROOT/scripts/verify_tile_kernels_patch.py"

say "done -- the reference path is now importable without TileLang"
echo "  check it with: DEEPINFER_SIMULATE_NO_TILELANG=1 python -c 'import tile_kernels.torch'"
echo "  (that env var needs the sitecustomize from scripts/simulate_no_tilelang/)"