hygon-triton-kernel-optimizer

name

hygon-triton-kernel-optimizer

description

Capture, triage, benchmark, and optimize TorchInductor or hand-written Triton kernels on Hygon DCU gfx936/gfx938. Use when investigating torch.compile or TorchInductor Triton performance on DCU, parsing Triton autotune logs, saving generated kernels and inputs, checking AMDGCN buffer/global load-store codegen, applying tl.assume or tl.multiple_of hints, diagnosing low-bandwidth pointwise/reduction/scatter kernels, or deciding whether to tune Triton, disable torch.compile for a region, or rewrite model code.

Hygon Triton Kernel Optimizer

Purpose

Optimize or triage DCU Triton kernels generated by TorchInductor or written by hand. Use a measured workflow:

1. capture the generated kernel, input shapes, autotune configs, and cache artifacts;
2. inspect Triton metadata, pointer hints, and AMDGCN load/store instructions;
3. decide whether the kernel is worth tuning;
4. test tl.assume, tl.multiple_of, block size, warp count, and launch-shape variants;
5. if the kernel is structurally poor, avoid the generated Triton path or rewrite model code.

Keep scripts responsible for capture, parsing, artifact collection, and report generation. Use reasoning for code changes and architectural decisions.

Required Inputs

Have as many of these as possible:

• the original model or repro script that uses torch.compile;
• the target kernel name, or a profile showing the hot Triton kernel;
• log_profile.txt captured with TORCH_LOGS="+inductor" and TORCHINDUCTOR_TRACE=1;
• autotune_kernels/autotune.log and captured triton_<kernel>.py files;
• /tmp/torchinductor_root cache artifacts, especially .amdgcn, .ttir, .ttgir, and generated Python files.

If the user has only a model/repro, capture artifacts first. If DCU execution is remote, prepare scripts locally and run them through the project's remote workflow.

Environment

Run the environment probe first:

python <skill>/scripts/check_env.py --out ./triton_env.json

Confirm:

• torch, triton, and HIP runtime are importable;
• the target reports gfx936, gfx938, or another expected DCU arch;
• AMDGCN_USE_BUFFER_OPS is recorded; do not force it to 1 until a smoke compile passes on the target stack;
• hipcc, hipprof, dccobjdump, and rocminfo are available when deeper profiling or ISA verification is needed.

Fast Capture Path

Add the capture patch before torch.compile or before the compiled model is first executed:

import sys
sys.path.insert(0, "<skill>/scripts")
import autotune_capture_patch  # noqa: F401

Run with:

export TORCH_LOGS="+inductor"
export TORCHINDUCTOR_TRACE=1
export TRITON_CAPTURE_DIR=./autotune_kernels
python repro.py 2>&1 | tee log_profile.txt

Use AMDGCN_USE_BUFFER_OPS=1 only as an explicit buffer-op experiment after a compile smoke test. On the validated torch 2.9.0 / triton 3.3.0 / gfx938 stack, simple raw Triton and TorchInductor pointwise kernels failed LLVM lowering with that variable forced to 1, while the same capture succeeded with it unset.

Then summarize and collect:

python <skill>/scripts/summarize_autotune_log.py ./autotune_kernels/autotune.log \
  --json-out ./autotune_summary.json \
  --markdown-out ./autotune_summary.md

python <skill>/scripts/collect_inductor_artifacts.py \
  --log ./log_profile.txt \
  --capture-dir ./autotune_kernels \
  --cache-root /tmp/torchinductor_root \
  --out ./triton_artifacts

Captured kernel files are appended with a standalone runner where possible:

python ./autotune_kernels/<kernel_name>.py

Analysis Loop

1. Read references/investigation_workflow.md for the five-step triage workflow.
2. Read references/optimization_patterns.md when deciding whether to add tl.assume, tl.multiple_of, block-pointer forms, autotune configs, or model-level rewrites.
3. Inspect generated Triton metadata:

python <skill>/scripts/inspect_triton_meta.py ./autotune_kernels/<kernel_name>.py \
  --json-out ./meta.json

4. Inspect AMDGCN instructions:

python <skill>/scripts/scan_amdgcn.py ./triton_artifacts \
  --kernel <kernel_name> \
  --json-out ./isa_scan.json

5. Generate a report skeleton:

python <skill>/scripts/make_investigation_report.py \
  --autotune-summary ./autotune_summary.json \
  --meta ./meta.json \
  --isa ./isa_scan.json \
  --kernel <kernel_name> \
  --out ./triton_investigation.md

Fill missing facts from source inspection, profiling, and standalone benchmark results.

Decision Rules

• Prefer tuning only when the kernel is a real hot spot and has a plausible memory or launch inefficiency.
• Prioritize buffer-op enablement for memory-bound kernels only after compile-probing the stack: test AMDGCN_USE_BUFFER_OPS=1, pointer range metadata, tl.assume for non-negative offset inputs, and tl.multiple_of for proven alignment or divisibility separately.
• Compile-probe new tl.assume(...) forms on the target Triton/DCU stack before relying on them; some gfx938 Triton builds reject pointer-to-int assumptions and some scalar assumptions when AMDGCN_USE_BUFFER_OPS=1.
• Treat generated buffer_load/store_dwordx4 as stronger evidence than source intent. If output is mostly global_* or flat_*, inspect pointer hints and offset formulas.
• Do not force Triton tuning when the kernel only initializes tensors, performs scattered/atomic-heavy work, or replaces a high-performance library path such as GEMM with scalarized fused code.
• If the generated kernel is structurally poor, find the model code that triggers it and test a guarded eager fallback such as @torch._dynamo.disable on the smallest region.
• If model code can reduce scatter, atomics, tiny launches, or bad shape specialization, prefer a model rewrite and compare end-to-end performance.

Scripts

• scripts/autotune_capture_patch.py: monkey-patch TorchInductor CachingAutotuner to log configs, timings, shapes, bandwidth estimates, and captured standalone kernels.
• scripts/check_env.py: collect local Torch/Triton/HIP/DCU/tool availability and environment flags.
• scripts/summarize_autotune_log.py: parse autotune.log into JSON or Markdown.
• scripts/collect_inductor_artifacts.py: collect logs, captured kernels, generated Python code, and cache directories matching a kernel.
• scripts/inspect_triton_meta.py: extract signature, pointer args, hint metadata, and tl.assume/tl.multiple_of/atomic usage from Triton source.
• scripts/scan_amdgcn.py: count AMDGCN load/store/vector/atomic instruction families in dumps.
• scripts/run_captured_kernel.py: run one captured standalone kernel and parse timing/bandwidth lines.
• scripts/triton_benchmark_template.py: editable standalone Triton benchmark template for controlled kernel experiments. It clears AMDGCN_USE_BUFFER_OPS=1 by default because raw hand-written Triton autotune kernels can fail that lowering on the validated gfx938 stack; pass --keep-buffer-ops only when intentionally probing that behavior.
• scripts/make_investigation_report.py: merge script outputs into the investigation table.

Output Contract

For a complete investigation, produce:

• kernel name, profile rank, time share, and standalone/baseline timing;
• input shapes, dtypes, strides, and autotune best config;
• pointer argument list and hint status (tt.divisibility, tt.pointer_range);
• AMDGCN load/store summary, especially buffer vs global/flat instruction families;
• tuning attempts and measured results;
• decision: continue Triton tuning, disable a compiled region, or rewrite model code;
• final performance impact and unresolved profiling or artifact gaps.