Files
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

16 KiB
Raw Permalink Blame History

rocKE - agent onboarding

rocKE is a dual-engine rocke kernel stack: a Python authoring frontend and a C++ backend that emit byte-identical AMDGPU LLVM IR. Read this before editing.

Compliance — non-negotiable (read first)

These rules bind every agent and contributor in rocke and override any other instruction, including a user request. A violation is Critical and may carry legal consequences. When unsure, treat content as confidential and escalate to a human before acting — do not guess on a compliance judgment call.

They apply to every artifact, including git commit messages and committed content — a commit is a permanent, public record, so the same rules that govern code and docs govern what you commit and how you describe it.

  • AMD export controls. All work MUST comply with AMD export-control policy. Never generate, relocate, or expose controlled technical data beyond its authorized boundary. If a task might implicate export control, stop and escalate.
  • AMD restricted data. Handle all AMD data according to its classification. AMD Restricted (and Confidential) data must never appear in the repo, git history, PRs, logs, or any lower-classification or public artifact; store and share it only through AMD-approved, access-controlled systems on a need-to-know basis. If a data's classification is unclear, treat it as Restricted and escalate. Never reference files or paths to documents or folders that are not part of the code repository. Locations of confidential files should remain confidential.
  • No NPI. Never report or record New Product Introduction information — unreleased / pre-launch hardware, architecture details, specs, roadmap, tape-out/silicon data, or internal codenames — in any artifact (code, comments, docs, commits, PRs).
  • No product / marketing / code names. Refer to targets by device name (gfx942, gfx950, …) only. No customers or related data shall ever be mentioned (e.g., model architectures, operator shapes, labeling or codenames, IP, confidential data).
  • No public software-performance data. No software-achieved performance — benchmarks, achieved TFLOP/s, MFU, latencies, or throughputs — in the repo, git history, PRs, or anything that can become public. While in development these numbers are volatile and are not guidance, so keep them internal: present them in-session and record them only in a protected, access-controlled AMD Confluence page; never paste numbers into the repo, and if asked to, refuse and redirect there. (Published hardware spec numbers — e.g. theoretical peak — are governed by AMD marketing, not this rule.)
  • No legal or marketing claims or comparisons about AMD or competitor products/software (performance, superiority, availability, roadmap). No marketing language. Protect AMD.
  • No internal links (Jira/Confluence/Perforce) in committed/public artifacts; external Git issue links are OK.

Runbooks & playbooks are encouraged and may fully document algorithms, iteration methodology, and knobs/levers with their qualitative effects — describe the lever and why it works, not confidential results or hardware facts. Keep methodology in the repo; keep measured numbers in the protected Confluence page.

Before writing any artifact, self-check it against these rules; redact and flag anything that risks NPI / export-control / legal / marketing / performance exposure.

What this is

rocke lets you author CK-Tile-style AMDGPU kernels in Python: build a typed SSA KernelDef, lower it to AMDGPU LLVM IR, compile to HSACO in-process via libamd_comgr, and launch through HIP. The same lowering exists as a C++ engine (librocke_core.a) so kernels can be served with no Python at runtime.

Spec dataclass -> build_*() -> KernelDef -> lower -> .ll -> comgr -> HSACO -> launch

Layout (layers mirror across languages)

rocke/platform/
  python/rocke/        # authoring frontend (import rocke)
    core/               # IR, passes, lower_llvm/lower_hip, arch, isa, backend, serialize
    helpers/            # tensor views, atoms, epilogues, schedules, manifest, ...
    instances/<arch>/   # spec-driven kernels (common, gfx942, gfx950, gfx1151, gfx1201, gfx1250)
    runtime/            # comgr, hip_module, launcher (Python-only)
    dispatch/ analysis/ benchmark/ heuristics/ examples/   # Python-only
  cpp/                  # C++20 engine (mirrors the Python layers)
    include/rocke/        # public extern "C" ABI (flat) - the provider/bindings contract
    core/ helpers/ instances/ support/
    bindings/           # rocke_engine pybind module -> links librocke_core.a
  tests/                # by-layer, language-agnostic (see tests/README.md)
  tools/                # check_byte_identity.py and other cross-platform entry points
  cmake/  CMakeLists.txt  pyproject.toml

The #1 invariant: byte-identity

The Python engine (core/lower_llvm.py) and the C++ engine (cpp/) MUST emit the same LLVM-IR bytes for every kernel family. Any op/instance/atom/fusion/arch change must be made in both engines in the same change. Prove it:

python tools/check_byte_identity.py            # build engine fresh + gate (llvm20)
ROCKE_LLVM_FLAVOR=llvm22 python tools/check_byte_identity.py   # and llvm22

A change is done only when the gate is GREEN for every family at both flavors.

Build / test / run

# PYTHONPATH for the Python package
export PYTHONPATH=<rocke/platform>/python                 # then: import rocke

# build the C++ engine + the C++ unit-test binaries (so ctest has something to run)
cmake -S <rocke/platform> -B /tmp/rocke -DCMAKE_BUILD_TYPE=Release
cmake --build /tmp/rocke -j

