polish the viz readme (#17236)

* polish the viz readme

* style

* no epilog=

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

tinygrad/viz/README

@@ -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

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
+```

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, "
+  parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli")
-                                   "5=rewrite steps and kernel graph, 6=all UOp graphs, 7=all rewrites")
   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")