diff --git a/AGENTS.md b/AGENTS.md index 7a7be64f5b..962039e68a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,4 +3,4 @@ - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`) - Run `python -m mypy tinygrad/` to typecheck - Run `python -m ruff check .` to lint -- Read `./tinygrad/viz/README` for profiling +- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules diff --git a/tinygrad/viz/README b/tinygrad/viz/README deleted file mode 100644 index 8c6a80e19c..0000000000 --- a/tinygrad/viz/README +++ /dev/null @@ -1,61 +0,0 @@ -VIZ is a tool for inspecting tinygrad's compilation process and performance profiling. - -to use: -1. Run tinygrad with VIZ=1 (this saves the pkls and launches the server in interactive shells) -2. That's it! - -This can: -1. See all schedules -2. See all graphs and how they were rewritten -3. See generated code -4. See profile -5. (AMD only) see instruction level SQTT profile - -VIZ pkls can be viewed in two ways: - -1. Web browser: python -m tinygrad.viz.serve -2. Command line: python -m tinygrad.viz.cli (add --json for srcipting) - -By default, VIZ UIs automatically load the latest files. - -user story: viewing profiling data -* tinygrad ran 32 LLM decode steps: -web: click "profiler", view the timeline of all python codegen and GPU kernels. -cli: Run `DEBUG=3 python -m tinygrad.viz.cli --json` to extract kernel timing info and ASTs in JSON format. - - note: Make sure to add NO_COLOR=1 to disable colored output. - -user story: viewing code -* tinygrad ran 3 schedules: init the model + first train step, train step, test step - * schedule 1 (123) = main.py:97 - * schedule 2 (97) = main.py:97 - * schedule 3 (10) = main.py:145 -* web: click "schedule 1", get list of kernels (like DEBUG=2) -* cli: `python -m tinygrad.viz.cli -s TINY "Schedule 3 Kernels n1"` - * kernel 1 "E_34_34" -- 'sin' - * kernel 2 "R_4545" -* web: click "E_34_34" -* cli: `python -m tinygrad.viz.cli -s TINY "do_to_program for E_34_34" "initial symbolic"` - * pre-rewritten UOp graph (step through rewrite here) - * post-rewritten UOp graph - * UOp list - * generated code - -user story: debugging scheduler -* tinygrad ran 3 schedules: init the model + first train step, train step, test step - * ... -* click "schedule 1 graph", get a graph of the schedule in UOps - * step through rewrite rules - * see how things are broken into kernels - * see why two kernels didn't fuse - -user story: SQTT / PMC profiling - -note: SQTT has additional overhead, to enable it, set VIZ=2. - -* tinygrad ran custom assembly GEMM kernel. -* web: click "SQTT gemm SE:1 PKTS", see wave instruction scheduling and CU execution unit occupancy at every clock cycle. -* cli: python -m tinygrad.viz.cli -s "kernel SQTT SE:0 PKTS" - -* get bank conflicts: -* web: click "gemm PMC" -* cli: python -m tinygrad.viz.cli -s "gemm PMC" | rg -A 16 SQC_LDS_BANK_CONFLICT diff --git a/tinygrad/viz/README.md b/tinygrad/viz/README.md new file mode 100644 index 0000000000..5b0c37bcc1 --- /dev/null +++ b/tinygrad/viz/README.md @@ -0,0 +1,93 @@ +VIZ is a tool for inspecting tinygrad's rewrites and runtime profiling. + +to use: +1. Run tinygrad with `VIZ=1` (this saves the pkls and launches the server in interactive shells) +2. That's it! + +# VIZ in the command line + +Use `python -m tinygrad.viz.cli` (add --json for scripting) to view the full timeline of events. + +### Environment variables + +Setting `DEBUG` includes the following data in the output stream. These show up as raw `{"value": "..."}` lines when using `--json`. + +| DEBUG | Includes | +|------:|----------| +| 3 | Base AST | +| 4 | Generated source | +| 5 | Rewrite steps and kernel graph | +| 6 | All UOp graphs | +| 7 | All rewrites | + + +VIZ defaults to colored output. Set `NO_COLOR=1` to disable colors. + +### Profiling examples + +Get kernel times and ASTs + +```bash +DEBUG=3 python -m tinygrad.viz.cli --json > /tmp/events.jsonl +``` + +Select events between two markers + +Markers are set using the `profile_marker` helper in user code. To list them: +``` +python -m tinygrad.viz.cli | rg MARKER +``` +Then: +``` +python -m tinygrad.viz.cli --interval "train @ 2" "train @ 3" +``` + +Set `-t` to aggregate events. + +### Rewrites Debugging example + +First, find the rewrite you are looking for. This can be a schedule or kernel: + +```bash +python -m tinygrad.viz.cli -s TINY | rg Schedule +python -m tinygrad.viz.cli -s TINY | rg E_3 +``` + +List all rewrite passes: + +Rewrite pass names come from `graph_rewrite(..., name="...")` in user code. +```bash +python -m tinygrad.viz.cli -s TINY "Schedule 6 Kernels n1" --ls +``` + +Show the input graph for each pass + +```bash +DEBUG=6 python -m tinygrad.viz.cli -s TINY "Schedule 6 Kernels n1" +``` + +Show all rewrites + +```bash +# for the entire scheduler +DEBUG=7 python -m tinygrad.viz.cli -s TINY "Schedule 6 Kernels n1" + +# or for a specific pass +DEBUG=7 python -m tinygrad.viz.cli -s TINY "Schedule 6 Kernels n1" "earliest rewrites" +``` + +# SQTT / PMC profiling (DEV=AMD only) + +SQTT has additional overhead. Set VIZ=2 to include it in pkls. + +Examples: + +Get all SQTT packets +```bash +python -m tinygrad.viz.cli -s "kernel SQTT SE:0 PKTS" --json +``` + +Get bank conflicts: +```bash +python -m tinygrad.viz.cli -s "gemm PMC" | rg -A 16 SQC_LDS_BANK_CONFLICT +``` diff --git a/tinygrad/viz/cli.py b/tinygrad/viz/cli.py index fd5b622124..50449843be 100755 --- a/tinygrad/viz/cli.py +++ b/tinygrad/viz/cli.py @@ -218,8 +218,7 @@ def main(args) -> None: for k in (produce_top_kernels if args.t else produce_all_kernels)(): render_event(k) def get_arg_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli", epilog="DEBUG modes (cumulative): 3=base AST, 4=generated source, " - "5=rewrite steps and kernel graph, 6=all UOp graphs, 7=all rewrites") + parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli") parser.add_argument("-s", "--src", nargs="+", default=[], metavar="NAME", help="Select a data source (default: all)") parser.add_argument("--list", "--ls", dest="list", action="store_true", help="List sources") parser.add_argument("--interval", nargs="+", metavar=("START", "END"), help="Optional start and end marker")