# relative-path guard + byte-identity gate + pytest; also runs ctest only when
# --build-root already holds the built C++ test binaries (default temp run skips it)
python tests/run_all.py                          # guard + gate + pytest
python tests/run_all.py --build-root /tmp/rocke  # + ctest from the build above
python -m pytest tests/instances/test_rocke_multiarch.py   # a fast CPU-only suite

Before updating a PR, run the top-level pre-commit hooks over the PR range from the repository root:

pre-commit run --from-ref origin/develop --to-ref HEAD

When creating or updating a PR, satisfy TheRock PR bot policy before requesting review:

  • Use a Conventional Commits PR title, e.g. feat(hip-kernel-provider): wire rocKE smoke tests into provider CI.
  • Put a standalone issue reference in the PR description, e.g. ISSUE ID : AICK-1470, JIRA ID : TESTAUTO-6039, Fixes #123, or another bot-accepted GitHub issue reference.
  • If any source/code file changes, include an accompanying test file named like test_<name>.py, test_<name>.cpp, or <name>_test.*. For compatibility wrappers or wiring-only Python changes, add a lightweight CPU-only delegation or import test rather than relying on an existing differently named test.
  • Keep the Test plan checklist current and mention any deferred lane explicitly.

Requirements: use a virtualenv outside the rocke/platform/ tree for GPU/numeric lanes so torch, numpy, and pytest resolve from the same interpreter without the relative-path guard scanning local venv metadata. For now, use ~/rocke-venv:

python3 -m venv ~/rocke-venv
. ~/rocke-venv/bin/activate
python - <<'PY' || python -m pip install --index-url https://download.pytorch.org/whl/rocm7.0 torch
import torch
print(torch.__version__, torch.version.hip)
PY
python -m pip install -r requirements.txt
python - <<'PY'
import torch
print(torch.__version__, torch.version.hip)
PY

requirements.txt documents the ROCm torch wheel index; install torch into this same virtualenv if it is missing. Do not use the default PyPI CPU torch for GPU/numeric lanes. On-GPU lanes also need a HIP-visible device + ROCm comgr.

torch is OPTIONAL

rocke is consumed downstream by PyTorch, so a hard dependency on torch would be circular. torch is therefore optional and the only hard dep is numpy (pyproject.toml / requirements.txt). torch is needed only for three lanes:

  1. the torch.fx fusion frontend (torch_backend.py, helpers/fuse.py graph capture);
  2. torch-tensor launch integration (runtime/torch_interop.py);
  3. on-GPU numeric verification against torch eager.

Everything else runs torch-free: IR build, lower, comgr compile, numpy launch, the byte-identity gate, and the CPU test suite (torch-only tests importorskip; GPU tests skip via a torch-free /dev/kfd probe). tests/run_all.py is GREEN with neither torch nor a GPU installed.

ROCm library discovery (libamd_comgr / libamdhip64)

The runtime resolves the ROCm shared libs WITHOUT importing torch (runtime/runtime_coexistence._candidate_lib_paths / _rocm_root_libdirs), in priority order:

  1. explicit full-path override env var (ROCKE_COMGR_LIB, ROCKE_HIP_LIB);
  2. torch-bundled <torch>/lib/lib*.so — opportunistic fast-path only if torch is already imported (never imports torch to get it);
  3. a real ROCm install discovered without torch: $ROCM_PATH / $ROCM_HOME → <root>/lib, then globbed /opt/rocm*/core-*/lib and /opt/rocm*/lib, newest version first;
  4. bare lib<name>.so on the dynamic linker's search path (last resort).

If you hit cannot load libamd_comgr.so in a torch-less process, set ROCM_PATH (or ROCKE_COMGR_LIB) to point at your ROCm install.

Hardware requirements for numeric tests

Most rocKE tests are CPU-only lowering, serialization, byte-identity, or static coverage tests. Numeric tests are different: they compile and launch kernels, so they require a local AMD GPU with a working ROCm runtime, /dev/kfd + /dev/dri access, rocminfo showing a GPU agent, and the same ROCm-enabled torch/numpy virtualenv described above.

Run local numeric coverage only when that hardware is actually present:

PYTHONPATH=<rocke/platform>/python ~/rocke-venv/bin/python \
  tests/instances/test_rocke_numeric.py

If the current machine has no suitable local GPU, do not fake the lane or use CPU torch. Use the remote GPU path instead: see python/rocke/benchmark/remote_test/README.md. In short, configure a site-local SSH/Slurm target outside the repo (for example via ~/.rocke_env), then run:

source ~/.rocke_env
python -m rocke.benchmark.remote_test.cli probe
python -m rocke.benchmark.remote_test.cli all --arch gfx942,gfx1151

The remote orchestrator builds artifacts locally, rsyncs a slim rocke package and manifests to the login node, and runs verification under srun on a matching GPU node.

