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
151 lines
6.4 KiB
Markdown
151 lines
6.4 KiB
Markdown
# rocKE
|
|
|
|
**rocKE** is a dual-engine rocke kernel stack for AMDGPU: a **Python authoring
|
|
frontend** (`rocke`) and a **C++ engine** (`cpp/` → `librocke_core.a`) that emit
|
|
**byte-identical** AMDGPU LLVM IR. You author kernels in Python (build a typed
|
|
SSA `KernelDef`), lower to LLVM IR, compile to HSACO in-process via
|
|
`libamd_comgr`, and launch through HIP. The same lowering exists in C++ so
|
|
kernels can be served with **no Python at runtime**.
|
|
|
|
```text
|
|
Spec dataclass -> build_*() -> KernelDef -> lower -> .ll -> comgr -> HSACO -> launch
|
|
```
|
|
|
|
> Deeper docs: agent/onboarding notes in [`AGENTS.md`](AGENTS.md), the canonical
|
|
> build + artifact-hygiene reference in [`BUILD.md`](BUILD.md), the test-tree map
|
|
> in [`tests/README.md`](tests/README.md), and architecture/runtime/development
|
|
> docs under [`dsl_docs/`](dsl_docs).
|
|
|
|
## Layout
|
|
|
|
```text
|
|
rocke/platform/
|
|
python/rocke/ # authoring frontend (import rocke): core, helpers, instances/<arch>,
|
|
# runtime, dispatch, analysis, benchmark, heuristics, examples
|
|
cpp/ # C++20 engine (mirrors the Python layers): core, helpers, instances,
|
|
# support; include/ckc (public extern "C" ABI); bindings (rocke_engine)
|
|
tests/ # by-layer, language-agnostic; run_all.py is the entrypoint
|
|
dsl_docs/ # merged architecture / runtime / development docs
|
|
tools/ # cross-platform Python tooling (check_byte_identity.py, ...)
|
|
cmake/ CMakeLists.txt pyproject.toml requirements.txt
|
|
```
|
|
|
|
`python/rocke/<layer>` and `cpp/<layer>` mirror each other by layer name. The
|
|
two engines must stay byte-identical (see [`dsl_docs/development/engine_parity.md`](dsl_docs/development/engine_parity.md)).
|
|
|
|
## Prerequisites
|
|
|
|
- **ROCm 7.x** with `libamd_comgr` and `libamdhip64` on the dynamic-linker path
|
|
(7.2 recommended → `ROCKE_LLVM_FLAVOR=llvm22`; older → `llvm20`). ROCm is a
|
|
system dependency, not a pip package.
|
|
- **Python 3.10+**, a C++20 compiler (`amdclang++`/`clang++`/`g++`), CMake 3.16+
|
|
(3.18+ for the `cpp/bindings` pybind module).
|
|
- Python deps: `pip install -r requirements.txt` (adds `numpy`, plus `pytest` /
|
|
`pybind11` in the `dev` extra). **torch** must be the ROCm build for your
|
|
system, installed separately from the ROCm wheel index (not via this repo),
|
|
for example: `pip install --index-url https://download.pytorch.org/whl/rocm7.0 torch`.
|
|
- GPU lanes need a HIP-visible device; on a shared box that may require running
|
|
under `sudo -E` (or membership in the `video`/`render` groups).
|
|
|
|
```bash
|
|
export ROCKE=$(pwd) # run from this rocke/platform/ directory
|
|
export PYTHONPATH=$ROCKE/Python
|
|
```
|
|
|
|
## Build
|
|
|
|
### The engine (`librocke_core.a`)
|
|
|
|
```bash
|
|
cmake -S "$ROCKE" -B /tmp/rocke -DCMAKE_BUILD_TYPE=Release
|
|
cmake --build /tmp/rocke --target rocke_core -j$(nproc)
|
|
# -> /tmp/rocke/librocke_core.a
|
|
```
|
|
|
|
The top-level `CMakeLists.txt` globs `cpp/**/*.cpp` (excluding `cpp/bindings/`)
|
|
into `rocke_core`, with the public ABI headers at `cpp/include`. Optional
|
|
diagnostic build: add `-DROCKE_SANITIZE=ON` (ASan/UBSan; not for shipping).
|
|
|
|
### The `rocke_engine` Python binding (optional, enables the `cpp` backend)
|
|
|
|
```bash
|
|
cmake -S "$ROCKE/cpp/bindings" -B /tmp/rocke_pybind -DCMAKE_BUILD_TYPE=Release \
|
|
-DROCKE_ENGINE_ARCHIVE=/tmp/rocke/librocke_core.a \
|
|
-Dpybind11_DIR="$(python -m pybind11 --cmakedir)" \
|
|
-DPYTHON_EXECUTABLE="$(which python)"
|
|
cmake --build /tmp/rocke_pybind -j$(nproc)
|
|
export PYTHONPATH="$ROCKE/Python:/tmp/rocke_pybind" # so `import rocke_engine` resolves
|
|
```
|
|
|
|
If `rocke_engine` is not importable, the default `cpp` backend transparently falls
|
|
back to the native Python lowerer (set `ROCKE_CPP_STRICT=1` to make that an
|
|
error). See [`dsl_docs/reference/env_flags.md`](dsl_docs/reference/env_flags.md).
|
|
|
|
## Run
|
|
|
|
Author + lower a kernel in Python (no GPU needed to build/lower):
|
|
|
|
```bash
|
|
python -c "
|
|
from rocke.instances.common.gemm_universal import (
|
|
UniversalGemmSpec, TileSpec, TraitSpec, DataSpec, build_universal_gemm)
|
|
from rocke.core.lower_llvm import lower_kernel_to_llvm
|
|
spec = UniversalGemmSpec(name='demo',
|
|
tile=TileSpec(tile_m=128, tile_n=128, tile_k=32, warp_m=2, warp_n=2, warp_k=1,
|
|
warp_tile_m=16, warp_tile_n=16, warp_tile_k=16),
|
|
trait=TraitSpec(pipeline='compv3', epilogue='default'),
|
|
data=DataSpec(dtype_a='fp16', dtype_b='fp16', dtype_c='fp16', dtype_acc='fp32'),
|
|
wave_size=64, block_size=256, batched=False)
|
|
print(lower_kernel_to_llvm(build_universal_gemm(spec, arch='gfx950'), arch='gfx950')[:400])
|
|
"
|
|
```
|
|
|
|
Compile + launch on a GPU uses `rocke.runtime` (comgr + HIP). The default arch
|
|
is `gfx950`; on-GPU paths can target the local device via
|
|
`rocke.runtime.hip_module.get_device_arch()`.
|
|
|
|
## Test
|
|
|
|
One cross-platform entrypoint runs the relative-path guard → byte-identity gate
|
|
→ pytest. It **also** runs `ctest`, but only when the `--build-root` already
|
|
contains the built C++ test binaries — the gate itself builds just `rocke_core`,
|
|
so a default run (fresh temp build-root) does the guard + gate + pytest and
|
|
skips ctest. To include the C++ unit tests, point `--build-root` at a full
|
|
build (see step 3 below):
|
|
|
|
```bash
|
|
python "$ROCKE/tests/run_all.py" # guard + gate + pytest
|
|
python "$ROCKE/tests/run_all.py" --only gemm # scope the gate to some families
|
|
python "$ROCKE/tests/run_all.py" --no-gate # skip the engine build/gate
|
|
python "$ROCKE/tests/run_all.py" --build-root /tmp/rocke # also run ctest from this build
|
|
```
|
|
|
|
Or run the pieces individually:
|
|
|
|
```bash
|
|
# 1) Byte-identity gate: build the engine fresh + prove C-engine == Python-engine .ll
|
|
python "$ROCKE/tools/check_byte_identity.py"
|
|
ROCKE_LLVM_FLAVOR=llvm22 python "$ROCKE/tools/check_byte_identity.py" # also the llvm22 flavor
|
|
|
|
# 2) Python unit/integration suite
|
|
python -m pytest "$ROCKE/tests"
|
|
|
|
# 3) C++ unit tests (after a build that includes the test targets)
|
|
cmake -S "$ROCKE" -B /tmp/rocke && cmake --build /tmp/rocke -j$(nproc)
|
|
ctest --test-dir /tmp/rocke --output-on-failure
|
|
|
|
# 4) On-GPU numeric (needs a HIP device; use sudo -E on a shared box)
|
|
sudo -n -E env PYTHONPATH="$ROCKE/Python" python \
|
|
"$ROCKE/tests/instances/differential/numeric.py" --arch gfx950
|
|
```
|
|
|
|
Arch coverage (gfx942 / gfx1151 / gfx1201 in addition to gfx950) rides on the
|
|
emitter config set the byte-identity gate enumerates; there is no separate arch
|
|
sweep. See [`tests/README.md`](tests/README.md) for the per-layer test map and
|
|
the dedup/coverage notes.
|
|
|
|
## License
|
|
|
|
Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
SPDX-License-Identifier: MIT
|