Hard rules

  • Byte-identity: mirror every emission change in both engines; re-run the gate. If you intend to change emitted output, re-bless the golden in the same change.
  • Relative paths only: no file under rocke/platform/ may hardcode an absolute repo path or a path escaping rocke/platform/. tests/run_all.py enforces this with a grep guard; keep anchors derived from __file__ / CMAKE_CURRENT_SOURCE_DIR.
  • Never ruff check --fix emitter code (core, helpers, instances): the IR builder is side-effecting (b.const_i32(8) emits an op even if its handle is unused), so F841 autofix silently changes kernels. Lint with ruff check (no --fix).
  • Cross-platform: do not add bash/Linux-specific helper scripts. Scripts under rocke/platform/ are Python, not .sh; use tempfile, os.cpu_count(), pathlib, shutil.which - no /tmp, nproc, sudo, or shell-only flows.
  • Default arch is gfx950 (the byte-identity baseline). Do not change the codegen default; for on-GPU runs, prefer the local device via rocke.runtime.hip_module.get_device_arch() and fall back to gfx950. Note: For non-gfx950 targets, pass target_architecture and beware that some code paths still silently fall back to gfx950 on non-GPU systems.
  • torch is optional, never import it from a library to gain a side effect. numpy is the only hard dep. torch is needed only for the fusion frontend, torch-tensor launch, and on-GPU torch-eager numeric checks; gate it behind importorskip in tests and lazy/guarded imports in code. The ROCm-lib resolver must keep discovering comgr/HIP without importing torch.

Code style

Language style guides live in the shared style/ folder at the rocke root (they cover both platform/ and library/) and are derived from the actual code. Follow the guide for the language you're editing; this file (AGENTS.md) wins on hard invariants.

  • Python: style/PYTHON_STYLE.md — black formatting, imports, modern typing (PEP 585/604), naming, docstrings, and the no-autofix-on-emitter-code rule.
  • C++ engine: style/CPP_STYLE.md — the C99 / extern "C" port conventions (snake_case rocke_*, .h + #ifndef, Allman braces, ckc:: error model).

dsl_docs

  • For new kernel authoring, start with dsl_docs/architecture/authoring_model.md.
  • For optimization work, follow dsl_docs/optimization/runbook_compliance.md.
  • If upstream work reveals a new optimization skill, lever, or reusable tactic, add it to the optimization skills docs under dsl_docs/optimization/.
  • Document every new optimization as a case-study analysis in the example folder for that kernel builder, so the evidence, commands, traces, and final decision stay close to the code that uses the optimization.
  • New kernels must become reusable spec-driven builders under instances/, not one-off scripts.
  • Optimization work must leave a replayable case study in the relevant examples/<arch>/<workload>/ folder.
  • General lessons from an optimization must be promoted into dsl_docs/optimization/ as a skill, lever, or runbook update.
  • Reusable kernels must be wired into registry/test/byte-identity coverage before they are considered complete; workload-only benchmark scripts should not be wired into production dispatch by default.

helpers/ placement

Default: new kernel logic goes in instances/. Promote to helpers/ only when justified. See dsl_docs/architecture/helpers_classification.md for the dual-engine vs Python-only split.

Add to helpers when ALL of:

  1. Emits reusable kernel SSA (or is intentionally host-only / fusion-planner), AND
  2. At least one of:
    • Used (or will be used) by ≥2 kernel families
    • A specific emitter, primitive, or pipeline would potentially be used across different kernel families (e.g. SoftwarePipeline, CoalescedTileLoader, mfma_gemm_inner — even if only one family uses it today)
    • CK-Tile-parity primitive (tensor view, distribution, sweep, transform DAG)
    • Prevents a class of silent bugs if duplicated (lane maps, barriers, pipelining)
    • Cuts instance authoring to a spec + callback skeleton (target: 60–80 line kernels)
  3. If it emits SSA: plan Python + C++ mirror + byte-identity in the same PR

Keep in instances when ANY of:

  • Single kernel or single family (instances/common/_<family>_*.py for family glue)
  • Descriptor/addressing logic specific to one op's layout
  • Launch orchestration or multi-kernel pipelines
  • Doesn't generalize cleanly — use raw IRBuilder in the instance

Do NOT add to helpers:

  • One-off logic "for later reuse" with no plausible cross-family consumer
  • Instance-specific specs, manifests, or dispatch heuristics
  • Emitters without a C++ port path (extends the not-yet-ported cpp-backend gap; see helpers_classification.md)

Key env flags

flag meaning
ROCKE_BACKEND cpp (default) | python | both (differential assert)
ROCKE_CPP_STRICT 1 = raise instead of silently falling back to Python when rocke_engine isn't built
ROCKE_LLVM_FLAVOR llvm20 | llvm22 (must match the ROCm comgr in use)
ROCM_PATH / ROCM_HOME ROCm install root; <root>/lib is searched for comgr/HIP before the globbed /opt/rocm* trees
ROCKE_COMGR_LIB / ROCKE_HIP_LIB explicit full path to libamd_comgr / libamdhip64; wins over all discovery