ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-11 09:40:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/develop' into ginolu/sparge_attention

Browse Source
This commit is contained in:
Gino Lu
2026-04-13 03:27:43 -04:00
parent 9317fc4a85 fa4473fde6
commit 643ad35de2
660 changed files with 40115 additions and 22333 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
.gitignore vendored
View File

@@ -112,4 +112,30 @@ test_data/*
experimental/grouped_convolution_tile_instances/instances/*
!experimental/grouped_convolution_tile_instances/instances/*.in
!experimental/grouped_convolution_tile_instances/instances/*.inc
!experimental/grouped_convolution_tile_instances/instances/*.hpp
experimental/grouped_convolution_tile_instances/*.inc
# Heuristics: benchmark data (never in git)
dispatcher/heuristics/data/
# Heuristics: experimental/training artifacts (exclude from git)
dispatcher/heuristics/models/**/oof_predictions.parquet
dispatcher/heuristics/models/**/cv_metrics_*.json
dispatcher/heuristics/models/**/eval_report.json
dispatcher/heuristics/models/**/feature_importances_*.json
dispatcher/heuristics/models/**/model_tflops_ihem.lgbm
dispatcher/heuristics/models/**/model_tflops_log.lgbm
dispatcher/heuristics/models/**/model_tflops_log_big.lgbm
# Heuristics: keep in git (production model files):
# models/{op}_{dtype}_{arch}/model_tflops.lgbm
# models/{op}_{dtype}_{arch}/model_latency.lgbm
# models/{op}_{dtype}_{arch}/model_bandwidth.lgbm
# models/{op}_{dtype}_{arch}/feature_spec.json
# models/{op}_{dtype}_{arch}/train_manifest.json
# Heuristics: logs and caches
dispatcher/heuristics/*.log
dispatcher/heuristics/__pycache__/
dispatcher/heuristics/tests/__pycache__/
dispatcher/heuristics/.pytest_cache/

4
.readthedocs.yaml
View File

@@ -4,13 +4,13 @@
version: 2
sphinx:
configuration: docs/conf.py
configuration: projects/composablekernel/docs/conf.py
formats: [htmlzip, pdf, epub]
python:
install:
- requirements: docs/sphinx/requirements.txt
- requirements: projects/composablekernel/docs/sphinx/requirements.txt
build:
os: ubuntu-22.04

1
CHANGELOG.md
View File

@@ -22,6 +22,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj
* Added FP8 block scale quantization for FMHA forward kernel.
* Added gfx11 support for FMHA.
* Added microscaling (MX) FP8/FP4 support on gfx950 for FMHA forward kernel ("qr" pipeline only).
* Added FP8 per-tensor quantization support for FMHA forward V3 pipeline on gfx950.
### Changed

51
Dockerfile
View File

@@ -2,22 +2,33 @@
FROM ubuntu:24.04
ARG DEBIAN_FRONTEND=noninteractive
ARG ROCMVERSION=7.1.1
ARG DEB_ROCM_REPO=http://repo.radeon.com/rocm/apt/.apt_$ROCMVERSION/
ARG TARBALL_URL=https://rocm.nightlies.amd.com/tarball/therock-dist-linux-gfx90X-dcgpu-7.12.0a20260218.tar.gz
ARG compiler_version=""
ARG compiler_commit=""
ARG CK_SCCACHE=""
ARG DEB_ROCM_REPO=http://repo.radeon.com/rocm/apt/.apt_$ROCMVERSION/
ENV APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=DontWarn
ENV DEBIAN_FRONTEND=noninteractive
ENV PATH=$PATH:/opt/rocm/bin
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/lib
ENV HIP_PLATFORM=amd
# Add rocm repository
RUN set -xe && \
apt-get update && apt-get install -y --allow-unauthenticated apt-utils wget gnupg2 curl
RUN wget https://repo.radeon.com/amdgpu-install/7.1.1/ubuntu/noble/amdgpu-install_7.1.1.70101-1_all.deb && \
apt install ./amdgpu-install_7.1.1.70101-1_all.deb -y && \
apt update && \
apt install python3-setuptools python3-wheel -y && \
apt install rocm-dev -y
RUN if [ "$compiler_version" = "therock" ]; then \
rm -rf /opt/rocm && mkdir /opt/rocm && \
echo "Downloading ROCm tarball from $TARBALL_URL..." && \
wget -q -O /tmp/rocm.tar.gz "$TARBALL_URL" && \
echo "Extracting tarball to /opt/rocm..." && \
tar -xzf /tmp/rocm.tar.gz -C /opt/rocm --strip-components=1 ; \
else echo "using the release compiler" && \
wget https://repo.radeon.com/amdgpu-install/7.1.1/ubuntu/noble/amdgpu-install_7.1.1.70101-1_all.deb && \
apt install ./amdgpu-install_7.1.1.70101-1_all.deb -y && \
apt update && \
apt install python3-setuptools python3-wheel -y && \
apt install rocm-dev -y; \
fi
# Install SCCACHE
ENV SCCACHE_VERSION="0.14.0"
@@ -34,7 +45,6 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-
build-essential \
cmake \
git \
hip-rocclr \
iputils-ping \
jq \
libelf-dev \
@@ -44,8 +54,8 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-
net-tools \
pkg-config \
python3-full \
python3-pip \
redis \
rocm-llvm-dev \
sshpass \
stunnel \
software-properties-common \
@@ -88,26 +98,3 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-
git clone -b master https://github.com/ROCm/rocm-cmake.git && \
cd rocm-cmake && mkdir build && cd build && \
cmake .. && cmake --build . && cmake --build . --target install
WORKDIR /
# Add alternative compilers, if necessary
ENV compiler_version=$compiler_version
ENV compiler_commit=$compiler_commit
RUN sh -c "echo compiler version = '$compiler_version'" && \
sh -c "echo compiler commit = '$compiler_commit'"
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-mainline" ] ) && [ "$compiler_commit" = "" ]; then \
git clone -b "$compiler_version" https://github.com/ROCm/llvm-project.git && \
cd llvm-project && mkdir build && cd build && \
cmake -DCMAKE_INSTALL_PREFIX=/opt/rocm/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm && \
make -j 8 ; \
else echo "using the release compiler"; \
fi
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-mainline" ] ) && [ "$compiler_commit" != "" ]; then \
git clone -b "$compiler_version" https://github.com/ROCm/llvm-project.git && \
cd llvm-project && git checkout "$compiler_commit" && echo "checking out commit $compiler_commit" && mkdir build && cd build && \
cmake -DCMAKE_INSTALL_PREFIX=/opt/rocm/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm && \
make -j 8 ; \
else echo "using the release compiler"; \
fi

2
Dockerfile.aiter
View File

@@ -4,7 +4,7 @@ ARG AITER_BRANCH="main"
ARG CK_AITER_BRANCH="develop"
# CK_FROM_ROCM_LIBRARIES - 1: CK from rocm-libraries sparse-checkout; 0: direct clone from ROCm/composable_kernel
ARG CK_FROM_ROCM_LIBRARIES=1
RUN pip install pandas zmq einops ninja tabulate && \
RUN pip install pandas zmq einops ninja tabulate vcs_versioning && \
pip install numpy==1.26.2 && \
sudo mkdir /home/jenkins && \
sudo mkdir /home/jenkins/workspace && \

66
Dockerfile.compiler
View File

@@ -9,18 +9,70 @@ ENV compiler_commit=$compiler_commit
RUN sh -c "echo compiler version = '$compiler_version'" && \
sh -c "echo compiler commit = '$compiler_commit'"
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-mainline" ] ) && [ "$compiler_commit" = "" ]; then \
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-staging" ] ) && [ "$compiler_commit" = "" ]; then \
git clone -b "$compiler_version" https://github.com/ROCm/llvm-project.git && \
cd llvm-project && mkdir build && cd build && \
cmake -DCMAKE_INSTALL_PREFIX=/opt/rocm/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm && \
make -j 16 ; \
cd llvm-project && git log -1 && mkdir build && cd build && \
cmake -G Ninja \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_ENABLE_PROJECTS="clang;lld;clang-tools-extra;flang;mlir" \
-DLLVM_LIT_ARGS="-vv --show-unsupported --show-xfail -j 32" \
-DPACKAGE_VENDOR="AMD" \
-DCMAKE_INSTALL_PREFIX=/home/$USER/rocm/pure_llvm_1.0 \
-DLLVM_ENABLE_ASSERTIONS=ON \
-DLLVM_BUILD_DOCS=ON \
-DLLVM_TARGETS_TO_BUILD=all \
-DLIBOMPTARGET_ENABLE_DEBUG=ON \
-DOFFLOAD_ENABLE_EMISSARY_APIS=OFF \
-DCLANG_DEFAULT_LINKER=lld \
-DCLANG_DEFAULT_PIE_ON_LINUX=0 \
-DLLVM_ENABLE_RUNTIMES="libcxx;libcxxabi;openmp;compiler-rt;libunwind;flang-rt" \
-DLIBCXX_ENABLE_SHARED=OFF \
-DLIBCXX_ENABLE_STATIC=ON \
-DLIBCXX_INSTALL_LIBRARY=OFF \
-DLIBCXX_INSTALL_HEADERS=OFF \
-DLIBCXXABI_ENABLE_SHARED=OFF \
-DLIBCXXABI_ENABLE_STATIC=ON \
-DLIBCXXABI_INSTALL_STATIC_LIBRARY=OFF \
-DLLVM_ENABLE_ASSERTIONS=1 \
-DLLVM_ENABLE_Z3_SOLVER=OFF \
-DLLVM_ENABLE_ZLIB=ON \
-DLLVM_LINK_LLVM_DYLIB=OFF \
-DCLANG_LINK_CLANG_DYLIB=OFF \
../llvm && \
ninja -j16 ; \
else echo "using the release compiler"; \
fi
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-mainline" ] ) && [ "$compiler_commit" != "" ]; then \
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-staging" ] ) && [ "$compiler_commit" != "" ]; then \
git clone -b "$compiler_version" https://github.com/ROCm/llvm-project.git && \
cd llvm-project && git checkout "$compiler_commit" && echo "checking out commit $compiler_commit" && mkdir build && cd build && \
cmake -DCMAKE_INSTALL_PREFIX=/opt/rocm/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm && \
make -j 16 ; \
cmake -G Ninja \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_ENABLE_PROJECTS="clang;lld;clang-tools-extra;flang;mlir" \
-DLLVM_LIT_ARGS="-vv --show-unsupported --show-xfail -j 32" \
-DPACKAGE_VENDOR="AMD" \
-DCMAKE_INSTALL_PREFIX=/home/$USER/rocm/pure_llvm_1.0 \
-DLLVM_ENABLE_ASSERTIONS=ON \
-DLLVM_BUILD_DOCS=ON \
-DLLVM_TARGETS_TO_BUILD=all \
-DLIBOMPTARGET_ENABLE_DEBUG=ON \
-DOFFLOAD_ENABLE_EMISSARY_APIS=OFF \
-DCLANG_DEFAULT_LINKER=lld \
-DCLANG_DEFAULT_PIE_ON_LINUX=0 \
-DLLVM_ENABLE_RUNTIMES="libcxx;libcxxabi;openmp;compiler-rt;libunwind;flang-rt" \
-DLIBCXX_ENABLE_SHARED=OFF \
-DLIBCXX_ENABLE_STATIC=ON \
-DLIBCXX_INSTALL_LIBRARY=OFF \
-DLIBCXX_INSTALL_HEADERS=OFF \
-DLIBCXXABI_ENABLE_SHARED=OFF \
-DLIBCXXABI_ENABLE_STATIC=ON \
-DLIBCXXABI_INSTALL_STATIC_LIBRARY=OFF \
-DLLVM_ENABLE_ASSERTIONS=1 \
-DLLVM_ENABLE_Z3_SOLVER=OFF \
-DLLVM_ENABLE_ZLIB=ON \
-DLLVM_LINK_LLVM_DYLIB=OFF \
-DCLANG_LINK_CLANG_DYLIB=OFF \
../llvm && \
ninja -j16 ; \
else echo "using the release compiler"; \
fi

43
Dockerfile.fa Normal file
View File

@@ -0,0 +1,43 @@
ARG BASE_DOCKER="rocm/pytorch:latest"
FROM $BASE_DOCKER
ARG FA_ORIGIN="ROCm"
ARG FA_BRANCH="tridao"
ARG CK_FA_ORIGIN="ROCm"
ARG CK_FA_BRANCH="develop"
# CK_FROM_ROCM_LIBRARIES - 1: CK from rocm-libraries sparse-checkout; 0: direct clone from ROCm/composable_kernel
ARG CK_FROM_ROCM_LIBRARIES=1
ARG GPU_ARCHS="gfx90a;gfx942;gfx950"
RUN set -x ; \
sudo mkdir /home/jenkins && \
sudo mkdir /home/jenkins/workspace && \
cd /home/jenkins/workspace && rm -rf rocm-libraries ck && \
if [ "$CK_FROM_ROCM_LIBRARIES" = "1" ]; then \
git clone --depth 1 -b "$CK_FA_BRANCH" --no-checkout --filter=blob:none https://github.com/$CK_FA_ORIGIN/rocm-libraries.git && \
cd rocm-libraries && \
git sparse-checkout init --cone && \
git sparse-checkout set projects/composablekernel && \
git checkout "$CK_FA_BRANCH" && \
ROCM_LIBRARIES_SHA=$(git rev-parse --short HEAD) && \
mv projects/composablekernel ../ck && \
cd ../ck && rm -rf ../rocm-libraries && \
git init && \
git config user.name "assistant-librarian[bot]" && \
git config user.email "assistant-librarian[bot]@users.noreply.github.com" && \
git branch -m "$CK_FA_BRANCH" && git add -A && \
git commit -m "import from ROCm/rocm-libraries@$ROCM_LIBRARIES_SHA" > /dev/null ; \
else \
git clone --depth 1 -b "$CK_FA_BRANCH" https://github.com/$CK_FA_ORIGIN/composable_kernel.git ck ; \
fi && \
cd /home/jenkins/workspace && rm -rf flash-attention && \
git clone --depth 1 -b "$FA_BRANCH" --recursive "https://github.com/$FA_ORIGIN/flash-attention.git" && \
cd flash-attention && \
rm -rf csrc/composable_kernel/ && \
git clone -b "$CK_FA_BRANCH" ../ck csrc/composable_kernel/ && git add csrc/composable_kernel && \
MAX_JOBS=$(nproc) GPU_ARCHS="$GPU_ARCHS" /opt/venv/bin/python3 -u -m pip install --no-build-isolation -v . && \
groupadd -g 1001 jenkins && \
useradd -u 1001 -g 1001 -m -s /bin/bash jenkins && \
chown -R jenkins:jenkins /home/jenkins && \
chmod -R a+rwx /home/jenkins && \
chown -R jenkins:jenkins /tmp && \
chmod -R a+rwx /tmp && \
sudo usermod -aG irc jenkins

80
Dockerfile.manylinux
View File

@@ -3,7 +3,6 @@ ARG DEBIAN_FRONTEND=noninteractive
ARG ROCMVERSION=7.2
ARG compiler_version=""
ARG compiler_commit=""
ARG CK_SCCACHE=""
ARG DEB_ROCM_REPO=http://repo.radeon.com/rocm/apt/.apt_$ROCMVERSION/
ENV APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=DontWarn
ENV DEBIAN_FRONTEND=noninteractive
@@ -19,16 +18,15 @@ RUN wget https://repo.radeon.com/amdgpu-install/7.2/rhel/8.10/amdgpu-install-7.2
dnf install python3-setuptools python3-wheel -y && \
dnf install rocm-dev -y
## Sccache binary built from source for ROCm, only install if CK_SCCACHE is defined
ARG SCCACHE_REPO_URL=http://compute-artifactory.amd.com/artifactory/rocm-generic-experimental/rocm-sccache
# Install SCCACHE
ENV SCCACHE_VERSION="0.14.0"
ENV SCCACHE_INSTALL_LOCATION=/usr/local/.cargo/bin
ENV PATH=$PATH:${SCCACHE_INSTALL_LOCATION}
ENV CK_SCCACHE=$CK_SCCACHE
RUN if [ "$CK_SCCACHE" != "" ]; then \
mkdir -p ${SCCACHE_INSTALL_LOCATION} && \
curl ${SCCACHE_REPO_URL}/portable/0.2.16/sccache-0.2.16-alpha.1-rocm --output ${SCCACHE_INSTALL_LOCATION}/sccache && \
chmod +x ${SCCACHE_INSTALL_LOCATION}/sccache; \
fi
RUN set -x && \
mkdir -p ${SCCACHE_INSTALL_LOCATION} && \
wget -qO sccache.tar.gz https://github.com/mozilla/sccache/releases/latest/download/sccache-v$SCCACHE_VERSION-x86_64-unknown-linux-musl.tar.gz && \
tar -xzf sccache.tar.gz --strip-components=1 -C ${SCCACHE_INSTALL_LOCATION} && \
chmod +x ${SCCACHE_INSTALL_LOCATION}/sccache
# Install dependencies
RUN dnf update -y && DEBIAN_FRONTEND=noninteractive dnf install -y \
@@ -83,19 +81,71 @@ ENV compiler_commit=$compiler_commit
RUN sh -c "echo compiler version = '$compiler_version'" && \
sh -c "echo compiler commit = '$compiler_commit'"
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-mainline" ] ) && [ "$compiler_commit" = "" ]; then \
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-staging" ] ) && [ "$compiler_commit" = "" ]; then \
git clone -b "$compiler_version" https://github.com/ROCm/llvm-project.git && \
cd llvm-project && mkdir build && cd build && \
cmake -DCMAKE_INSTALL_PREFIX=/opt/rocm/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm && \
make -j 8 ; \
cmake -G Ninja \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_ENABLE_PROJECTS="clang;lld;clang-tools-extra;flang;mlir" \
-DLLVM_LIT_ARGS="-vv --show-unsupported --show-xfail -j 32" \
-DPACKAGE_VENDOR="AMD" \
-DCMAKE_INSTALL_PREFIX=/home/$USER/rocm/pure_llvm_1.0 \
-DLLVM_ENABLE_ASSERTIONS=ON \
-DLLVM_BUILD_DOCS=ON \
-DLLVM_TARGETS_TO_BUILD=all \
-DLIBOMPTARGET_ENABLE_DEBUG=ON \
-DOFFLOAD_ENABLE_EMISSARY_APIS=OFF \
-DCLANG_DEFAULT_LINKER=lld \
-DCLANG_DEFAULT_PIE_ON_LINUX=0 \
-DLLVM_ENABLE_RUNTIMES="libcxx;libcxxabi;openmp;compiler-rt;libunwind;flang-rt" \
-DLIBCXX_ENABLE_SHARED=OFF \
-DLIBCXX_ENABLE_STATIC=ON \
-DLIBCXX_INSTALL_LIBRARY=OFF \
-DLIBCXX_INSTALL_HEADERS=OFF \
-DLIBCXXABI_ENABLE_SHARED=OFF \
-DLIBCXXABI_ENABLE_STATIC=ON \
-DLIBCXXABI_INSTALL_STATIC_LIBRARY=OFF \
-DLLVM_ENABLE_ASSERTIONS=1 \
-DLLVM_ENABLE_Z3_SOLVER=OFF \
-DLLVM_ENABLE_ZLIB=ON \
-DLLVM_LINK_LLVM_DYLIB=OFF \
-DCLANG_LINK_CLANG_DYLIB=OFF \
../llvm && \
ninja -j16 ; \
else echo "using the release compiler"; \
fi
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-mainline" ] ) && [ "$compiler_commit" != "" ]; then \
RUN if ( [ "$compiler_version" = "develop" ] || [ "$compiler_version" = "amd-staging" ] ) && [ "$compiler_commit" != "" ]; then \
git clone -b "$compiler_version" https://github.com/ROCm/llvm-project.git && \
cd llvm-project && git checkout "$compiler_commit" && echo "checking out commit $compiler_commit" && mkdir build && cd build && \
cmake -DCMAKE_INSTALL_PREFIX=/opt/rocm/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm && \
make -j 8 ; \
cmake -G Ninja \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_ENABLE_PROJECTS="clang;lld;clang-tools-extra;flang;mlir" \
-DLLVM_LIT_ARGS="-vv --show-unsupported --show-xfail -j 32" \
-DPACKAGE_VENDOR="AMD" \
-DCMAKE_INSTALL_PREFIX=/home/$USER/rocm/pure_llvm_1.0 \
-DLLVM_ENABLE_ASSERTIONS=ON \
-DLLVM_BUILD_DOCS=ON \
-DLLVM_TARGETS_TO_BUILD=all \
-DLIBOMPTARGET_ENABLE_DEBUG=ON \
-DOFFLOAD_ENABLE_EMISSARY_APIS=OFF \
-DCLANG_DEFAULT_LINKER=lld \
-DCLANG_DEFAULT_PIE_ON_LINUX=0 \
-DLLVM_ENABLE_RUNTIMES="libcxx;libcxxabi;openmp;compiler-rt;libunwind;flang-rt" \
-DLIBCXX_ENABLE_SHARED=OFF \
-DLIBCXX_ENABLE_STATIC=ON \
-DLIBCXX_INSTALL_LIBRARY=OFF \
-DLIBCXX_INSTALL_HEADERS=OFF \
-DLIBCXXABI_ENABLE_SHARED=OFF \
-DLIBCXXABI_ENABLE_STATIC=ON \
-DLIBCXXABI_INSTALL_STATIC_LIBRARY=OFF \
-DLLVM_ENABLE_ASSERTIONS=1 \
-DLLVM_ENABLE_Z3_SOLVER=OFF \
-DLLVM_ENABLE_ZLIB=ON \
-DLLVM_LINK_LLVM_DYLIB=OFF \
-DCLANG_LINK_CLANG_DYLIB=OFF \
../llvm && \
ninja -j16 ; \
else echo "using the release compiler"; \
fi

370
Jenkinsfile vendored
View File

@@ -81,71 +81,6 @@ def checkoutComposableKernel()
checkout scm
}
// Given a pattern, check if the log contains the pattern and return the context.
def checkForPattern(pattern, log) {
def lines = log.split('\n')
for (int i = 0; i < lines.size(); i++) {
if (lines[i] =~ pattern) {
echo "Found pattern match in log for ${pattern}"
// Get the two lines before and after failure.
def contextStart = Math.max(0, i - 2)
def contextEnd = Math.min(lines.size() - 1, i + 2)
def contextLines = []
for (int j = contextStart; j <= contextEnd; j++) {
contextLines.add(lines[j])
}
return [found: true, matchedLine: lines[i], context: contextLines.join('\n')]
}
}
echo "No pattern match found in log for ${pattern}"
return [found: false, matchedLine: "", context: ""]
}
// Scan the build logs for failures and send notifications.
def sendFailureNotifications() {
// Error patterns to scan build logs for specific failure types and send detailed notifications.
def failurePatterns = [
[pattern: /login attempt to .* failed with status: 401 Unauthorized/, description: "Docker registry authentication failed"],
[pattern: /.*docker login failed.*/, description: "Docker login failed"],
[pattern: /HTTP request sent .* 404 Not Found/, description: "HTTP request failed with 404"],
[pattern: /cat: .* No such file or directory/, description: "GPU not found"],
[pattern: /.*GPU not found.*/, description: "GPU not found"],
[pattern: /Could not connect to Redis at .* Connection timed out/, description: "Redis connection timed out"],
[pattern: /.*unauthorized: your account must log in with a Personal Access Token.*/, description: "Docker login failed"],
[pattern: /.*sccache: error: Server startup failed: Address in use.*/, description: "Sccache Error"]
]
// Get the build log.
def buildLog = sh(script: 'wget -q --no-check-certificate -O - ' + BUILD_URL + 'consoleText', returnStdout: true)
echo "Checking for failure patterns..."
// Check for patterns in the log.
// def foundPatterns = []
// for (patternMap in failurePatterns) {
// def result = checkForPattern(patternMap.pattern, buildLog)
// if (result.found) {
// foundPatterns.add([
// description: patternMap.description,
// matchedLine: result.matchedLine,
// context: result.context
// ])
// }
// }
echo "Done checking for failure patterns..."
// Send a notification for each matched failure pattern.
for (patternMap in foundPatterns) {
withCredentials([string(credentialsId: 'ck_ci_errors_webhook_url', variable: 'WEBHOOK_URL')]) {
sh '''
curl -X POST "${WEBHOOK_URL}" \
-H 'Content-Type: application/json' \
-d '{"text": "\\n\\n**Build Failed**\\n\\n**Issues detected:** ''' + patternMap.description + '''\\n\\n**Log context:**\\n```\\n''' + patternMap.context.replace("'", "\\'") + '''\\n```\\n\\n**Job:** ''' + env.JOB_NAME + '''\\n\\n**Build:** #''' + env.BUILD_NUMBER + '''\\n\\n**URL:** ''' + env.RUN_DISPLAY_URL + '''"}'
'''
}
}
echo "Done failure pattern checking and notifications"
}
def generateAndArchiveBuildTraceVisualization(String buildTraceFileName) {
try {
checkoutComposableKernel()
@@ -479,51 +414,86 @@ def getDockerImage(Map conf=[:]){
return [retimage, image]
}
def buildDocker(install_prefix){
// Build and push a docker image, capturing its digest into the specified env var.
// If forceBuild is false, will skip building if the image already exists in the registry.
def buildAndPushDockerImage(String install_prefix, String image_name, String dockerExtraArgs, boolean forceBuild){
show_node_info()
env.DOCKER_BUILDKIT=1
checkoutComposableKernel()
def image_name = getDockerImageName()
def base_image_name = getBaseDockerImageName()
echo "Building Docker for ${image_name}"
def dockerArgs = "--build-arg PREFIX=${install_prefix} --build-arg CK_SCCACHE='${env.CK_SCCACHE}' --build-arg compiler_version='${params.COMPILER_VERSION}' --build-arg compiler_commit='${params.COMPILER_COMMIT}' --build-arg ROCMVERSION='${params.ROCMVERSION}' "
if(params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-mainline" || params.COMPILER_COMMIT != ""){
dockerArgs = dockerArgs + " --no-cache --build-arg BASE_DOCKER='${base_image_name}' -f projects/composablekernel/Dockerfile.compiler . "
}
else if(params.RUN_AITER_TESTS){
image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_aiter"
dockerArgs = dockerArgs + " --no-cache -f projects/composablekernel/Dockerfile.aiter --build-arg AITER_BRANCH='${params.aiter_branch}' --build-arg CK_AITER_BRANCH='${params.ck_aiter_branch}' . "
}
else if(params.RUN_PYTORCH_TESTS){
image_name = "${env.CK_DOCKERHUB}:ck_pytorch"
dockerArgs = dockerArgs + " --no-cache -f projects/composablekernel/Dockerfile.pytorch --build-arg CK_PYTORCH_BRANCH='${params.ck_pytorch_branch}' . "
}
else{
dockerArgs = dockerArgs + " -f projects/composablekernel/Dockerfile . "
}
echo "Build Args: ${dockerArgs}"
try{
if(params.BUILD_DOCKER || params.RUN_AITER_TESTS || params.RUN_PYTORCH_TESTS){
//force building the new docker if that parameter is true
echo "Building image: ${image_name}"
retimage = docker.build("${image_name}", dockerArgs)
withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
retimage.push()
}
sh 'docker images -q -f dangling=true | xargs --no-run-if-empty docker rmi'
}
else{
def dockerArgs = "--build-arg PREFIX=${install_prefix} --build-arg compiler_version='${params.COMPILER_VERSION}' --build-arg compiler_commit='${params.COMPILER_COMMIT}' --build-arg ROCMVERSION='${params.ROCMVERSION}' "
dockerArgs += " " + dockerExtraArgs
if(!forceBuild){
try{
echo "Checking for image: ${image_name}"
sh "docker manifest inspect --insecure ${image_name}"
echo "Image: ${image_name} found! Skipping building image"
return image_name
}
catch(Exception ex){
echo "Unable to locate image: ${image_name}. Will attempt to build image now."
}
}
catch(Exception ex){
echo "Unable to locate image: ${image_name}. Building image now"
retimage = docker.build("${image_name}", dockerArgs)
withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
retimage.push()
}
echo "Building image: ${image_name} with args: ${dockerArgs}"
def retimage = docker.build("${image_name}", dockerArgs)
withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
retimage.push()
}
def digest = sh(returnStdout: true, script: "docker inspect --format='{{index .RepoDigests 0}}' ${image_name}").trim()
echo "Built image digest: ${digest}"
echo "Pruning dangling Docker images to free disk space on CI agent"
sh "docker image prune -f --filter 'dangling=true' || true"
return digest
}
def buildDockerBase(install_prefix){
def image_name = getDockerImageName()
def base_image_name = getBaseDockerImageName()
echo "Building Docker for ${image_name}"
def dockerExtraArgs = " -f projects/composablekernel/Dockerfile . "
if(params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-staging" || params.COMPILER_COMMIT != ""){
dockerExtraArgs = " --no-cache --build-arg BASE_DOCKER='${base_image_name}' -f projects/composablekernel/Dockerfile.compiler . "
}
else if(params.COMPILER_VERSION == "therock"){
dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile . "
}
env.CK_BASE_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, params.BUILD_DOCKER.toBoolean())
}
def buildDockerPytorch(install_prefix){
def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_pytorch"
def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.pytorch --build-arg CK_PYTORCH_BRANCH='${params.ck_pytorch_branch}' . "
env.CK_PYTORCH_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true)
}
def buildDockerAiter(install_prefix){
def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_aiter"
def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.aiter --build-arg AITER_BRANCH='${params.aiter_branch}' --build-arg CK_AITER_BRANCH='${params.ck_aiter_branch}' . "
env.CK_AITER_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true)
}
def buildDockerFa(install_prefix){
def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_fa"
def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.fa"
dockerExtraArgs += " --build-arg BASE_DOCKER='${params.fa_base_docker}'"
dockerExtraArgs += " --build-arg FA_BRANCH='${params.fa_branch}'"
dockerExtraArgs += " --build-arg CK_FA_BRANCH='${params.ck_fa_branch}'"
dockerExtraArgs += " --build-arg GPU_ARCHS='gfx942;gfx950'"
dockerExtraArgs += " . "
env.CK_FA_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true)
}
def buildDocker(install_prefix){
buildDockerBase(install_prefix)
if (params.RUN_PYTORCH_TESTS.toBoolean()) {
buildDockerPytorch(install_prefix)
}
if (params.RUN_AITER_TESTS.toBoolean()) {
buildDockerAiter(install_prefix)
}
if (params.RUN_FA_TESTS.toBoolean()) {
buildDockerFa(install_prefix)
}
}
@@ -535,10 +505,10 @@ def get_docker_options(){
else{ //only add kfd and dri paths if you actually going to run somthing on GPUs
dockerOpts = "--network=host --device=/dev/kfd --device=/dev/dri --group-add video --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined"
}
if (params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-mainline" || params.COMPILER_COMMIT != ""){
if (params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-staging" || params.COMPILER_VERSION == "therock" || params.COMPILER_COMMIT != ""){
// the --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 env variable is required when building code with offload-compress flag with
// newer clang22 compilers and running with older hip runtima libraries
dockerOpts = dockerOpts + " --env HIP_CLANG_PATH='/llvm-project/build/bin' --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 "
dockerOpts = dockerOpts + " --env HIP_CLANG_PATH='/llvm-project/build/bin' --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 --env HIP_PLATFORM=amd "
}
// on some machines the group ids for video and render groups may not be the same as in the docker image!
def video_id = sh(returnStdout: true, script: 'getent group video | cut -d: -f3')
@@ -1148,99 +1118,73 @@ def process_results(Map conf=[:]){
}
}
def run_aiter_tests(Map conf=[:]){
def run_downstream_tests(Map conf=[:]){
show_node_info()
checkoutComposableKernel()
//use the latest pytorch image
def image = "${env.CK_DOCKERHUB_PRIVATE}:ck_aiter"
def dockerOpts=get_docker_options() + ' --group-add irc '
def dockerOpts = get_docker_options() + ' --group-add irc '
gitStatusWrapper(credentialsId: "${env.ck_git_creds}", gitHubContext: "${env.STAGE_NAME}", account: 'ROCm', repo: 'rocm-libraries') {
try
{
echo "Pulling image: ${image}"
retimage = docker.image("${image}")
echo "Pulling image: ${conf.image}"
retimage = docker.image("${conf.image}")
withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
retimage.pull()
}
}
catch(Exception ex)
{
error "Unable to locate image: ${image}"
error "Unable to locate image: ${conf.image}"
}
}
withDockerContainer(image: image, args: dockerOpts) {
timeout(time: 5, unit: 'HOURS'){
withDockerContainer(image: conf.image, args: dockerOpts) {
timeout(time: conf.get("timeoutHours", 2), unit: 'HOURS'){
try{
sh "rocminfo"
sh "python3 --version"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_mha.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_mha_varlen.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_batch_prefill.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_2stage.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_blockscale.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_ep.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting_mxfp4.py"
sh "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_tkw1.py"
for (cmd in conf.execute_cmds) {
sh "${cmd}"
}
}
catch(e){
echo "Throwing error exception while running AITER tests"
echo "Throwing error exception while running ${env.STAGE_NAME}"
echo 'Exception occurred: ' + e.toString()
throw e
}
finally{
echo "Finished running AITER tests"
echo "Finished running ${env.STAGE_NAME}"
}
}
}
}
def run_pytorch_tests(Map conf=[:]){
show_node_info()
checkoutComposableKernel()
//use the latest pytorch-nightly image
def image = "${env.CK_DOCKERHUB}:ck_pytorch"
def dockerOpts=get_docker_options() + ' --group-add irc '
gitStatusWrapper(credentialsId: "${env.ck_git_creds}", gitHubContext: "${env.STAGE_NAME}", account: 'ROCm', repo: 'rocm-libraries') {
try
{
echo "Pulling image: ${image}"
retimage = docker.image("${image}")
withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
retimage.pull()
}
}
catch(Exception ex)
{
error "Unable to locate image: ${image}"
}
}
withDockerContainer(image: image, args: dockerOpts) {
timeout(time: 2, unit: 'HOURS'){
try{
sh "rocminfo"
sh "python3 --version"
sh "python3 /tmp/pytorch/tools/amd_build/build_amd.py"
sh "USE_ROCM_CK_SDPA=1 PYTORCH_ROCM_ARCH=gfx942 python /tmp/pytorch/setup.py develop"
}
catch(e){
echo "Throwing error exception while building Pytorch"
echo 'Exception occurred: ' + e.toString()
throw e
}
finally{
echo "Finished building Pytorch"
}
}
}
def getPytorchTestsCmds() {
return [
"python3 /tmp/pytorch/tools/amd_build/build_amd.py",
"USE_ROCM_CK_SDPA=1 PYTORCH_ROCM_ARCH=gfx942 python /tmp/pytorch/setup.py develop"
]
}
def getAiterTestsCmds() {
return [
"python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_mha.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_mha_varlen.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_batch_prefill.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe_2stage.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe_blockscale.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe_ep.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting_mxfp4.py",
"python3 /home/jenkins/workspace/aiter/op_tests/test_moe_tkw1.py"
]
}
def getFaTestsCmds() {
return [
"python3 -u -m pytest /home/jenkins/workspace/flash-attention/tests/test_flash_attn_ck.py"
]
}
//launch develop branch daily jobs
@@ -1248,15 +1192,20 @@ CRON_SETTINGS = BRANCH_NAME == "develop" ? '''0 23 * * * % RUN_FULL_QA=true;RUN_
0 22 * * * % RUN_FULL_QA=true;DISABLE_DL_KERNELS=true;RUN_TILE_ENGINE_BASIC_TESTS=true;RUN_TILE_ENGINE_GEMM_TESTS=true;RUN_PERFORMANCE_TESTS=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true
0 21 * * * % RUN_GROUPED_CONV_LARGE_CASES_TESTS=true;hipTensor_test=true;BUILD_GFX101=false;BUILD_GFX908=false;BUILD_GFX942=true;BUILD_GFX950=true;RUN_PERFORMANCE_TESTS=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true;BUILD_PACKAGES=true
0 19 * * * % BUILD_DOCKER=true;COMPILER_VERSION=develop;BUILD_COMPILER=/llvm-project/build/bin/clang++;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true
0 17 * * * % BUILD_DOCKER=true;COMPILER_VERSION=amd-mainline;BUILD_COMPILER=/llvm-project/build/bin/clang++;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true
0 15 * * * % BUILD_INSTANCES_ONLY=true;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;FORCE_CI=true
0 13 * * * % RUN_FULL_CONV_TILE_TESTS=true;RUN_AITER_TESTS=true;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;FORCE_CI=true
0 11 * * * % RUN_PYTORCH_TESTS=true;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;BUILD_GFX101=false;BUILD_GFX103=false;BUILD_GFX11=false;BUILD_GFX12=false;BUILD_GFX90A=false;FORCE_CI=true''' : ""
0 17 * * * % BUILD_DOCKER=true;COMPILER_VERSION=therock;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true
0 15 * * * % BUILD_DOCKER=true;COMPILER_VERSION=amd-staging;BUILD_COMPILER=/llvm-project/build/bin/clang++;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true
0 13 * * * % BUILD_INSTANCES_ONLY=true;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;FORCE_CI=true
0 11 * * * % RUN_FULL_CONV_TILE_TESTS=true;RUN_AITER_TESTS=true;RUN_FA_TESTS=true;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;FORCE_CI=true
0 9 * * * % RUN_PYTORCH_TESTS=true;USE_SCCACHE=false;RUN_PERFORMANCE_TESTS=false;BUILD_GFX101=false;BUILD_GFX103=false;BUILD_GFX11=false;BUILD_GFX12=false;BUILD_GFX90A=false;FORCE_CI=true''' : ""
CURRENT_BRANCH_NAME = env.CHANGE_BRANCH ? env.CHANGE_BRANCH : env.BRANCH_NAME
POLL_SPEC = BRANCH_NAME == "develop" ? 'H H/6 * * *' : ''
pipeline {
agent none
triggers {
parameterizedCron(CRON_SETTINGS)
pollSCM(POLL_SPEC)
}
options {
skipDefaultCheckout()
@@ -1278,7 +1227,7 @@ pipeline {
string(
name: 'COMPILER_VERSION',
defaultValue: '',
description: 'Specify which version of compiler to use: release, develop, amd-mainline, or leave blank (default).')
description: 'Specify which version of compiler to use: develop, amd-staging, therock, or leave blank (default).')
string(
name: 'COMPILER_COMMIT',
defaultValue: '',
@@ -1381,8 +1330,8 @@ pipeline {
description: "Build CK and run tests on gfx12 (default: ON)")
booleanParam(
name: "NINJA_BUILD_TRACE",
defaultValue: false,
description: "Generate a ninja build trace (default: OFF)")
defaultValue: true,
description: "Generate a ninja build trace (default: ON)")
booleanParam(
name: "NINJA_FTIME_TRACE",
defaultValue: false,
@@ -1409,8 +1358,8 @@ pipeline {
description: "Try building PYTORCH with latest CK develop branch (default: OFF)")
string(
name: 'ck_pytorch_branch',
defaultValue: 'develop',
description: 'Specify which branch of CK to test with Pytorch (default: develop)')
defaultValue: CURRENT_BRANCH_NAME,
description: 'Specify which branch of CK to test with Pytorch (default: current branch)')
booleanParam(
name: "RUN_AITER_TESTS",
defaultValue: false,
@@ -1425,8 +1374,24 @@ pipeline {
description: 'Specify which branch of AITER to use (default: main)')
string(
name: 'ck_aiter_branch',
defaultValue: 'develop',
description: 'Specify which branch of CK to test with AITER (default: develop)')
defaultValue: CURRENT_BRANCH_NAME,
description: 'Specify which branch of CK to test with AITER (default: current branch)')
booleanParam(
name: "RUN_FA_TESTS",
defaultValue: false,
description: "Run Flash Attention tests with latest CK develop branch (default: OFF)")
string(
name: 'fa_base_docker',
defaultValue: 'rocm/pytorch:rocm7.1.1_ubuntu24.04_py3.12_pytorch_release_2.9.1',
description: 'Specify which base docker image to use for flash-attention tests')
string(
name: 'fa_branch',
defaultValue: 'ck_improve_main',
description: 'Specify which branch of flash-attention to use (default: ck_improve_main)')
string(
name: 'ck_fa_branch',
defaultValue: CURRENT_BRANCH_NAME,
description: 'Specify which branch of CK to test with flash-attention (default: current branch)')
booleanParam(
name: "FORCE_CI",
defaultValue: false,
@@ -1519,7 +1484,7 @@ pipeline {
}
}
}
stage("Run Pytorch Tests")
stage("Run Downstream Tests")
{
when {
beforeAgent true
@@ -1535,20 +1500,10 @@ pipeline {
}
agent{ label rocmnode("gfx942")}
steps{
run_pytorch_tests()
run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: getPytorchTestsCmds())
cleanWs()
}
}
}
}
stage("Run AITER Tests")
{
when {
beforeAgent true
expression { env.SHOULD_RUN_CI.toBoolean() }
}
parallel
{
stage("Run AITER Tests on gfx942")
{
when {
@@ -1557,7 +1512,7 @@ pipeline {
}
agent{ label rocmnode("gfx942")}
steps{
run_aiter_tests()
run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: getAiterTestsCmds())
cleanWs()
}
}
@@ -1569,7 +1524,31 @@ pipeline {
}
agent{ label rocmnode("gfx950")}
steps{
run_aiter_tests()
run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: getAiterTestsCmds())
cleanWs()
}
}
stage("Run FA Tests on gfx942")
{
when {
beforeAgent true
expression { params.RUN_FA_TESTS.toBoolean() }
}
agent{ label rocmnode("gfx942")}
steps{
run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: getFaTestsCmds())
cleanWs()
}
}
stage("Run FA Tests on gfx950")
{
when {
beforeAgent true
expression { params.RUN_FA_TESTS.toBoolean() }
}
agent{ label rocmnode("gfx950")}
steps{
run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: getFaTestsCmds())
cleanWs()
}
}
@@ -2109,7 +2088,10 @@ pipeline {
description: 'Some checks have failed'
node(rocmnode("nogpu")) {
script {
sendFailureNotifications()
checkoutComposableKernel()
}
withCredentials([string(credentialsId: 'ck_ci_errors_webhook_url', variable: 'WEBHOOK_URL')]) {
sh 'bash projects/composablekernel/script/infra_helper/send_failure_notifications.sh'
}
}
}

4
codegen/test/grouped_conv_fwd_multiple_d_v1.cpp
View File

@@ -198,10 +198,6 @@ struct Epilogue
input_left_pads,
input_right_pads);
// auto res = rtc::from_gpu(out_dev);
// pass &= ck::utils::check_err(res, out_host, "Error: incorrect results!", 1e-5f, 1e-4f);
// assert(pass);
// Simple check: this checks that the output from each instance matches the output from the
// first instance
CHECK(report(solution, check(rtc::from_gpu(out_dev))));

4
codegen/test/grouped_conv_fwd_multiple_d_v2.cpp
View File

@@ -198,10 +198,6 @@ struct Epilogue
input_left_pads,
input_right_pads);
// auto res = rtc::from_gpu(out_dev);
// pass &= ck::utils::check_err(res, out_host, "Error: incorrect results!", 1e-5f, 1e-4f);
// assert(pass);
// Simple check: this checks that the output from each instance matches the output from the
// first instance
CHECK(report(solution, check(rtc::from_gpu(out_dev))));

4
codegen/test/grouped_conv_fwd_multiple_d_v3.cpp
View File

@@ -198,10 +198,6 @@ struct Epilogue
input_left_pads,
input_right_pads);
// auto res = rtc::from_gpu(out_dev);
// pass &= ck::utils::check_err(res, out_host, "Error: incorrect results!", 1e-5f, 1e-4f);
// assert(pass);
// Simple check: this checks that the output from each instance matches the output from the
// first instance
CHECK(report(solution, check(rtc::from_gpu(out_dev))));

4
codegen/test/grouped_conv_fwd_multiple_d_v4.cpp
View File

@@ -198,10 +198,6 @@ struct Epilogue
input_left_pads,
input_right_pads);
// auto res = rtc::from_gpu(out_dev);
// pass &= ck::utils::check_err(res, out_host, "Error: incorrect results!", 1e-5f, 1e-4f);
// assert(pass);
// Simple check: this checks that the output from each instance matches the output from the
// first instance
CHECK(report(solution, check(rtc::from_gpu(out_dev))));

210
dispatcher/README.md
View File

@@ -1,6 +1,6 @@
# CK Tile Dispatcher
A unified kernel dispatch system for AMD GPUs with C++ and Python frontends.
A unified kernel dispatch system for AMD GPUs with C++ and Python frontends, supporting GEMM and Grouped Convolution operations.
**Validated Platform:** AMD Instinct MI300 series (gfx942)
@@ -154,6 +154,8 @@ rocminfo | grep -i "gfx"
### Install Python Dependencies
#### Core Dependencies (Required)
NumPy is required for Python examples and kernel generation. We recommend using a virtual environment:
**Option 1: Using standard venv**
@@ -165,8 +167,8 @@ python3 -m venv .venv
source .venv/bin/activate # Linux/macOS
# .venv\Scripts\activate # Windows
# Install NumPy
pip install numpy
# Install core dependencies
pip install -r python/requirements.txt
```
**Option 2: Using uv (faster alternative)**
@@ -179,17 +181,38 @@ uv venv .venv
source .venv/bin/activate # Linux/macOS
# .venv\Scripts\activate # Windows
# Install NumPy
uv pip install numpy
# Install core dependencies
uv pip install -r python/requirements.txt
```
**Option 3: System-wide install (not recommended)**
```bash
pip install numpy
pip install -r python/requirements.txt
```
> **Note:** Always activate your virtual environment before running CMake or Python examples.
#### ML Heuristics Dependencies (Optional)
For ML-based kernel selection (examples 09-11), install additional dependencies:
```bash
# Activate your virtual environment first
source .venv/bin/activate
# Install ML dependencies (LightGBM, pandas, pyarrow, scikit-learn)
pip install -r requirements-ml.txt
```
**Why separate?** ML dependencies are large (especially pyarrow) and not needed for basic dispatcher usage. Install only if you need:
- ML-based kernel selection (`examples/gemm/python/09_ml_heuristic.py`)
- Model training (`heuristics/train.py`)
- Model evaluation (`heuristics/evaluate.py`)
- Automated benchmark analysis
**Core dependencies:** ~50 MB (NumPy only)
**With ML dependencies:** ~500 MB (includes LightGBM, pandas, pyarrow, scikit-learn)
### Supported Data Types
CK Tile supports a wide range of data types for GEMM operations:
@@ -319,8 +342,8 @@ ls examples/libdispatcher_gemm_lib.so
| `CMAKE_PREFIX_PATH` | - | ROCm installation path |
| `CMAKE_CXX_COMPILER` | - | Path to hipcc compiler |
⚠️ **Important:** Always use `-DCMAKE_BUILD_TYPE=Release` for benchmarking. Debug builds are slower.
⚠️ **Important:** Note that the current system provides single GPU target support for architecture-based kernel filtering, please do not use multiple GPU targets at a time (if necessary, please compile into different build directories).
WARNING: **Important:** Always use `-DCMAKE_BUILD_TYPE=Release` for benchmarking. Debug builds are slower.
WARNING: **Important:** Note that the current system provides single GPU target support for architecture-based kernel filtering, please do not use multiple GPU targets at a time (if necessary, please compile into different build directories).
---
@@ -340,6 +363,15 @@ cd build/examples
./gemm_04_heuristics # Heuristic kernel selection
./gemm_05_json_export # Registry JSON export
./gemm_06_multi_registry # Multiple registries
# Grouped Convolution Examples
./grouped_conv_01_basic # Declaration patterns + GPU execution
./grouped_conv_02_all_dirs # Forward/BwdData/BwdWeight with GPU
./grouped_conv_03_bench_val # Benchmark + CPU reference validation
./grouped_conv_04_registry_json # Heuristic selection + JSON export
./grouped_conv_05_bwd_data # Backward data + CPU validation
./grouped_conv_06_bwd_weight # Backward weight + CPU validation
./grouped_conv_07_benchmark # Multi-tile ResNet benchmark
```
### Python Examples
@@ -352,8 +384,16 @@ cd /path/to/composable_kernel/dispatcher
# GEMM Examples
python3 examples/gemm/python/01_basic_gemm.py # Basic multi-kernel GEMM
python3 examples/gemm/python/04_validation.py # CPU reference validation
python3 examples/gemm/python/07_stress_test.py # Stress test (48 kernels)
python3 examples/gemm/python/07_stress_test.py # Stress test
python3 examples/gemm/python/08_heuristics.py # Heuristic selection
# Grouped Convolution Examples
python3 examples/grouped_conv/python/01_basic_grouped_conv.py # Config patterns + registry + GPU
python3 examples/grouped_conv/python/02_forward.py # Forward 2D/3D + CPU ref
python3 examples/grouped_conv/python/03_bwd_data.py # Backward data + CPU ref
python3 examples/grouped_conv/python/04_bwd_weight.py # Backward weight + CPU ref
python3 examples/grouped_conv/python/05_benchmark.py # Multi-problem benchmark
python3 examples/grouped_conv/python/06_registry_json.py # Heuristic selection + JSON
```
### Example Output
@@ -470,6 +510,42 @@ python3 examples/gemm/python/10_advanced_benchmark.py \
---
## ML-Based Kernel Selection (Optional)
The dispatcher includes ML heuristics for automated kernel selection using trained LightGBM models.
**Prerequisites:** Install ML dependencies first:
```bash
pip install -r requirements-ml.txt # ~500 MB (LightGBM, pandas, pyarrow, scikit-learn)
```
**Documentation:** See [heuristics/README.md](heuristics/README.md) for:
- Training and evaluating models
- Feature engineering (72 features)
- Using pre-trained models
- Python API reference
**Examples:**
```bash
python3 examples/gemm/python/09_ml_heuristic.py # ML-based kernel selection
python3 examples/gemm/python/10_rank_kernels.py # Kernel ranking
```
**Model Compression:** Trained models are stored in compressed `.lgbm.gz` format to save space (~67% size reduction). Python tools automatically decompress models on first use. For C++ examples, decompress manually:
```bash
# If you have compressed models
cd heuristics/models/gemm_universal_fp16_gfx950
gunzip model_tflops.lgbm.gz
# Then use in C++ example
cd ../../../build
./gemm_09_ml_heuristic --model ../heuristics/models/gemm_universal_fp16_gfx950/model_tflops.lgbm
```
---
## External Integration
### Using Dispatcher in Your Own Project
@@ -588,7 +664,7 @@ lib = DispatcherLib.load("/absolute/path/to/libdispatcher_gemm_lib.so")
### Data Flow
```
KernelConfig → Registry → Dispatcher → GPU Execution
KernelConfig -> Registry -> Dispatcher -> GPU Execution
```
1. **KernelConfig**: Defines kernel parameters (tile sizes, data types, layouts)
@@ -784,31 +860,49 @@ make -j$(nproc)
```
dispatcher/
├── README.md # This file
├── CMakeLists.txt # Build configuration
│
├── include/ck_tile/dispatcher/ # C++ headers
│ ├── dispatcher.hpp # GEMM dispatcher
│ ├── registry.hpp # Kernel registry
│ └── kernel_key.hpp # Kernel configuration
│
├── src/ # C++ implementation
│
├── codegen/ # Kernel generation
│ ├── unified_gemm_codegen.py # GEMM kernel generator
│ └── arch_specs.json # GPU specifications
│
├── bindings/ctypes/ # Python ctypes interface
│ └── gemm_ctypes_lib.cpp # GEMM Python library
│
├── examples/ # Examples
│ └── gemm/
│ ├── cpp/ # C++ GEMM examples (01-06)
│ └── python/ # Python GEMM examples (01-11)
│
├── scripts/ # Build scripts
│
└── tests/ # Unit tests
|---- README.md # This file
|---- CMakeLists.txt # Build configuration
|
|---- include/ck_tile/dispatcher/ # C++ headers
| |---- dispatcher.hpp # Main dispatcher include
| |---- registry.hpp # GEMM kernel registry
| |---- kernel_key.hpp # Kernel configuration
| |---- grouped_conv_config.hpp # Grouped conv configuration
| |---- grouped_conv_problem.hpp # Grouped conv problem (with builder)
| |---- grouped_conv_kernel_decl.hpp # Grouped conv kernel declarations
| |---- grouped_conv_registry.hpp # Grouped conv registry (thread-safe)
| +---- grouped_conv_utils.hpp # Grouped conv utilities
|
|---- src/ # C++ implementation
|
|---- codegen/ # Kernel generation
| |---- codegen_common.py # Shared: TileConfig, TraitConfigBase, type mappings
| |---- unified_gemm_codegen.py # GEMM kernel generator
| |---- unified_grouped_conv_codegen.py # Grouped conv kernel generator
| +---- arch_specs.json # GPU specifications
|
|---- python/ # Python utilities
| |---- dispatcher_common.py # Shared: paths, validation, Colors, phased output
| |---- ctypes_utils.py # GEMM ctypes utilities
| +---- grouped_conv_utils.py # Grouped conv utilities
|
|---- scripts/ # Build scripts
| |---- compile_gemm_examples.py # GEMM build script
| +---- compile_grouped_conv_examples.py # Grouped conv build script
|
|---- bindings/ctypes/ # Python ctypes interface
| |---- gemm_ctypes_lib.cpp # GEMM Python library
| +---- conv_ctypes_lib.cpp # Grouped conv Python library
|
|---- examples/ # Examples
| |---- gemm/
| | |---- cpp/ # C++ GEMM examples (01-07)
| | +---- python/ # Python GEMM examples (01-11)
| +---- grouped_conv/
| |---- cpp/ # C++ Grouped Conv examples (01-07)
| +---- python/ # Python Grouped Conv examples (01-06)
|
+---- tests/ # Unit tests (C++ and Python)
```
---
@@ -820,17 +914,49 @@ dispatcher/
| GEMM C++ | [examples/gemm/cpp/README.md](examples/gemm/cpp/README.md) |
| GEMM Python | [examples/gemm/python/README.md](examples/gemm/python/README.md) |
| Codegen | [codegen/README.md](codegen/README.md) |
| Python Utils | [python/README.md](python/README.md) |
| C++ Headers | [include/ck_tile/dispatcher/README.md](include/ck_tile/dispatcher/README.md) |
---
## Archived Content
## Grouped Convolution Support
Convolution examples and utilities have been archived to `ck-2/conv_archive/dispatcher/`:
- `examples/conv/cpp/` - 11 C++ convolution examples
- `examples/conv/python/` - 14 Python convolution examples
- `codegen/unified_conv_codegen.py` - Conv kernel generator
- `include/ck_tile/dispatcher/conv_*.hpp` - Conv headers
- `python/conv_utils.py` - Conv Python utilities
Grouped convolution is fully supported alongside GEMM, with shared infrastructure to eliminate duplication.
### Python
```bash
# Generate grouped conv kernels
python3 codegen/unified_grouped_conv_codegen.py \
--output-dir build/generated_kernels \
--datatype fp16 --variant forward --ndim-spatial 2
# Build grouped conv examples
python3 scripts/compile_grouped_conv_examples.py examples/grouped_conv/cpp/01_basic_grouped_conv.cpp
```
### Key Files
| Component | File |
|-----------|------|
| C++ Headers | `include/ck_tile/dispatcher/grouped_conv_*.hpp` |
| Python Codegen | `codegen/unified_grouped_conv_codegen.py` |
| Python Utils | `python/grouped_conv_utils.py` |
| Build Script | `scripts/compile_grouped_conv_examples.py` |
| Shared Codegen | `codegen/codegen_common.py` |
| Shared Utils | `python/dispatcher_common.py` |
### Variants
- **Forward** (`grouped_conv_fwd`) - Standard grouped convolution
- **Backward Data** (`grouped_conv_bwd_data`) - Gradient w.r.t. input
- **Backward Weight** (`grouped_conv_bwd_weight`) - Gradient w.r.t. weights
### Shared Infrastructure
GEMM and grouped convolution share common code to avoid duplication:
- `codegen/codegen_common.py` - TileConfig, TraitConfigBase, type mappings, parallel generation, arch-aware expansion
- `python/dispatcher_common.py` - Path helpers, validation, auto-correction, Colors, phased output
---

24
dispatcher/bindings/README.md
View File

@@ -6,13 +6,13 @@ This directory contains language bindings for the CK Tile Dispatcher.
```
bindings/
├── ctypes/ # Python ctypes bindings (C API)
│ ├── gemm_ctypes_lib.cpp # GEMM dispatcher C API
│ ├── conv_ctypes_lib.cpp # Convolution dispatcher C API (fwd + bwd_data)
│ ├── conv_bwdw_ctypes_lib.cpp # Convolution backward weight C API
│ ├── gpu_helper.cpp # CLI helper for Python
│ └── CMakeLists.txt
└── README.md
|---- ctypes/ # Python ctypes bindings (C API)
| |---- gemm_ctypes_lib.cpp # GEMM dispatcher C API
| |---- conv_ctypes_lib.cpp # Grouped conv dispatcher C API (fwd + bwd_data)
| |---- conv_bwdw_ctypes_lib.cpp # Grouped conv backward weight C API (separate library)
| |---- gpu_helper.cpp # CLI helper for Python
| +---- CMakeLists.txt
+---- README.md
```
## ctypes Bindings
@@ -65,7 +65,7 @@ lib.dispatcher_cleanup()
| `dispatcher_export_registry_json()` | Export registry as JSON |
| `dispatcher_cleanup()` | Release resources |
### Convolution API
### Grouped Convolution API
| Function | Description |
|----------|-------------|
@@ -105,5 +105,11 @@ Output is JSON for easy parsing:
See the examples that use these bindings:
- **GEMM**: `dispatcher/examples/gemm/python/`
- **Conv**: `dispatcher/examples/conv/python/`
### Grouped Convolution
Grouped convolution C++ headers and Python utilities are in:
- **C++ Headers**: `dispatcher/include/ck_tile/dispatcher/grouped_conv_*.hpp`
- **Python Utils**: `dispatcher/python/grouped_conv_utils.py`
- **Build Script**: `dispatcher/scripts/compile_grouped_conv_examples.py`

4
dispatcher/bindings/ctypes/CMakeLists.txt
View File

@@ -78,7 +78,7 @@ endif()
# Look for forward kernels
file(GLOB CONV_FWD_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/conv_fwd_*.hpp")
# Look for backward data kernels
file(GLOB CONV_BWDD_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/conv_bwdd_*.hpp")
file(GLOB CONV_BWDD_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/conv_bwd_data_*.hpp")
# Fallback: any conv kernel (for backwards compatibility)
file(GLOB CONV_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/conv_*.hpp")
@@ -112,7 +112,7 @@ endif()
# Add backward data kernel if available
if(CONV_BWDD_KERNEL_HEADERS)
list(GET CONV_BWDD_KERNEL_HEADERS 0 CONV_BWDD_KERNEL_HEADER)
message(STATUS "Found Conv BWD_DATA kernel for ctypes lib: ${CONV_BWDD_KERNEL_HEADER}")
message(STATUS "Found Conv BWD_DATA kernel for ctypes lib: ${CONV_BWD_DATA_KERNEL_HEADER}")
target_compile_options(dispatcher_conv_lib PRIVATE -include ${CONV_BWDD_KERNEL_HEADER})
target_compile_definitions(dispatcher_conv_lib PRIVATE CONV_BWD_DATA_AVAILABLE)
endif()

4
dispatcher/bindings/ctypes/conv_bwdw_ctypes_lib.cpp
View File

@@ -53,6 +53,7 @@ struct ConvBwdwProblemC
int stride_d, stride_h, stride_w;
int pad_d, pad_h, pad_w;
int dilation_d, dilation_h, dilation_w;
int split_k;
};
// =============================================================================
@@ -108,8 +109,7 @@ static float run_bwd_weight_impl(const void* input_ptr,
grad_weight_ptr, // wei_ptr = grad_weight (output)
{}, // ds_ptr
grad_output_ptr, // out_ptr = grad_output
1 // k_batch
);
(prob->split_k > 1) ? prob->split_k : 1);
ck_tile::stream_config stream_cfg{static_cast<hipStream_t>(stream), true, 1, 3, 10};

643
dispatcher/bindings/ctypes/conv_ctypes_lib.cpp
View File

@@ -1,128 +1,46 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Convolution Dispatcher ctypes Library
*
* Provides C API for Python ctypes integration.
* Supports forward convolution. Backward operations require additional headers.
*
* REQUIRED: Forward kernel header must be force-included via -include flag.
* OPTIONAL: Backward kernels can be added with CONV_BWD_DATA_AVAILABLE/CONV_BWD_WEIGHT_AVAILABLE
*
* Usage from Python:
* lib = ctypes.CDLL("libdispatcher_conv.so")
* lib.conv_dispatcher_init()
* lib.conv_dispatcher_run(...)
*/
//
// Multi-kernel grouped convolution dispatcher for Python ctypes.
//
// Supports: forward / backward-data / backward-weight x 2D / 3D
//
// The dispatch header (conv_python_dispatch.hpp) is force-included via
// -include and brings in ALL compiled kernels with these aliases:
//
// 2D launchers (from include_all headers):
// SelectedConvKernelLauncher (forward 2D)
// SelectedConvBwdDataLauncher (backward-data 2D)
// SelectedConvBwdWeightLauncher (backward-weight 2D)
//
// 3D launchers (from dispatch header):
// ConvFwd3dLauncher (forward 3D)
// ConvBwdData3dLauncher (backward-data 3D)
// ConvBwdWeight3dLauncher (backward-weight 3D)
//
// Usage from Python:
// lib = ctypes.CDLL("libdispatcher_conv_lib.so")
// lib.conv_dispatcher_init()
// lib.conv_dispatcher_run(input, weight, output, &problem, stream)
#include <cstring>
#include <memory>
#include <vector>
#include <stdexcept>
#include <hip/hip_runtime.h>
#include "ck_tile/dispatcher/conv_utils.hpp"
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
using namespace ck_tile::dispatcher;
// Global state (using shared_ptr for safe memory management)
static std::shared_ptr<ConvRegistry> g_registry = nullptr;
static std::shared_ptr<ConvDispatcher> g_dispatcher = nullptr;
static std::vector<const ConvKernelInstance*> g_kernels;
extern "C" {
// =============================================================================
// Initialization
// =============================================================================
int conv_dispatcher_init()
// =========================================================================
// Problem definition (matches Python ctypes struct exactly)
// =========================================================================
enum ConvDirection
{
if(g_registry)
return 0; // Already initialized
g_registry = std::make_shared<ConvRegistry>();
g_dispatcher = std::make_shared<ConvDispatcher>(g_registry.get());
// Register kernel configurations using simple ConvKernelSet
// (actual kernel launch uses the force-included SelectedConvKernelLauncher)
using namespace ck_tile::dispatcher::conv_decl;
// Forward kernels (required - must be force-included)
// Must match: conv_fwd_fp16_nhwgc_2d_compv4_cshuffle_intrawave_128x128x64_2x2x1_32x32x16_dsb
ConvKernelSet fwd_set;
fwd_set.add(ConvSignature().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
ConvAlgorithm()
.tile(128, 128, 64) // tile_m x tile_n x tile_k
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv4")
.scheduler("intrawave"),
"gfx942");
g_registry->register_set(fwd_set, ConvRegistry::Priority::High);
#ifdef CONV_BWD_DATA_AVAILABLE
// Backward data kernels
// Must match: conv_bwdd_fp16_nhwgc_2d_compv3_cshuffle_intrawave_128x128x64_2x2x1_32x32x16
ConvKernelSet bwd_data_set;
bwd_data_set.add(ConvSignature().dtype("fp16").layout("nhwgc").conv_type("bwd_data").dims(2),
ConvAlgorithm()
.tile(128, 128, 64) // tile_m x tile_n x tile_k
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave"),
"gfx942");
g_registry->register_set(bwd_data_set, ConvRegistry::Priority::High);
#endif
return 0;
}
int conv_dispatcher_cleanup()
{
// shared_ptr automatically handles cleanup when reset
g_dispatcher.reset();
g_registry.reset();
g_kernels.clear();
return 0;
}
// =============================================================================
// Registry Management
// =============================================================================
int conv_dispatcher_get_kernel_count()
{
if(!g_registry)
return 0;
return static_cast<int>(g_registry->size());
}
int conv_dispatcher_get_kernel_name(int index, char* buffer, int buffer_size)
{
if(index < 0 || !buffer || buffer_size <= 0)
return -1;
if(!g_registry)
return -1;
// Use registry to get kernel names (they are registered with full names)
const auto& kernels = g_registry->all_kernels();
if(static_cast<size_t>(index) >= kernels.size())
return -1;
const auto* kernel = kernels[index];
std::strncpy(buffer, kernel->name().c_str(), buffer_size - 1);
buffer[buffer_size - 1] = '\0';
return 0;
}
// =============================================================================
// Problem Definition
// =============================================================================
CONV_FORWARD = 0,
CONV_BWD_DATA = 1,
CONV_BWD_WEIGHT = 2
};
struct ConvProblemC
{
@@ -132,267 +50,33 @@ struct ConvProblemC
int stride_d, stride_h, stride_w;
int pad_d, pad_h, pad_w;
int dilation_d, dilation_h, dilation_w;
int direction; // 0=forward, 1=bwd_data, 2=bwd_weight
int direction;
int split_k;
};
// =============================================================================
// Kernel Selection
// =============================================================================
// =========================================================================
// Initialization / lifecycle
// =========================================================================
int conv_dispatcher_init() { return 0; }
int conv_dispatcher_cleanup() { return 0; }
int conv_dispatcher_is_supported(const ConvProblemC* prob)
{
if(!g_registry || !prob)
return 0;
ConvProblem problem;
problem.N = prob->N;
problem.G = prob->G;
problem.C = prob->C;
problem.K = prob->K;
problem.input_spatial = {prob->input_d, prob->input_h, prob->input_w};
problem.filter_spatial = {prob->filter_z, prob->filter_y, prob->filter_x};
problem.stride = {prob->stride_d, prob->stride_h, prob->stride_w};
problem.padding = {prob->pad_d, prob->pad_h, prob->pad_w};
problem.dilation = {prob->dilation_d, prob->dilation_h, prob->dilation_w};
problem.op = static_cast<ConvOp>(prob->direction);
problem.compute_output_size();
const auto* kernel = g_dispatcher->select(problem);
return kernel ? 1 : 0;
}
int conv_dispatcher_select_kernel(const ConvProblemC* prob, char* kernel_name, int buffer_size)
{
if(!g_registry || !prob || !kernel_name || buffer_size <= 0)
return -1;
ConvProblem problem;
problem.N = prob->N;
problem.G = prob->G;
problem.C = prob->C;
problem.K = prob->K;
problem.input_spatial = {prob->input_d, prob->input_h, prob->input_w};
problem.filter_spatial = {prob->filter_z, prob->filter_y, prob->filter_x};
problem.stride = {prob->stride_d, prob->stride_h, prob->stride_w};
problem.padding = {prob->pad_d, prob->pad_h, prob->pad_w};
problem.dilation = {prob->dilation_d, prob->dilation_h, prob->dilation_w};
problem.op = static_cast<ConvOp>(prob->direction);
problem.compute_output_size();
const auto* kernel = g_dispatcher->select(problem);
if(!kernel)
return -1;
std::strncpy(kernel_name, kernel->name().c_str(), buffer_size - 1);
kernel_name[buffer_size - 1] = '\0';
return 0;
}
// =============================================================================
// Convolution Execution
// =============================================================================
// Helper to build ConvParam
static ck_tile::conv::ConvParam build_conv_param(const ConvProblemC* prob)
{
// Determine if this is 2D or 3D convolution
const bool is_3d = (prob->input_d > 1 || prob->filter_z > 1);
if(is_3d)
{
// 3D convolution: use all spatial dimensions
return ck_tile::conv::ConvParam{3,
prob->G,
prob->N,
prob->K,
prob->C,
{prob->filter_z, prob->filter_y, prob->filter_x},
{prob->input_d, prob->input_h, prob->input_w},
{prob->stride_d, prob->stride_h, prob->stride_w},
{prob->dilation_d, prob->dilation_h, prob->dilation_w},
{prob->pad_d, prob->pad_h, prob->pad_w},
{prob->pad_d, prob->pad_h, prob->pad_w}};
}
else
{
// 2D convolution: only use H, W dimensions
return ck_tile::conv::ConvParam{2,
prob->G,
prob->N,
prob->K,
prob->C,
{prob->filter_y, prob->filter_x},
{prob->input_h, prob->input_w},
{prob->stride_h, prob->stride_w},
{prob->dilation_h, prob->dilation_w},
{prob->pad_h, prob->pad_w},
{prob->pad_h, prob->pad_w}};
}
}
// Forward convolution (required - kernel header must be force-included)
static float run_forward(const void* input_ptr,
const void* weight_ptr,
void* output_ptr,
const ConvProblemC* prob,
void* stream)
{
auto conv_param = build_conv_param(prob);
ck_tile::GroupedConvFwdHostArgs<> args(conv_param, input_ptr, weight_ptr, {}, output_ptr, 1);
ck_tile::stream_config stream_cfg{static_cast<hipStream_t>(stream), true, 1, 3, 10};
// SelectedConvKernelLauncher is defined in the force-included forward kernel header
return SelectedConvKernelLauncher::launch(args, stream_cfg);
}
#ifdef CONV_BWD_DATA_AVAILABLE
// Backward data convolution (optional)
// Computes: grad_input = conv_bwd_data(weight, grad_output)
//
// Parameters:
// grad_output_ptr: dY - gradient from next layer (const, read-only INPUT)
// weight_ptr: W - frozen weights (const, read-only INPUT)
// grad_input_ptr: dX - gradient for input (writable, OUTPUT)
static float run_bwd_data(const void* grad_output_ptr,
const void* weight_ptr,
void* grad_input_ptr,
const ConvProblemC* prob,
void* stream)
{
auto conv_param = build_conv_param(prob);
// CK Tile API uses tensor POSITION names (from forward pass), not data flow:
// in_ptr = input tensor position = grad_input_ptr (dX, OUTPUT of bwd_data)
// wei_ptr = weight tensor = weight_ptr (W, const)
// out_ptr = output tensor position = grad_output_ptr (dY, INPUT to bwd_data)
ck_tile::GroupedConvBwdDataHostArgs args(
conv_param, grad_input_ptr, weight_ptr, {}, grad_output_ptr, 1);
ck_tile::stream_config stream_cfg{static_cast<hipStream_t>(stream), true, 1, 3, 10};
return SelectedConvBwdDataLauncher::launch(args, stream_cfg);
}
#endif
#ifdef CONV_BWD_WEIGHT_AVAILABLE
// Backward weight convolution (optional)
// Parameters:
// input_ptr: original forward input X (const, read-only)
// grad_output_ptr: gradient from next layer dY (const, read-only)
// grad_weight_ptr: gradient of weights dW (writable, OUTPUT)
static float run_bwd_weight(const void* input_ptr,
const void* grad_output_ptr,
void* grad_weight_ptr,
const ConvProblemC* prob,
void* stream)
{
auto conv_param = build_conv_param(prob);
// GroupedConvBwdWeightHostArgs constructor order:
// (param, in=X, wei=dW (output), ds, out=dY (input), k_batch)
// Note: wei_ptr is the OUTPUT (grad_weight), out_ptr is the INPUT (grad_output)
ck_tile::GroupedConvBwdWeightHostArgs args(
conv_param, input_ptr, grad_weight_ptr, {}, grad_output_ptr, 1);
ck_tile::stream_config stream_cfg{static_cast<hipStream_t>(stream), true, 1, 3, 10};
return SelectedConvBwdWeightLauncher::launch(args, stream_cfg);
}
#endif
/**
* @brief Execute convolution based on direction specified in prob
*
* Parameter mapping varies by direction:
* Forward (direction=0):
* input_ptr = X (input tensor)
* weight_ptr = W (weight tensor)
* output_ptr = Y (output buffer)
*
* Backward Data (direction=1):
* input_ptr = dY (grad_output - gradient from next layer)
* weight_ptr = W (weight tensor, frozen)
* output_ptr = dX (grad_input buffer)
*
* Backward Weight (direction=2):
* input_ptr = X (forward input tensor)
* weight_ptr = dY (grad_output - gradient from next layer)
* output_ptr = dW (grad_weight buffer)
*/
float conv_dispatcher_run(const void* input_ptr,
const void* weight_ptr,
void* output_ptr,
const ConvProblemC* prob,
void* stream)
{
// Validate all required pointers before kernel launch
if(!g_dispatcher || !prob)
return -1.0f;
if(!input_ptr || !weight_ptr || !output_ptr)
return -1.0f; // Null data pointer would cause kernel crash
// Build problem for kernel selection
ConvProblem problem;
problem.N = prob->N;
problem.G = prob->G;
problem.C = prob->C;
problem.K = prob->K;
problem.input_spatial = {prob->input_d, prob->input_h, prob->input_w};
problem.filter_spatial = {prob->filter_z, prob->filter_y, prob->filter_x};
problem.stride = {prob->stride_d, prob->stride_h, prob->stride_w};
problem.padding = {prob->pad_d, prob->pad_h, prob->pad_w};
problem.dilation = {prob->dilation_d, prob->dilation_h, prob->dilation_w};
problem.op = static_cast<ConvOp>(prob->direction);
problem.compute_output_size();
// Select kernel
const auto* kernel = g_dispatcher->select(problem);
if(!kernel)
return -1.0f;
// Dispatch based on direction
switch(prob->direction)
{
case 0: // Forward (always available)
return run_forward(input_ptr, weight_ptr, output_ptr, prob, stream);
#ifdef CONV_BWD_DATA_AVAILABLE
case 1: // Backward data
// Convention: caller passes (grad_output, weight, grad_input_buffer)
// in the (input_ptr, weight_ptr, output_ptr) slots respectively.
// run_bwd_data expects: (grad_output, weight, grad_input)
return run_bwd_data(input_ptr, weight_ptr, output_ptr, prob, stream);
#endif
#ifdef CONV_BWD_WEIGHT_AVAILABLE
case 2: // Backward weight
// Convention: caller passes (input, grad_output, grad_weight_buffer)
// in the (input_ptr, weight_ptr, output_ptr) slots respectively.
// run_bwd_weight expects: (input, grad_output, grad_weight)
return run_bwd_weight(input_ptr, weight_ptr, output_ptr, prob, stream);
#endif
default: return -1.0f;
}
}
// =============================================================================
// Info
// =============================================================================
const char* conv_dispatcher_version() { return "1.0.0"; }
// =========================================================================
// Library info
// =========================================================================
const char* conv_dispatcher_version() { return "2.0.0"; }
int conv_dispatcher_has_kernels()
{
return 1; // Forward kernel is required
#if defined(CONV_FWD_2D_AVAILABLE) || defined(CONV_FWD_3D_AVAILABLE)
return 1;
#else
return 0;
#endif
}
int conv_dispatcher_has_bwd_data()
{
#ifdef CONV_BWD_DATA_AVAILABLE
#if defined(CONV_BWD_DATA_2D_AVAILABLE) || defined(CONV_BWD_DATA_3D_AVAILABLE)
return 1;
#else
return 0;
@@ -401,11 +85,240 @@ int conv_dispatcher_has_bwd_data()
int conv_dispatcher_has_bwd_weight()
{
#ifdef CONV_BWD_WEIGHT_AVAILABLE
#if defined(CONV_BWD_WEIGHT_2D_AVAILABLE) || defined(CONV_BWD_WEIGHT_3D_AVAILABLE)
return 1;
#else
return 0;
#endif
}
int conv_dispatcher_get_kernel_count()
{
return CONV_KERNEL_COUNT; // defined in conv_python_dispatch.hpp
}
int conv_dispatcher_get_kernel_name(int index, char* buffer, int buffer_size)
{
if(!buffer || buffer_size <= 0 || index < 0 || index >= CONV_KERNEL_COUNT)
return -1;
std::strncpy(buffer, CONV_KERNEL_NAMES[index], buffer_size - 1);
buffer[buffer_size - 1] = '\0';
return 0;
}
// =========================================================================
// Support query
// =========================================================================
bool conv_dispatcher_is_supported(const ConvProblemC* prob)
{
if(!prob)
return false;
const bool is_3d = (prob->input_d > 1 || prob->filter_z > 1);
switch(prob->direction)
{
case CONV_FORWARD:
#if defined(CONV_FWD_3D_AVAILABLE)
if(is_3d)
return true;
#endif
#if defined(CONV_FWD_2D_AVAILABLE)
if(!is_3d)
return true;
#endif
return false;
case CONV_BWD_DATA:
#if defined(CONV_BWD_DATA_3D_AVAILABLE)
if(is_3d)
return true;
#endif
#if defined(CONV_BWD_DATA_2D_AVAILABLE)
if(!is_3d)
return true;
#endif
return false;
case CONV_BWD_WEIGHT:
#if defined(CONV_BWD_WEIGHT_3D_AVAILABLE)
if(is_3d)
return true;
#endif
#if defined(CONV_BWD_WEIGHT_2D_AVAILABLE)
if(!is_3d)
return true;
#endif
return false;
default: return false;
}
}
// =========================================================================
// ConvParam builders
// =========================================================================
static ck_tile::conv::ConvParam make_param_2d(const ConvProblemC* p)
{
return ck_tile::conv::ConvParam{2,
p->G,
p->N,
p->K,
p->C,
{p->filter_y, p->filter_x},
{p->input_h, p->input_w},
{p->stride_h, p->stride_w},
{p->dilation_h, p->dilation_w},
{p->pad_h, p->pad_w},
{p->pad_h, p->pad_w}};
}
static ck_tile::conv::ConvParam make_param_3d(const ConvProblemC* p)
{
return ck_tile::conv::ConvParam{3,
p->G,
p->N,
p->K,
p->C,
{p->filter_z, p->filter_y, p->filter_x},
{p->input_d, p->input_h, p->input_w},
{p->stride_d, p->stride_h, p->stride_w},
{p->dilation_d, p->dilation_h, p->dilation_w},
{p->pad_d, p->pad_h, p->pad_w},
{p->pad_d, p->pad_h, p->pad_w}};
}
// =========================================================================
// Kernel launch helpers
// =========================================================================
#ifdef CONV_FWD_2D_AVAILABLE
static float
launch_fwd_2d(const void* in, const void* wei, void* out, const ConvProblemC* p, hipStream_t stream)
{
auto param = make_param_2d(p);
ck_tile::GroupedConvFwdHostArgs<> args(param, in, wei, {}, out, 1);
ck_tile::stream_config sc{stream, true, 1, 3, 10};
return SelectedConvKernelLauncher::launch(args, sc);
}
#endif
#ifdef CONV_FWD_3D_AVAILABLE
static float
launch_fwd_3d(const void* in, const void* wei, void* out, const ConvProblemC* p, hipStream_t stream)
{
auto param = make_param_3d(p);
ck_tile::GroupedConvFwdHostArgs<> args(param, in, wei, {}, out, 1);
ck_tile::stream_config sc{stream, true, 1, 3, 10};
return ConvFwd3dLauncher::launch(args, sc);
}
#endif
#ifdef CONV_BWD_DATA_2D_AVAILABLE
static float launch_bwd_data_2d(
const void* dy, const void* wei, void* dx, const ConvProblemC* p, hipStream_t stream)
{
auto param = make_param_2d(p);
ck_tile::GroupedConvBwdDataHostArgs args(param, dx, wei, {}, dy, 1);
ck_tile::stream_config sc{stream, true, 1, 3, 10};
return SelectedConvBwdDataLauncher::launch(args, sc);
}
#endif
#ifdef CONV_BWD_DATA_3D_AVAILABLE
static float launch_bwd_data_3d(
const void* dy, const void* wei, void* dx, const ConvProblemC* p, hipStream_t stream)
{
auto param = make_param_3d(p);
ck_tile::GroupedConvBwdDataHostArgs args(param, dx, wei, {}, dy, 1);
ck_tile::stream_config sc{stream, true, 1, 3, 10};
return ConvBwdData3dLauncher::launch(args, sc);
}
#endif
#ifdef CONV_BWD_WEIGHT_2D_AVAILABLE
static float launch_bwd_weight_2d(
const void* x, const void* dy, void* dw, const ConvProblemC* p, hipStream_t stream)
{
auto param = make_param_2d(p);
const int k_batch = (p->split_k > 1) ? p->split_k : 1;
ck_tile::GroupedConvBwdWeightHostArgs args(param, x, dw, {}, dy, k_batch);
ck_tile::stream_config sc{stream, true, 1, 3, 10};
return SelectedConvBwdWeightLauncher::launch(args, sc);
}
#endif
#ifdef CONV_BWD_WEIGHT_3D_AVAILABLE
static float launch_bwd_weight_3d(
const void* x, const void* dy, void* dw, const ConvProblemC* p, hipStream_t stream)
{
auto param = make_param_3d(p);
const int k_batch = (p->split_k > 1) ? p->split_k : 1;
ck_tile::GroupedConvBwdWeightHostArgs args(param, x, dw, {}, dy, k_batch);
ck_tile::stream_config sc{stream, true, 1, 3, 10};
return ConvBwdWeight3dLauncher::launch(args, sc);
}
#endif
// =========================================================================
// Main dispatch
//
// direction=0 (forward): a=X(input), b=W(weight), c=Y(output)
// direction=1 (bwd_data): a=dY(grad_out), b=W(weight), c=dX(grad_in)
// direction=2 (bwd_weight): a=X(input), b=dY(grad_out), c=dW(grad_wei)
// =========================================================================
float conv_dispatcher_run(
const void* a_ptr, const void* b_ptr, void* c_ptr, const ConvProblemC* prob, void* stream)
{
if(!prob || !a_ptr || !b_ptr || !c_ptr)
return -1.0f;
const bool is_3d = (prob->input_d > 1 || prob->filter_z > 1);
hipStream_t hip_stream = static_cast<hipStream_t>(stream);
try
{
switch(prob->direction)
{
case CONV_FORWARD:
#ifdef CONV_FWD_3D_AVAILABLE
if(is_3d)
return launch_fwd_3d(a_ptr, b_ptr, c_ptr, prob, hip_stream);
#endif
#ifdef CONV_FWD_2D_AVAILABLE
if(!is_3d)
return launch_fwd_2d(a_ptr, b_ptr, c_ptr, prob, hip_stream);
#endif
return -2.0f;
case CONV_BWD_DATA:
#ifdef CONV_BWD_DATA_3D_AVAILABLE
if(is_3d)
return launch_bwd_data_3d(a_ptr, b_ptr, c_ptr, prob, hip_stream);
#endif
#ifdef CONV_BWD_DATA_2D_AVAILABLE
if(!is_3d)
return launch_bwd_data_2d(a_ptr, b_ptr, c_ptr, prob, hip_stream);
#endif
return -2.0f;
case CONV_BWD_WEIGHT:
#ifdef CONV_BWD_WEIGHT_3D_AVAILABLE
if(is_3d)
return launch_bwd_weight_3d(a_ptr, b_ptr, c_ptr, prob, hip_stream);
#endif
#ifdef CONV_BWD_WEIGHT_2D_AVAILABLE
if(!is_3d)
return launch_bwd_weight_2d(a_ptr, b_ptr, c_ptr, prob, hip_stream);
#endif
return -2.0f;
default: return -1.0f;
}
}
catch(const std::exception&)
{
return -3.0f; // Kernel rejected args (e.g. unsupported tile/channel combo)
}
catch(...)
{
return -3.0f;
}
}
} // extern "C"

16
dispatcher/codegen/ADDING_NEW_GPU.md
View File

@@ -9,8 +9,8 @@ Guide for adding support for a new AMD GPU architecture to the CK Tile Dispatche
The dispatcher uses `arch_specs.json` as the **single source of truth** for GPU specifications:
```
arch_specs.json → generate_arch_specs.py → arch_specs_generated.py (Python)
→ arch_specs_generated.hpp (C++)
arch_specs.json -> generate_arch_specs.py -> arch_specs_generated.py (Python)
-> arch_specs_generated.hpp (C++)
```
## Quick Start
@@ -175,14 +175,14 @@ for error in result.errors:
```
codegen/
├── arch_specs.json # Single source of truth (EDIT THIS)
├── generate_arch_specs.py # Generator script
├── arch_specs_generated.py # Generated Python module
└── ADDING_NEW_GPU.md # This file
|---- arch_specs.json # Single source of truth (EDIT THIS)
|---- generate_arch_specs.py # Generator script
|---- arch_specs_generated.py # Generated Python module
+---- ADDING_NEW_GPU.md # This file
include/ck_tile/dispatcher/
├── arch_specs_generated.hpp # Generated C++ header
└── arch_filter.hpp # C++ filter
|---- arch_specs_generated.hpp # Generated C++ header
+---- arch_filter.hpp # C++ filter
```
## Best Practices

47
dispatcher/codegen/README.md
View File

@@ -1,11 +1,22 @@
# CK Tile GEMM Unified Code Generator
# CK Tile Unified Code Generators
Single source of truth for all GEMM kernel generation.
Single source of truth for GEMM and Grouped Convolution kernel generation.
> **See also:** [Main Dispatcher README](../README.md) for installation and core concepts.
## Shared Infrastructure
Both GEMM and Grouped Conv generators share common code via `codegen_common.py`:
- `TileConfig` - Dataclass for tile dimensions
- `TraitConfigBase` - Base for kernel trait configurations with arch-aware validation
- `CommonTypeMappings` - Dtype-to-C++ type mappings
- `parallel_generate()` - Parallel kernel generation with per-kernel progress logging
- Arch-aware expansion helpers (`valid_wave_configs`, `valid_warp_configs`, etc.)
## Quick Start
### GEMM
```bash
cd dispatcher/codegen
@@ -22,6 +33,25 @@ python3 unified_gemm_codegen.py \
--variants standard preshuffle multi_d
```
### Grouped Convolution
```bash
cd dispatcher/codegen
# Generate forward FP16 grouped conv kernels
python3 unified_grouped_conv_codegen.py \
--output-dir ../build/generated_kernels \
--datatype fp16 \
--variant forward \
--ndim-spatial 2
# Generate backward data kernels
python3 unified_grouped_conv_codegen.py \
--output-dir ../build/generated_kernels \
--variant backward_data \
--ndim-spatial 2
```
## Using from Python
```python
@@ -58,13 +88,13 @@ results = codegen.generate_all()
## Variants
### Standard
Basic GEMM: `C = A × B`
Basic GEMM: `C = A x B`
### PreShuffle
Optimized weight access with LDS pre-shuffling. Best for large matrices.
### Multi-D
Element-wise fusion: `C = op(A × B + D0 + D1 + ...)`
Element-wise fusion: `C = op(A x B + D0 + D1 + ...)`
Supported ops: `PassThrough`, `MultiDAdd`, `Relu`, `Gelu`, `Sigmoid`, `Tanh`
@@ -72,10 +102,11 @@ Supported ops: `PassThrough`, `MultiDAdd`, `Relu`, `Gelu`, `Sigmoid`, `Tanh`
```
generated_kernels/
├── gemm_fp16_rcr_compv4_..._128x128x32_....hpp
├── gemm_fp16_rcr_compv4_..._preshuffle.hpp
├── gemm_fp16_rcr_compv4_..._multid_Relu_d1.hpp
└── ...
|---- gemm_fp16_rcr_compv4_..._128x128x32_....hpp # GEMM kernels
|---- gemm_fp16_rcr_compv4_..._preshuffle.hpp
|---- gemm_fp16_rcr_compv4_..._multid_Relu_d1.hpp
|---- grouped_conv_fwd_fp16_nhwgc_..._128x128x32_....hpp # Grouped conv kernels
+---- ...
```
## Configuration Files

350
dispatcher/codegen/codegen_common.py Normal file
View File

@@ -0,0 +1,350 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Shared codegen infrastructure for GEMM and grouped convolution code generators.
Extracted from unified_gemm_codegen.py + arch-aware expansion helpers from conv.
Both unified_gemm_codegen.py and unified_grouped_conv_codegen.py import from here
to eliminate duplication.
"""
import logging
import concurrent.futures
from dataclasses import dataclass
from typing import (
Callable,
ClassVar,
Dict,
FrozenSet,
List,
Optional,
Sequence,
Tuple,
TypeVar,
)
log = logging.getLogger(__name__)
T = TypeVar("T")
R = TypeVar("R")
ANY_INT = -1
# ============================================================================
# Tile and Trait Configuration (shared between GEMM and Conv)
# ============================================================================
@dataclass
class TileConfig:
"""Tile configuration parameters shared by GEMM and grouped conv."""
tile_m: int
tile_n: int
tile_k: int
warp_m: int
warp_n: int
warp_k: int
warp_tile_m: int
warp_tile_n: int
warp_tile_k: int
def is_valid(self) -> bool:
if self.tile_m <= 0 or self.tile_n <= 0 or self.tile_k <= 0:
return False
return (
self.tile_m % (self.warp_m * self.warp_tile_m) == 0
and self.tile_n % (self.warp_n * self.warp_tile_n) == 0
and self.tile_k % (self.warp_k * self.warp_tile_k) == 0
)
@dataclass
class TraitConfigBase:
"""
Base kernel trait configuration shared by GEMM and grouped conv.
GEMM extends this with ``persistent``; grouped conv extends with
``double_smem_buffer`` and ``num_groups_to_merge``.
"""
pipeline: str # mem, compv3, compv4, compv5, ...
epilogue: str # cshuffle, default
scheduler: str # intrawave, interwave
pad_m: bool
pad_n: bool
pad_k: bool
# Unsupported (pipeline, epilogue, scheduler) combinations.
# Only 'mem' and 'basic_v1' pipelines support interwave; all compute
# pipelines (compv3/v4/v5/v6/async) only support intrawave.
_UNSUPPORTED: ClassVar[FrozenSet] = frozenset(
{
("compv3", "cshuffle", "interwave"),
("compv3", "default", "interwave"),
("compv4", "cshuffle", "interwave"),
("compv4", "default", "interwave"),
("compv5", "cshuffle", "interwave"),
("compv5", "default", "interwave"),
("compv6", "cshuffle", "interwave"),
("compv6", "default", "interwave"),
("comp_async", "cshuffle", "interwave"),
("comp_async", "default", "interwave"),
("basic_async_v1", "cshuffle", "interwave"),
("basic_async_v1", "default", "interwave"),
}
)
def is_valid(self) -> bool:
return (self.pipeline, self.epilogue, self.scheduler) not in self._UNSUPPORTED
# ============================================================================
# Type Mappings (centralized for both GEMM and conv codegen)
# ============================================================================
class CommonTypeMappings:
"""Centralized type mappings shared by GEMM and grouped conv codegen."""
DTYPE_TO_CK = {
"fp16": "fp16_t",
"bf16": "bf16_t",
"fp32": "float",
"fp8": "fp8_t",
"bf8": "bf8_t",
"int8": "int8_t",
}
DTYPE_TO_CK_QUALIFIED = {
"fp16": "ck_tile::fp16_t",
"bf16": "ck_tile::bf16_t",
"fp32": "float",
"fp8": "ck_tile::fp8_t",
"bf8": "ck_tile::bf8_t",
"int8": "int8_t",
}
DTYPE_TO_DISPATCHER = {
"fp16": "DataType::FP16",
"bf16": "DataType::BF16",
"fp32": "DataType::FP32",
"fp8": "DataType::FP8",
"bf8": "DataType::BF8",
"int8": "DataType::INT8",
}
# GEMM-specific layout mappings ("r"/"c" for row/column major).
# Convolution layouts (NHWGC, GKYXC, etc.) are handled by
# unified_grouped_conv_codegen.py via GroupedConvLayout / GroupedConvTypeMappings.
GEMM_LAYOUT_TO_CK = {
"r": "tensor_layout::gemm::RowMajor",
"c": "tensor_layout::gemm::ColumnMajor",
}
LAYOUT_TO_CK = GEMM_LAYOUT_TO_CK # backward compat alias
GEMM_LAYOUT_TO_DISPATCHER = {
"r": "LayoutTag::RowMajor",
"c": "LayoutTag::ColMajor",
}
LAYOUT_TO_DISPATCHER = GEMM_LAYOUT_TO_DISPATCHER # backward compat alias
# GEMM-only pipeline mappings (used by unified_gemm_codegen.py).
# Convolution pipelines are in GroupedConvTypeMappings
# (unified_grouped_conv_codegen.py). CK Tile conv supports:
# BASIC_V1, Mem, CompV3, CompV4, CompV5, CompV6, ASYNC_V1, ASYNC_V4.
# The dispatcher currently generates: mem, compv3, compv4.
# preshufflev2 is GEMM-only (weight pre-shuffle for GEMM, not conv).
PIPELINE_TO_CK = {
"mem": "GemmPipelineAgBgCrMem",
"compv3": "GemmPipelineAgBgCrCompV3",
"compv4": "GemmPipelineAgBgCrCompV4",
"compv5": "GemmPipelineAgBgCrCompV5",
"preshufflev2": "WeightPreshufflePipelineAGmemBGmemCRegV2",
}
PIPELINE_TO_BASE = {
"mem": "BaseGemmPipelineAgBgCrMem",
"compv3": "BaseGemmPipelineAgBgCrCompV3",
"compv4": "BaseGemmPipelineAgBgCrCompV4",
"compv5": "BaseGemmPipelineAgBgCrCompV5",
"preshufflev2": "BaseWeightPreshufflePipelineAGmemBGmemCRegV2",
}
PIPELINE_TO_DISPATCHER = {
"mem": "Pipeline::Mem",
"compv3": "Pipeline::CompV3",
"compv4": "Pipeline::CompV4",
"compv5": "Pipeline::CompV5",
"preshufflev2": "Pipeline::PreShuffleV2",
}
SCHEDULER_TO_CK = {
"intrawave": "GemmPipelineScheduler::Intrawave",
"interwave": "GemmPipelineScheduler::Interwave",
"default": "GemmPipelineScheduler::Default",
}
SCHEDULER_TO_DISPATCHER = {
"intrawave": "Scheduler::Intrawave",
"interwave": "Scheduler::Interwave",
"default": "Scheduler::Auto",
}
EPILOGUE_TO_DISPATCHER = {
"cshuffle": "Epilogue::CShuffle",
"default": "Epilogue::Default",
}
@staticmethod
def get_output_dtype(dtype: str) -> str:
"""Get output datatype (fp8/bf8 -> fp16)."""
return "fp16" if dtype in ("fp8", "bf8") else dtype
# ============================================================================
# Code Generation Helpers
# ============================================================================
def generate_cpp_compilation_unit(kernel_name: str) -> str:
"""Generate a .cpp compilation unit that includes a kernel header.
This is the standard pattern: one .cpp per kernel that just includes
the generated .hpp header, causing template instantiation.
"""
return (
f"// Auto-generated compilation unit for {kernel_name}\n"
f'#include "{kernel_name}.hpp"\n'
)
def parallel_generate(
generate_fn: Callable[[T], R],
items: Sequence[T],
parallel: bool = True,
) -> List[R]:
"""Run ``generate_fn`` over ``items``, optionally in parallel.
Logs per-item progress (best-of-conv pattern).
Returns a flat list of results in completion order.
"""
results: List[R] = []
if not items:
return results
if parallel and len(items) > 1:
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = {executor.submit(generate_fn, item): item for item in items}
for future in concurrent.futures.as_completed(futures):
result = future.result()
results.append(result)
log.info("Generated: %s", futures[future])
else:
for item in items:
result = generate_fn(item)
results.append(result)
log.info("Generated: %s", item)
return results
# ============================================================================
# Arch-Aware Expansion Helpers (adopted from conv kernel_decl.hpp)
# ============================================================================
# These load from arch_specs_generated when available, falling back to
# hardcoded defaults that match the most common arch (gfx942).
_arch_data_cache: Optional[Dict] = None
def _get_arch_data() -> Dict:
"""Load arch filter data, with caching."""
global _arch_data_cache
if _arch_data_cache is not None:
return _arch_data_cache
try:
from arch_specs_generated import (
WARP_SUPPORTED_COMBINATIONS,
WARP_TILE_SUPPORTED_COMBINATIONS,
TRAIT_UNSUPPORTED_COMBINATIONS,
get_supported_archs,
)
_arch_data_cache = {
"warp_combos": WARP_SUPPORTED_COMBINATIONS,
"warp_tile_combos": WARP_TILE_SUPPORTED_COMBINATIONS,
"trait_unsupported": TRAIT_UNSUPPORTED_COMBINATIONS,
"supported_archs": get_supported_archs(),
}
except ImportError:
_arch_data_cache = {
"warp_combos": {
"gfx942": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
"gfx90a": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
},
"warp_tile_combos": {
"gfx942": {"fp16_fp16_fp32": [[16, 16, 16], [32, 32, 16]]},
"gfx90a": {"fp16_fp16_fp32": [[16, 16, 16], [32, 32, 16]]},
},
"trait_unsupported": {
("compv3", "cshuffle", "interwave"),
("compv4", "cshuffle", "interwave"),
},
"supported_archs": ["gfx90a", "gfx942", "gfx950"],
}
return _arch_data_cache
def valid_wave_configs(arch: str) -> List[List[int]]:
"""Return valid [wave_m, wave_n, wave_k] combos for *arch*."""
data = _get_arch_data()
return data["warp_combos"].get(arch, [[2, 2, 1]])
def valid_warp_configs(arch: str, dtype: str) -> List[List[int]]:
"""Return valid [warp_tile_m, warp_tile_n, warp_tile_k] combos for *arch*/*dtype*.
The dtype key is constructed as ``{dtype}_{dtype}_{acc}`` where acc is
fp32 for float types and int32 for int8.
"""
data = _get_arch_data()
acc = "int32" if dtype == "int8" else "fp32"
dtype_key = f"{dtype}_{dtype}_{acc}"
arch_tiles = data["warp_tile_combos"].get(arch, {})
return arch_tiles.get(dtype_key, [[32, 32, 16]])
def valid_trait_configs() -> List[Tuple[str, str]]:
"""Return valid (pipeline, scheduler) pairs.
Compute pipelines only support intrawave; mem supports both.
"""
return [
("compv3", "intrawave"),
("compv4", "intrawave"),
("compv5", "intrawave"),
("mem", "intrawave"),
("mem", "interwave"),
]
def needs_wave_expansion(config: dict) -> bool:
"""True if wave_m or wave_n is a wildcard (ANY_INT = -1)."""
return config.get("wave_m", 2) == ANY_INT or config.get("wave_n", 2) == ANY_INT
def needs_warp_expansion(config: dict) -> bool:
"""True if warp_m or warp_n is a wildcard (ANY_INT = -1)."""
return config.get("warp_m", 32) == ANY_INT or config.get("warp_n", 32) == ANY_INT
def needs_pipeline_expansion(config: dict) -> bool:
"""True if pipeline is a wildcard (\"*\")."""
return config.get("pipeline", "compv4") == "*"

8
dispatcher/codegen/generate_dispatcher_registration.py
View File

@@ -109,7 +109,7 @@ inline void register_all_kernels()
"""
output_file.write_text(content)
print(f"✓ Generated registration header: {output_file}")
print(f"OK Generated registration header: {output_file}")
def generate_registration_cpp(kernels: List[KernelConfig], output_file: Path):
@@ -143,7 +143,7 @@ namespace generated {
"""
output_file.write_text(content)
print(f"✓ Generated registration implementation: {output_file}")
print(f"OK Generated registration implementation: {output_file}")
def generate_kernel_wrapper_header(kernel: KernelConfig, output_dir: Path):
@@ -414,8 +414,8 @@ def main():
with open(manifest_output, "w") as f:
json.dump(manifest_data, f, indent=2)
print(f"✓ Generated manifest: {manifest_output}")
print("\n✓ Registration code generation complete!")
print(f"OK Generated manifest: {manifest_output}")
print("\nOK Registration code generation complete!")
print(f" Total kernels: {len(kernels)}")
print(" Output files:")
print(f" - {registration_header}")

8
dispatcher/codegen/generate_kernel_wrappers.py
View File

@@ -17,10 +17,10 @@ Usage:
Output structure:
build/kernel_wrappers/
├── gemm_fp16_rcr_128x128x32.cpp
├── gemm_fp16_rcr_256x256x64.cpp
├── conv_fwd_fp16_2d_128x128.cpp
└── ...
|---- gemm_fp16_rcr_128x128x32.cpp
|---- gemm_fp16_rcr_256x256x64.cpp
|---- conv_fwd_fp16_2d_128x128.cpp
+---- ...
Each .cpp simply includes its corresponding .hpp and forces symbol emission.
"""

32
dispatcher/codegen/kernel_config_loader.py
View File

@@ -359,8 +359,8 @@ class ConvTraitConfig:
@dataclass
class ConvKernelConfig:
"""Complete convolution kernel configuration"""
class GroupedConvKernelConfig:
"""Complete grouped convolution kernel configuration"""
tile: ConvTileConfig = field(default_factory=ConvTileConfig)
trait: ConvTraitConfig = field(default_factory=ConvTraitConfig)
@@ -419,7 +419,11 @@ class ConvKernelConfig:
def kernel_name(self) -> str:
"""Generate kernel name from config"""
variant_map = {"forward": "fwd", "bwd_data": "bwdd", "bwd_weight": "bwdw"}
variant_map = {
"forward": "fwd",
"bwd_data": "bwd_data",
"bwd_weight": "bwd_weight",
}
var_str = variant_map.get(self.variant, self.variant)
name = f"conv_{var_str}_{self.dtype_input}_{self.ndim}d"
@@ -433,11 +437,11 @@ class ConvKernelConfig:
@dataclass
class ConvKernelConfigSet:
class GroupedConvKernelConfigSet:
"""A set of convolution kernel configurations loaded from JSON"""
name: str = "default"
configs: List[ConvKernelConfig] = field(default_factory=list)
configs: List[GroupedConvKernelConfig] = field(default_factory=list)
# Tile parameter ranges
tile_m_values: List[int] = field(default_factory=lambda: [128])
@@ -481,7 +485,7 @@ class ConvKernelConfigSet:
layout: str = "nhwgc"
gpu_targets: List[str] = field(default_factory=lambda: ["gfx942"])
def generate_configs(self) -> Iterator[ConvKernelConfig]:
def generate_configs(self) -> Iterator[GroupedConvKernelConfig]:
"""Generate all kernel configurations (cartesian product)"""
# Tile parameters
tile_params = itertools.product(
@@ -548,7 +552,7 @@ class ConvKernelConfigSet:
double_smem_buffer=trait[6],
num_groups_to_merge=trait[7],
)
yield ConvKernelConfig(
yield GroupedConvKernelConfig(
tile=tile_cfg,
trait=trait_cfg,
dtype_input=self.dtype_input,
@@ -599,7 +603,9 @@ class ConvKernelConfigSet:
return tile_count * trait_count * extra_count * len(self.gpu_targets)
def load_conv_kernel_configs(json_path: str | Path) -> ConvKernelConfigSet:
def load_grouped_conv_kernel_configs(
json_path: str | Path,
) -> GroupedConvKernelConfigSet:
"""
Load convolution kernel configurations from a JSON file.
@@ -607,14 +613,14 @@ def load_conv_kernel_configs(json_path: str | Path) -> ConvKernelConfigSet:
json_path: Path to JSON configuration file
Returns:
ConvKernelConfigSet with all parameter values loaded
GroupedConvKernelConfigSet with all parameter values loaded
"""
json_path = Path(json_path)
with open(json_path) as f:
data = json.load(f)
config_set = ConvKernelConfigSet()
config_set = GroupedConvKernelConfigSet()
# Name
config_set.name = data.get("kernel_set_name", json_path.stem)
@@ -680,15 +686,15 @@ def load_conv_kernel_configs(json_path: str | Path) -> ConvKernelConfigSet:
def generate_cpp_conv_kernel_set_declaration(
config_set: ConvKernelConfigSet,
config_set: GroupedConvKernelConfigSet,
set_name: Optional[str] = None,
) -> str:
"""
Generate C++ DECL_CONV_KERNEL_SET code from a ConvKernelConfigSet.
Generate C++ DECL_GROUPED_CONV_KERNEL_SET code from a GroupedConvKernelConfigSet.
"""
name = set_name or config_set.name
lines = [f"DECL_CONV_KERNEL_SET({name},"]
lines = [f"DECL_GROUPED_CONV_KERNEL_SET({name},"]
for config in config_set.generate_configs():
line = f' .add("{config.dtype_input}", "{config.variant}", {config.ndim}, '

236
dispatcher/codegen/unified_gemm_codegen.py
View File

@@ -7,7 +7,7 @@
Unified GEMM Code Generator - Single Source of Truth
This is THE unified code generator for all GEMM kernel variants:
- Standard GEMM (C = A × B)
- Standard GEMM (C = A x B)
- Preshuffle GEMM (optimized weight access)
- Multi-D GEMM (element-wise fusion)
@@ -25,6 +25,12 @@ from dataclasses import dataclass, asdict
from enum import Enum
import concurrent.futures
from codegen_common import (
TileConfig,
TraitConfigBase,
CommonTypeMappings as TypeMappings,
)
# Import architecture filter for GPU-specific validation
try:
from arch_filter import ArchFilter, KernelConfig as ArchKernelConfig, OperatorType
@@ -194,62 +200,14 @@ class GemmVariant(Enum):
MULTI_D = "multi_d"
@dataclass
class TileConfig:
"""Tile configuration parameters"""
tile_m: int
tile_n: int
tile_k: int
warp_m: int
warp_n: int
warp_k: int
warp_tile_m: int
warp_tile_n: int
warp_tile_k: int
def is_valid(self) -> bool:
"""Validate tile configuration"""
return (
self.tile_m % (self.warp_m * self.warp_tile_m) == 0
and self.tile_n % (self.warp_n * self.warp_tile_n) == 0
and self.tile_k % (self.warp_k * self.warp_tile_k) == 0
and self.tile_m > 0
and self.tile_n > 0
and self.tile_k > 0
)
# TileConfig imported from codegen_common
@dataclass
class TraitConfig:
"""Kernel trait configuration"""
class TraitConfig(TraitConfigBase):
"""GEMM-specific trait configuration extending TraitConfigBase with persistent mode."""
pipeline: str # mem, compv3, compv4
epilogue: str # default, cshuffle
scheduler: str # intrawave, interwave
pad_m: bool
pad_n: bool
pad_k: bool
persistent: bool
def is_valid(self) -> bool:
"""Check if trait combination is valid"""
# Unsupported combinations
# Only 'mem' pipeline supports interwave scheduler.
# All compute pipelines (compv3/v4/v5/v6/async) only support intrawave.
unsupported = {
("compv3", "cshuffle", "interwave"),
("compv3", "default", "interwave"),
("compv4", "cshuffle", "interwave"),
("compv4", "default", "interwave"),
("compv5", "cshuffle", "interwave"),
("compv5", "default", "interwave"),
("compv6", "cshuffle", "interwave"),
("compv6", "default", "interwave"),
("comp_async", "cshuffle", "interwave"),
("comp_async", "default", "interwave"),
}
return (self.pipeline, self.epilogue, self.scheduler) not in unsupported
persistent: bool = False
@dataclass
@@ -345,89 +303,7 @@ class KernelConfig:
# ============================================================================
class TypeMappings:
"""Centralized type mappings for code generation"""
DTYPE_TO_CK = {
"fp16": "fp16_t",
"bf16": "bf16_t",
"fp32": "float",
"fp8": "fp8_t",
"bf8": "bf8_t",
"int8": "int8_t",
}
# Fully-qualified types for use outside of 'using namespace ck_tile' scope
DTYPE_TO_CK_QUALIFIED = {
"fp16": "ck_tile::fp16_t",
"bf16": "ck_tile::bf16_t",
"fp32": "float", # Built-in type, no namespace
"fp8": "ck_tile::fp8_t",
"bf8": "ck_tile::bf8_t",
"int8": "int8_t", # Built-in type
}
DTYPE_TO_DISPATCHER = {
"fp16": "DataType::FP16",
"bf16": "DataType::BF16",
"fp32": "DataType::FP32",
"fp8": "DataType::FP8",
"bf8": "DataType::BF8",
"int8": "DataType::INT8",
}
LAYOUT_TO_CK = {
"r": "tensor_layout::gemm::RowMajor",
"c": "tensor_layout::gemm::ColumnMajor",
}
LAYOUT_TO_DISPATCHER = {
"r": "LayoutTag::RowMajor",
"c": "LayoutTag::ColMajor",
}
PIPELINE_TO_CK = {
"mem": "GemmPipelineAgBgCrMem",
"compv3": "GemmPipelineAgBgCrCompV3",
"compv4": "GemmPipelineAgBgCrCompV4",
"preshufflev2": "WeightPreshufflePipelineAGmemBGmemCRegV2",
}
PIPELINE_TO_BASE = {
"mem": "BaseGemmPipelineAgBgCrMem",
"compv3": "BaseGemmPipelineAgBgCrCompV3",
"compv4": "BaseGemmPipelineAgBgCrCompV4",
"preshufflev2": "BaseWeightPreshufflePipelineAGmemBGmemCRegV2",
}
PIPELINE_TO_DISPATCHER = {
"mem": "Pipeline::Mem",
"compv3": "Pipeline::CompV3",
"compv4": "Pipeline::CompV4",
"preshufflev2": "Pipeline::PreShuffleV2",
}
SCHEDULER_TO_CK = {
"intrawave": "GemmPipelineScheduler::Intrawave",
"interwave": "GemmPipelineScheduler::Interwave",
"default": "GemmPipelineScheduler::Default",
}
SCHEDULER_TO_DISPATCHER = {
"intrawave": "Scheduler::Intrawave",
"interwave": "Scheduler::Interwave",
"default": "Scheduler::Auto",
}
EPILOGUE_TO_DISPATCHER = {
"cshuffle": "Epilogue::CShuffle",
"default": "Epilogue::Default",
}
@staticmethod
def get_output_dtype(dtype: str) -> str:
"""Get output datatype (fp8/bf8 -> fp16)"""
return "fp16" if dtype in ["fp8", "bf8"] else dtype
# TypeMappings imported from codegen_common as CommonTypeMappings -> TypeMappings alias
# ============================================================================
@@ -1068,7 +944,11 @@ class UnifiedGemmCodegen:
}
def generate_all(self, parallel: bool = True) -> Dict:
"""Generate all kernels"""
"""Generate all kernels.
When parallel=True, all configs across all variants are collected first,
then generated concurrently in a single thread pool for maximum throughput.
"""
log.info("Generating GEMM kernels:")
log.info(f" Datatype: {self.datatype}")
log.info(f" Layout: {self.layout}")
@@ -1078,49 +958,24 @@ class UnifiedGemmCodegen:
results = {"kernels": [], "wrappers": [], "failed": []}
# Get configurations
# Collect ALL configs across all variants/preselected sets upfront
all_configs = []
if self.use_preselected:
configs = self._get_preselected_configs()
log.info(f" Total configurations: {len(configs)}")
all_configs = self._get_preselected_configs()
log.info(f" Total configurations: {len(all_configs)}")
else:
for variant in self.variants:
log.info(f"\nGenerating {variant.value} kernels...")
configs = self._get_configs_for_variant(variant)
log.info(f" Configurations: {len(configs)}")
log.info(f" {variant.value}: {len(configs)} configurations")
all_configs.extend(configs)
log.info(f" Total across all variants: {len(all_configs)}")
if parallel:
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = [
executor.submit(self._generate_one, cfg) for cfg in configs
]
for future in concurrent.futures.as_completed(futures):
try:
k, w = future.result()
results["kernels"].append(k)
results["wrappers"].append(w)
except Exception as e:
results["failed"].append(str(e))
log.error(f"Failed: {e}")
else:
for cfg in configs:
try:
k, w = self._generate_one(cfg)
results["kernels"].append(k)
results["wrappers"].append(w)
except Exception as e:
results["failed"].append(str(e))
log.error(f"Failed: {e}")
# Generate registration header
if results["wrappers"]:
self._generate_registration_header(results["wrappers"])
return results
# Generate from preselected set
if parallel:
# Generate all configs in a single parallel pass
if parallel and all_configs:
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = [executor.submit(self._generate_one, cfg) for cfg in configs]
futures = [
executor.submit(self._generate_one, cfg) for cfg in all_configs
]
for future in concurrent.futures.as_completed(futures):
try:
k, w = future.result()
@@ -1130,7 +985,7 @@ class UnifiedGemmCodegen:
results["failed"].append(str(e))
log.error(f"Failed: {e}")
else:
for cfg in configs:
for cfg in all_configs:
try:
k, w = self._generate_one(cfg)
results["kernels"].append(k)
@@ -1139,7 +994,6 @@ class UnifiedGemmCodegen:
results["failed"].append(str(e))
log.error(f"Failed: {e}")
# Generate registration header
if results["wrappers"]:
self._generate_registration_header(results["wrappers"])
@@ -1638,12 +1492,19 @@ def main():
# Write to temp file and use as config
import tempfile
import os as _os
with tempfile.NamedTemporaryFile(
_tmp_config = tempfile.NamedTemporaryFile(
mode="w", suffix=".json", delete=False
) as f:
json.dump(full_config, f)
args.config = Path(f.name)
)
try:
json.dump(full_config, _tmp_config)
_tmp_config.close()
args.config = Path(_tmp_config.name)
except Exception:
_tmp_config.close()
_os.unlink(_tmp_config.name)
raise
except json.JSONDecodeError as e:
logging.error(f"Invalid tile-config-json: {e}")
return 1
@@ -1672,7 +1533,7 @@ def main():
results = codegen.generate_all(parallel=not args.no_parallel)
logging.info("\n✅ Generation complete!")
logging.info("\nGeneration complete.")
logging.info(f" Kernels: {len(results['kernels'])}")
logging.info(f" Wrappers: {len(results['wrappers'])}")
logging.info(f" Failed: {len(results['failed'])}")
@@ -1684,7 +1545,7 @@ def main():
# Generate dispatcher registration if requested
if args.register:
logging.info("\n📝 Generating dispatcher registration code...")
logging.info("\nGenerating dispatcher registration code...")
try:
from generate_dispatcher_registration import (
scan_generated_headers,
@@ -1701,11 +1562,20 @@ def main():
)
generate_registration_cpp(kernels, reg_dir / "dispatcher_registration.cpp")
logging.info(f"✓ Generated registration code for {len(kernels)} kernels")
logging.info(f"Generated registration code for {len(kernels)} kernels")
except Exception as e:
logging.error(f"Failed to generate registration code: {e}")
return 1
# Clean up temp config file if we created one
if args.tile_config_json and args.config and args.config.exists():
try:
import os as _os
_os.unlink(args.config)
except OSError:
pass
return 0 if not results["failed"] else 1

1757
dispatcher/codegen/unified_grouped_conv_codegen.py Normal file
View File

File diff suppressed because it is too large Load Diff

119
dispatcher/examples/CMakeLists.txt
View File

@@ -187,7 +187,6 @@ function(add_gpu_example NAME SOURCE KERNEL_HEADER)
if(HEADER_NAME STREQUAL "register_all_kernels.hpp")
# Registration header - examples include it directly
target_compile_options(${NAME} PRIVATE
-DGEMM_KERNEL_AVAILABLE=1
-mllvm -enable-noalias-to-md-conversion=0
-Wno-undefined-func-template
-Wno-float-equal
@@ -315,6 +314,7 @@ function(add_declarative_gpu_example NAME SOURCE)
target_include_directories(${NAME} PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/../../include
${CMAKE_CURRENT_SOURCE_DIR}/../include
${CMAKE_CURRENT_SOURCE_DIR}/../..
${EXAMPLE_KERNEL_DIR}
${EXAMPLE_KERNEL_DIR}/dispatcher_wrappers
)
@@ -322,7 +322,6 @@ function(add_declarative_gpu_example NAME SOURCE)
# Force-include the generated registration header
target_compile_options(${NAME} PRIVATE
-include ${EXAMPLE_HEADER}
-DGEMM_KERNEL_AVAILABLE=1
-mllvm -enable-noalias-to-md-conversion=0
-Wno-undefined-func-template
-Wno-float-equal
@@ -345,6 +344,56 @@ add_declarative_gpu_example(gemm_03_benchmark_validation gemm/cpp/03_benchmark_v
add_declarative_gpu_example(gemm_04_heuristics gemm/cpp/04_heuristics.cpp)
add_declarative_gpu_example(gemm_05_json_export gemm/cpp/05_json_export.cpp)
add_declarative_gpu_example(gemm_06_multi_registry gemm/cpp/06_multi_registry.cpp)
add_declarative_gpu_example(gemm_07_gfx950_minimal gemm/cpp/07_gfx950_minimal.cpp)
# ML Heuristic example -- requires LightGBM shared library
# Derive site-packages from active Python interpreter (respects virtualenvs)
find_package(Python3 COMPONENTS Interpreter)
set(LIGHTGBM_SEARCH_PATHS)
if(Python3_FOUND AND Python3_EXECUTABLE)
execute_process(
COMMAND ${Python3_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_path('purelib'))"
OUTPUT_VARIABLE PYTHON_SITE_PACKAGES
OUTPUT_STRIP_TRAILING_WHITESPACE
ERROR_QUIET
)
if(PYTHON_SITE_PACKAGES)
list(APPEND LIGHTGBM_SEARCH_PATHS "${PYTHON_SITE_PACKAGES}/lightgbm/lib")
endif()
endif()
# Fallback to common Python 3.x site-packages if auto-detection failed
if(NOT PYTHON_SITE_PACKAGES)
list(APPEND LIGHTGBM_SEARCH_PATHS
"$ENV{HOME}/.local/lib/python3.12/site-packages/lightgbm/lib"
)
endif()
find_library(LIGHTGBM_LIB NAMES LightGBM lib_lightgbm _lightgbm
HINTS ${CMAKE_PREFIX_PATH}
PATHS ${LIGHTGBM_SEARCH_PATHS}
NO_DEFAULT_PATH
DOC "LightGBM shared library for ML heuristics"
)
# Fallback: search default paths (respects LightGBM_DIR if set by user)
if(NOT LIGHTGBM_LIB)
find_library(LIGHTGBM_LIB NAMES LightGBM lib_lightgbm)
endif()
if(LIGHTGBM_LIB)
add_declarative_gpu_example(gemm_09_ml_heuristic gemm/cpp/09_ml_heuristic.cpp)
target_link_libraries(gemm_09_ml_heuristic PRIVATE ${LIGHTGBM_LIB})
message(STATUS "LightGBM found: ${LIGHTGBM_LIB} -- building gemm_09_ml_heuristic")
else()
message(STATUS "LightGBM not found -- skipping gemm_09_ml_heuristic")
message(STATUS " To enable ML heuristic example:")
message(STATUS " 1. Activate virtualenv: source .venv/bin/activate")
message(STATUS " 2. Install: pip install -r ../requirements-ml.txt")
message(STATUS " 3. Reconfigure: cmake ..")
message(STATUS " Or set CMAKE_PREFIX_PATH or LightGBM_DIR to LightGBM location")
endif()
# =============================================================================
# GEMM Python Library - Single Fallback Kernel
@@ -394,19 +443,79 @@ if(hip_FOUND)
endif()
add_dependencies(dispatcher_gemm_lib generate_gemm_fallback_kernel)
# =============================================================================
# Grouped Convolution C++ Examples
# =============================================================================
add_declarative_gpu_example(grouped_conv_01_basic grouped_conv/cpp/01_basic_grouped_conv.cpp)
add_declarative_gpu_example(grouped_conv_02_all_dirs grouped_conv/cpp/02_all_directions.cpp)
add_declarative_gpu_example(grouped_conv_03_bench_val grouped_conv/cpp/03_benchmark_validation.cpp)
add_declarative_gpu_example(grouped_conv_04_registry_json grouped_conv/cpp/04_registry_json.cpp)
add_declarative_gpu_example(grouped_conv_05_bwd_data grouped_conv/cpp/05_bwd_data.cpp)
add_declarative_gpu_example(grouped_conv_06_bwd_weight grouped_conv/cpp/06_bwd_weight.cpp)
add_declarative_gpu_example(grouped_conv_07_benchmark grouped_conv/cpp/07_multi_tile_benchmark.cpp)
# =============================================================================
# Grouped Convolution Python Library - Multi-Kernel (fwd/bwd_data/bwd_weight x 2D/3D)
# =============================================================================
# Kernel output directory for the Python conv library
set(CONV_FALLBACK_KERNEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/conv_python_fallback")
set(CONV_DISPATCH_HEADER "${CONV_FALLBACK_KERNEL_DIR}/conv_python_dispatch.hpp")
# Generate ALL conv kernels (fwd/bwd_data/bwd_weight x 2D/3D x multiple tile configs)
# then create the dispatch header with 2D/3D aliases
add_custom_command(
OUTPUT ${CONV_DISPATCH_HEADER}
COMMAND ${CMAKE_COMMAND} -E make_directory ${CONV_FALLBACK_KERNEL_DIR}
COMMAND python3 ${CMAKE_CURRENT_SOURCE_DIR}/../codegen/unified_grouped_conv_codegen.py
--variant forward bwd_data bwd_weight --ndim 2 3
--datatype fp16 --arch ${GPU_TARGET}
--output ${CONV_FALLBACK_KERNEL_DIR}
COMMAND python3 ${CMAKE_CURRENT_SOURCE_DIR}/../scripts/generate_conv_dispatch_header.py
--kernel-dir ${CONV_FALLBACK_KERNEL_DIR}
--output ${CONV_DISPATCH_HEADER}
COMMENT "Generating conv kernels (fwd/bwd_data/bwd_weight x 2D/3D) for Python library..."
VERBATIM
)
add_custom_target(generate_conv_fallback_kernels DEPENDS ${CONV_DISPATCH_HEADER})
# Conv dynamic library for Python (all 6 kernel variants)
add_library(dispatcher_conv_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/../bindings/ctypes/conv_ctypes_lib.cpp)
target_link_libraries(dispatcher_conv_lib PRIVATE ck_tile_dispatcher)
target_include_directories(dispatcher_conv_lib PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/../../include
${CMAKE_CURRENT_SOURCE_DIR}/../include
${CONV_FALLBACK_KERNEL_DIR}
)
target_compile_options(dispatcher_conv_lib PRIVATE
-include ${CONV_DISPATCH_HEADER}
-DGFX_ARCH="${GPU_TARGET}"
-mllvm -enable-noalias-to-md-conversion=0
-Wno-undefined-func-template
-Wno-float-equal
--offload-compress
)
if(hip_FOUND)
target_link_libraries(dispatcher_conv_lib PRIVATE hip::device hip::host)
endif()
add_dependencies(dispatcher_conv_lib generate_conv_fallback_kernels)
message(STATUS "GEMM examples configured - kernels will be generated during 'make'")
message(STATUS "Grouped Conv examples configured - kernels will be generated during 'make'")
# Convenience target to build all Python ctypes libraries
add_custom_target(python_libs
DEPENDS dispatcher_gemm_lib
COMMENT "Building Python ctypes libraries (GEMM)"
DEPENDS dispatcher_gemm_lib dispatcher_conv_lib
COMMENT "Building Python ctypes libraries (GEMM + Conv)"
)
# =============================================================================
# Per-Architecture Kernel Generation Targets
# =============================================================================
set(SUPPORTED_GPU_ARCHS gfx942 gfx90a gfx1100 gfx1030)
set(SUPPORTED_GPU_ARCHS gfx942 gfx950 gfx90a gfx1100 gfx1030)
foreach(ARCH ${SUPPORTED_GPU_ARCHS})
# GEMM kernels for this arch

45
dispatcher/examples/README.md
View File

@@ -1,8 +1,6 @@
# CK Tile Dispatcher Examples
Comprehensive examples for GEMM operations with GPU execution.
> **Note**: Convolution examples have been moved to `ck-2/conv_archive/` for reference.
Comprehensive examples for GEMM and Grouped Convolution operations with GPU execution.
---
@@ -60,11 +58,11 @@ python3 examples/gemm/python/08_heuristics.py
```
examples/
├── gemm/
│ ├── cpp/ # 6 C++ GEMM examples
│ └── python/ # 11 Python GEMM examples
│
└── README.md
|---- gemm/
| |---- cpp/ # 6 C++ GEMM examples
| +---- python/ # 11 Python GEMM examples
|
+---- README.md
```
---
@@ -201,10 +199,31 @@ rocminfo | grep "Name:"
---
## Archived Examples
## Grouped Convolution
Convolution examples have been archived to `ck-2/conv_archive/dispatcher/`:
- `examples/conv/cpp/` - 11 C++ convolution examples
- `examples/conv/python/` - 14 Python convolution examples
Grouped convolution support has been re-introduced with a unified infrastructure shared with GEMM.
See the archive for convolution functionality reference.
### Infrastructure
The grouped convolution code generation, utilities, and build scripts are available:
| Component | Location |
|-----------|----------|
| C++ Headers | `include/ck_tile/dispatcher/grouped_conv_*.hpp` |
| Python Codegen | `codegen/unified_grouped_conv_codegen.py` |
| Python Utils | `python/grouped_conv_utils.py` |
| Build Script | `scripts/compile_grouped_conv_examples.py` |
### Building Grouped Conv Kernels
```bash
# Generate grouped conv kernels
python3 codegen/unified_grouped_conv_codegen.py \
--output-dir build/generated_kernels \
--datatype fp16 --variant forward --ndim-spatial 2
# Compile a grouped conv example
python3 scripts/compile_grouped_conv_examples.py my_grouped_conv_example.cpp
```
See the [main README](../README.md#grouped-convolution-support) for more details.

20
dispatcher/examples/gemm/cpp/02_multi_size.cpp
View File

@@ -21,9 +21,9 @@
* - pipeline: "compv3" -> 1 option (compv4 requires special handling)
* - scheduler: "intrawave" -> 1 option
*
* Raw expansion: 3 × 2 = 6 configs, but arch filter validates each:
* - tile_m must be divisible by (warp_m × warp_tile_m)
* - tile_n must be divisible by (warp_n × warp_tile_n)
* Raw expansion: 3 x 2 = 6 configs, but arch filter validates each:
* - tile_m must be divisible by (warp_m x warp_tile_m)
* - tile_n must be divisible by (warp_n x warp_tile_n)
* - Some wave/warp combos invalid: (4,1,1)+(32,32,16), (1,4,1)+(32,32,16)
* Result: 4 valid wildcard kernels + 1 explicit = 5 total
*
@@ -70,13 +70,13 @@ DECL_KERNEL_SET(multi_size_kernels,
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(64, 64, 64)
.wave(ANY_INT, ANY_INT, 1) // ANY_INT → (1,4,1), (2,2,1), (4,1,1)
.warp(-1, -1, -1) // -1 same as ANY_INT → (16,16,32), (32,32,16)
.pipeline("*") // "*" → valid pipelines
.scheduler("*") // "*" → valid schedulers
.wave(ANY_INT, ANY_INT, 1) // ANY_INT -> (1,4,1), (2,2,1), (4,1,1)
.warp(-1, -1, -1) // -1 same as ANY_INT -> (16,16,32), (32,32,16)
.pipeline("*") // "*" -> valid pipelines
.scheduler("*") // "*" -> valid schedulers
.epilogue("cshuffle"),
"gfx942"));
// Raw: 3×2=6, arch filter removes 2 invalid → 4 valid kernels
// Raw: 3x2=6, arch filter removes 2 invalid -> 4 valid kernels
// =============================================================================
// MAIN
@@ -116,8 +116,8 @@ int main(int argc, char* argv[])
.pipeline("*") -> expands to valid pipelines = 1
.scheduler("*") -> expands to valid schedulers = 1
Expanded: 3 × 2 = 6 configs, but arch filter validates each:
- wave×warp must divide tile: (4,1,1)×(32,32,16) invalid for 64x64
Expanded: 3 x 2 = 6 configs, but arch filter validates each:
- wave x warp must divide tile: (4,1,1)x(32,32,16) invalid for 64x64
- Result: 4 valid kernels from wildcard + 1 explicit = 5 total
)";

191
dispatcher/examples/gemm/cpp/07_gfx950_minimal.cpp Normal file
View File

@@ -0,0 +1,191 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Example 07: Minimal gfx950 (CDNA4 / MI350) GEMM
*
* Demonstrates the dispatcher working with gfx950-specific kernels:
*
* - fp16 GEMM with standard tile configs
* - fp8 GEMM with gfx950-extended warp tiles (16x16x128)
* - 160KB LDS: gfx950 doubles the LDS from 64KB to 160KB
*
* Build: cd dispatcher/build && cmake .. -DGPU_TARGETS=gfx950 && make gemm_07_gfx950_minimal
*/
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include "ck_tile/dispatcher.hpp"
#include "ck_tile/dispatcher/kernel_decl.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::backends;
using namespace ck_tile::dispatcher::utils;
using Signature = decl::Signature;
using Algorithm = decl::Algorithm;
// =============================================================================
// gfx950-targeted kernel declarations
// =============================================================================
DECL_KERNEL_SET(gfx950_gemm_kernels,
// fp16 128x128x32 -- bread-and-butter config, works on all CDNA
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 32)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx950")
// fp16 128x128x64 -- deeper K tile using more LDS
// LDS usage: 128*64*2 + 128*64*2 = 32768 bytes (fits 64KB, gfx950 has 160KB)
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 64)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx950")
// fp16 64x64x32 -- small-tile variant for small problems
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(64, 64, 32)
.wave(2, 2, 1)
.warp(16, 16, 32)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx950"));
// =============================================================================
// MAIN
// =============================================================================
int main(int argc, char* argv[])
{
ExampleArgs args("Example 07: gfx950 Minimal GEMM",
"Demonstrates gfx950 (CDNA4 / MI350) dispatcher");
args.add_flag("--list", "List registered kernels");
args.add_flag("--list-verbose", "List registered kernels with full details");
args.add_option("--M", "1024", "Problem M dimension");
args.add_option("--N", "1024", "Problem N dimension");
args.add_option("--K", "1024", "Problem K dimension");
args.add_option("--arch", "gfx950", "GPU architecture (default: gfx950)");
if(!args.parse(argc, argv))
return 0;
std::string gfx_arch = args.get("--arch", "gfx950");
print_header("Example 07: gfx950 (CDNA4) Minimal GEMM");
// =========================================================================
// Architecture info
// =========================================================================
std::cout << "\ngfx950 (CDNA4 / MI350) highlights:\n";
std::cout << " - 160KB LDS (up from 64KB on gfx942)\n";
std::cout << " - Extended FP8 warp tiles: 16x16x128, 32x32x64\n";
std::cout << " - Packed FP4 support (pk_fp4)\n";
std::cout << " - Same warp configs as gfx942: [1,4,1], [2,2,1], [4,1,1]\n\n";
// =========================================================================
// Register kernels
// =========================================================================
std::cout << "Registering kernels for " << gfx_arch << "...\n";
Registry registry;
registry.set_name("gfx950_gemm");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)\n";
if(args.has("--list") || args.has("--list-verbose"))
{
std::cout << "\n";
print_registered_kernels(registry, std::cout, args.has("--list-verbose"));
return 0;
}
if(registry.size() == 0)
{
std::cerr << "ERROR: No kernels registered for " << gfx_arch << "!\n";
std::cerr << " Did you build with -DGPU_TARGETS=gfx950?\n";
return 1;
}
// =========================================================================
// Create Dispatcher
// =========================================================================
Dispatcher dispatcher(&registry);
// =========================================================================
// Setup Problem
// =========================================================================
const int M = args.get_int("--M", 1024);
const int N = args.get_int("--N", 1024);
const int K = args.get_int("--K", 1024);
std::cout << "\nProblem: " << M << " x " << N << " x " << K << "\n";
Problem problem(M, N, K);
using DataType = ck_tile::fp16_t;
GpuBuffer<DataType> a_dev(M * K);
GpuBuffer<DataType> b_dev(K * N);
GpuBuffer<DataType> c_dev(M * N);
std::vector<DataType> a_host(M * K, DataType(1.0f));
std::vector<DataType> b_host(K * N, DataType(1.0f));
a_dev.copy_from_host(a_host.data());
b_dev.copy_from_host(b_host.data());
c_dev.zero();
// =========================================================================
// Select and Run
// =========================================================================
auto selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cerr << "ERROR: No suitable kernel found for " << M << "x" << N << "x" << K << "\n";
return 1;
}
std::cout << " Selected: " << selected->get_name() << "\n";
float time_ms = dispatcher.run(a_dev.get(), b_dev.get(), c_dev.get(), problem, nullptr);
std::cout << " Time: " << std::fixed << std::setprecision(4) << time_ms << " ms\n";
std::cout << " TFLOPS: " << std::setprecision(2) << calculate_tflops(M, N, K, time_ms) << "\n";
// =========================================================================
// Verify
// =========================================================================
std::cout << "\nVerification:\n";
std::vector<DataType> c_host(M * N);
c_dev.copy_to_host(c_host.data());
const float expected = static_cast<float>(K);
int errors = 0;
for(int i = 0; i < std::min(M * N, 1024); ++i)
{
if(std::abs(static_cast<float>(c_host[i]) - expected) > 0.01f * expected + 1.0f)
++errors;
}
bool passed = (errors == 0);
std::cout << " Expected value: " << expected << "\n";
std::cout << " Errors (first 1024 elements): " << errors << "\n";
std::cout << " Status: " << (passed ? "PASS" : "FAIL") << "\n";
print_separator();
return passed ? 0 : 1;
}

211
dispatcher/examples/gemm/cpp/09_ml_heuristic.cpp Normal file
View File

@@ -0,0 +1,211 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Example 09: ML-Based Kernel Selection (Native C++)
*
* Uses a trained LightGBM model loaded via the C API to predict TFLOPS
* for each kernel in the registry and select the best one. The kernels
* are JIT-compiled at build time via DECL_KERNEL_SET (same as other examples).
*
* Build: cd dispatcher/build && cmake .. && make gemm_09_ml_heuristic
* Run: ./gemm_09_ml_heuristic --model <path_to_model.lgbm>
*/
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <chrono>
#include "ck_tile/dispatcher.hpp"
#include "ck_tile/dispatcher/kernel_decl.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
#include "ck_tile/dispatcher/ml_heuristic.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::utils;
using Signature = decl::Signature;
using Algorithm = decl::Algorithm;
// Multiple kernel configs for ML to choose from
DECL_KERNEL_SET(ml_kernels,
// Small tiles
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(64, 64, 32)
.wave(2, 2, 1)
.warp(16, 16, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(64, 64, 64)
.wave(2, 2, 1)
.warp(16, 16, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
// Medium tiles
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 32)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 64)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 64)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv4")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
// Large tiles
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(256, 256, 32)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(256, 128, 32)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942")
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 256, 32)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx942"));
int main(int argc, char* argv[])
{
ExampleArgs args("Example 09: ML-Based Kernel Selection",
"Uses trained LightGBM model for kernel selection");
args.add_option("--arch", "gfx942", "GPU architecture");
args.add_option("--model", "", "Path to LightGBM model file (.lgbm)");
args.add_option("--log_transform", "false", "Model uses log1p transform");
if(!args.parse(argc, argv))
return 0;
print_header("Example 09: ML-Based Kernel Selection");
std::string gfx_arch = args.get("--arch", "gfx942");
std::string model_path = args.get("--model", "");
bool log_transform = (args.get("--log_transform", "false") == "true");
if(model_path.empty())
{
std::cerr << "Error: --model <path> is required" << std::endl;
std::cerr << "Usage: ./gemm_09_ml_heuristic --model path/to/model_tflops.lgbm" << std::endl;
return 1;
}
// Setup Registry (kernels are JIT compiled from DECL_KERNEL_SET above)
Registry registry;
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << "Registry: " << registry.size() << " kernel(s)" << std::endl;
// Load ML model and create heuristic
HardwareProfile hw;
MLHeuristic ml_heuristic(model_path, &registry, hw, log_transform);
if(!ml_heuristic.is_loaded())
{
std::cerr << "Failed to load model. Exiting." << std::endl;
return 1;
}
// Wire ML heuristic into dispatcher
Dispatcher dispatcher(&registry);
dispatcher.set_strategy(Dispatcher::SelectionStrategy::Heuristic);
dispatcher.set_heuristic([&ml_heuristic](const Problem& p) { return ml_heuristic(p); });
std::cout << "Strategy: ML Heuristic (LightGBM)" << std::endl;
// Test with different problem sizes
using DataType = ck_tile::fp16_t;
std::vector<std::tuple<int, int, int>> sizes = {
{128, 128, 64},
{512, 512, 256},
{1024, 1024, 512},
{2048, 2048, 1024},
};
std::cout << std::endl
<< std::setw(20) << "Shape" << std::setw(30) << "Selected Kernel" << std::setw(15)
<< "Pred TFLOPS" << std::setw(12) << "Select ms" << std::setw(10) << "Status"
<< std::endl;
std::cout << std::string(87, '-') << std::endl;
bool all_passed = true;
for(const auto& [M, N, K] : sizes)
{
Problem problem;
problem.M = M;
problem.N = N;
problem.K = K;
problem.k_batch = 1;
auto t0 = std::chrono::high_resolution_clock::now();
auto kernel = dispatcher.select_kernel(problem);
auto t1 = std::chrono::high_resolution_clock::now();
double select_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
std::string size_str =
std::to_string(M) + "x" + std::to_string(N) + "x" + std::to_string(K);
if(!kernel)
{
std::cout << std::setw(20) << size_str << std::setw(30) << "NONE" << std::setw(15)
<< "N/A" << std::setw(12) << std::fixed << std::setprecision(2) << select_ms
<< std::setw(10) << "FAIL" << std::endl;
all_passed = false;
continue;
}
double pred = ml_heuristic.predict_tflops(problem, kernel->get_key());
std::string name = kernel->get_key().encode_identifier();
if(name.length() > 27)
name = name.substr(0, 27) + "..";
std::cout << std::setw(20) << size_str << std::setw(30) << name << std::setw(15)
<< std::fixed << std::setprecision(2) << pred << std::setw(12)
<< std::setprecision(2) << select_ms << std::setw(10) << "OK" << std::endl;
}
std::cout << std::endl
<< (all_passed ? "*** ALL TESTS PASSED ***" : "*** SOME TESTS FAILED ***")
<< std::endl;
return all_passed ? 0 : 1;
}

18
dispatcher/examples/gemm/cpp/README.md
View File

@@ -29,14 +29,14 @@ cd examples
## Examples
| Example | Description | Complexity |
|---------|-------------|------------|
| [01_basic_gemm.cpp](01_basic_gemm.cpp) | Basic GEMM with declarative API, autofill, autocorrect | ★☆☆☆☆ |
| [02_multi_size.cpp](02_multi_size.cpp) | Wildcard expansion for multiple configurations | ★★☆☆☆ |
| [03_benchmark_validation.cpp](03_benchmark_validation.cpp) | Performance benchmarking with CPU reference validation | ★★☆☆☆ |
| [04_heuristics.cpp](04_heuristics.cpp) | Heuristic-based kernel selection | ★★★☆☆ |
| [05_json_export.cpp](05_json_export.cpp) | Registry JSON export for external tools | ★★☆☆☆ |
| [06_multi_registry.cpp](06_multi_registry.cpp) | Multiple registries with named kernel sets | ★★★☆☆ |
| Example | Description |
|---------|-------------|
| [01_basic_gemm.cpp](01_basic_gemm.cpp) | Basic GEMM with declarative API, autofill, autocorrect |
| [02_multi_size.cpp](02_multi_size.cpp) | Wildcard expansion for multiple configurations |
| [03_benchmark_validation.cpp](03_benchmark_validation.cpp) | Performance benchmarking with CPU reference validation |
| [04_heuristics.cpp](04_heuristics.cpp) | Heuristic-based kernel selection |
| [05_json_export.cpp](05_json_export.cpp) | Registry JSON export for external tools |
| [06_multi_registry.cpp](06_multi_registry.cpp) | Multiple registries with named kernel sets |
## Example Details
@@ -225,5 +225,5 @@ DECL_KERNEL_SET(my_kernels,
## Related Documentation
- [Python GEMM Examples](../python/README.md)
- [Convolution Examples](../../conv/cpp/README.md)
- [C++ Headers (GEMM + Grouped Conv)](../../../include/ck_tile/dispatcher/README.md)
- [Main Dispatcher README](../../../README.md)

291
dispatcher/examples/gemm/python/01_basic_gemm.py
View File

@@ -7,41 +7,37 @@
Example 01: Basic GEMM with Multiple Kernels
Demonstrates:
1. Declaring multiple kernel configurations
2. Printing all registered kernels
3. Running each kernel and validating output
1. Building a Registry with multiple kernel configurations
2. Parallel JIT compilation via registry.build()
3. Running each kernel and validating output against NumPy reference
4. Comparing performance across kernels
Complexity: ★★☆☆☆
Usage:
python3 01_basic_gemm.py
python3 01_basic_gemm.py --help
python3 01_basic_gemm.py --dtype bf16
python3 01_basic_gemm.py --size 2048
python3 01_basic_gemm.py --num-kernels 4
python3 01_basic_gemm.py --workers 4
"""
import sys
import time
import argparse
from pathlib import Path
from dataclasses import dataclass
from typing import List
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
import numpy as np
from ctypes_utils import (
KernelConfig,
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
Registry,
detect_gpu_arch,
)
@dataclass
class KernelSpec:
"""Specification for a kernel configuration"""
name: str
tile_m: int
tile_n: int
@@ -50,80 +46,37 @@ class KernelSpec:
scheduler: str = "intrawave"
# Define multiple kernel configurations to test (50+ kernels)
KERNEL_SPECS = [
# Small tiles - compv3
# Small tiles
KernelSpec("small_64x64_k32", 64, 64, 32, "compv3"),
KernelSpec("small_64x64_k64", 64, 64, 64, "compv3"),
# Small tiles - compv4
KernelSpec("small_64x64_v4_k32", 64, 64, 32, "compv4"),
KernelSpec("small_64x64_v4_k64", 64, 64, 64, "compv4"),
# Medium tiles - compv3
# Medium tiles
KernelSpec("med_128x128_k32", 128, 128, 32, "compv3"),
KernelSpec("med_128x128_k64", 128, 128, 64, "compv3"),
KernelSpec("med_128x128_k128", 128, 128, 128, "compv3"),
# Medium tiles - compv4
KernelSpec("med_128x128_v4_k32", 128, 128, 32, "compv4"),
KernelSpec("med_128x128_v4_k64", 128, 128, 64, "compv4"),
KernelSpec("med_128x128_v4_k128", 128, 128, 128, "compv4"),
# Rectangular tiles - compv3
# Rectangular tiles
KernelSpec("rect_64x128_k32", 64, 128, 32, "compv3"),
KernelSpec("rect_64x128_k64", 64, 128, 64, "compv3"),
KernelSpec("rect_128x64_k32", 128, 64, 32, "compv3"),
KernelSpec("rect_128x64_k64", 128, 64, 64, "compv3"),
# Rectangular tiles - compv4
KernelSpec("rect_64x128_v4_k32", 64, 128, 32, "compv4"),
KernelSpec("rect_64x128_v4_k64", 64, 128, 64, "compv4"),
KernelSpec("rect_128x64_v4_k32", 128, 64, 32, "compv4"),
KernelSpec("rect_128x64_v4_k64", 128, 64, 64, "compv4"),
# Large tiles - compv3
# Large tiles
KernelSpec("large_256x128_k32", 256, 128, 32, "compv3"),
KernelSpec("large_256x128_k64", 256, 128, 64, "compv3"),
KernelSpec("large_128x256_k32", 128, 256, 32, "compv3"),
KernelSpec("large_128x256_k64", 128, 256, 64, "compv3"),
KernelSpec("large_256x256_k32", 256, 256, 32, "compv3"),
KernelSpec("large_256x256_k64", 256, 256, 64, "compv3"),
# Large tiles - compv4
KernelSpec("large_256x128_v4_k32", 256, 128, 32, "compv4"),
KernelSpec("large_256x128_v4_k64", 256, 128, 64, "compv4"),
KernelSpec("large_128x256_v4_k32", 128, 256, 32, "compv4"),
KernelSpec("large_128x256_v4_k64", 128, 256, 64, "compv4"),
KernelSpec("large_256x256_v4_k32", 256, 256, 32, "compv4"),
KernelSpec("large_256x256_v4_k64", 256, 256, 64, "compv4"),
# Interwave scheduler variants
KernelSpec("int_64x64_k32", 64, 64, 32, "compv3", "interwave"),
# Interwave scheduler
KernelSpec("int_128x128_k32", 128, 128, 32, "compv3", "interwave"),
KernelSpec("int_128x128_k64", 128, 128, 64, "compv3", "interwave"),
KernelSpec("int_256x128_k32", 256, 128, 32, "compv3", "interwave"),
# More tile_k variations - compv3
KernelSpec("med_128x128_k16", 128, 128, 16, "compv3"),
KernelSpec("rect_64x128_k16", 64, 128, 16, "compv3"),
KernelSpec("rect_128x64_k16", 128, 64, 16, "compv3"),
# More tile_k variations - compv4
KernelSpec("med_128x128_v4_k16", 128, 128, 16, "compv4"),
KernelSpec("rect_64x128_v4_k16", 64, 128, 16, "compv4"),
KernelSpec("rect_128x64_v4_k16", 128, 64, 16, "compv4"),
# Additional rectangular
KernelSpec("rect_32x64_k32", 32, 64, 32, "compv3"),
KernelSpec("rect_64x32_k32", 64, 32, 32, "compv3"),
KernelSpec("rect_32x128_k32", 32, 128, 32, "compv3"),
KernelSpec("rect_128x32_k32", 128, 32, 32, "compv3"),
# Additional compv4 variants
KernelSpec("rect_32x64_v4_k32", 32, 64, 32, "compv4"),
KernelSpec("rect_64x32_v4_k32", 64, 32, 32, "compv4"),
KernelSpec("rect_32x128_v4_k32", 32, 128, 32, "compv4"),
KernelSpec("rect_128x32_v4_k32", 128, 32, 32, "compv4"),
]
def create_kernel_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfig:
"""Create a KernelConfig from a spec"""
# Adjust warp tiles based on tile size
if spec.tile_m <= 64:
warp_m, warp_n = 16, 16
else:
warp_m, warp_n = 32, 32
def spec_to_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfig:
warp_m, warp_n = (16, 16) if spec.tile_m <= 64 else (32, 32)
return KernelConfig(
dtype_a=dtype,
dtype_b=dtype,
@@ -148,180 +101,118 @@ def create_kernel_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfi
)
def print_kernel_table(specs: List[KernelSpec], dtype: str):
"""Print a formatted table of kernel configurations"""
print("\n" + "=" * 70)
print(f" DECLARED KERNEL CONFIGURATIONS ({len(specs)} kernels)")
print("=" * 70)
print(f"\n {'#':<3} {'Name':<18} {'Tile':<14} {'Pipeline':<10} {'Scheduler':<12}")
print(" " + "-" * 68)
for i, spec in enumerate(specs, 1):
tile = f"{spec.tile_m}x{spec.tile_n}x{spec.tile_k}"
print(
f" {i:<3} {spec.name:<18} {tile:<14} {spec.pipeline:<10} {spec.scheduler:<12}"
)
print(" " + "-" * 68)
print(f" Data type: {dtype}")
def main():
parser = argparse.ArgumentParser(
description="Basic GEMM Example with Multiple Kernels",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
python3 01_basic_gemm.py # Default FP16 with 4 kernels
python3 01_basic_gemm.py --dtype bf16 # BF16 mode
python3 01_basic_gemm.py --size 2048 # Larger problem size
python3 01_basic_gemm.py --num-kernels 2 # Test only 2 kernels
""",
)
parser = argparse.ArgumentParser(description="Basic GEMM with Multiple Kernels")
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument("--size", type=int, default=512, help="Problem size MxNxK")
parser.add_argument("--num-kernels", type=int, default=0, help="0 = all")
parser.add_argument(
"--dtype",
default="fp16",
choices=["fp16", "bf16", "fp32"],
help="Data type (default: fp16)",
)
parser.add_argument(
"--arch",
default="gfx942",
help="Target architecture (default: gfx942)",
)
parser.add_argument(
"--size",
type=int,
default=512,
help="Problem size MxNxK (default: 512)",
)
parser.add_argument(
"--num-kernels",
type=int,
default=0,
help="Number of kernels to test (0 = all)",
"--workers", type=int, default=0, help="Max parallel JIT workers (0 = auto)"
)
args = parser.parse_args()
reset_for_example()
print("=" * 70)
print("Example 01: Basic GEMM with Multiple Kernels")
print("=" * 70)
# Select kernels to test
specs = KERNEL_SPECS[: args.num_kernels] if args.num_kernels > 0 else KERNEL_SPECS
# =========================================================================
# Step 1: Print all kernel configurations
# =========================================================================
print_kernel_table(specs, args.dtype)
# =========================================================================
# Step 2: Setup and test each kernel
# =========================================================================
print("\n" + "=" * 70)
print(" RUNNING KERNELS")
print("=" * 70)
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
M, N, K = args.size, args.size, args.size
results = []
print(f"\n Problem size: {M}x{N}x{K}\n")
# Step 1: Build registry
print(
f" {'#':<3} {'Name':<18} {'Tile':<14} {'Time (ms)':>10} {'TFLOPS':>10} {'Max Err':>10} {'Status':<8}"
f"\n {len(specs)} kernel configurations, dtype={args.dtype}, arch={args.arch}"
)
print(" " + "-" * 78)
for i, spec in enumerate(specs, 1):
# Create unique test data per kernel
np.random.seed(42 + i * 1000)
A = (np.random.randn(M, K) * 0.1).astype(np_dtype)
B = (np.random.randn(K, N) * 0.1).astype(np_dtype)
# Create config and setup dispatcher
config = create_kernel_config(spec, args.dtype, args.arch)
setup = setup_gemm_dispatcher(
config=config,
registry_name=f"kernel_{spec.name}",
verbose=False,
auto_rebuild=True,
print(f"\n {'#':<3} {'Name':<22} {'Tile':<14} {'Pipeline':<10} {'Scheduler':<12}")
print(" " + "-" * 64)
for i, s in enumerate(specs, 1):
print(
f" {i:<3} {s.name:<22} {s.tile_m}x{s.tile_n}x{s.tile_k:<6} {s.pipeline:<10} {s.scheduler:<12}"
)
reg = Registry(name="basic_gemm")
for s in specs:
reg.register_kernel(spec_to_config(s, args.dtype, args.arch))
# Step 2: Parallel JIT build via registry.build()
workers = args.workers if args.workers > 0 else None
print(
f"\n--- Parallel JIT Build ({len(specs)} kernels{f', workers={workers}' if workers else ''}) ---"
)
t0 = time.perf_counter()
setups = reg.build(verbose=False, max_workers=workers)
jit_build_s = time.perf_counter() - t0
built = sum(1 for s in setups if s.success)
print(f" Built: {built}/{len(specs)} kernels in {jit_build_s:.1f} s")
if built == 0:
print(" ERROR: No kernels built")
return 1
# Step 3: Run each kernel and validate
print(f"\n--- Running Kernels (problem {args.size}x{args.size}x{args.size}) ---")
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
M = N = K = args.size
np.random.seed(42)
A = (np.random.randn(M, K) * 0.1).astype(np_dtype)
B = (np.random.randn(K, N) * 0.1).astype(np_dtype)
C_ref = np.matmul(A.astype(np.float32), B.astype(np.float32)).astype(np_dtype)
print(
f"\n {'#':<3} {'Name':<22} {'Tile':<14} {'Time(ms)':>10} {'TFLOPS':>10} {'MaxErr':>10} {'Status':<6}"
)
print(" " + "-" * 80)
results = []
for i, (spec, setup) in enumerate(zip(specs, setups), 1):
tile = f"{spec.tile_m}x{spec.tile_n}x{spec.tile_k}"
if not setup.success:
print(
f" {i:<3} {spec.name:<18} {tile:<14} {'N/A':>10} {'N/A':>10} {'N/A':>10} {'FAIL':<8}"
f" {i:<3} {spec.name:<22} {tile:<14} {'---':>10} {'---':>10} {'---':>10} {'SKIP':<6}"
)
results.append((spec.name, False, 0, 0, 0))
cleanup_gemm()
results.append((spec.name, False, 0.0, 0.0, 0.0))
continue
dispatcher = setup.dispatcher
# Check if size is supported
if not dispatcher.is_supported(M, N, K):
disp = setup.dispatcher
if not disp.is_supported(M, N, K):
print(
f" {i:<3} {spec.name:<18} {tile:<14} {'N/A':>10} {'N/A':>10} {'N/A':>10} {'SKIP':<8}"
f" {i:<3} {spec.name:<22} {tile:<14} {'---':>10} {'---':>10} {'---':>10} {'SKIP':<6}"
)
results.append((spec.name, False, 0, 0, 0))
cleanup_gemm()
results.append((spec.name, False, 0.0, 0.0, 0.0))
continue
# Run GEMM
result = dispatcher.run(A, B, M, N, K)
if not result.success:
res = disp.run(A, B, M, N, K)
if not res.success:
print(
f" {i:<3} {spec.name:<18} {tile:<14} {'N/A':>10} {'N/A':>10} {'N/A':>10} {'FAIL':<8}"
f" {i:<3} {spec.name:<22} {tile:<14} {'---':>10} {'---':>10} {'---':>10} {'FAIL':<6}"
)
results.append((spec.name, False, 0, 0, 0))
cleanup_gemm()
results.append((spec.name, False, 0.0, 0.0, 0.0))
continue
# Validate against NumPy reference
C_ref = np.matmul(A.astype(np.float32), B.astype(np.float32)).astype(np_dtype)
max_err = np.max(np.abs(result.output - C_ref))
# Check if within tolerance
passed = max_err < 1e-2
status = "PASS" if passed else "FAIL"
max_err = float(np.max(np.abs(res.output - C_ref)))
ok = max_err < 1e-2
tag = "PASS" if ok else "FAIL"
print(
f" {i:<3} {spec.name:<18} {tile:<14} {result.time_ms:>10.4f} {result.tflops:>10.2f} {max_err:>10.2e} {status:<8}"
f" {i:<3} {spec.name:<22} {tile:<14} {res.time_ms:>10.4f} {res.tflops:>10.2f} {max_err:>10.2e} {tag:<6}"
)
results.append((spec.name, passed, result.time_ms, result.tflops, max_err))
cleanup_gemm()
# =========================================================================
# Step 3: Summary
# =========================================================================
print("\n" + "=" * 70)
print(" SUMMARY")
print("=" * 70)
results.append((spec.name, ok, res.time_ms, res.tflops, max_err))
# Step 4: Summary
passed = sum(1 for r in results if r[1])
failed = len(results) - passed
valid = [r for r in results if r[1]]
print(f"\n Results: {passed}/{len(results)} kernels passed")
print(f" Problem: {M}x{N}x{K}, dtype={args.dtype}")
if results:
valid_results = [r for r in results if r[1]]
if valid_results:
best = max(valid_results, key=lambda x: x[3])
print(f"\n Best kernel: {best[0]} ({best[3]:.2f} TFLOPS)")
if failed == 0:
print("\n *** ALL KERNELS PASSED ***")
else:
print(f"\n *** {failed} KERNELS FAILED ***")
print("\n" + "=" * 70)
print(f" Results: {passed}/{len(results)} passed")
print(f" Problem: {M}x{N}x{K}, dtype={args.dtype}")
print(f" JIT time: {jit_build_s:.1f} s (parallel)")
if valid:
best = max(valid, key=lambda x: x[3])
print(f" Best: {best[0]} ({best[3]:.2f} TFLOPS)")
print(f" Status: {'PASS' if failed == 0 else 'FAIL'}")
print("=" * 70)
return 0 if failed == 0 else 1

35
dispatcher/examples/gemm/python/02_batch_gemm.py
View File

@@ -6,9 +6,7 @@
"""
Example 02: Batch GEMM
Runs multiple GEMM operations with different sizes.
Complexity: ★★☆☆☆
Runs multiple GEMM operations with different sizes using JIT compilation.
Usage:
python3 02_batch_gemm.py
@@ -25,9 +23,8 @@ import numpy as np
from ctypes_utils import (
KernelConfig,
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
Registry,
detect_gpu_arch,
)
@@ -55,20 +52,20 @@ Examples:
help="Maximum problem size (default: 4096)",
)
parser.add_argument(
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
"--arch",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo)",
)
args = parser.parse_args()
reset_for_example()
print("=" * 60)
print("Example 02: Batch GEMM")
print("=" * 60)
# =========================================================================
# Step 1: Setup dispatcher
# Step 1: JIT build dispatcher
# =========================================================================
print("\nStep 1: Setup Dispatcher")
print("\nStep 1: JIT Build Dispatcher")
config = KernelConfig(
dtype_a=args.dtype,
@@ -80,19 +77,22 @@ Examples:
gfx_arch=args.arch,
)
setup = setup_gemm_dispatcher(config, registry_name="batch_gemm", verbose=True)
if not setup.success:
print(f" ERROR: {setup.error}")
reg = Registry(name="batch_gemm")
reg.register_kernel(config)
setups = reg.build(verbose=True)
if not setups or not setups[0].success:
error = setups[0].error if setups else "No kernels built"
print(f" ERROR: {error}")
return 1
dispatcher = setup.dispatcher
dispatcher = setups[0].dispatcher
# =========================================================================
# Step 2: Run batch of different sizes
# =========================================================================
print("\nStep 2: Run Batch")
# Generate sizes up to max_size
all_sizes = [
(256, 256, 256),
(512, 512, 512),
@@ -135,9 +135,6 @@ Examples:
avg_tflops = (total_ops / 1e12) / (total_time / 1000)
print(f"\n Total: {total_time:.2f} ms, Average: {avg_tflops:.2f} TFLOPS")
# Cleanup
cleanup_gemm()
print("\n" + "=" * 60)
print("Batch GEMM complete!")
print("=" * 60)

37
dispatcher/examples/gemm/python/03_benchmark.py
View File

@@ -6,9 +6,8 @@
"""
Example 03: Benchmark
Performance benchmarking with compute-optimized kernel configuration.
Complexity: ★★★☆☆
Performance benchmarking with compute-optimized kernel configuration
using JIT compilation.
Usage:
python3 03_benchmark.py
@@ -26,9 +25,8 @@ import numpy as np
from ctypes_utils import (
KernelConfig,
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
Registry,
detect_gpu_arch,
)
@@ -63,20 +61,20 @@ Examples:
"--iterations", type=int, default=10, help="Benchmark iterations (default: 10)"
)
parser.add_argument(
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
"--arch",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo)",
)
args = parser.parse_args()
reset_for_example()
print("=" * 60)
print("Example 03: Benchmark")
print("=" * 60)
# =========================================================================
# Step 1: Setup dispatcher with compute-optimized config
# Step 1: JIT build dispatcher with compute-optimized config
# =========================================================================
print("\nStep 1: Setup Dispatcher")
print("\nStep 1: JIT Build Dispatcher")
config = KernelConfig(
dtype_a=args.dtype,
@@ -90,12 +88,16 @@ Examples:
gfx_arch=args.arch,
)
setup = setup_gemm_dispatcher(config, registry_name="benchmark", verbose=True)
if not setup.success:
print(f" ERROR: {setup.error}")
reg = Registry(name="benchmark")
reg.register_kernel(config)
setups = reg.build(verbose=True)
if not setups or not setups[0].success:
error = setups[0].error if setups else "No kernels built"
print(f" ERROR: {error}")
return 1
dispatcher = setup.dispatcher
dispatcher = setups[0].dispatcher
# =========================================================================
# Step 2: Benchmark
@@ -130,11 +132,9 @@ Examples:
A = np.random.randn(M, K).astype(np_dtype) * 0.1
B = np.random.randn(K, N).astype(np_dtype) * 0.1
# Warmup
for _ in range(args.warmup):
dispatcher.run(A, B, M, N, K)
# Benchmark
times = []
for _ in range(args.iterations):
result = dispatcher.run(A, B, M, N, K)
@@ -150,9 +150,6 @@ Examples:
f" {M:>4}x{N:>4}x{K:<4} | {min_time:>10.4f} | {avg_time:>10.4f} | {tflops:>10.2f}"
)
# Cleanup
cleanup_gemm()
# Summary
print("\n" + "=" * 60)
print("Summary")

34
dispatcher/examples/gemm/python/04_validation.py
View File

@@ -6,9 +6,7 @@
"""
Example 04: Validation
Validates GPU GEMM against NumPy reference.
Complexity: ★★★☆☆
Validates GPU GEMM against NumPy reference using JIT compilation.
Usage:
python3 04_validation.py
@@ -26,9 +24,8 @@ import numpy as np
from ctypes_utils import (
KernelConfig,
Validator,
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
Registry,
detect_gpu_arch,
)
@@ -56,20 +53,20 @@ Examples:
"--atol", type=float, default=1e-2, help="Absolute tolerance (default: 1e-2)"
)
parser.add_argument(
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
"--arch",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo)",
)
args = parser.parse_args()
reset_for_example()
print("=" * 60)
print("Example 04: Validation")
print("=" * 60)
# =========================================================================
# Step 1: Setup dispatcher
# Step 1: JIT build dispatcher
# =========================================================================
print("\nStep 1: Setup Dispatcher")
print("\nStep 1: JIT Build Dispatcher")
config = KernelConfig(
dtype_a=args.dtype,
@@ -81,12 +78,16 @@ Examples:
gfx_arch=args.arch,
)
setup = setup_gemm_dispatcher(config, registry_name="validation", verbose=True)
if not setup.success:
print(f" ERROR: {setup.error}")
reg = Registry(name="validation")
reg.register_kernel(config)
setups = reg.build(verbose=True)
if not setups or not setups[0].success:
error = setups[0].error if setups else "No kernels built"
print(f" ERROR: {error}")
return 1
dispatcher = setup.dispatcher
dispatcher = setups[0].dispatcher
# =========================================================================
# Step 2: Run validation tests
@@ -139,9 +140,6 @@ Examples:
print(f" {name:<15} | {M}x{N}x{K:<5} | {max_err:>10.2e} | FAILED")
failed += 1
# Cleanup
cleanup_gemm()
# Summary
print("\n" + "=" * 60)
total = passed + failed

6
dispatcher/examples/gemm/python/05_numpy_integration.py
View File

@@ -8,7 +8,6 @@ Example 05: NumPy Integration
Shows how to create a GPU-accelerated matmul wrapper.
Complexity: ★★☆☆☆
Usage:
python3 05_numpy_integration.py
@@ -29,6 +28,7 @@ from ctypes_utils import (
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
detect_gpu_arch,
)
@@ -70,7 +70,9 @@ Examples:
help="Data type (default: fp16)",
)
parser.add_argument(
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
"--arch",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo)",
)
args = parser.parse_args()

6
dispatcher/examples/gemm/python/06_json_export.py
View File

@@ -8,7 +8,6 @@ Example 06: JSON Export
Exports registry configuration to JSON.
Complexity: ★★☆☆☆
Usage:
python3 06_json_export.py
@@ -28,6 +27,7 @@ from ctypes_utils import (
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
detect_gpu_arch,
)
@@ -54,7 +54,9 @@ Examples:
help="Data type (default: fp16)",
)
parser.add_argument(
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
"--arch",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo)",
)
args = parser.parse_args()

6
dispatcher/examples/gemm/python/07_stress_test.py
View File

@@ -18,7 +18,6 @@ This tests:
- Multiple data types (fp16, bf16)
- Different schedulers (intrawave, interwave)
Complexity: ★★★★☆
Usage:
python3 07_stress_test.py
@@ -43,6 +42,7 @@ from ctypes_utils import (
cleanup_gemm,
reset_for_example,
Validator,
detect_gpu_arch,
)
@@ -413,8 +413,8 @@ Examples:
)
parser.add_argument(
"--arch",
default="gfx942",
help="Target architecture (default: gfx942)",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo, override with --arch gfxNNN)",
)
args = parser.parse_args()

6
dispatcher/examples/gemm/python/08_heuristics.py
View File

@@ -19,7 +19,6 @@ Heuristic strategies:
- Memory-bound: Optimize memory access for bandwidth-limited cases
- Latency-focused: Minimize kernel launch overhead for small problems
Complexity: ★★★★☆
Usage:
python3 08_heuristics.py
@@ -43,6 +42,7 @@ from ctypes_utils import (
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
detect_gpu_arch,
)
@@ -561,8 +561,8 @@ Examples:
)
parser.add_argument(
"--arch",
default="gfx942",
help="Target architecture (default: gfx942)",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo, override with --arch gfxNNN)",
)
args = parser.parse_args()

305
dispatcher/examples/gemm/python/09_ml_heuristic.py Normal file
View File

@@ -0,0 +1,305 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 09: ML-Based Kernel Selection
Uses a trained LightGBM model to select the optimal kernel for each problem
size. The model predicts TFLOPS for every candidate in the kernel pool and
picks the highest-scoring one, which is then JIT-compiled and run.
This replaces the hand-crafted rules in 08_heuristics.py with a data-driven
approach achieving 97-98% of oracle-best TFLOPS efficiency.
Complexity: *****
Prerequisites:
- Trained model in dispatcher/heuristics/models/gemm_universal_fp8_gfx950/
- lightgbm, pandas, numpy, pyarrow installed
Usage:
python3 09_ml_heuristic.py
python3 09_ml_heuristic.py --dtype fp16 --arch gfx942
"""
import sys
import argparse
import time
from pathlib import Path
from dataclasses import dataclass
from typing import List
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "heuristics"))
import numpy as np
from ctypes_utils import (
KernelConfig,
setup_gemm_dispatcher,
cleanup_gemm,
)
from predict import Predictor
@dataclass
class KernelSpec:
"""Kernel specification -- same structure as 08_heuristics.py"""
name: str
tile_m: int
tile_n: int
tile_k: int
pipeline: str = "compv3"
scheduler: str = "intrawave"
wave_m: int = 2
wave_n: int = 2
wave_k: int = 1
warp_m: int = 32
warp_n: int = 32
warp_k: int = 16
# Kernel pool: representative configs spanning small to large tiles,
# compv3/compv4/mem pipelines, and intrawave/interwave schedulers.
KERNEL_POOL = [
# Small tiles
KernelSpec("s_64x64_k32_v3", 64, 64, 32, "compv3", warp_m=16, warp_n=16),
KernelSpec("s_64x64_k64_v3", 64, 64, 64, "compv3", warp_m=16, warp_n=16),
KernelSpec("s_64x64_k128_v3", 64, 64, 128, "compv3", warp_m=16, warp_n=16),
KernelSpec("s_64x64_k32_v4", 64, 64, 32, "compv4", warp_m=16, warp_n=16),
KernelSpec("s_64x64_k64_mem", 64, 64, 64, "mem", warp_m=16, warp_n=16),
KernelSpec("s_64x64_k128_mem", 64, 64, 128, "mem", warp_m=16, warp_n=16),
# Medium tiles
KernelSpec("m_128x128_k32_v3", 128, 128, 32, "compv3"),
KernelSpec("m_128x128_k64_v3", 128, 128, 64, "compv3"),
KernelSpec("m_128x128_k128_v3", 128, 128, 128, "compv3"),
KernelSpec("m_128x128_k32_v4", 128, 128, 32, "compv4"),
KernelSpec("m_128x128_k64_v4", 128, 128, 64, "compv4"),
KernelSpec("m_128x128_k64_mem", 128, 128, 64, "mem"),
KernelSpec("m_128x128_k128_mem", 128, 128, 128, "mem"),
# Rectangular medium
KernelSpec("r_64x128_k32", 64, 128, 32, "compv3", warp_m=16),
KernelSpec("r_128x64_k32", 128, 64, 32, "compv3", warp_n=16),
KernelSpec("r_64x128_k64", 64, 128, 64, "compv3", warp_m=16),
KernelSpec("r_128x64_k64", 128, 64, 64, "compv3", warp_n=16),
# Large tiles
KernelSpec("l_256x128_k32", 256, 128, 32, "compv3"),
KernelSpec("l_128x256_k32", 128, 256, 32, "compv3"),
KernelSpec("l_256x256_k32", 256, 256, 32, "compv3"),
KernelSpec("l_256x256_k64", 256, 256, 64, "compv3"),
# Interwave variants
KernelSpec("m_128x128_k64_iw", 128, 128, 64, "compv3", "interwave"),
KernelSpec("m_128x128_k64_mem_iw", 128, 128, 64, "mem", "interwave"),
]
def spec_to_feature_dict(spec: KernelSpec, dtype: str, layout: str) -> dict:
"""Convert a KernelSpec to the dict format the feature engine expects.
Note: pad_m/n/k default to True to match KernelConfig defaults and actual
compiled kernels. This ensures the ML model receives the correct padding
flags that will be used during JIT compilation.
"""
return {
"kernel_name": spec.name,
"tile_m": spec.tile_m,
"tile_n": spec.tile_n,
"tile_k": spec.tile_k,
"warp_m": spec.wave_m,
"warp_n": spec.wave_n,
"warp_k": spec.wave_k,
"warp_tile_m": spec.warp_m,
"warp_tile_n": spec.warp_n,
"warp_tile_k": spec.warp_k,
"pipeline": spec.pipeline,
"scheduler": spec.scheduler,
"epilogue": "cshuffle",
"pad_m": True, # Match KernelConfig default
"pad_n": True, # Match KernelConfig default
"pad_k": True, # Match KernelConfig default
"persistent": False,
"dtype": dtype,
"layout": layout,
}
def spec_to_kernel_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfig:
"""Convert a KernelSpec to the dispatcher's KernelConfig for JIT compilation."""
return KernelConfig(
dtype_a=dtype,
dtype_b=dtype,
dtype_c=dtype,
dtype_acc="fp32",
layout_a="row",
layout_b="col",
layout_c="row",
tile_m=spec.tile_m,
tile_n=spec.tile_n,
tile_k=spec.tile_k,
wave_m=spec.wave_m,
wave_n=spec.wave_n,
wave_k=spec.wave_k,
warp_m=spec.warp_m,
warp_n=spec.warp_n,
warp_k=spec.warp_k,
pipeline=spec.pipeline,
scheduler=spec.scheduler,
epilogue="cshuffle",
gfx_arch=arch,
)
def ml_select_kernel(
predictor: Predictor,
pool: List[KernelSpec],
M: int,
N: int,
K: int,
dtype: str,
layout: str,
) -> tuple:
"""Score all kernels in the pool and return (best_spec, predicted_tflops)."""
problem = {"m": M, "n": N, "k": K, "dtype": dtype, "layout": layout, "split_k": 1}
kernel_dicts = [spec_to_feature_dict(s, dtype, layout) for s in pool]
ranked = predictor.rank_kernels(problem, kernel_dicts)
if not ranked:
return pool[0], 0.0
best_name, best_tflops = ranked[0]
best_spec = next((s for s in pool if s.name == best_name), pool[0])
return best_spec, best_tflops
def main():
parser = argparse.ArgumentParser(description="ML-based kernel selection for GEMM")
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16", "fp8"])
parser.add_argument("--arch", default="gfx942")
parser.add_argument(
"--model_dir",
default=str(
Path(__file__).parent.parent.parent.parent
/ "heuristics"
/ "models"
/ "gemm_universal_fp8_gfx950"
),
)
parser.add_argument(
"--no_run", action="store_true", help="Only predict, don't run GEMMs"
)
args = parser.parse_args()
print("=" * 75)
print(" Example 09: ML-Based Kernel Selection")
print("=" * 75)
print(f"\n Model: {args.model_dir}")
print(f" Dtype: {args.dtype}")
print(f" Arch: {args.arch}")
print(f" Pool: {len(KERNEL_POOL)} kernels")
predictor = Predictor(args.model_dir)
print(" Model loaded successfully")
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float16
test_sizes = [
(128, 128, 64),
(256, 256, 128),
(512, 512, 256),
(1024, 1024, 512),
(2048, 2048, 1024),
]
header = f"{'Shape':<20} {'Selected Kernel':<25} {'Pred TFLOPS':>12}"
if not args.no_run:
header += f" {'Time (ms)':>10} {'TFLOPS':>10} {'Status':<8}"
print(f"\n {header}")
print(" " + "-" * len(header))
results = []
for M, N, K in test_sizes:
t0 = time.time()
best_spec, pred_tflops = ml_select_kernel(
predictor, KERNEL_POOL, M, N, K, args.dtype, "rcr"
)
_ = (time.time() - t0) * 1000 # ML selection time (unused)
size_str = f"{M}x{N}x{K}"
line = f" {size_str:<20} {best_spec.name:<25} {pred_tflops:>12.2f}"
if args.no_run:
print(line)
results.append((size_str, best_spec.name, True, 0, pred_tflops))
continue
config = spec_to_kernel_config(best_spec, args.dtype, args.arch)
setup = setup_gemm_dispatcher(
config=config,
registry_name=f"ml_{best_spec.name}",
verbose=False,
auto_rebuild=True,
)
if not setup.success:
line += f" {'N/A':>10} {'N/A':>10} {'BUILD':>8}"
print(line)
results.append((size_str, best_spec.name, False, 0, 0))
cleanup_gemm()
continue
dispatcher = setup.dispatcher
if not dispatcher.is_supported(M, N, K):
line += f" {'N/A':>10} {'N/A':>10} {'UNSUP':>8}"
print(line)
results.append((size_str, best_spec.name, False, 0, 0))
cleanup_gemm()
continue
np.random.seed(42)
A = (np.random.randn(M, K) * 0.1).astype(np_dtype)
B = (np.random.randn(K, N) * 0.1).astype(np_dtype)
result = dispatcher.run(A, B, M, N, K)
if result.success:
C_ref = np.matmul(A.astype(np.float32), B.astype(np.float32)).astype(
np_dtype
)
max_err = np.max(np.abs(result.output - C_ref))
passed = max_err < 1e-2
status = "PASS" if passed else "FAIL"
line += f" {result.time_ms:>10.4f} {result.tflops:>10.2f} {status:<8}"
results.append(
(size_str, best_spec.name, passed, result.time_ms, result.tflops)
)
else:
line += f" {'N/A':>10} {'N/A':>10} {'FAIL':<8}"
results.append((size_str, best_spec.name, False, 0, 0))
print(line)
cleanup_gemm()
# Summary
print("\n" + "=" * 75)
print(" SUMMARY")
print("=" * 75)
passed = sum(1 for r in results if r[2])
print(f"\n Results: {passed}/{len(results)} tests passed")
valid = [r for r in results if r[2] and r[4] > 0]
if valid:
avg = sum(r[4] for r in valid) / len(valid)
print(f" Average TFLOPS: {avg:.2f}")
if passed == len(results):
print("\n *** ALL TESTS PASSED ***")
print("=" * 75)
return 0 if passed == len(results) else 1
if __name__ == "__main__":
sys.exit(main())

6
dispatcher/examples/gemm/python/09_multi_registry.py
View File

@@ -8,7 +8,6 @@ Example 09: Multiple Registries
Demonstrates multiple registries for different optimization targets.
Complexity: ★★★★★
Usage:
python3 09_multi_registry.py
@@ -30,6 +29,7 @@ from ctypes_utils import (
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
detect_gpu_arch,
)
@@ -50,7 +50,9 @@ Examples:
help="Data type (default: fp16)",
)
parser.add_argument(
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
"--arch",
default=detect_gpu_arch(),
help="Target architecture (auto-detected from rocminfo)",
)
args = parser.parse_args()

7
dispatcher/examples/gemm/python/10_advanced_benchmark.py
View File

@@ -33,6 +33,7 @@ from ctypes_utils import (
setup_gemm_dispatcher,
cleanup_gemm,
reset_for_example,
detect_gpu_arch,
)
@@ -69,7 +70,11 @@ def parse_args():
# Kernel configuration
parser.add_argument("--dtype", default="fp16", help="Data type")
parser.add_argument("--pipeline", default="compv4", help="Pipeline type")
parser.add_argument("--arch", default="gfx942", help="GPU architecture")
parser.add_argument(
"--arch",
default=detect_gpu_arch(),
help="GPU architecture (auto-detected from rocminfo)",
)
return parser.parse_args()

16
dispatcher/examples/gemm/python/11_json_import.py
View File

@@ -15,7 +15,6 @@ Key Features:
- Use arch_filter validation on loaded configs
- Export to C++ DECL_KERNEL_SET format
Complexity: ★★★☆☆
Usage:
python3 11_json_import.py
@@ -45,6 +44,7 @@ from ctypes_utils import ( # noqa: E402
cleanup_gemm,
reset_for_example,
validate_kernel_config,
detect_gpu_arch,
)
# Sample JSON configuration (embedded for demonstration)
@@ -141,8 +141,8 @@ Examples:
)
parser.add_argument(
"--arch",
default="gfx942",
help="Target GPU architecture (default: gfx942)",
default=detect_gpu_arch(),
help="Target GPU architecture (auto-detected from rocminfo, override with --arch gfxNNN)",
)
args = parser.parse_args()
@@ -236,13 +236,13 @@ Examples:
else:
invalid_count += 1
if invalid_count <= 3: # Show first 3 invalid
print(f"\n ✗ Invalid: {config.kernel_name()}")
print(f"\n FAIL Invalid: {config.kernel_name()}")
for error in result.errors:
print(f" Error: {error}")
print("\n Validation Summary:")
print(f" ✓ Valid: {valid_count}")
print(f" ✗ Invalid: {invalid_count}")
print(f" OK Valid: {valid_count}")
print(f" FAIL Invalid: {invalid_count}")
print(f" Total: {len(configs)}")
# =========================================================================
@@ -275,12 +275,12 @@ Examples:
disp_config, registry_name="json_import", verbose=False
)
if setup.success:
print(" ✓ Dispatcher setup successful")
print(" OK Dispatcher setup successful")
print(
f" Kernel header: {setup.kernel_header.name if setup.kernel_header else 'N/A'}"
)
else:
print(f" ⚠ Dispatcher setup: {setup.error}")
print(f" WARNING Dispatcher setup: {setup.error}")
print(" (This is expected if kernels aren't generated)")
# =========================================================================

2
dispatcher/examples/gemm/python/README.md
View File

@@ -295,5 +295,5 @@ Compilation time scales roughly linearly with kernel count.
## Related Documentation
- [C++ GEMM Examples](../cpp/README.md)
- [Python Conv Examples](../../conv/python/README.md)
- [Python Utilities](../../../python/README.md)
- [Main Dispatcher README](../../../README.md)

203
dispatcher/examples/grouped_conv/cpp/01_basic_grouped_conv.cpp Normal file
View File

@@ -0,0 +1,203 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 01: Basic Grouped Convolution
//
// Demonstrates three declaration patterns (mirrors GEMM 01):
// 1. AUTOFILL - tile + pipeline only, wave/warp auto-filled
// 2. AUTOCORRECT - invalid wave(1,1,1) corrected to valid config
// 3. FULL - all parameters explicit (matches validated gfx942 config)
//
// Then runs the forward convolution on GPU and verifies output.
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_01_basic
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
// Three declaration patterns -- codegen auto-fills/auto-corrects as needed
DECL_GROUPED_CONV_KERNEL_SET(
basic_conv_kernels,
// Pattern 1: AUTOFILL - only tile + pipeline, rest auto-filled
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo().tile(1, 128, 128).pipeline("compv4").scheduler("intrawave"),
"gfx950")
// Pattern 2: AUTOCORRECT - wave(1,1,1) invalid, corrected to (1,4,1)
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo()
.tile(1, 64, 64)
.wave(1, 1, 1)
.warp(16, 16, 32)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle")
.vector_sizes(4, 8, 8),
"gfx950")
// Pattern 3: FULL - all parameters explicit (validated config)
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo()
.tile(1, 128, 128)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle")
.vector_sizes(4, 8, 8)
.block_per_cu(1),
"gfx950"));
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 01: Basic Grouped Convolution",
"Declaration patterns + GPU execution");
args.add_option("--arch", "gfx950", "GPU architecture");
args.add_option("--size", "14", "Spatial size (H=W)");
args.add_option("-n", "1", "Batch size");
args.add_option("-g", "1", "Groups");
args.add_option("-c", "64", "Input channels C");
args.add_option("-k", "128", "Output channels K");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 01: Basic Grouped Convolution");
std::string gfx_arch = args.get("--arch", "gfx950");
int N = args.get_int("-n", 1);
int G = args.get_int("-g", 1);
int C = args.get_int("-c", 64);
int K = args.get_int("-k", 128);
int HW = args.get_int("--size", 14);
int Y = 3, X = 3;
// Step 1: Show declared kernel sets
std::cout << "\nStep 1: Declared Kernel Sets\n";
GroupedConvKernelSetRegistry::instance().print();
// Step 2: Register kernels
std::cout << "\nStep 2: Register Kernels\n";
GroupedConvRegistry registry;
registry.set_name("basic_conv");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)\n";
// Step 3: Create dispatcher
std::cout << "\nStep 3: Create Dispatcher\n";
GroupedConvDispatcher dispatcher(&registry);
// Step 4: Build problem using CK Tile ConvParam
std::cout << "\nStep 4: Problem\n";
auto problem = create_grouped_conv2d_problem(N, C, K, HW, HW, Y, X, 1, 1);
problem.op = GroupedConvOp::Forward;
print_grouped_conv_problem(problem);
ck_tile::conv::ConvParam conv_param{
2,
static_cast<ck_tile::index_t>(G),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(K),
static_cast<ck_tile::index_t>(C),
{static_cast<ck_tile::index_t>(Y), static_cast<ck_tile::index_t>(X)},
{static_cast<ck_tile::index_t>(HW), static_cast<ck_tile::index_t>(HW)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
auto in_desc =
ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
auto wei_desc =
ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(conv_param);
auto out_desc =
ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(conv_param);
ck_tile::HostTensor<InDataType> input_host(in_desc);
ck_tile::HostTensor<WeiDataType> weight_host(wei_desc);
ck_tile::HostTensor<OutDataType> output_host(out_desc);
ck_tile::FillUniformDistribution<InDataType>{-0.5f, 0.5f}(input_host);
ck_tile::FillUniformDistribution<WeiDataType>{-0.5f, 0.5f}(weight_host);
ck_tile::DeviceMem input_dev(input_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem weight_dev(weight_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem output_dev(output_host.get_element_space_size_in_bytes());
input_dev.ToDevice(input_host.data());
weight_dev.ToDevice(weight_host.data());
// Step 5: Select and run
std::cout << "\nStep 5: Select and Run\n";
auto* selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cerr << " ERROR: No kernel found for problem!\n";
return 1;
}
std::cout << " Selected: " << selected->name() << "\n";
float time_ms = dispatcher.run(input_dev.GetDeviceBuffer(),
weight_dev.GetDeviceBuffer(),
output_dev.GetDeviceBuffer(),
problem,
nullptr);
double tflops = calculate_conv_tflops(problem, time_ms);
std::cout << " Time: " << std::fixed << std::setprecision(4) << time_ms << " ms\n";
std::cout << " TFLOPS: " << std::setprecision(2) << tflops << "\n";
// Step 6: Verify
std::cout << "\nStep 6: Verify\n";
output_dev.FromDevice(output_host.data());
size_t total = output_host.get_element_space_size();
size_t nonzero = 0;
double checksum = 0.0;
for(size_t i = 0; i < total; ++i)
{
float v = static_cast<float>(output_host.data()[i]);
if(v != 0.0f)
++nonzero;
checksum += v;
}
bool passed = nonzero > 0;
std::cout << " Output elements: " << total << "\n";
std::cout << " Non-zero: " << nonzero << "/" << total
<< (nonzero > 0 ? " (kernel produced output)" : " WARNING: all zeros!") << "\n";
std::cout << " Checksum: " << std::fixed << std::setprecision(2) << checksum << "\n";
std::cout << " Status: " << (passed ? "PASS" : "FAIL") << "\n";
utils::print_separator();
std::cout << "DECLARATION PATTERNS:\n";
std::cout << " 1. AUTOFILL: tile + pipeline only, wave/warp auto-filled\n";
std::cout << " 2. AUTOCORRECT: invalid wave(1,1,1) corrected\n";
std::cout << " 3. FULL: all parameters explicit\n";
utils::print_separator();
return passed ? 0 : 1;
}

216
dispatcher/examples/grouped_conv/cpp/02_all_directions.cpp Normal file
View File

@@ -0,0 +1,216 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 02: All Convolution Directions
//
// Forward, backward-data, and backward-weight for 2D convolution,
// each executed on GPU with non-zero verification.
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_02_all_dirs
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
DECL_GROUPED_CONV_KERNEL_SET(
conv_fwd_2d,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo().tile(1, 128, 128).pipeline("compv4").vector_sizes(4, 8, 8),
"gfx950"));
DECL_GROUPED_CONV_KERNEL_SET(
conv_bwdd_2d,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("bwd_data").dims(2),
GroupedConvAlgo().tile(1, 128, 128).pipeline("compv3").vector_sizes(4, 8, 8),
"gfx950"));
DECL_GROUPED_CONV_KERNEL_SET(
conv_bwdw_2d,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("bwd_weight").dims(2),
GroupedConvAlgo()
.tile(1, 128, 128)
.pipeline("compv3")
.memory_op("atomic_add")
.vector_sizes(4, 8, 8),
"gfx950"));
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 02: All Convolution Directions",
"Forward/BwdData/BwdWeight with GPU execution and verification");
args.add_option("--arch", "gfx950", "GPU architecture");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 02: All Convolution Directions");
std::string gfx_arch = args.get("--arch", "gfx950");
GroupedConvRegistry registry;
registry.set_name("all_directions");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)\n";
GroupedConvDispatcher dispatcher(&registry);
const int N = 1, G = 1, C = 64, K = 128, Hi = 14, Wi = 14, Y = 3, X = 3;
ck_tile::conv::ConvParam conv_param{
2,
static_cast<ck_tile::index_t>(G),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(K),
static_cast<ck_tile::index_t>(C),
{static_cast<ck_tile::index_t>(Y), static_cast<ck_tile::index_t>(X)},
{static_cast<ck_tile::index_t>(Hi), static_cast<ck_tile::index_t>(Wi)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
auto in_desc =
ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
auto wei_desc =
ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(conv_param);
auto out_desc =
ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(conv_param);
ck_tile::HostTensor<InDataType> input(in_desc);
ck_tile::HostTensor<WeiDataType> weight(wei_desc);
ck_tile::HostTensor<OutDataType> output(out_desc);
ck_tile::FillUniformDistribution<InDataType>{-0.5f, 0.5f}(input);
ck_tile::FillUniformDistribution<WeiDataType>{-0.5f, 0.5f}(weight);
ck_tile::DeviceMem input_dev(input.get_element_space_size_in_bytes());
ck_tile::DeviceMem weight_dev(weight.get_element_space_size_in_bytes());
ck_tile::DeviceMem output_dev(output.get_element_space_size_in_bytes());
input_dev.ToDevice(input.data());
weight_dev.ToDevice(weight.data());
std::cout << "\n " << std::left << std::setw(12) << "Direction" << std::right << std::setw(10)
<< "Time(ms)" << std::setw(10) << "TFLOPS" << std::setw(14) << "NonZero"
<< std::setw(10) << "Status" << "\n";
std::cout << " " << std::string(56, '-') << "\n";
bool all_pass = true;
auto print_result =
[](const char* label, float time_ms, double tflops, size_t nz, size_t total, bool ok) {
std::cout << " " << std::left << std::setw(12) << label << std::right << std::fixed
<< std::setprecision(4) << std::setw(10) << time_ms << std::setprecision(2)
<< std::setw(10) << tflops << std::setw(14)
<< (std::to_string(nz) + "/" + std::to_string(total)) << std::setw(10)
<< (ok ? "OK" : "FAIL") << "\n";
};
// Forward: run(X, W, Y)
{
auto problem =
create_grouped_conv2d_problem(N, C, K, Hi, Wi, Y, X, 1, 1, GroupedConvOp::Forward);
float time_ms = dispatcher.run(input_dev.GetDeviceBuffer(),
weight_dev.GetDeviceBuffer(),
output_dev.GetDeviceBuffer(),
problem,
nullptr);
output_dev.FromDevice(output.data());
size_t nz = 0;
for(size_t i = 0; i < output.get_element_space_size(); ++i)
if(static_cast<float>(output.data()[i]) != 0.0f)
++nz;
bool ok = nz > 0;
print_result("forward",
time_ms,
calculate_conv_tflops(problem, time_ms),
nz,
output.get_element_space_size(),
ok);
if(!ok)
all_pass = false;
}
// Backward Data: run(dY, W, dX)
{
auto problem =
create_grouped_conv2d_problem(N, C, K, Hi, Wi, Y, X, 1, 1, GroupedConvOp::BackwardData);
ck_tile::HostTensor<InDataType> dx_host(in_desc);
ck_tile::DeviceMem dx_dev(dx_host.get_element_space_size_in_bytes());
float time_ms = dispatcher.run(output_dev.GetDeviceBuffer(), // dY (from forward pass)
weight_dev.GetDeviceBuffer(), // W
dx_dev.GetDeviceBuffer(), // dX (output)
problem,
nullptr);
dx_dev.FromDevice(dx_host.data());
size_t nz = 0;
for(size_t i = 0; i < dx_host.get_element_space_size(); ++i)
if(static_cast<float>(dx_host.data()[i]) != 0.0f)
++nz;
bool ok = nz > 0;
print_result("bwd_data",
time_ms,
calculate_conv_tflops(problem, time_ms),
nz,
dx_host.get_element_space_size(),
ok);
if(!ok)
all_pass = false;
}
// Backward Weight: run(X, dY, dW)
{
auto problem = create_grouped_conv2d_problem(
N, C, K, Hi, Wi, Y, X, 1, 1, GroupedConvOp::BackwardWeight);
ck_tile::HostTensor<WeiDataType> dw_host(wei_desc);
ck_tile::DeviceMem dw_dev(dw_host.get_element_space_size_in_bytes());
float time_ms = dispatcher.run(input_dev.GetDeviceBuffer(), // X
output_dev.GetDeviceBuffer(), // dY
dw_dev.GetDeviceBuffer(), // dW (output)
problem,
nullptr);
dw_dev.FromDevice(dw_host.data());
size_t nz = 0;
for(size_t i = 0; i < dw_host.get_element_space_size(); ++i)
if(static_cast<float>(dw_host.data()[i]) != 0.0f)
++nz;
bool ok = nz > 0;
print_result("bwd_weight",
time_ms,
calculate_conv_tflops(problem, time_ms),
nz,
dw_host.get_element_space_size(),
ok);
if(!ok)
all_pass = false;
}
utils::print_separator();
std::cout << " Status: " << (all_pass ? "PASS" : "FAIL") << "\n";
utils::print_separator();
return all_pass ? 0 : 1;
}

263
dispatcher/examples/grouped_conv/cpp/03_benchmark_validation.cpp Normal file
View File

@@ -0,0 +1,263 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 03: Benchmark and CPU-Reference Validation
//
// Runs a 2D grouped conv forward kernel on the GPU via dispatcher.run()
// and compares against the CK Tile host reference implementation.
// Exposes warmup/repeat/log_level as CLI args (matches example 20 pattern).
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_03_bench_val
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include <algorithm>
#include <numeric>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/host/reference/reference_grouped_conv_fwd.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
using AccDataType = float;
DECL_GROUPED_CONV_KERNEL_SET(
bench_kernels,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo().tile(1, 128, 128).pipeline("compv4").vector_sizes(4, 8, 8),
"gfx950")
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo().tile(1, 64, 64).pipeline("compv3").vector_sizes(4, 8, 8),
"gfx950"));
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 03: Benchmark & Validation",
"GPU execution with CPU reference validation");
args.add_option("-n", "1", "Batch size N");
args.add_option("-g", "1", "Groups G");
args.add_option("-c", "64", "Input channels C");
args.add_option("-k", "128", "Output channels K");
args.add_option("--size", "14", "Spatial size (H=W)");
args.add_option("--warmup", "3", "Warmup iterations");
args.add_option("--repeat", "10", "Benchmark iterations");
args.add_option("--arch", "gfx950", "GPU architecture");
args.add_flag("--no-verify", "Skip CPU validation");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 03: Grouped Conv Benchmark & Validation");
int N = args.get_int("-n", 1);
int G = args.get_int("-g", 1);
int C = args.get_int("-c", 64);
int K = args.get_int("-k", 128);
int Hi = args.get_int("--size", 14);
int Wi = Hi;
int Y = 3, X = 3;
int warmup = args.get_int("--warmup", 3);
int repeat = args.get_int("--repeat", 10);
bool verify = !args.has("--no-verify");
std::string gfx_arch = args.get("--arch", "gfx950");
std::cout << "\nProblem: N=" << N << " G=" << G << " C=" << C << " K=" << K << " Hi=" << Hi
<< " Wi=" << Wi << " Y=" << Y << " X=" << X << "\n";
std::cout << "Benchmark: warmup=" << warmup << " repeat=" << repeat << "\n";
// Step 1: Setup tensors using CK Tile descriptors
std::cout << "\nStep 1: Setup tensors\n";
ck_tile::conv::ConvParam conv_param{
2,
static_cast<ck_tile::index_t>(G),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(K),
static_cast<ck_tile::index_t>(C),
{static_cast<ck_tile::index_t>(Y), static_cast<ck_tile::index_t>(X)},
{static_cast<ck_tile::index_t>(Hi), static_cast<ck_tile::index_t>(Wi)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
auto in_desc =
ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
auto wei_desc =
ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(conv_param);
auto out_desc =
ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(conv_param);
ck_tile::HostTensor<InDataType> input(in_desc);
ck_tile::HostTensor<WeiDataType> weight(wei_desc);
ck_tile::HostTensor<OutDataType> output_gpu(out_desc);
ck_tile::HostTensor<OutDataType> output_cpu(out_desc);
ck_tile::FillUniformDistribution<InDataType>{-0.5f, 0.5f}(input);
ck_tile::FillUniformDistribution<WeiDataType>{-0.5f, 0.5f}(weight);
output_cpu.SetZero();
std::cout << " Input: " << input.get_element_space_size() << " elements\n";
std::cout << " Weight: " << weight.get_element_space_size() << " elements\n";
std::cout << " Output: " << output_gpu.get_element_space_size() << " elements\n";
// Step 2: CPU reference
if(verify)
{
std::cout << "\nStep 2: CPU Reference\n";
std::vector<ck_tile::long_index_t> strides_v = {1, 1};
std::vector<ck_tile::long_index_t> dilations_v = {1, 1};
std::vector<ck_tile::long_index_t> left_pads_v = {1, 1};
std::vector<ck_tile::long_index_t> right_pads_v = {1, 1};
ck_tile::reference_grouped_conv_fwd<2, InDataType, WeiDataType, OutDataType>(
input, weight, output_cpu, strides_v, dilations_v, left_pads_v, right_pads_v);
std::cout << " CPU ref[0..7]: ";
for(int i = 0; i < std::min(8, static_cast<int>(output_cpu.get_element_space_size())); ++i)
std::cout << std::fixed << std::setprecision(4)
<< static_cast<float>(output_cpu.data()[i]) << " ";
std::cout << "\n";
}
// Step 3: GPU execution via dispatcher
std::cout << "\nStep 3: GPU Execution (via dispatcher.run)\n";
GroupedConvRegistry registry;
registry.set_name("bench_val");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)\n";
GroupedConvDispatcher dispatcher(&registry);
auto problem = create_grouped_conv2d_problem(N, C, K, Hi, Wi, Y, X, 1, 1);
problem.op = GroupedConvOp::Forward;
auto* selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cerr << " ERROR: No kernel found!\n";
return 1;
}
std::cout << " Selected: " << selected->name() << "\n";
ck_tile::DeviceMem input_dev(input.get_element_space_size_in_bytes());
ck_tile::DeviceMem weight_dev(weight.get_element_space_size_in_bytes());
ck_tile::DeviceMem output_dev(output_gpu.get_element_space_size_in_bytes());
input_dev.ToDevice(input.data());
weight_dev.ToDevice(weight.data());
float elapsed_ms = dispatcher.run(input_dev.GetDeviceBuffer(),
weight_dev.GetDeviceBuffer(),
output_dev.GetDeviceBuffer(),
problem,
nullptr);
output_dev.FromDevice(output_gpu.data());
size_t total = output_gpu.get_element_space_size();
std::cout << " GPU out[0..7]: ";
for(int i = 0; i < std::min(8, static_cast<int>(total)); ++i)
std::cout << std::fixed << std::setprecision(4) << static_cast<float>(output_gpu.data()[i])
<< " ";
std::cout << "\n";
size_t nonzero_gpu = 0;
double gpu_sum = 0.0;
for(size_t i = 0; i < total; ++i)
{
float v = static_cast<float>(output_gpu.data()[i]);
if(v != 0.0f)
++nonzero_gpu;
gpu_sum += v;
}
std::cout << " GPU checksum: " << std::fixed << std::setprecision(6) << gpu_sum << "\n";
std::cout << " GPU non-zero: " << nonzero_gpu << "/" << total
<< (nonzero_gpu > 0 ? " (kernel produced output)" : " WARNING: all zeros!") << "\n";
int Ho = static_cast<int>(problem.Ho());
int Wo = static_cast<int>(problem.Wo());
double flops = 2.0 * G * N * K * C * Y * X * Ho * Wo;
double tflops = flops / (elapsed_ms * 1e9);
std::cout << " Time: " << std::fixed << std::setprecision(4) << elapsed_ms << " ms\n";
std::cout << " TFLOPS: " << std::setprecision(2) << tflops << "\n";
// Step 4: Validation
bool passed = true;
if(verify)
{
std::cout << "\nStep 4: Validation (GPU vs CPU)\n";
constexpr float rtol = 1e-2f;
constexpr float atol = 1e-2f;
float max_diff = 0.0f;
float max_rel = 0.0f;
size_t max_diff_idx = 0;
size_t num_elements = output_gpu.get_element_space_size();
size_t mismatches = 0;
for(size_t i = 0; i < num_elements; ++i)
{
float gpu_val = static_cast<float>(output_gpu.data()[i]);
float cpu_val = static_cast<float>(output_cpu.data()[i]);
float diff = std::abs(gpu_val - cpu_val);
float tol = atol + rtol * std::abs(cpu_val);
float rel = diff / (std::abs(cpu_val) + 1e-6f);
if(diff > max_diff)
{
max_diff = diff;
max_diff_idx = i;
}
max_rel = std::max(max_rel, rel);
if(diff > tol)
++mismatches;
}
passed = (mismatches == 0);
std::cout << " Side-by-side at worst element [" << max_diff_idx << "]:\n";
std::cout << " GPU: " << std::fixed << std::setprecision(6)
<< static_cast<float>(output_gpu.data()[max_diff_idx])
<< " CPU: " << static_cast<float>(output_cpu.data()[max_diff_idx])
<< " diff: " << std::scientific << max_diff << "\n";
std::cout << " Elements: " << num_elements << "\n";
std::cout << " Mismatches: " << mismatches << "/" << num_elements << "\n";
std::cout << " Max abs diff: " << std::scientific << max_diff << "\n";
std::cout << " Max rel diff: " << std::scientific << max_rel << "\n";
std::cout << " Status: " << (passed ? "PASSED" : "FAILED") << "\n";
}
utils::print_separator();
std::cout << "BENCHMARK & VALIDATION:\n";
std::cout << " GPU kernel: " << (selected ? selected->name() : "none") << "\n";
std::cout << " Performance: " << std::fixed << std::setprecision(2) << tflops
<< " TFLOPS\n";
std::cout << " CPU reference: reference_grouped_conv_fwd<2>()\n";
std::cout << " Validation: " << (passed ? "PASS" : "FAIL") << "\n";
utils::print_separator();
return passed ? 0 : 1;
}

154
dispatcher/examples/grouped_conv/cpp/04_registry_json.cpp Normal file
View File

@@ -0,0 +1,154 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 04: Heuristic Selection + JSON Export
//
// Demonstrates runtime kernel selection with heuristic ranking,
// GPU execution, and JSON registry export.
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_04_registry_json
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
// Two tile configs for heuristic selection
DECL_GROUPED_CONV_KERNEL_SET(
heuristic_kernels,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo().tile(1, 128, 128).pipeline("compv4").vector_sizes(4, 8, 8),
"gfx950")
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo().tile(1, 64, 64).pipeline("compv3").vector_sizes(4, 8, 8),
"gfx950"));
std::vector<std::string> conv_heuristic(const GroupedConvProblem& problem)
{
int64_t spatial = problem.Ho() * problem.Wo();
if(spatial > 400)
return {"128x128", "64x64"};
return {"64x64", "128x128"};
}
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 04: Heuristic + JSON",
"Runtime kernel selection and JSON export");
args.add_option("--arch", "gfx950", "GPU architecture");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 04: Heuristic Selection + JSON Export");
std::string gfx_arch = args.get("--arch", "gfx950");
// Step 1: Register
std::cout << "\nStep 1: Register Kernels" << std::endl;
GroupedConvRegistry registry;
registry.set_name("heuristic_conv");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)" << std::endl;
// Step 2: Heuristic dispatcher
std::cout << "\nStep 2: Heuristic Dispatcher" << std::endl;
GroupedConvDispatcher dispatcher(&registry);
dispatcher.set_strategy(GroupedConvDispatcher::SelectionStrategy::Heuristic);
dispatcher.set_heuristic(conv_heuristic);
// Step 3: Select kernels (no GPU yet)
std::cout << "\nStep 3: Kernel Selection" << std::endl;
auto problem = create_grouped_conv2d_problem(1, 64, 128, 14, 14, 3, 3, 1, 1);
auto* selected = dispatcher.select_kernel(problem);
std::cout << " Selected: " << (selected ? selected->name() : "none") << std::endl;
// Step 4: GPU execution
std::cout << "\nStep 4: GPU Execution" << std::endl;
ck_tile::conv::ConvParam cp{
2,
static_cast<ck_tile::index_t>(1),
static_cast<ck_tile::index_t>(1),
static_cast<ck_tile::index_t>(128),
static_cast<ck_tile::index_t>(64),
{static_cast<ck_tile::index_t>(3), static_cast<ck_tile::index_t>(3)},
{static_cast<ck_tile::index_t>(14), static_cast<ck_tile::index_t>(14)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
std::cout << " Creating tensors..." << std::endl;
auto in_d = ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(cp);
auto wei_d = ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(cp);
auto out_d = ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(cp);
ck_tile::HostTensor<InDataType> input(in_d);
ck_tile::HostTensor<WeiDataType> weight(wei_d);
ck_tile::HostTensor<OutDataType> output(out_d);
ck_tile::FillUniformDistribution<InDataType>{-0.5f, 0.5f}(input);
ck_tile::FillUniformDistribution<WeiDataType>{-0.5f, 0.5f}(weight);
std::cout << " Allocating device memory..." << std::endl;
ck_tile::DeviceMem in_dev(input.get_element_space_size_in_bytes());
ck_tile::DeviceMem wei_dev(weight.get_element_space_size_in_bytes());
ck_tile::DeviceMem out_dev(output.get_element_space_size_in_bytes());
in_dev.ToDevice(input.data());
wei_dev.ToDevice(weight.data());
std::cout << " Launching kernel..." << std::endl;
float time_ms = dispatcher.run(in_dev.GetDeviceBuffer(),
wei_dev.GetDeviceBuffer(),
out_dev.GetDeviceBuffer(),
problem,
nullptr);
std::cout << " Reading back..." << std::endl;
out_dev.FromDevice(output.data());
size_t nz = 0;
for(size_t i = 0; i < output.get_element_space_size(); ++i)
if(static_cast<float>(output.data()[i]) != 0.0f)
++nz;
std::cout << " Time: " << std::fixed << std::setprecision(4) << time_ms << " ms"
<< std::endl;
std::cout << " TFLOPS: " << std::setprecision(2) << calculate_conv_tflops(problem, time_ms)
<< std::endl;
std::cout << " NonZero: " << nz << "/" << output.get_element_space_size() << std::endl;
// Step 5: JSON export
std::cout << "\nStep 5: JSON Export" << std::endl;
std::string json = registry.export_json(false);
std::cout << " JSON size: " << json.size() << " bytes" << std::endl;
bool passed = nz > 0;
utils::print_separator();
std::cout << " Status: " << (passed ? "PASS" : "FAIL") << "\n";
utils::print_separator();
return passed ? 0 : 1;
}

183
dispatcher/examples/grouped_conv/cpp/05_bwd_data.cpp Normal file
View File

@@ -0,0 +1,183 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 05: Backward Data with CPU Reference Validation
//
// Computes dX = ConvBwdData(dY, W) on GPU via dispatcher.run()
// and validates against ck_tile::reference_grouped_conv_bwd_data.
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_05_bwd_data
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/host/reference/reference_grouped_conv_bwd_data.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
DECL_GROUPED_CONV_KERNEL_SET(
bwd_data_kernels,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("bwd_data").dims(2),
GroupedConvAlgo()
.tile(1, 128, 128)
.pipeline("compv3")
.scheduler("intrawave")
.vector_sizes(4, 8, 8),
"gfx950"));
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 05: Backward Data Validation",
"dX = ConvBwdData(dY, W) with CPU reference");
args.add_option("--arch", "gfx950", "GPU architecture");
args.add_option("-n", "1", "Batch size");
args.add_option("-c", "64", "Input channels");
args.add_option("-k", "128", "Output channels");
args.add_option("--size", "14", "Spatial size (H=W)");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 05: Backward Data with CPU Validation");
std::string gfx_arch = args.get("--arch", "gfx950");
int N = args.get_int("-n", 1), G = 1;
int C = args.get_int("-c", 64), K = args.get_int("-k", 128);
int Hi = args.get_int("--size", 14), Wi = Hi, Y = 3, X = 3;
// Setup
ck_tile::conv::ConvParam conv_param{
2,
static_cast<ck_tile::index_t>(G),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(K),
static_cast<ck_tile::index_t>(C),
{static_cast<ck_tile::index_t>(Y), static_cast<ck_tile::index_t>(X)},
{static_cast<ck_tile::index_t>(Hi), static_cast<ck_tile::index_t>(Wi)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
auto in_desc =
ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
auto wei_desc =
ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(conv_param);
auto out_desc =
ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(conv_param);
// dY (gradient from next layer) and W (weight) are inputs; dX is output
ck_tile::HostTensor<OutDataType> dy(out_desc);
ck_tile::HostTensor<WeiDataType> weight(wei_desc);
ck_tile::HostTensor<InDataType> dx_gpu(in_desc);
ck_tile::HostTensor<InDataType> dx_cpu(in_desc);
ck_tile::FillUniformDistribution<OutDataType>{-0.5f, 0.5f}(dy);
ck_tile::FillUniformDistribution<WeiDataType>{-0.5f, 0.5f}(weight);
dx_cpu.SetZero();
// CPU reference
std::cout << "\nStep 1: CPU Reference (bwd_data)\n";
std::vector<ck_tile::long_index_t> strides_v = {1, 1};
std::vector<ck_tile::long_index_t> dilations_v = {1, 1};
std::vector<ck_tile::long_index_t> left_pads_v = {1, 1};
std::vector<ck_tile::long_index_t> right_pads_v = {1, 1};
ck_tile::reference_grouped_conv_bwd_data<2, InDataType, WeiDataType, OutDataType>(
dx_cpu, weight, dy, strides_v, dilations_v, left_pads_v, right_pads_v);
std::cout << " CPU complete\n";
// GPU execution via dispatcher
std::cout << "\nStep 2: GPU Execution (via dispatcher.run)\n";
GroupedConvRegistry registry;
registry.set_name("bwd_data");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
GroupedConvDispatcher dispatcher(&registry);
auto problem =
create_grouped_conv2d_problem(N, C, K, Hi, Wi, Y, X, 1, 1, GroupedConvOp::BackwardData);
auto* selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cerr << " ERROR: No bwd_data kernel found!\n";
return 1;
}
std::cout << " Selected: " << selected->name() << "\n";
ck_tile::DeviceMem dy_dev(dy.get_element_space_size_in_bytes());
ck_tile::DeviceMem wei_dev(weight.get_element_space_size_in_bytes());
ck_tile::DeviceMem dx_dev(dx_gpu.get_element_space_size_in_bytes());
dy_dev.ToDevice(dy.data());
wei_dev.ToDevice(weight.data());
// dispatcher.run(dY, W, dX, problem) for bwd_data
float time_ms = dispatcher.run(dy_dev.GetDeviceBuffer(),
wei_dev.GetDeviceBuffer(),
dx_dev.GetDeviceBuffer(),
problem,
nullptr);
dx_dev.FromDevice(dx_gpu.data());
double tflops = (time_ms > 0) ? calculate_conv_tflops(problem, time_ms) : 0;
std::cout << " Time: " << std::fixed << std::setprecision(4) << time_ms << " ms\n";
std::cout << " TFLOPS: " << std::setprecision(2) << tflops << "\n";
// Validation
std::cout << "\nStep 3: Validation (GPU vs CPU)\n";
size_t num_elements = dx_gpu.get_element_space_size();
float max_abs = 0, max_rel = 0;
size_t mismatches = 0;
constexpr float rtol = 5e-2f, atol = 5e-2f;
for(size_t i = 0; i < num_elements; ++i)
{
float gv = static_cast<float>(dx_gpu.data()[i]);
float cv = static_cast<float>(dx_cpu.data()[i]);
float d = std::abs(gv - cv);
float r = d / (std::abs(cv) + 1e-6f);
max_abs = std::max(max_abs, d);
max_rel = std::max(max_rel, r);
if(d > atol + rtol * std::abs(cv))
++mismatches;
}
bool passed = (mismatches == 0);
std::cout << " Elements: " << num_elements << "\n";
std::cout << " Mismatches: " << mismatches << "\n";
std::cout << " Max abs diff: " << std::scientific << max_abs << "\n";
std::cout << " Max rel diff: " << std::scientific << max_rel << "\n";
utils::print_separator();
std::cout << " dX = ConvBwdData(dY, W)\n";
std::cout << " Status: " << (passed ? "PASS" : "FAIL") << "\n";
utils::print_separator();
return passed ? 0 : 1;
}

188
dispatcher/examples/grouped_conv/cpp/06_bwd_weight.cpp Normal file
View File

@@ -0,0 +1,188 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 06: Backward Weight with CPU Reference Validation
//
// Computes dW = ConvBwdWeight(X, dY) on GPU via dispatcher.run()
// and validates against ck_tile::reference_grouped_conv_bwd_weight.
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_06_bwd_weight
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/host/reference/reference_grouped_conv_bwd_weight.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
DECL_GROUPED_CONV_KERNEL_SET(
bwd_weight_kernels,
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("bwd_weight").dims(2),
GroupedConvAlgo()
.tile(1, 128, 128)
.pipeline("compv3")
.scheduler("intrawave")
.memory_op("atomic_add")
.vector_sizes(4, 8, 8),
"gfx950"));
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 06: Backward Weight Validation",
"dW = ConvBwdWeight(X, dY) with CPU reference");
args.add_option("--arch", "gfx950", "GPU architecture");
args.add_option("-n", "1", "Batch size");
args.add_option("-c", "64", "Input channels");
args.add_option("-k", "128", "Output channels");
args.add_option("--size", "14", "Spatial size (H=W)");
args.add_option("--split-k", "1", "Split-K factor for bwd_weight (k_batch)");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 06: Backward Weight with CPU Validation");
std::string gfx_arch = args.get("--arch", "gfx950");
int N = args.get_int("-n", 1), G = 1;
int C = args.get_int("-c", 64), K = args.get_int("-k", 128);
int Hi = args.get_int("--size", 14), Wi = Hi, Y = 3, X = 3;
// Setup
ck_tile::conv::ConvParam conv_param{
2,
static_cast<ck_tile::index_t>(G),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(K),
static_cast<ck_tile::index_t>(C),
{static_cast<ck_tile::index_t>(Y), static_cast<ck_tile::index_t>(X)},
{static_cast<ck_tile::index_t>(Hi), static_cast<ck_tile::index_t>(Wi)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
auto in_desc =
ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
auto wei_desc =
ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(conv_param);
auto out_desc =
ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(conv_param);
// X (input) and dY (gradient) are inputs; dW is output
ck_tile::HostTensor<InDataType> input(in_desc);
ck_tile::HostTensor<OutDataType> dy(out_desc);
ck_tile::HostTensor<WeiDataType> dw_gpu(wei_desc);
ck_tile::HostTensor<WeiDataType> dw_cpu(wei_desc);
ck_tile::FillUniformDistribution<InDataType>{-0.5f, 0.5f}(input);
ck_tile::FillUniformDistribution<OutDataType>{-0.5f, 0.5f}(dy);
dw_cpu.SetZero();
// CPU reference
std::cout << "\nStep 1: CPU Reference (bwd_weight)\n";
std::vector<ck_tile::long_index_t> strides_v = {1, 1};
std::vector<ck_tile::long_index_t> dilations_v = {1, 1};
std::vector<ck_tile::long_index_t> left_pads_v = {1, 1};
std::vector<ck_tile::long_index_t> right_pads_v = {1, 1};
ck_tile::reference_grouped_conv_bwd_weight<2, InDataType, WeiDataType, OutDataType>(
input, dw_cpu, dy, strides_v, dilations_v, left_pads_v, right_pads_v);
std::cout << " CPU complete\n";
// GPU execution
std::cout << "\nStep 2: GPU Execution (via dispatcher.run)\n";
GroupedConvRegistry registry;
registry.set_name("bwd_weight");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
GroupedConvDispatcher dispatcher(&registry);
auto problem =
create_grouped_conv2d_problem(N, C, K, Hi, Wi, Y, X, 1, 1, GroupedConvOp::BackwardWeight);
problem.split_k = args.get_int("--split-k", 1);
auto* selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cerr << " ERROR: No bwd_weight kernel found!\n";
return 1;
}
std::cout << " Selected: " << selected->name() << "\n";
ck_tile::DeviceMem in_dev(input.get_element_space_size_in_bytes());
ck_tile::DeviceMem dy_dev(dy.get_element_space_size_in_bytes());
ck_tile::DeviceMem dw_dev(dw_gpu.get_element_space_size_in_bytes());
in_dev.ToDevice(input.data());
dy_dev.ToDevice(dy.data());
if(problem.split_k > 1)
dw_dev.SetZero();
// dispatcher.run(X, dY, dW, problem) for bwd_weight
float time_ms = dispatcher.run(in_dev.GetDeviceBuffer(),
dy_dev.GetDeviceBuffer(),
dw_dev.GetDeviceBuffer(),
problem,
nullptr);
dw_dev.FromDevice(dw_gpu.data());
double tflops = (time_ms > 0) ? calculate_conv_tflops(problem, time_ms) : 0;
std::cout << " Time: " << std::fixed << std::setprecision(4) << time_ms << " ms\n";
std::cout << " TFLOPS: " << std::setprecision(2) << tflops << "\n";
// Validation
std::cout << "\nStep 3: Validation (GPU vs CPU)\n";
size_t num_elements = dw_gpu.get_element_space_size();
float max_abs = 0, max_rel = 0;
size_t mismatches = 0;
constexpr float rtol = 5e-2f, atol = 5e-2f;
for(size_t i = 0; i < num_elements; ++i)
{
float gv = static_cast<float>(dw_gpu.data()[i]);
float cv = static_cast<float>(dw_cpu.data()[i]);
float d = std::abs(gv - cv);
float r = d / (std::abs(cv) + 1e-6f);
max_abs = std::max(max_abs, d);
max_rel = std::max(max_rel, r);
if(d > atol + rtol * std::abs(cv))
++mismatches;
}
bool passed = (mismatches == 0);
std::cout << " Elements: " << num_elements << "\n";
std::cout << " Mismatches: " << mismatches << "\n";
std::cout << " Max abs diff: " << std::scientific << max_abs << "\n";
std::cout << " Max rel diff: " << std::scientific << max_rel << "\n";
utils::print_separator();
std::cout << " dW = ConvBwdWeight(X, dY)\n";
std::cout << " Status: " << (passed ? "PASS" : "FAIL") << "\n";
utils::print_separator();
return passed ? 0 : 1;
}

226
dispatcher/examples/grouped_conv/cpp/07_multi_tile_benchmark.cpp Normal file
View File

@@ -0,0 +1,226 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Example 07: Multi-Tile Benchmark
//
// Benchmarks multiple tile configurations across ResNet-like problem sizes.
// Exposes warmup, repeat, and init method as CLI args (matching CK Tile
// example 20 patterns).
//
// Build: cd dispatcher/build && cmake .. && make grouped_conv_07_benchmark
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::grouped_conv_utils;
using GroupedConvSig = grouped_conv_decl::GroupedConvSignature;
using GroupedConvAlgo = grouped_conv_decl::GroupedConvAlgorithm;
using InDataType = ck_tile::half_t;
using WeiDataType = ck_tile::half_t;
using OutDataType = ck_tile::half_t;
// Multiple tile configurations for benchmarking
DECL_GROUPED_CONV_KERNEL_SET(
benchmark_tiles,
// Small tile - compv3
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo()
.tile(1, 64, 64)
.wave(1, 4, 1)
.warp(16, 16, 32)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle")
.vector_sizes(4, 8, 8)
.block_per_cu(1),
"gfx950")
// Medium tile - compv3
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo()
.tile(1, 128, 128)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle")
.vector_sizes(4, 8, 8)
.block_per_cu(1),
"gfx950")
// Large tile - compv4 with double smem buffer
.add(GroupedConvSig().dtype("fp16").layout("nhwgc").conv_type("forward").dims(2),
GroupedConvAlgo()
.tile(1, 256, 256)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv4")
.scheduler("intrawave")
.epilogue("cshuffle")
.vector_sizes(4, 8, 8)
.block_per_cu(1),
"gfx950"));
int main(int argc, char* argv[])
{
utils::ExampleArgs args("Example 07: Multi-Tile Benchmark",
"Multiple tiles across ResNet-like problem sizes");
args.add_option("--arch", "gfx950", "GPU architecture");
args.add_option("--warmup", "5", "Warmup iterations (passed to stream_config)");
args.add_option("--repeat", "20", "Benchmark iterations (passed to stream_config)");
args.add_option("--init", "0", "Init method: 0=random, 1=linear, 2=constant(1)");
if(!args.parse(argc, argv))
return 0;
utils::print_header("Example 07: Multi-Tile Benchmark");
std::string gfx_arch = args.get("--arch", "gfx950");
int warmup = args.get_int("--warmup", 5);
int repeat = args.get_int("--repeat", 20);
int init_method = args.get_int("--init", 0);
std::cout << "\n Config: warmup=" << warmup << " repeat=" << repeat << " init=" << init_method
<< "\n";
GroupedConvRegistry registry;
registry.set_name("benchmark");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)\n";
GroupedConvDispatcher dispatcher(&registry);
// ResNet-like problem sizes
struct BenchProblem
{
const char* label;
int N, C, K, Hi, Wi, Y, X;
};
BenchProblem problems[] = {
{"ResNet-stage2", 1, 64, 64, 56, 56, 3, 3},
{"ResNet-stage3", 1, 128, 128, 28, 28, 3, 3},
{"ResNet-stage4", 1, 256, 256, 14, 14, 3, 3},
{"ResNet-stage5", 1, 512, 512, 7, 7, 3, 3},
{"Pointwise-1x1", 1, 256, 256, 56, 56, 1, 1},
{"Batch-8", 8, 64, 128, 56, 56, 3, 3},
};
std::cout << "\n " << std::left << std::setw(16) << "Problem" << std::right << std::setw(5)
<< "N" << std::setw(5) << "C" << std::setw(5) << "K" << std::setw(5) << "H"
<< std::setw(5) << "W" << std::setw(4) << "F" << std::setw(10) << "Time(ms)"
<< std::setw(10) << "TFLOPS" << std::setw(10) << "Status" << "\n";
std::cout << " " << std::string(74, '-') << "\n";
bool all_pass = true;
for(const auto& bp : problems)
{
auto problem =
create_grouped_conv2d_problem(bp.N, bp.C, bp.K, bp.Hi, bp.Wi, bp.Y, bp.X, 1, 1);
problem.op = GroupedConvOp::Forward;
ck_tile::conv::ConvParam conv_param{
2,
static_cast<ck_tile::index_t>(1),
static_cast<ck_tile::index_t>(bp.N),
static_cast<ck_tile::index_t>(bp.K),
static_cast<ck_tile::index_t>(bp.C),
{static_cast<ck_tile::index_t>(bp.Y), static_cast<ck_tile::index_t>(bp.X)},
{static_cast<ck_tile::index_t>(bp.Hi), static_cast<ck_tile::index_t>(bp.Wi)},
{1, 1},
{1, 1},
{1, 1},
{1, 1}};
using InLayout = ck_tile::tensor_layout::convolution::NHWGC;
using WeiLayout = ck_tile::tensor_layout::convolution::GKYXC;
using OutLayout = ck_tile::tensor_layout::convolution::NHWGK;
auto in_desc =
ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
auto wei_desc =
ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(
conv_param);
auto out_desc =
ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(
conv_param);
ck_tile::HostTensor<InDataType> input(in_desc);
ck_tile::HostTensor<WeiDataType> weight(wei_desc);
ck_tile::HostTensor<OutDataType> output(out_desc);
switch(init_method)
{
case 1:
ck_tile::FillMonotonicSeq<InDataType>{0.0f, 0.001f}(input);
ck_tile::FillMonotonicSeq<WeiDataType>{0.0f, 0.001f}(weight);
break;
case 2:
ck_tile::FillConstant<InDataType>{1.0f}(input);
ck_tile::FillConstant<WeiDataType>{1.0f}(weight);
break;
default:
ck_tile::FillUniformDistribution<InDataType>{-0.5f, 0.5f}(input);
ck_tile::FillUniformDistribution<WeiDataType>{-0.5f, 0.5f}(weight);
break;
}
ck_tile::DeviceMem in_dev(input.get_element_space_size_in_bytes());
ck_tile::DeviceMem wei_dev(weight.get_element_space_size_in_bytes());
ck_tile::DeviceMem out_dev(output.get_element_space_size_in_bytes());
in_dev.ToDevice(input.data());
wei_dev.ToDevice(weight.data());
float time_ms = 0;
bool ok = false;
try
{
time_ms = dispatcher.run(in_dev.GetDeviceBuffer(),
wei_dev.GetDeviceBuffer(),
out_dev.GetDeviceBuffer(),
problem,
nullptr);
out_dev.FromDevice(output.data());
size_t nz = 0;
for(size_t j = 0; j < output.get_element_space_size(); ++j)
if(static_cast<float>(output.data()[j]) != 0.0f)
++nz;
ok = nz > 0;
}
catch(const std::exception&)
{
ok = false;
}
double tflops = (time_ms > 0) ? calculate_conv_tflops(problem, time_ms) : 0;
std::string filter_str = std::to_string(bp.Y) + "x" + std::to_string(bp.X);
std::cout << " " << std::left << std::setw(16) << bp.label << std::right << std::setw(5)
<< bp.N << std::setw(5) << bp.C << std::setw(5) << bp.K << std::setw(5) << bp.Hi
<< std::setw(5) << bp.Wi << std::setw(4) << filter_str << std::fixed
<< std::setprecision(4) << std::setw(10) << time_ms << std::setprecision(2)
<< std::setw(10) << tflops << std::setw(10) << (ok ? "OK" : "FAIL") << "\n";
if(!ok)
all_pass = false;
}
utils::print_separator();
std::cout << " Warmup: " << warmup << ", Repeat: " << repeat << ", Init: " << init_method
<< "\n";
std::cout << " Status: " << (all_pass ? "PASS" : "FAIL") << "\n";
utils::print_separator();
return all_pass ? 0 : 1;
}

271
dispatcher/examples/grouped_conv/python/01_basic_grouped_conv.py Normal file
View File

@@ -0,0 +1,271 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 01: Basic Grouped Convolution
Demonstrates:
1. Three kernel configuration patterns (minimal, explicit, full ConvConfigBase)
2. Adding kernels to a registry
3. Validation and auto-correction
4. JIT compilation via registry.build()
5. GPU execution with CPU reference verification
Usage:
python3 01_basic_grouped_conv.py
python3 01_basic_grouped_conv.py --variant bwd_data
python3 01_basic_grouped_conv.py --arch gfx942
"""
import sys
import argparse
import time
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
from grouped_conv_utils import (
GroupedConvKernelConfig,
GroupedConvProblem,
GroupedConvRegistry,
validate_grouped_conv_config,
auto_correct_grouped_conv_config,
detect_gpu_arch,
)
def cpu_conv2d_fwd(inp, wei, prob):
"""Naive CPU reference: 2D forward, NHWGC layout."""
N, Hi, Wi, G, Cpg = inp.shape
_, Kpg, Y, X, _ = wei.shape
Ho, Wo = prob.Ho, prob.Wo
out = np.zeros((N, Ho, Wo, G, Kpg), dtype=np.float32)
for n in range(N):
for g in range(G):
for ho in range(Ho):
for wo in range(Wo):
for k in range(Kpg):
s = 0.0
for y in range(Y):
for x in range(X):
hi = (
ho * prob.stride_h
- prob.pad_h
+ y * prob.dilation_h
)
wi = (
wo * prob.stride_w
- prob.pad_w
+ x * prob.dilation_w
)
if 0 <= hi < Hi and 0 <= wi < Wi:
for c in range(Cpg):
s += float(inp[n, hi, wi, g, c]) * float(
wei[g, k, y, x, c]
)
out[n, ho, wo, g, k] = s
return out
def main():
parser = argparse.ArgumentParser(description="Basic Grouped Conv Example")
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument(
"--variant", default="forward", choices=["forward", "bwd_data", "bwd_weight"]
)
parser.add_argument("--ndim", type=int, default=2, choices=[2, 3])
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument(
"--workers", type=int, default=0, help="Max JIT workers (0=auto)"
)
args = parser.parse_args()
print("=" * 70)
print("Example 01: Basic Grouped Convolution")
print("=" * 70)
# =========================================================================
# Step 1: Three kernel configuration patterns
# =========================================================================
print("\n--- Step 1: Kernel Configuration Patterns ---")
# Pattern 1: MINIMAL -- only variant/dtype/arch, everything else auto-filled
config_minimal = GroupedConvKernelConfig(
variant=args.variant,
ndim_spatial=args.ndim,
arch=args.arch,
dtype=args.dtype,
)
print("\n Pattern 1: MINIMAL (defaults auto-filled)")
config_minimal.print_config(indent=" ")
# Pattern 2: EXPLICIT tile/wave/warp -- user controls tiling strategy
config_explicit = GroupedConvKernelConfig(
variant=args.variant,
ndim_spatial=args.ndim,
arch=args.arch,
dtype=args.dtype,
tile_m=1,
tile_n=64,
tile_k=64,
wave_m=1,
wave_n=4,
wave_k=1,
warp_tile_m=16,
warp_tile_n=16,
warp_tile_k=32,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
)
print("\n Pattern 2: EXPLICIT tile/wave/warp")
config_explicit.print_config(indent=" ")
# Pattern 3: FULL ConvConfigBase -- every parameter specified
config_full = GroupedConvKernelConfig(
variant=args.variant,
ndim_spatial=args.ndim,
arch=args.arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
num_wave_groups=1,
num_groups_to_merge=1,
)
print("\n Pattern 3: FULL (all ConvConfigBase fields)")
config_full.print_config(indent=" ")
# =========================================================================
# Step 2: Build a registry with multiple configs
# =========================================================================
print("\n--- Step 2: Build Registry ---")
registry = GroupedConvRegistry("basic_conv")
registry.add(config_minimal)
registry.add(config_explicit)
registry.add(config_full)
registry.print_registry()
# =========================================================================
# Step 3: Validate and auto-correct
# =========================================================================
print("\n--- Step 3: Validate & Auto-Correct ---")
for i, cfg in enumerate(registry.kernels):
result = validate_grouped_conv_config(cfg.to_dict())
if result.is_valid:
print(f" Config [{i}] {cfg.tile_str}: VALID")
else:
print(f" Config [{i}] {cfg.tile_str}: needs correction")
corrected, result = auto_correct_grouped_conv_config(cfg.to_dict())
print(f" After correction: valid={result.is_valid}")
# =========================================================================
# Step 4: JIT compile via registry.build()
# =========================================================================
print("\n--- Step 4: JIT Build (via registry.build()) ---")
# Use only the first config for the actual GPU run
jit_reg = GroupedConvRegistry("jit")
jit_reg.add(config_minimal)
workers = args.workers if args.workers > 0 else None
t0 = time.perf_counter()
runners = jit_reg.build(verbose=False, max_workers=workers)
jit_build_s = time.perf_counter() - t0
key = (args.variant, args.ndim)
if key not in runners:
print(" JIT build failed")
return 1
runner = runners[key]
print(f" JIT build: {jit_build_s:.3f} s")
print(f" Library: {runner.library_path}")
print(f" Kernels: {runner.lib.kernel_names()}")
# =========================================================================
# Step 5: Define problem + GPU execution
# =========================================================================
print("\n--- Step 5: GPU Execution ---")
prob = GroupedConvProblem(
N=1,
C=64,
K=128,
Hi=16,
Wi=16,
Y=3,
X=3,
stride_h=1,
stride_w=1,
pad_h=1,
pad_w=1,
direction=args.variant,
)
prob.print_problem()
inp = np.random.uniform(-0.5, 0.5, prob.input_shape()).astype(np.float16)
wei = np.random.uniform(-0.5, 0.5, prob.weight_shape()).astype(np.float16)
res = runner.run(inp, wei, prob)
if not res.success:
print(f" GPU execution failed: {res.error}")
runner.cleanup()
return 1
print(f" Time: {res.time_ms:.4f} ms")
print(f" TFLOPS: {res.tflops:.2f}")
print(
f" Output: shape={res.output.shape}, range=[{res.output.min():.3f}, {res.output.max():.3f}]"
)
# =========================================================================
# Step 6: CPU reference (forward 2D only)
# =========================================================================
verified = False
if args.variant == "forward" and args.ndim == 2:
print("\n--- Step 6: CPU Reference Verification ---")
ref = cpu_conv2d_fwd(inp, wei, prob)
gpu_f32 = res.output.astype(np.float32)
diff = np.abs(gpu_f32 - ref)
max_abs = diff.max()
max_rel = (diff / (np.abs(ref) + 1e-6)).max()
match = np.allclose(gpu_f32, ref, atol=0.05, rtol=0.05)
print(f" max_abs_diff: {max_abs:.6f}")
print(f" max_rel_diff: {max_rel:.6f}")
print(f" Match: {match}")
verified = match
runner.cleanup()
# Summary
print("\n" + "=" * 70)
status = (
"PASS" if res.success and (verified or args.variant != "forward") else "FAIL"
)
print(f" Status: {status}")
print(
f" {config_minimal.name} | {prob.gflops:.2f} GFLOPs | {res.tflops:.2f} TFLOPS"
)
print(f" JIT build time: {jit_build_s:.3f} s")
print(f" Registry: {len(registry)} configs (3 patterns demonstrated)")
print("=" * 70)
return 0 if status == "PASS" else 1
if __name__ == "__main__":
sys.exit(main())

222
dispatcher/examples/grouped_conv/python/02_forward.py Normal file
View File

@@ -0,0 +1,222 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 02: Forward Convolution (2D + 3D)
Declares forward kernels with explicit tile/wave/warp/pipeline parameters,
builds a registry, JIT compiles, runs on GPU, and validates against CPU reference.
Usage:
python3 02_forward.py
python3 02_forward.py --arch gfx942
"""
import sys
import argparse
import time
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
from grouped_conv_utils import (
GroupedConvKernelConfig,
GroupedConvProblem,
GroupedConvRegistry,
detect_gpu_arch,
)
def cpu_conv2d_fwd(inp, wei, prob):
"""Naive CPU reference: 2D forward, NHWGC layout."""
N, Hi, Wi, G, C = inp.shape
_, Kpg, Y, X, _ = wei.shape
Ho, Wo = prob.Ho, prob.Wo
out = np.zeros((N, Ho, Wo, G, Kpg), dtype=np.float32)
for n in range(N):
for g in range(G):
for ho in range(Ho):
for wo in range(Wo):
for k in range(Kpg):
s = 0.0
for y in range(Y):
for x in range(X):
hi = ho * prob.stride_h - prob.pad_h + y
wi = wo * prob.stride_w - prob.pad_w + x
if 0 <= hi < Hi and 0 <= wi < Wi:
for c in range(C):
s += float(inp[n, hi, wi, g, c]) * float(
wei[g, k, y, x, c]
)
out[n, ho, wo, g, k] = s
return out
def main():
parser = argparse.ArgumentParser(description="Forward Convolution (2D + 3D)")
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument(
"--workers", type=int, default=0, help="Max JIT workers (0=auto)"
)
args = parser.parse_args()
arch = args.arch
print("=" * 70)
print("Example 02: Forward Convolution (2D + 3D)")
print("=" * 70)
print(f"\n Arch: {arch}, Dtype: {args.dtype}")
# =========================================================================
# Step 1: Declare forward kernels with explicit parameters
# =========================================================================
print("\n--- Step 1: Declare Forward Kernels ---")
reg = GroupedConvRegistry("forward_conv")
# Forward 2D: compv4, 128x128 tile, wave 2x2x1, warp 32x32x16
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv4",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
# Forward 3D: compv3, 64x64 tile, wave 1x4x1, warp 16x16x32
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=3,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=64,
tile_k=64,
wave_m=1,
wave_n=4,
wave_k=1,
warp_tile_m=16,
warp_tile_n=16,
warp_tile_k=32,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
reg.print_registry()
# =========================================================================
# Step 2: JIT build via registry
# =========================================================================
print("\n--- Step 2: JIT Build ---")
workers = args.workers if args.workers > 0 else None
t0 = time.perf_counter()
runners = reg.build(verbose=False, max_workers=workers)
jit_s = time.perf_counter() - t0
print(f" Built {len(runners)} runners in {jit_s:.1f}s")
for key in [("forward", 2), ("forward", 3)]:
tag = "OK" if key in runners else "FAILED"
print(f" {key[0]} {key[1]}D: {tag}")
if ("forward", 2) not in runners:
print(" ERROR: forward 2D JIT failed")
return 1
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
# =========================================================================
# Step 3: Forward 2D -- GPU + CPU reference
# =========================================================================
print("\n--- Step 3: Forward 2D ---")
prob_2d = GroupedConvProblem(
N=1, C=64, K=64, Hi=8, Wi=8, Y=3, X=3, pad_h=1, pad_w=1, direction="forward"
)
prob_2d.print_problem()
x = np.random.uniform(-0.5, 0.5, prob_2d.input_shape()).astype(np_dtype)
w = np.random.uniform(-0.5, 0.5, prob_2d.weight_shape()).astype(np_dtype)
res = runners[("forward", 2)].run(x, w, prob_2d)
print(f" Time: {res.time_ms:.4f} ms")
print(f" TFLOPS: {res.tflops:.2f}")
print(
f" Output: shape={res.output.shape}, nonzero={np.count_nonzero(res.output)}/{res.output.size}"
)
ref = cpu_conv2d_fwd(x, w, prob_2d)
diff = np.abs(res.output.astype(np.float32) - ref)
match_2d = np.allclose(res.output.astype(np.float32), ref, atol=0.05)
print(f" CPU ref: max_abs={diff.max():.6f}, match={match_2d}")
# =========================================================================
# Step 4: Forward 3D -- GPU + non-zero check
# =========================================================================
ok_3d = True
if ("forward", 3) in runners:
print("\n--- Step 4: Forward 3D ---")
prob_3d = GroupedConvProblem(
N=1,
C=64,
K=64,
Di=8,
Hi=8,
Wi=8,
Z=3,
Y=3,
X=3,
pad_d=1,
pad_h=1,
pad_w=1,
direction="forward",
)
prob_3d.print_problem()
x3 = np.random.uniform(-0.5, 0.5, prob_3d.input_shape()).astype(np_dtype)
w3 = np.random.uniform(-0.5, 0.5, prob_3d.weight_shape()).astype(np_dtype)
res3 = runners[("forward", 3)].run(x3, w3, prob_3d)
nz = np.count_nonzero(res3.output)
ok_3d = res3.success and nz > 0
print(f" Time: {res3.time_ms:.4f} ms")
print(f" TFLOPS: {res3.tflops:.2f}")
print(f" NonZero: {nz}/{res3.output.size}")
for r in runners.values():
r.cleanup()
passed = res.success and match_2d and ok_3d
print("\n" + "=" * 70)
print(f" Forward 2D: {'PASS' if match_2d else 'FAIL'} (CPU validated)")
print(f" Forward 3D: {'PASS' if ok_3d else 'FAIL'} (non-zero check)")
print(f" JIT build: {jit_s:.1f}s")
print(f" Status: {'PASS' if passed else 'FAIL'}")
print("=" * 70)
return 0 if passed else 1
if __name__ == "__main__":
sys.exit(main())

214
dispatcher/examples/grouped_conv/python/03_bwd_data.py Normal file
View File

@@ -0,0 +1,214 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 03: Backward Data Convolution (2D + 3D)
dX = ConvBwdData(dY, W)
Declares backward-data kernels with explicit parameters,
builds a registry, JIT compiles, runs on GPU, and validates
against a CPU reference.
Usage:
python3 03_bwd_data.py
"""
import sys
import argparse
import time
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
from grouped_conv_utils import (
GroupedConvKernelConfig,
GroupedConvProblem,
GroupedConvRegistry,
detect_gpu_arch,
)
def cpu_conv2d_bwd_data(dy, wei, prob):
"""CPU ref: compute dX from dY and W."""
N, Ho, Wo, G, Kpg = dy.shape
_, _, Y, X, C = wei.shape
Hi, Wi = prob.Hi, prob.Wi
dx = np.zeros((N, Hi, Wi, G, C), dtype=np.float32)
for n in range(N):
for g in range(G):
for hi in range(Hi):
for wi in range(Wi):
for c in range(C):
s = 0.0
for y in range(Y):
for x in range(X):
ho = hi + prob.pad_h - y
wo = wi + prob.pad_w - x
if ho % prob.stride_h == 0 and wo % prob.stride_w == 0:
ho //= prob.stride_h
wo //= prob.stride_w
if 0 <= ho < Ho and 0 <= wo < Wo:
for k in range(Kpg):
s += float(dy[n, ho, wo, g, k]) * float(
wei[g, k, y, x, c]
)
dx[n, hi, wi, g, c] = s
return dx
def main():
parser = argparse.ArgumentParser(description="Backward Data (2D + 3D)")
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument("--workers", type=int, default=0)
args = parser.parse_args()
arch = args.arch
print("=" * 70)
print("Example 03: Backward Data Convolution (2D + 3D)")
print("=" * 70)
print(f"\n Arch: {arch}, Dtype: {args.dtype}")
print(" dX = ConvBwdData(dY, W)")
# =========================================================================
# Step 1: Declare bwd_data kernels
# =========================================================================
print("\n--- Step 1: Declare BwdData Kernels ---")
reg = GroupedConvRegistry("bwd_data_conv")
# BwdData 2D: compv3, 128x128 tile
reg.add(
GroupedConvKernelConfig(
variant="bwd_data",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
# BwdData 3D: compv3, 64x64 tile
reg.add(
GroupedConvKernelConfig(
variant="bwd_data",
ndim_spatial=3,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=64,
tile_k=64,
wave_m=1,
wave_n=4,
wave_k=1,
warp_tile_m=16,
warp_tile_n=16,
warp_tile_k=32,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
reg.print_registry()
# =========================================================================
# Step 2: JIT build
# =========================================================================
print("\n--- Step 2: JIT Build ---")
workers = args.workers if args.workers > 0 else None
t0 = time.perf_counter()
runners = reg.build(verbose=False, max_workers=workers)
jit_s = time.perf_counter() - t0
print(f" Built {len(runners)} runners in {jit_s:.1f}s")
if ("bwd_data", 2) not in runners:
print(" ERROR: bwd_data 2D JIT failed")
return 1
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
# =========================================================================
# Step 3: BwdData 2D -- GPU + CPU reference
# =========================================================================
print("\n--- Step 3: Backward Data 2D ---")
prob = GroupedConvProblem(
N=1, C=32, K=32, Hi=8, Wi=8, Y=3, X=3, pad_h=1, pad_w=1, direction="bwd_data"
)
prob.print_problem()
dy = np.random.uniform(-0.5, 0.5, prob.output_shape()).astype(np_dtype)
w = np.random.uniform(-0.5, 0.5, prob.weight_shape()).astype(np_dtype)
res = runners[("bwd_data", 2)].run(dy, w, prob)
print(f" Time: {res.time_ms:.4f} ms")
print(f" TFLOPS: {res.tflops:.2f}")
print(f" NonZero: {np.count_nonzero(res.output)}/{res.output.size}")
ref = cpu_conv2d_bwd_data(dy, w, prob)
diff = np.abs(res.output.astype(np.float32) - ref)
match_2d = np.allclose(res.output.astype(np.float32), ref, atol=0.1)
print(f" CPU ref: max_abs={diff.max():.6f}, match={match_2d}")
# =========================================================================
# Step 4: BwdData 3D -- GPU + non-zero check
# =========================================================================
ok_3d = True
if ("bwd_data", 3) in runners:
print("\n--- Step 4: Backward Data 3D ---")
prob3 = GroupedConvProblem(
N=1,
C=32,
K=32,
Di=6,
Hi=6,
Wi=6,
Z=3,
Y=3,
X=3,
pad_d=1,
pad_h=1,
pad_w=1,
direction="bwd_data",
)
dy3 = np.random.uniform(-0.5, 0.5, prob3.output_shape()).astype(np_dtype)
w3 = np.random.uniform(-0.5, 0.5, prob3.weight_shape()).astype(np_dtype)
res3 = runners[("bwd_data", 3)].run(dy3, w3, prob3)
nz = np.count_nonzero(res3.output)
ok_3d = res3.success and nz > 0
print(f" Time: {res3.time_ms:.4f} ms, NonZero: {nz}/{res3.output.size}")
for r in runners.values():
r.cleanup()
passed = res.success and match_2d and ok_3d
print("\n" + "=" * 70)
print(f" BwdData 2D: {'PASS' if match_2d else 'FAIL'} (CPU validated)")
print(f" BwdData 3D: {'PASS' if ok_3d else 'FAIL'}")
print(f" Status: {'PASS' if passed else 'FAIL'}")
print("=" * 70)
return 0 if passed else 1
if __name__ == "__main__":
sys.exit(main())

224
dispatcher/examples/grouped_conv/python/04_bwd_weight.py Normal file
View File

@@ -0,0 +1,224 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 04: Backward Weight Convolution (2D + 3D)
dW = ConvBwdWeight(X, dY)
Declares backward-weight kernels with explicit parameters,
builds a registry, JIT compiles, runs on GPU, and validates
against a CPU reference.
Usage:
python3 04_bwd_weight.py
"""
import sys
import argparse
import time
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
from grouped_conv_utils import (
GroupedConvKernelConfig,
GroupedConvProblem,
GroupedConvRegistry,
detect_gpu_arch,
)
def cpu_conv2d_bwd_weight(x, dy, prob):
"""CPU ref: compute dW from X and dY."""
N, Hi, Wi, G, C = x.shape
_, Ho, Wo, _, Kpg = dy.shape
Y, X_ = prob.Y, prob.X
dw = np.zeros((G, Kpg, Y, X_, C), dtype=np.float32)
for g in range(G):
for k in range(Kpg):
for y in range(Y):
for xf in range(X_):
for c in range(C):
s = 0.0
for n in range(N):
for ho in range(Ho):
for wo in range(Wo):
hi = ho * prob.stride_h - prob.pad_h + y
wi = wo * prob.stride_w - prob.pad_w + xf
if 0 <= hi < Hi and 0 <= wi < Wi:
s += float(x[n, hi, wi, g, c]) * float(
dy[n, ho, wo, g, k]
)
dw[g, k, y, xf, c] = s
return dw
def main():
parser = argparse.ArgumentParser(description="Backward Weight (2D + 3D)")
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument("--workers", type=int, default=0)
parser.add_argument(
"--split-k", type=int, default=1, help="Split-K factor for bwd_weight (k_batch)"
)
args = parser.parse_args()
arch = args.arch
print("=" * 70)
print("Example 04: Backward Weight Convolution (2D + 3D)")
print("=" * 70)
print(f"\n Arch: {arch}, Dtype: {args.dtype}")
print(" dW = ConvBwdWeight(X, dY)")
# =========================================================================
# Step 1: Declare bwd_weight kernels
# =========================================================================
print("\n--- Step 1: Declare BwdWeight Kernels ---")
reg = GroupedConvRegistry("bwd_weight_conv")
# BwdWeight 2D: compv3, 128x128 tile
reg.add(
GroupedConvKernelConfig(
variant="bwd_weight",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
# BwdWeight 3D: compv3, 64x64 tile
reg.add(
GroupedConvKernelConfig(
variant="bwd_weight",
ndim_spatial=3,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=64,
tile_k=64,
wave_m=1,
wave_n=4,
wave_k=1,
warp_tile_m=16,
warp_tile_n=16,
warp_tile_k=32,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
reg.print_registry()
# =========================================================================
# Step 2: JIT build
# =========================================================================
print("\n--- Step 2: JIT Build ---")
workers = args.workers if args.workers > 0 else None
t0 = time.perf_counter()
runners = reg.build(verbose=False, max_workers=workers)
jit_s = time.perf_counter() - t0
print(f" Built {len(runners)} runners in {jit_s:.1f}s")
if ("bwd_weight", 2) not in runners:
print(" ERROR: bwd_weight 2D JIT failed")
return 1
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
# =========================================================================
# Step 3: BwdWeight 2D -- GPU + CPU reference
# =========================================================================
print("\n--- Step 3: Backward Weight 2D ---")
prob = GroupedConvProblem(
N=1,
C=32,
K=32,
Hi=8,
Wi=8,
Y=3,
X=3,
pad_h=1,
pad_w=1,
direction="bwd_weight",
split_k=args.split_k,
)
prob.print_problem()
x = np.random.uniform(-0.5, 0.5, prob.input_shape()).astype(np_dtype)
dy = np.random.uniform(-0.5, 0.5, prob.output_shape()).astype(np_dtype)
res = runners[("bwd_weight", 2)].run(x, dy, prob)
print(f" Time: {res.time_ms:.4f} ms")
print(f" TFLOPS: {res.tflops:.2f}")
print(f" NonZero: {np.count_nonzero(res.output)}/{res.output.size}")
ref = cpu_conv2d_bwd_weight(x, dy, prob)
diff = np.abs(res.output.astype(np.float32) - ref)
match_2d = np.allclose(res.output.astype(np.float32), ref, atol=0.5)
print(f" CPU ref: max_abs={diff.max():.6f}, match={match_2d}")
# =========================================================================
# Step 4: BwdWeight 3D -- GPU + non-zero check
# =========================================================================
ok_3d = True
if ("bwd_weight", 3) in runners:
print("\n--- Step 4: Backward Weight 3D ---")
prob3 = GroupedConvProblem(
N=1,
C=32,
K=32,
Di=6,
Hi=6,
Wi=6,
Z=3,
Y=3,
X=3,
pad_d=1,
pad_h=1,
pad_w=1,
direction="bwd_weight",
)
x3 = np.random.uniform(-0.5, 0.5, prob3.input_shape()).astype(np_dtype)
dy3 = np.random.uniform(-0.5, 0.5, prob3.output_shape()).astype(np_dtype)
res3 = runners[("bwd_weight", 3)].run(x3, dy3, prob3)
nz = np.count_nonzero(res3.output)
ok_3d = res3.success and nz > 0
print(f" Time: {res3.time_ms:.4f} ms, NonZero: {nz}/{res3.output.size}")
for r in runners.values():
r.cleanup()
passed = res.success and match_2d and ok_3d
print("\n" + "=" * 70)
print(f" BwdWeight 2D: {'PASS' if match_2d else 'FAIL'} (CPU validated)")
print(f" BwdWeight 3D: {'PASS' if ok_3d else 'FAIL'}")
print(f" Status: {'PASS' if passed else 'FAIL'}")
print("=" * 70)
return 0 if passed else 1
if __name__ == "__main__":
sys.exit(main())

318
dispatcher/examples/grouped_conv/python/05_benchmark.py Normal file
View File

@@ -0,0 +1,318 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 05: Multi-Problem GPU Benchmark
Declares kernels with explicit tile/wave/warp/pipeline parameters for
all directions, builds registries, JIT compiles, and benchmarks across
ResNet-like problem sizes with configurable warmup/repeat.
Usage:
python3 05_benchmark.py
python3 05_benchmark.py --warmup 3 --repeat 10
python3 05_benchmark.py --workers 4
"""
import sys
import argparse
import time
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
from grouped_conv_utils import (
GroupedConvKernelConfig,
GroupedConvProblem,
GroupedConvRegistry,
detect_gpu_arch,
)
def compute_bytes(prob, dtype_bytes=2):
in_elems = 1
for d in prob.input_shape():
in_elems *= d
wei_elems = 1
for d in prob.weight_shape():
wei_elems *= d
out_elems = 1
for d in prob.output_shape():
out_elems *= d
return (in_elems + wei_elems + out_elems) * dtype_bytes
def main():
parser = argparse.ArgumentParser(description="Multi-Problem GPU Benchmark")
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument("--warmup", type=int, default=3, help="Warmup iterations")
parser.add_argument("--repeat", type=int, default=5, help="Benchmark iterations")
parser.add_argument(
"--workers", type=int, default=0, help="Max JIT workers (0=auto)"
)
args = parser.parse_args()
print("=" * 70)
print("Example 05: Multi-Problem GPU Benchmark")
print("=" * 70)
print(f"\n Arch: {args.arch}, Dtype: {args.dtype}")
print(f" Warmup: {args.warmup}, Repeat: {args.repeat}")
# =========================================================================
# Step 1: Declare all kernels with explicit parameters
# =========================================================================
print("\n--- Step 1: Declare Kernels ---")
reg = GroupedConvRegistry("benchmark")
# Forward 2D: compv4, 128x128 tile
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=2,
arch=args.arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv4",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
# Forward 3D: compv3, 64x64 tile
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=3,
arch=args.arch,
dtype=args.dtype,
tile_m=1,
tile_n=64,
tile_k=64,
wave_m=1,
wave_n=4,
wave_k=1,
warp_tile_m=16,
warp_tile_n=16,
warp_tile_k=32,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
# BwdData 2D: compv3, 128x128 tile
reg.add(
GroupedConvKernelConfig(
variant="bwd_data",
ndim_spatial=2,
arch=args.arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
# BwdWeight 2D: compv3, 128x128 tile
reg.add(
GroupedConvKernelConfig(
variant="bwd_weight",
ndim_spatial=2,
arch=args.arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
)
)
reg.print_registry()
# =========================================================================
# Step 2: JIT build
# =========================================================================
print("\n--- Step 2: JIT Build ---")
workers = args.workers if args.workers > 0 else None
t0 = time.perf_counter()
runner_by_key = reg.build(verbose=False, max_workers=workers)
jit_s = time.perf_counter() - t0
for key in [("forward", 2), ("forward", 3), ("bwd_data", 2), ("bwd_weight", 2)]:
tag = "OK" if key in runner_by_key else "FAILED"
print(f" {key[0]:12s} {key[1]}D: {tag}")
print(f" JIT build time: {jit_s:.3f} s")
missing = [
k
for k in [("forward", 2), ("forward", 3), ("bwd_data", 2), ("bwd_weight", 2)]
if k not in runner_by_key
]
if missing:
print(f"\n ERROR: missing {missing}")
return 1
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
def bench_run(runner, inp, wei, prob):
for _ in range(args.warmup):
runner.run(inp, wei, prob)
times = []
for _ in range(args.repeat):
r = runner.run(inp, wei, prob)
if r.success:
times.append(r.time_ms)
if not times:
return 0.0, 0.0
return min(times), sum(times) / len(times)
# =========================================================================
# Step 3: 2D Forward benchmark
# =========================================================================
print("\n--- Step 3: Forward 2D Benchmark ---")
print(
f"{'Problem':<18} {'N':>3} {'C':>4} {'K':>4} {'H':>3} {'W':>3} "
f"{'F':>3} {'Min(ms)':>9} {'Avg(ms)':>9} {'TFLOPS':>8} {'GB/s':>8}"
)
print("-" * 85)
all_ok = True
for label, n, c, k, h, w, y, x, s, p in [
("ResNet-stage2", 1, 64, 64, 56, 56, 3, 3, 1, 1),
("ResNet-stage3", 1, 128, 128, 28, 28, 3, 3, 1, 1),
("ResNet-stage4", 1, 256, 256, 14, 14, 3, 3, 1, 1),
("ResNet-stage5", 1, 512, 512, 7, 7, 3, 3, 1, 1),
("Pointwise-1x1", 1, 256, 256, 56, 56, 1, 1, 1, 0),
("Batch-8", 8, 64, 128, 56, 56, 3, 3, 1, 1),
("Batch-32", 32, 64, 128, 56, 56, 3, 3, 1, 1),
]:
prob = GroupedConvProblem(
N=n,
C=c,
K=k,
Hi=h,
Wi=w,
Y=y,
X=x,
stride_h=s,
stride_w=s,
pad_h=p,
pad_w=p,
direction="forward",
)
inp = np.random.uniform(-0.3, 0.3, prob.input_shape()).astype(np_dtype)
wei = np.random.uniform(-0.3, 0.3, prob.weight_shape()).astype(np_dtype)
min_ms, avg_ms = bench_run(runner_by_key[("forward", 2)], inp, wei, prob)
if avg_ms > 0:
tflops = prob.flops / (avg_ms * 1e9)
bw = compute_bytes(prob) / (avg_ms * 1e6)
print(
f"{label:<18} {n:>3} {c:>4} {k:>4} {h:>3} {w:>3} "
f"{y}x{x} {min_ms:>9.4f} {avg_ms:>9.4f} {tflops:>8.2f} {bw:>8.1f}"
)
else:
all_ok = False
# =========================================================================
# Step 4: 3D Forward
# =========================================================================
print("\n--- Step 4: Forward 3D ---")
for label, n, c, k, d, h, w, z, y, x in [
("3D-small", 1, 64, 64, 8, 16, 16, 3, 3, 3),
("3D-medium", 1, 64, 128, 16, 32, 32, 3, 3, 3),
]:
prob = GroupedConvProblem(
N=n, C=c, K=k, Di=d, Hi=h, Wi=w, Z=z, Y=y, X=x, direction="forward"
)
inp = np.random.uniform(-0.3, 0.3, prob.input_shape()).astype(np_dtype)
wei = np.random.uniform(-0.3, 0.3, prob.weight_shape()).astype(np_dtype)
min_ms, avg_ms = bench_run(runner_by_key[("forward", 3)], inp, wei, prob)
if avg_ms > 0:
tflops = prob.flops / (avg_ms * 1e9)
print(f" {label:<14} {min_ms:.4f} / {avg_ms:.4f} ms {tflops:.2f} TFLOPS")
# =========================================================================
# Step 5: Backward directions
# =========================================================================
print("\n--- Step 5: Backward Directions ---")
for label, direction in [
("bwd_data ResNet-s3", "bwd_data"),
("bwd_weight ResNet-s3", "bwd_weight"),
]:
prob = GroupedConvProblem(
N=1,
C=128,
K=128,
Hi=28,
Wi=28,
Y=3,
X=3,
stride_h=1,
stride_w=1,
pad_h=1,
pad_w=1,
direction=direction,
)
inp = np.random.uniform(-0.3, 0.3, prob.input_shape()).astype(np_dtype)
wei = np.random.uniform(-0.3, 0.3, prob.weight_shape()).astype(np_dtype)
min_ms, avg_ms = bench_run(runner_by_key[(direction, 2)], inp, wei, prob)
if avg_ms > 0:
tflops = prob.flops / (avg_ms * 1e9)
print(
f" {label:<14} {direction:>12} {min_ms:.4f} / {avg_ms:.4f} ms {tflops:.2f} TFLOPS"
)
for runner in runner_by_key.values():
runner.cleanup()
print("\n" + "=" * 70)
print(f" JIT build: {jit_s:.3f} s")
print(f" Warmup: {args.warmup}, Repeat: {args.repeat}")
print(f" Status: {'PASS' if all_ok else 'FAIL'}")
print("=" * 70)
return 0 if all_ok else 1
if __name__ == "__main__":
sys.exit(main())

274
dispatcher/examples/grouped_conv/python/06_registry_json.py Normal file
View File

@@ -0,0 +1,274 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Example 06: Registry, Heuristic Selection & JSON Export
Declares multiple kernel configurations with different tile sizes,
builds a registry, demonstrates heuristic runtime kernel selection,
JSON round-trip, and GPU execution.
Usage:
python3 06_registry_json.py
python3 06_registry_json.py --workers 4
"""
import sys
import time
import argparse
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
from grouped_conv_utils import (
GroupedConvKernelConfig,
GroupedConvProblem,
GroupedConvRegistry,
detect_gpu_arch,
)
def conv_heuristic(problem):
spatial = problem.Ho * problem.Wo
if spatial > 400:
return ["256", "128", "64"]
return ["64", "128", "256"]
def main():
parser = argparse.ArgumentParser(description="Registry, Heuristic & JSON")
parser.add_argument("--arch", default=detect_gpu_arch())
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
parser.add_argument("--workers", type=int, default=0)
args = parser.parse_args()
arch = args.arch
print("=" * 70)
print("Example 06: Registry, Heuristic Selection & JSON Export")
print("=" * 70)
print(f"\n Arch: {arch}, Dtype: {args.dtype}")
# Step 1: Declare kernels with full explicit parameters
print("\n--- Step 1: Declare Kernels + Build Registry ---")
reg = GroupedConvRegistry("conv_tiles")
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=256,
tile_k=256,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
num_wave_groups=1,
num_groups_to_merge=1,
)
)
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv4",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
num_wave_groups=1,
num_groups_to_merge=1,
)
)
reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=64,
tile_k=64,
wave_m=1,
wave_n=4,
wave_k=1,
warp_tile_m=16,
warp_tile_n=16,
warp_tile_k=32,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
block_per_cu=1,
num_wave_groups=1,
num_groups_to_merge=1,
)
)
reg.print_registry()
# Step 2: Heuristic kernel selection
print("\n--- Step 2: Heuristic Kernel Selection ---")
problems = [
(
"small_7x7",
GroupedConvProblem(
N=1,
C=512,
K=512,
Hi=7,
Wi=7,
Y=3,
X=3,
pad_h=1,
pad_w=1,
direction="forward",
),
),
(
"medium_14x14",
GroupedConvProblem(
N=1,
C=256,
K=256,
Hi=14,
Wi=14,
Y=3,
X=3,
pad_h=1,
pad_w=1,
direction="forward",
),
),
(
"large_56x56",
GroupedConvProblem(
N=1,
C=64,
K=128,
Hi=56,
Wi=56,
Y=3,
X=3,
pad_h=1,
pad_w=1,
direction="forward",
),
),
]
print(f" {'Problem':<16} {'Spatial':>8} {'Selected Kernel':<50}")
print(f" {'-' * 74}")
for label, prob in problems:
selected = reg.select(prob, heuristic=conv_heuristic)
spatial = prob.Ho * prob.Wo
sel_name = selected.name if selected else "none"
print(f" {label:<16} {spatial:>8} {sel_name:<50}")
# Step 3: JSON round-trip
print("\n--- Step 3: JSON Round-Trip ---")
json_str = reg.to_json()
print(f" Exported: {len(json_str)} bytes, {len(reg)} kernels")
imported = GroupedConvRegistry.from_json(json_str)
print(f" Imported: {len(imported)} kernels")
orig = reg.kernels[0]
imp = imported.kernels[0]
rt_ok = (
orig.vector_size_a == imp.vector_size_a
and orig.block_per_cu == imp.block_per_cu
and orig.tile_n == imp.tile_n
)
print(f" Full fields round-trip: {'OK' if rt_ok else 'FAIL'}")
# Step 4: JIT build + GPU execution
print("\n--- Step 4: JIT Build + GPU Execution ---")
workers = args.workers if args.workers > 0 else None
jit_reg = GroupedConvRegistry("jit_conv")
jit_reg.add(
GroupedConvKernelConfig(
variant="forward",
ndim_spatial=2,
arch=arch,
dtype=args.dtype,
tile_m=1,
tile_n=128,
tile_k=128,
wave_m=2,
wave_n=2,
wave_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv4",
scheduler="intrawave",
epilogue="cshuffle",
vector_size_a=4,
vector_size_b=8,
vector_size_c=8,
)
)
t0 = time.perf_counter()
runners = jit_reg.build(verbose=False, max_workers=workers)
jit_s = time.perf_counter() - t0
if ("forward", 2) not in runners:
print(" JIT build failed")
return 1
runner = runners[("forward", 2)]
print(f" JIT build: {jit_s:.3f} s")
print(f" Library: {runner.library_path}")
prob = GroupedConvProblem(
N=1, C=128, K=128, Hi=16, Wi=16, Y=3, X=3, pad_h=1, pad_w=1, direction="forward"
)
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
inp = np.random.uniform(-0.3, 0.3, prob.input_shape()).astype(np_dtype)
wei = np.random.uniform(-0.3, 0.3, prob.weight_shape()).astype(np_dtype)
res = runner.run(inp, wei, prob)
runner.cleanup()
if res.success:
print(f" Time: {res.time_ms:.4f} ms")
print(f" TFLOPS: {res.tflops:.2f}")
print(f" NonZero: {np.count_nonzero(res.output)}/{res.output.size}")
gpu_ok = res.success
print("\n" + "=" * 70)
print(f" Registry: {len(reg)} kernels (3 tile configs)")
print(" Heuristic: spatial-based selection demonstrated")
print(f" JSON: round-trip {'OK' if rt_ok else 'FAIL'}")
print(f" GPU: {'OK' if gpu_ok else 'FAIL'}")
print(f" Status: {'PASS' if gpu_ok and rt_ok else 'FAIL'}")
print("=" * 70)
return 0 if gpu_ok and rt_ok else 1
if __name__ == "__main__":
sys.exit(main())

60
dispatcher/heuristics/.gitignore vendored Normal file
View File

@@ -0,0 +1,60 @@
# Python bytecode and caches
__pycache__/
*.pyc
*.pyo
*.pyd
.Python
# Jupyter notebooks
*.ipynb
.ipynb_checkpoints/
# Virtual environments
.venv/
venv/
ENV/
# IDE and editor files
.vscode/
.idea/
*.swp
*.swo
*~
# Test output and logs
*.log
test_output.log
custom_shapes_gpu_test.log
# Benchmark and analysis output files
*.csv
*.json
!models/*/feature_spec.json
!models/*/train_manifest.json
# Data files (parquet, arrow)
*.parquet
*.arrow
# Temporary and NFS files
.nfs*
*.tmp
*.bak
# Decompressed model files (compressed .lgbm.gz versions are tracked)
models/**/*.lgbm
# User-specific test and analysis scripts
test_*.py
!tests/test_*.py
find_*.py
oracle_*.json
validation_results_*.csv
custom_shapes_*.csv
fp16_bf16_*.csv
# Ignore all markdown files except tracked documentation
*.md
!DATA_GENERATION.md
!LEARNINGS.md
!README.md

412
dispatcher/heuristics/DATA_GENERATION.md Normal file
View File

@@ -0,0 +1,412 @@
# Data Generation Guide
This document explains how to build benchmark binaries from the CK Tile engine,
generate benchmark datasets, and manage them for the ML kernel performance
prediction system.
## Overview
The ML heuristic needs benchmark data: measured TFLOPS, latency, and bandwidth
for every (problem shape, kernel config) pair. The tile engine builds one
executable per kernel configuration. Each executable benchmarks a single kernel
on a given problem size and outputs JSON with performance metrics.
```
CK source --> CMake configure --> ninja build --> benchmark binaries
(4608 per op/dtype/layout)
benchmark binaries --> run on GPU --> streaming log --> parquet dataset
(per shape) (JSON blocks) (canonical schema)
```
## Prerequisites
- **ROCm**: HIP >= 6.0.3 (for gfx950: HIP >= 6.0.4)
- **Build tools**: CMake >= 3.21, Ninja, HIP-aware clang compiler
- **Python**: 3.10+ with `pandas`, `pyarrow`
- **GPU**: ROCm-capable AMD GPU (MI250X, MI300X, MI355X, etc.)
---
## Part 1: Building Benchmark Binaries from the Tile Engine
If you already have pre-built binaries (e.g., in `/workspace/ck_tile/bin/`),
skip to Part 2. This section explains how to build them from source.
### Step 1: CMake Configure
From the CK repository root:
```bash
cmake -S /workspace/rocm-libraries/projects/composablekernel \
-B build \
-DCMAKE_BUILD_TYPE=Release \
-DGPU_TARGETS="gfx950" \
-DGEMM_UNIVERSAL_DATATYPE="fp8" \
-DGEMM_UNIVERSAL_LAYOUT="rcr" \
-G Ninja
```
**Key CMake variables:**
| Variable | Default | Description |
|---|---|---|
| `GPU_TARGETS` | (required) | Target GPU architectures. Supported: `gfx90a`, `gfx942`, `gfx950`, `gfx1201`. Semicolon-separated for multiple. |
| `GEMM_UNIVERSAL_DATATYPE` | `"fp8;fp16"` | Data types to build. Options: `fp8`, `fp16`, `bf16`, `bf8`. Semicolon-separated. |
| `GEMM_UNIVERSAL_LAYOUT` | `"rcr;rrr;crr;ccr"` | Layouts to build. Semicolon-separated. |
| `GEMM_UNIVERSAL_CONFIG_FILE` | `"default_config.json"` | Kernel config file (in the `configs/` directory). Controls which tile sizes, warp configs, pipelines, etc. are enumerated. |
| `ENABLE_CCACHE_GEMM_UNIVERSAL` | `OFF` | Enable ccache for faster rebuilds. |
**Example: build only fp8 RCR for gfx950 (fastest, ~4608 kernels):**
```bash
cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \
-DGPU_TARGETS="gfx950" \
-DGEMM_UNIVERSAL_DATATYPE="fp8" \
-DGEMM_UNIVERSAL_LAYOUT="rcr" \
-G Ninja
```
**Example: build all dtypes and layouts (slow, ~4608 * 4 * 4 = ~73K kernels):**
```bash
cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \
-DGPU_TARGETS="gfx950" \
-DGEMM_UNIVERSAL_DATATYPE="fp8;fp16;bf16;bf8" \
-DGEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \
-G Ninja
```
### What happens during configure
1. CMake calls `gemm_universal_instance_builder.py --list_kernels` to enumerate
all valid kernel configurations from the config JSON.
2. It writes `gemm_universal_kernel_list.txt` (one kernel per line) and
`gemm_universal_kernel_count.txt` to the build directory.
3. For each kernel, it creates a ninja build target.
### Step 2: Build
```bash
# Build all benchmarks for the configured dtypes/layouts
ninja -C build benchmark_gemm_universal_all
# Or build a specific dtype/layout combo
ninja -C build benchmark_gemm_universal_fp8_rcr
# Or build by pipeline type
ninja -C build benchmark_gemm_universal_compv4_pipeline
ninja -C build benchmark_gemm_universal_mem_pipeline
# Or build a single specific kernel
ninja -C build benchmark_gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128
```
**Build time estimates:**
- ~4608 kernels (one dtype, one layout): 1-4 hours depending on CPU cores
- Use `-j <N>` to control parallelism: `ninja -C build -j 32 benchmark_gemm_universal_fp8_rcr`
### Step 3: Verify binaries
Binaries are placed in `build/bin/`:
```bash
ls build/bin/benchmark_gemm_universal_fp8_rcr_* | wc -l
# Expected: 4608 (for default config)
# Test one binary
./build/bin/benchmark_gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128 \
-m=1024 -n=1024 -k=1024 -warmup=3 -repeat=10 -verify=0
```
### Kernel config files
The config files live in:
```
tile_engine/ops/gemm/gemm_universal/configs/
default_config.json # Default: full enumeration
default_ci_config.json # CI: reduced set for fast testing
user_provided_config.json # Custom: your own subset
```
To use a custom config:
```bash
cmake ... -DGEMM_UNIVERSAL_CONFIG_FILE="user_provided_config.json"
```
The config controls which tile sizes (e.g., 128x128x64, 256x256x32), warp
configurations (e.g., 2x2x1, 1x4x1), pipelines (compv3, compv4, mem),
schedulers, and other parameters are included in the kernel enumeration.
### Building StreamK / other ops
The same pattern applies to other tile engine ops:
```bash
# StreamK
ninja -C build benchmark_gemm_streamk_fp8_rcr
# Grouped convolution
ninja -C build benchmark_grouped_conv_fwd_fp16_nhwgc
```
Each op has its own instance builder and config directory.
---
## Part 2: Running Benchmarks and Generating Data
## Quick Start
### 1. Run benchmarks for a set of shapes
Each binary accepts `-m=`, `-n=`, `-k=`, `-warmup=`, `-repeat=`, `-verify=` flags
and outputs JSON to stdout:
```bash
/workspace/ck_tile/bin/benchmark_gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128 \
-m=1024 -n=1024 -k=1024 -warmup=3 -repeat=10 -verify=0
```
Output:
```json
{
"name": "gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_...",
"problem": {
"split_k": 1, "m": 1024, "n": 1024, "k": 1024,
"dtype_a": "fp8", "dtype_b": "fp8", ...
},
"perf_result": {
"latency(ms)": 0.04,
"tflops(TFlops)": 204.60,
"bandwidth(GB/s)": 624.39
}
}
```
### 2. Batch generation using provided scripts
**Wide coverage (diverse shapes across all regimes):**
```bash
python3 generate_wide_coverage.py \
--bin_dir /workspace/ck_tile/bin \
--out_dir data/wide_coverage \
--batch_size 25 \
--warmup 3 --repeat 10
```
**Edge-case dimensions (N=1, K=1, small N/K):**
```bash
python3 generate_edge_dims.py
```
Both scripts write streaming log files that `data_pipeline.py` can parse.
### 3. Parse logs into parquet
```bash
python3 data_pipeline.py <log_file> \
-o data/my_dataset.parquet \
--arch gfx950 \
--capture_hw
```
The `--capture_hw` flag runs `rocminfo` once and injects the GPU hardware
profile (CU count, clock speed, cache sizes, etc.) into every row.
## Canonical Data Schema
Every parquet file follows this schema:
| Column | Type | Description |
|---|---|---|
| `op_type` | str | `gemm_universal`, `gemm_streamk`, etc. |
| `dtype` | str | `fp8`, `fp16`, `bf16`, `bf8` |
| `layout` | str | `rcr`, `rrr`, `crr`, `ccr` |
| `arch` | str | `gfx942`, `gfx950`, etc. |
| `kernel_name` | str | Full kernel identifier |
| `m`, `n`, `k` | int | Problem dimensions |
| `split_k` | int | Split-K factor (1 = standard) |
| `measured_tflops` | float | Ground-truth TFLOPS |
| `latency_ms` | float | Measured latency |
| `bandwidth_gb_s` | float | Measured bandwidth |
| `is_valid` | bool | True if tflops > 0 and latency > 0 |
| `tile_m`, `tile_n`, `tile_k` | int | Tile dimensions |
| `warp_m`, `warp_n`, `warp_k` | int | Warp config |
| `warp_tile_m/n/k` | int | Warp tile dimensions |
| `pipeline` | str | `compv3`, `compv4`, `mem`, etc. |
| `scheduler` | str | `intrawave`, `interwave` |
| `epilogue` | str | `cshuffle`, `default` |
| `pad_m`, `pad_n`, `pad_k` | bool | Padding flags |
| `persistent` | bool | Persistent kernel flag |
| `run_id` | str | Unique collection run identifier |
## Shape Selection Guidelines
Good training data requires diverse shapes. Cover all of these regimes:
### By M dimension (batch size / output rows)
- **M=1**: single-token inference (hardest case for tiling)
- **Tiny M (2-16)**: small batch inference
- **Small M (32-128)**: medium batch
- **Medium M (256-2048)**: large batch / training
- **Large M (4096-20480)**: very large batch
### By N and K dimension
- **N=1**: vector-matrix multiply (degenerate)
- **K=1**: rank-1 update / outer product (degenerate)
- **Small N or K (2-16)**: stress tile efficiency
- **Deep K (K > 4096)**: compute-bound regime
- **Shallow K (K < 256)**: memory-bound regime
### By shape family
- **Square**: M ~ N ~ K (powers of 2)
- **Tall**: M >> N (tall output matrix)
- **Wide**: N >> M (wide output matrix)
- **Deep-K**: K >> M and K >> N
### Special cases
- **Prime dimensions**: 17, 31, 127, 251, 509, 1021, 2039, 4093
(worst-case for tile alignment, tests padding logic)
- **Non-power-of-2**: 48, 96, 192, 384, 576, 768, 1536, 3072, 4608
(common in LLM architectures)
- **LLM inference shapes**: DeepSeek, LLaMA-7B, LLaMA-70B MLP/attention dims
### Minimum recommended coverage
For a production-quality model, aim for:
- At least 200 unique (M, N, K) shapes
- At least 10 shapes per shape family
- All kernel configs (4608 for fp8 RCR) run against every shape
- Multiple layouts if training a cross-layout model
## Benchmark Quality Guidelines
### Warmup and repeat
- Minimum `warmup=3`, `repeat=10` for fast iteration
- Production quality: `warmup=5`, `repeat=20` for stable measurements
- The `perf_result` values are averaged over `repeat` iterations
### Noise handling
- Use **median** latency when aggregating multiple runs of the same benchmark
- Flag measurements where coefficient of variation exceeds 10%
- Avoid benchmarking under thermal throttling (check GPU temperature)
- Lock GPU clocks if possible for reproducibility
### Environment metadata
Store with every dataset:
- GPU model and architecture (from `rocminfo`)
- ROCm driver version
- Clock mode (default / locked)
- Git hash of the CK tile engine build (if available)
- Timestamp
## Adding Data for a New Op
To generate benchmark data for a new operation (e.g., `gemm_streamk`):
1. **Build the binaries** using the tile engine:
```bash
ninja -C build benchmark_gemm_streamk_fp8_rcr
```
2. **Write a generation script** (or modify `generate_wide_coverage.py`):
- Change the executable glob pattern to match the new op
- Add any op-specific CLI flags the binaries need
3. **Run and parse**:
```bash
python3 data_pipeline.py my_streamk_run.log \
-o data/gemm_streamk_fp8_gfx950.parquet --arch gfx950
```
4. **Train**:
```bash
python3 train.py --op gemm_streamk --dtype fp8 --arch gfx950 \
--data_dir data/ --out_dir models/gemm_streamk_fp8_gfx950
```
## Adding Data for a New Layout
Same binaries, same shapes -- just change the layout filter:
```bash
# Build rrr binaries
ninja -C build benchmark_gemm_universal_fp8_rrr
# Generate and parse
# ... (same flow, different bin_dir or executable glob)
# Train a cross-layout model by putting all layouts in the same data_dir
python3 train.py --data_dir data/ --out_dir models/gemm_universal_fp8_gfx950_all_layouts
```
The feature engine includes `layout` as a categorical feature, so one model
can handle all layouts.
## Incremental Data Collection
When you have a trained model and want to add more data:
1. Generate new data (new shapes, new layouts, etc.)
2. Parse into parquet alongside existing data
3. Warm-start from the previous model:
```bash
python3 train.py --data_dir data/ --out_dir models/v2 \
--warm_start models/v1 \
--warm_start_n_estimators 200
```
This adds 200 new trees on top of the existing model. The feature schema
must match exactly (enforced automatically).
## File Organization
Recommended directory structure:
```
heuristics/
data/
gemm_universal_fp8_rcr_gfx950.parquet # original 108 shapes
wide_coverage/ # batch log files
wide_coverage_batch_001.log
wide_coverage_batch_002.log
...
edge_dims/ # N=1, K=1 edge cases
edge_dims_batch_001.log
...
models/
gemm_universal_fp8_gfx950/ # trained model artifacts
model_tflops.lgbm
model_latency.lgbm
model_bandwidth.lgbm
feature_spec.json
train_manifest.json
cv_metrics_tflops.json
eval_report.json
...
```
## Troubleshooting
### Benchmark binary exits with non-zero code
Some kernel configs are invalid for certain problem sizes (e.g., tile_m=256
with M=16). The data pipeline marks these as `is_valid=False` and they are
filtered out during training. This is expected.
### Edge dims produce very few results
N=1 and K=1 shapes are degenerate -- most kernel configurations have minimum
dimension requirements and will fail or produce zero TFLOPS. The small number
of valid results is still useful (it tells the model which configs work for
these shapes).
### Benchmarks are slow
Each shape requires running all 4608 kernel executables sequentially. At
~0.01s per kernel, that is ~46 seconds per shape. For 700 shapes, expect
~9 hours. Tips:
- Run on a dedicated GPU (no other workloads)
- Use `--batch_size 25` to get incremental output
- Parse and train on partial data while generation continues
### Data from different GPUs / driver versions
Store `run_id` and hardware metadata with each dataset. Training on mixed
data is allowed but not recommended for production models. Filter to a
single `run_id` or `arch` for clean experiments.

151
dispatcher/heuristics/LEARNINGS.md Normal file
View File

@@ -0,0 +1,151 @@
# Learnings and Design Decisions
Empirical findings from building the CK Tile kernel performance prediction system.
These inform the current defaults and explain why certain approaches were chosen.
## 1. Log-Transform is Essential for Cross-Scale Accuracy
**Problem**: GEMM TFLOPS spans 5 orders of magnitude across different problem
sizes. When training on raw TFLOPS, the regression loss (RMSE) is dominated by
large shapes where absolute errors are biggest. The model learns to predict
large shapes accurately but ignores tiny shapes where the TFLOPS values are
much lower.
**Evidence** (168 shapes, 626K rows, 5-fold GroupKFold CV):
| Model | Mean Eff | P10 Eff | tiny_m Eff | Min Eff |
| ----------------------------- | ---------- | ---------- | ---------- | ---------- |
| Raw TFLOPS (500 trees) | 92.73% | 80.24% | 84.55% | 4.26% |
| **log1p(TFLOPS)** (500 trees) | **96.92%** | **94.34%** | **94.89%** | **60.27%** |
| log1p(TFLOPS) (2000 trees) | 97.51% | 93.89% | 96.04% | 63.56% |
**Solution**: Train on `log1p(measured_tflops)` and apply `expm1()` to
predictions. This is now the default in `train.py`. Pass `--no_log_transform`
to revert to raw regression (not recommended).
**Why log1p, not log**: `log1p(x) = log(1 + x)` handles zero and near-zero
TFLOPS gracefully, whereas `log(x)` produces -inf for x=0.
## 2. Tiny-M Shapes are the Hardest Case
M=1 (single-token inference) shapes are fundamentally different from batch shapes:
- Most kernel configurations produce very low TFLOPS
- The "best" kernel is often only marginally better than the rest
- The oracle performance itself is very low, so any prediction error tanks efficiency
- Many kernels fail outright (tile_m=128 with M=1 wastes 127/128 of the tile)
The bottom shapes in our evaluation are all M=1, with efficiencies in the
63-70% range. These shapes have such low absolute performance that the model's
noise floor exceeds the performance difference between kernels.
**Mitigation**: Log-transform helps significantly (tiny_m improved from 84% to
96%). For production use with M=1, consider a dedicated fallback (e.g.,
hardcoded kernel selection for M < 4 based on known-good configs).
## 3. IHEM (Hard Example Mining) Hurts When Scale is the Issue
We tried Iterative Hard Example Mining with sample reweighting (2x-5x weight
on hard shapes). Result: it made things **worse**, degrading mean efficiency
from 94.31% to 92.90% over 3 iterations.
**Why**: The hard shapes are hard because of scale mismatch, not because the
model lacks capacity. Reweighting amplifies the small-TFLOPS rows, which
distorts the learned relationship between features and performance for the
majority of shapes. The log-transform was the correct fix -- it addresses the
root cause (scale) rather than the symptom (bad predictions on tiny shapes).
**Lesson**: IHEM is useful when the model has capacity gaps (e.g., certain
pipeline types are underrepresented). It is counterproductive when the issue
is target-variable scale. Always try target transforms before reweighting.
## 4. GroupKFold Key = (M, N, K) Forces Generalization
The validation uses `GroupKFold` where the group key is `(M, N, K)` -- all
kernels for the same shape go to the same fold. This means:
- The model is always evaluated on shapes it has **never seen** during training
- Layout is excluded from the key, forcing the model to generalize across layouts
- Since models are per-arch, `arch` is implicit (constant within one training run)
This is much stricter than random row splitting, where the model would see some
kernels for each shape during training. Our efficiency numbers are conservative
estimates of real-world performance on unseen shapes.
## 5. Model Size vs Accuracy Tradeoff
| Config | Trees | Leaves | LR | Mean Eff | P10 Eff | Train Time |
| ------------------ | -------- | ------- | -------- | ---------- | ---------- | ------------- |
| Small (default v1) | 500 | 127 | 0.05 | 96.92% | 94.34% | ~20s |
| **Big (current)** | **2000** | **255** | **0.02** | **97.51%** | **93.89%** | **~25s/fold** |
The bigger model improved mean efficiency by 0.6% but P10 didn't improve
(actually slightly worse). The extra capacity helps on medium shapes but
doesn't crack the tiny-M floor. This suggests the feature set, not model
capacity, is the limiting factor for the hardest shapes.
For C++ deployment, the bigger model (2000 trees, 255 leaves) is still fast
enough -- LightGBM inference is O(trees * log(leaves)) per sample, which is
~microseconds even at 2000 trees.
## 6. N=1 and K=1 Shapes are Degenerate
We generated benchmark data for 546 edge-case shapes (N=1, K=1, small N/K).
Result: **zero valid kernel results** across 94 shapes. All 4608 kernels either
fail or produce 0 TFLOPS for these degenerate dimensions.
This means:
- The tile engine kernels have hard minimum dimension requirements
- N=1 / K=1 shapes cannot be handled by the current kernel set
- These shapes need dedicated kernels (e.g., BLAS-1/BLAS-2 fallbacks)
- The ML model should not be expected to handle them -- they should be filtered
out before reaching the heuristic
## 7. Feature Engineering Insights
From LightGBM feature importances on the log-target model:
**Top features** (by split count):
- `M, N, K` -- raw dimensions are always the most important
- `tile_m, tile_n, tile_k` -- the tile shape is the primary kernel differentiator
- `overall_tile_efficiency` -- how well the shape fits the tile (the interaction)
- `num_tiles_m, total_output_tiles` -- work decomposition
- `arithmetic_intensity` -- compute vs memory bound regime
- `pipeline` -- pipeline type (compv3 vs compv4 vs mem) significantly affects perf
**Low-importance features**:
- Hardware constants (CUs, clock, caches) -- they're constant within one arch
model, so they provide no discriminative signal. They'll become important when
training cross-arch models.
- `split_k` -- always 1 in current data
- `persistent` -- rarely True in current kernel set
## 8. Warm-Start Works for Incremental Updates
LightGBM's `init_model` parameter successfully continues training from an
existing model. New trees are added on top of existing ones. Key considerations:
- Feature schema must match exactly (enforced by `check_feature_compatibility`)
- Use fewer new trees (200-500) since we're refining, not starting fresh
- The `train_manifest.json` tracks the full lineage (total trees, data sizes)
- Quality should be at least as good as the base model (tested)
## 9. Data Volume Matters More Than Model Complexity
| Dataset | Shapes | Rows | Mean Eff (log, 500 trees) |
| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------ | ---- | ----------------------------- |
| Original (DeepSeek only) | 108 | 418K | 98.28% (on seen distribution) |
| + Wide coverage M=1 distribution. Adding 60 diverse shapes (many M=1) exposed the model's weakness on tiny shapes. More diverse training data is always better than a bigger model on narrow data.Summary of DefaultsBased on these findings, the current defaults in `train.py` are:- **Target transform**: `log1p` for TFLOPS and bandwidth (scale normalization)- **Model**: 2000 trees, 255 leaves, max depth 15, LR 0.02- **Validation**: 5-fold GroupKFold, key = (M, N, K)- **Early stopping**: patience 100 (let trees fully converge)- **Warm start**: 500 new trees (was 200, increased for bigger base model) | 168 | 626K | 96.92% (harder distribution) |
The original 108-shape model looked great (98.28%) but was overfitting to the
DeepSeek LLM inference

271
dispatcher/heuristics/README.md Normal file
View File

@@ -0,0 +1,271 @@
# CK Tile Heuristics: ML-Based Kernel Selection
Fast, accurate kernel selection for CK Tile operations using LightGBM regression
with Origami-augmented feature engineering.
## What This Does
Instead of running all 4608+ kernel configurations on the GPU to find the best
one (exhaustive search taking ~46 seconds per shape), this system trains an ML
model that predicts TFLOPS for any (problem, kernel) pair in microseconds. It
scores all candidates instantly and picks the best kernel -- achieving 98.28%
of oracle-best TFLOPS efficiency across 108 tested shapes.
## Quick Start
### 1. Generate and convert benchmark data
**Step 1: Generate benchmark data**
```bash
python3 generate_benchmark_data.py \
--build_dir /path/to/build \
--output_dir data/fp16_original \
--dtype fp16 \
--layout rcr \
--num_build_jobs 4 \
--warmup 10 \
--repeat 50
```
This outputs JSON with all benchmark results.
**Step 2: Convert JSON to parquet training format**
```bash
python3 convert_json_to_parquet.py \
--input data/fp16_original/benchmark_results_fp16_rcr.json \
--output data/fp16_original/fp16_training_data.parquet \
--arch gfx950
```
The converter automatically fixes pad flags for `_mem` kernels and validates data.
**Alternative: Parse existing logs**
If you have raw benchmark logs from CK Tile:
```bash
python3 data_pipeline.py ck_tile_testrun_2.log \
-o data/gemm_universal_fp8_rcr_gfx950.parquet \
--arch gfx950 --capture_hw
```
### 2. Train a model
```bash
python3 train.py \
--data_dir data/ \
--out_dir models/gemm_universal_fp8_gfx950 \
--op gemm_universal --dtype fp8 --arch gfx950
```
**Note**: Trained models are automatically compressed to `.lgbm.gz` format to save space (~67% reduction). The Python tools automatically decompress them on first use and cache the decompressed version. For warm-start training, decompression happens automatically.
### 3. Evaluate
```bash
python3 evaluate.py \
--model_dir models/gemm_universal_fp8_gfx950 \
--data_dir data/ --op gemm_universal --dtype fp8
```
### 4. Predict the best kernel for a problem
```bash
python3 predict.py \
--model_dir models/gemm_universal_fp8_gfx950 \
--m 128 --n 1536 --k 7168 --layout rcr
```
### 5. Search for optimal configs (optional)
```bash
python3 search.py \
--model_dir models/gemm_universal_fp8_gfx950 \
--m 128 --n 1536 --k 7168 \
--strategy random --budget 500 --top_k 10
```
### 6. Using models in C++ (requires decompression)
C++ code uses the LightGBM C API which requires uncompressed `.lgbm` files. If you have compressed models (`.lgbm.gz`), decompress them first:
```bash
cd models/gemm_universal_fp16_gfx950
gunzip model_tflops.lgbm.gz
```
Then use in C++ examples:
```bash
cd dispatcher/build
./gemm_09_ml_heuristic --model ../heuristics/models/gemm_universal_fp16_gfx950/model_tflops.lgbm
```
**Note**: Python tools automatically decompress `.lgbm.gz` files on first use, so you can run Python scripts first to trigger decompression, then use the same models in C++.
## Architecture
```
Problem (M, N, K, dtype, layout)
|
v
FeatureEngine.extract_batch() <-- 55 features: problem, kernel, interaction, hardware
|
v
LGBMRegressor.predict() <-- predicts TFLOPS for each candidate kernel
|
v
Sort by predicted TFLOPS <-- rank all candidates
|
v
Select Top-1 kernel <-- 98.28% mean efficiency, <1ms inference
```
Three models are trained per (op, dtype, arch):
- **TFLOPS model** (primary): used for kernel ranking
- **Latency model** (auxiliary): for latency-sensitive workloads
- **Bandwidth model** (auxiliary): for memory-bound analysis
## File Inventory
| File | Purpose |
|---|---|
| `generate_benchmark_data.py` | Build and run benchmarks across ~25 diverse problem sizes, output JSON |
| `convert_json_to_parquet.py` | Convert benchmark JSON to parquet training format, fix `_mem` pad flags |
| `data_pipeline.py` | Parse raw benchmark logs into canonical parquet datasets |
| `feature_engine.py` | 55-feature extraction: problem, kernel, interaction, hardware profile |
| `train.py` | Multi-target LGBMRegressor training with GroupKFold CV, IHEM, warm-start |
| `predict.py` | Predictor class: predict TFLOPS/latency/bandwidth, rank kernels |
| `evaluate.py` | Full evaluation: global metrics, per-shape/layout/pipeline slices |
| `search.py` | Surrogate search: discrete DE, random top-K |
| `generate_wide_coverage.py` | Generate benchmark data across 706 diverse shapes |
| `generate_edge_dims.py` | Generate N=1, K=1, and other edge-case shapes |
| `DATA_GENERATION.md` | Detailed guide for building binaries and generating data |
| `plan.md` | Full design plan with architecture, milestones, and rationale |
## Features Used (55 total)
### Problem features (13)
`M, N, K, split_k, log2(M), log2(N), log2(K), log2(MNK),
arithmetic_intensity, aspect_ratio_mn, aspect_ratio_mk, aspect_ratio_nk, layout`
### Kernel features (17)
`tile_m, tile_n, tile_k, warp_m, warp_n, warp_k, warp_tile_m, warp_tile_n,
warp_tile_k, pipeline, scheduler, epilogue, pad_m, pad_n, pad_k, persistent,
num_warps, tile_volume, tile_mn, lds_usage_estimate, lds_usage_ratio`
### Interaction features (9)
`num_tiles_m, num_tiles_n, num_tiles_k, total_output_tiles,
tile_eff_m, tile_eff_n, tile_eff_k, overall_tile_efficiency, cu_utilization`
### Hardware profile features (12)
`hw_num_cus, hw_simds_per_cu, hw_total_simds, hw_shader_engines,
hw_max_clock_mhz, hw_max_waves_per_cu, hw_wavefront_size, hw_lds_capacity,
hw_l1_cache_kb, hw_l2_cache_kb, hw_l3_cache_kb, hw_num_xcd`
## Model Performance
### fp8 RCR, gfx950
| Metric | 108 shapes (original) | 168 shapes (wide coverage) |
|---|---|---|
| Mean TFLOPS Efficiency | 98.28% | 97.51% |
| P10 TFLOPS Efficiency | 94.64% | 93.89% |
| tiny_m (M=1) Efficiency | 95.57% | 96.04% |
| R2 (TFLOPS) | 0.997 | 0.993 |
### fp16 RCR, gfx950
Trained on 25 shapes, 1,024 kernels, 21,920 valid benchmarks.
| Metric | Value |
|---|---|
| Mean TFLOPS Efficiency | 99.36% |
| P10 TFLOPS Efficiency | 98.05% |
| P50 TFLOPS Efficiency | 100.00% |
| Min Efficiency | 95.45% |
| NDCG@1 | 64.00% |
| Top-5 Hit Rate | 88.00% |
**Shape Family Breakdown:**
| Shape Family | Mean Eff | P10 Eff | Shapes |
|---|---|---|---|
| Large M (M≥1024) | 99.54% | 99.07% | 4 |
| Medium M (128≤M<1024) | 99.62% | 98.74% | 7 |
| Small M (8≤M<128) | 98.82% | 96.22% | 8 |
| Tiny M (M<8) | 99.65% | 98.96% | 6 |
**Pipeline Breakdown:**
| Pipeline | Mean Eff | P10 Eff |
|---|---|---|
| compv3 | 99.75% | 99.09% |
| compv4 | 99.40% | 98.54% |
| mem | 99.08% | 96.59% |
Training uses `log1p(TFLOPS)` as the target by default, which normalizes the
scale across shapes spanning 0.02 to 2230 TFLOPS. This was the key finding
that improved tiny-M shapes from 84% to 96% efficiency. See
[LEARNINGS.md](LEARNINGS.md) for details.
## Validation
Training uses `GroupKFold(n_splits=5)` with group key `(M, N, K)` to ensure
the model is evaluated on shapes it has never seen during training. Layout is
excluded from the group key to force cross-layout generalization.
## Incremental Training (Warm Start)
When new benchmark data arrives, update the model without retraining from scratch:
```bash
python3 train.py \
--data_dir data/ \
--out_dir models/v2 \
--warm_start models/gemm_universal_fp8_gfx950 \
--warm_start_n_estimators 200
```
This adds 200 new trees on top of the existing model. Feature schemas must
match exactly (automatically enforced).
## Extending to New Ops
Adding support for a new operation (e.g., `gemm_streamk`, `grouped_conv`):
1. **Build binaries**: `ninja -C build benchmark_gemm_streamk_fp8_rcr`
2. **Subclass `FeatureEngine`**: add op-specific features (e.g., StreamK split factor)
3. **Generate data**: run benchmarks across diverse shapes
4. **Train**: `python3 train.py --op gemm_streamk --dtype fp8 --data_dir data/ --out_dir models/`
The training, evaluation, prediction, and search infrastructure is fully
op-agnostic -- only the feature engine needs a new subclass.
## Tests
102 tests covering all modules:
```bash
python3 -m pytest tests/ -v
```
Test coverage includes:
- Log parsing with malformed JSON, empty logs, single-kernel shapes
- Feature formula correctness (tile efficiency, LDS usage, arithmetic intensity)
- Corner-case shapes: M=1, N=1, K=1, prime dimensions, 20480x7168x256
- Batch vs single extraction parity
- Parameter space validation and projection
- Predictor: single/batch prediction, ranking, missing models, empty inputs
- Training: group keys, efficiency computation, warm-start, feature compatibility
- Search: random, DE, config validity, determinism
## Documentation
- **[README.md](README.md)**: This file -- quick start, architecture, performance
- **[DATA_GENERATION.md](DATA_GENERATION.md)**: Complete guide for building tile engine
binaries, running benchmarks, managing datasets, and troubleshooting
- **[LEARNINGS.md](LEARNINGS.md)**: Empirical findings and design decisions (log-transform,
IHEM results, tiny-M analysis, feature importance, N=1/K=1 edge cases)

4
dispatcher/heuristics/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
# CK Tile Heuristics: ML-based kernel selection

67
dispatcher/heuristics/collect_additional.sh Executable file
View File

@@ -0,0 +1,67 @@
#!/bin/bash
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
# Generate additional benchmark data for shapes NOT in the original log.
# Runs in background; outputs streaming JSON that can be parsed by data_pipeline.py.
BIN_DIR="/workspace/ck_tile/bin"
OUT_LOG="data/additional_shapes.log"
WARMUP=3
REPEAT=10
mkdir -p data
# Additional shapes: square powers-of-2 and common ML sizes not in original DeepSeek set
SHAPES=(
"64,64,64"
"128,128,128"
"256,256,256"
"512,512,512"
"1024,1024,1024"
"2048,2048,2048"
"4096,4096,4096"
"1,4096,4096"
"8,4096,4096"
"32,4096,4096"
"128,4096,4096"
"1,4096,11008"
"32,4096,11008"
"1,8192,8192"
"32,8192,8192"
"1,8192,28672"
"32,8192,28672"
"256,256,8192"
"8192,8192,256"
"1024,4096,1024"
"4096,1024,4096"
"2048,8192,2048"
)
echo "CK Tile Additional Shapes Benchmark" > "$OUT_LOG"
echo "GPU ID: 0" >> "$OUT_LOG"
echo "Implementation: gemm_universal" >> "$OUT_LOG"
echo "" >> "$OUT_LOG"
SHAPE_IDX=0
for SHAPE in "${SHAPES[@]}"; do
IFS=',' read -r M N K <<< "$SHAPE"
SHAPE_IDX=$((SHAPE_IDX + 1))
echo "========================================" >> "$OUT_LOG"
echo "Shape $SHAPE_IDX: M=$M N=$N K=$K dtype=fp8 layout=rcr" >> "$OUT_LOG"
echo "========================================" >> "$OUT_LOG"
KERNEL_COUNT=0
for EXE in "$BIN_DIR"/benchmark_gemm_universal_fp8_rcr_*; do
KERNEL_COUNT=$((KERNEL_COUNT + 1))
OUTPUT=$("$EXE" -m="$M" -n="$N" -k="$K" -warmup=$WARMUP -repeat=$REPEAT -verify=0 2>/dev/null)
# Extract just the JSON block
echo "$OUTPUT" | sed -n '/{/,/^}/p' >> "$OUT_LOG"
done
echo "Found $KERNEL_COUNT kernels" >> "$OUT_LOG"
echo "Completed shape $SHAPE_IDX: M=$M N=$N K=$K ($KERNEL_COUNT kernels)" >&2
done
echo "Done generating additional data" >&2

233
dispatcher/heuristics/convert_json_to_parquet.py Normal file
View File

@@ -0,0 +1,233 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Convert benchmark JSON results to parquet format for training.
Usage:
python convert_json_to_parquet.py \
--input benchmark_results_fp16_rcr.json \
--output fp16_training_data.parquet
Features:
- Converts JSON benchmark results to flat row format
- Automatically fixes pad flags for _mem kernels
- Captures both successes and failures
- Compatible with existing training data format
"""
import argparse
import json
import pandas as pd
from pathlib import Path
def convert_json_to_parquet(json_file: Path, output_file: Path, arch: str = "gfx950"):
"""Convert benchmark JSON to parquet training data format."""
print(f"Loading {json_file}...")
with open(json_file) as f:
data = json.load(f)
metadata = data.get("metadata", {})
dtype = metadata.get("dtype", "fp16")
layout = metadata.get("layout", "rcr")
print(f" Data type: {dtype}")
print(f" Layout: {layout}")
print(f" Kernels: {metadata.get('num_kernels', 0)}")
print(f" Problem sizes: {metadata.get('num_problems', 0)}")
print()
rows = []
for kernel_result in data["results"]:
kernel_config = kernel_result["kernel_config"]
for benchmark in kernel_result["benchmarks"]:
# Common fields for both valid and invalid runs
row = {
"op_type": "gemm_universal",
"dtype": dtype,
"layout": layout,
"arch": arch,
"kernel_name": kernel_config["name"],
"m": benchmark["m"],
"n": benchmark["n"],
"k": benchmark["k"],
"split_k": 1,
"is_valid": benchmark["is_valid"],
"run_id": 0,
"pipeline": kernel_config["pipeline"],
"epilogue": kernel_config["epilogue"],
"scheduler": kernel_config["scheduler"],
"pad_m": kernel_config["pad_m"],
"pad_n": kernel_config["pad_n"],
"pad_k": kernel_config["pad_k"],
"persistent": kernel_config["persistent"],
"tile_m": kernel_config["tile_m"],
"tile_n": kernel_config["tile_n"],
"tile_k": kernel_config["tile_k"],
"warp_m": kernel_config["warp_m"],
"warp_n": kernel_config["warp_n"],
"warp_k": kernel_config["warp_k"],
"warp_tile_m": kernel_config["warp_tile_m"],
"warp_tile_n": kernel_config["warp_tile_n"],
"warp_tile_k": kernel_config["warp_tile_k"],
}
if benchmark["is_valid"]:
# Valid run - include performance metrics
row["measured_tflops"] = benchmark["tflops"]
row["latency_ms"] = benchmark["avg_time_ms"]
# Calculate bandwidth if needed
m, n, k = benchmark["m"], benchmark["n"], benchmark["k"]
bytes_transferred = (m * k + k * n + m * n) * 2 # FP16 = 2 bytes
if benchmark["avg_time_ms"] > 0:
row["bandwidth_gb_s"] = (bytes_transferred / 1e9) / (
benchmark["avg_time_ms"] / 1000
)
else:
row["bandwidth_gb_s"] = 0.0
else:
# Failed run - zero metrics
row["measured_tflops"] = 0.0
row["latency_ms"] = 0.0
row["bandwidth_gb_s"] = 0.0
rows.append(row)
df = pd.DataFrame(rows)
print(f"Converted {len(df):,} benchmark results")
print(f" Valid: {df['is_valid'].sum():,}")
print(f" Failed: {(~df['is_valid']).sum():,}")
print()
# Fix pad flags for _mem kernels (critical for P1 features!)
print("Fixing pad flags for _mem kernels...")
mem_mask = df["pipeline"] == "mem"
mem_count = mem_mask.sum()
if mem_count > 0:
df.loc[mem_mask, "pad_m"] = True
df.loc[mem_mask, "pad_n"] = True
df.loc[mem_mask, "pad_k"] = True
print(f" ✓ Fixed {mem_count:,} _mem kernel rows")
print()
# Save to parquet
df.to_parquet(output_file, index=False)
print(f"✓ Saved to {output_file}")
print()
# Show statistics
print("=" * 80)
print("STATISTICS")
print("=" * 80)
print()
print("Dimension ranges:")
print(f" M: {df['m'].min():,} - {df['m'].max():,}")
print(f" N: {df['n'].min():,} - {df['n'].max():,}")
print(f" K: {df['k'].min():,} - {df['k'].max():,}")
print()
print("Pipeline distribution:")
print(df["pipeline"].value_counts())
print()
print("Pad flag distribution:")
pad_combos = df[["pad_m", "pad_n", "pad_k"]].value_counts()
print(pad_combos)
print()
if (~df["is_valid"]).sum() > 0:
print("Failure analysis:")
failed = df[~df["is_valid"]]
print(f" Total failures: {len(failed):,}")
# Group by pipeline
print("\n By pipeline:")
for pipeline, count in failed["pipeline"].value_counts().items():
print(f" {pipeline}: {count:,}")
# Show sample failures
print("\n Sample failures:")
for _, row in failed.head(5).iterrows():
print(
f" {row['kernel_name'][:60]:60s} M={row['m']:4d} N={row['n']:4d} K={row['k']:4d}"
)
return df
def merge_datasets(parquet_files: list[Path], output_file: Path):
"""Merge multiple parquet files into one."""
print("=" * 80)
print("MERGING DATASETS")
print("=" * 80)
print()
dfs = []
for pq_file in parquet_files:
if pq_file.exists():
df = pd.read_parquet(pq_file)
print(f" {pq_file.name}: {len(df):,} rows")
dfs.append(df)
else:
print(f" ✗ {pq_file} not found, skipping")
if not dfs:
print("No files to merge!")
return
combined = pd.concat(dfs, ignore_index=True)
combined.to_parquet(output_file, index=False)
print()
print(f"✓ Merged {len(combined):,} total rows to {output_file}")
print()
# Show dtype distribution
print("Data type distribution:")
print(combined["dtype"].value_counts())
print()
print("Layout distribution:")
print(combined["layout"].value_counts())
def main():
parser = argparse.ArgumentParser(
description="Convert benchmark JSON to parquet training data",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__,
)
parser.add_argument(
"--input", type=str, required=True, help="Input JSON file from benchmark"
)
parser.add_argument("--output", type=str, required=True, help="Output parquet file")
parser.add_argument("--arch", type=str, default="gfx950", help="GPU architecture")
parser.add_argument(
"--merge_with", type=str, nargs="*", help="Additional parquet files to merge"
)
args = parser.parse_args()
input_file = Path(args.input)
output_file = Path(args.output)
# Convert JSON to parquet
df = convert_json_to_parquet(input_file, output_file, args.arch)
# Merge if requested
if args.merge_with:
merge_files = [output_file] + [Path(f) for f in args.merge_with]
merged_output = output_file.parent / f"{output_file.stem}_merged.parquet"
merge_datasets(merge_files, merged_output)
if __name__ == "__main__":
main()

394
dispatcher/heuristics/data_pipeline.py Normal file
View File

@@ -0,0 +1,394 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Data pipeline for CK Tile heuristics.
Parses benchmark logs and structured JSON into a canonical parquet dataset.
Supports:
- Streaming log format (Shape N: headers + inline JSON) from ck_tile profiling runs
- Structured JSON from generate_benchmark_data.py
- Direct parquet passthrough
"""
import json
import re
import subprocess
import hashlib
from pathlib import Path
from typing import Optional
import pandas as pd
CANONICAL_COLUMNS = [
"op_type",
"dtype",
"layout",
"arch",
"kernel_name",
"m",
"n",
"k",
"split_k",
"measured_tflops",
"latency_ms",
"bandwidth_gb_s",
"is_valid",
"tile_m",
"tile_n",
"tile_k",
"warp_m",
"warp_n",
"warp_k",
"warp_tile_m",
"warp_tile_n",
"warp_tile_k",
"pipeline",
"scheduler",
"epilogue",
"pad_m",
"pad_n",
"pad_k",
"persistent",
"run_id",
]
def parse_kernel_name(name: str) -> dict:
"""Extract kernel config fields from a gemm_universal kernel name.
Name format:
gemm_universal_{dtype}_{layout}_{pipeline}_{epilogue}_{scheduler}
_{padM}_{padN}_{padK}_{persistent}_{tileM}x{tileN}x{tileK}
_{warpM}x{warpN}x{warpK}_{warpTileM}x{warpTileN}x{warpTileK}
"""
result = {}
try:
prefix_match = re.match(
r"gemm_universal_(\w+?)_((?:rcr|rrr|crr|ccr))_(.*)", name
)
if not prefix_match:
return result
result["dtype"] = prefix_match.group(1)
result["layout"] = prefix_match.group(2)
remainder = prefix_match.group(3)
parts = remainder.split("_")
if len(parts) < 10:
return result
result["pipeline"] = parts[0]
result["epilogue"] = parts[1]
result["scheduler"] = parts[2]
result["pad_m"] = parts[3] == "True"
result["pad_n"] = parts[4] == "True"
result["pad_k"] = parts[5] == "True"
result["persistent"] = parts[6] == "True"
tile_dims = parts[7].split("x")
warp_dims = parts[8].split("x")
warp_tile_dims = parts[9].split("x")
result["tile_m"] = int(tile_dims[0])
result["tile_n"] = int(tile_dims[1])
result["tile_k"] = int(tile_dims[2])
result["warp_m"] = int(warp_dims[0])
result["warp_n"] = int(warp_dims[1])
result["warp_k"] = int(warp_dims[2])
result["warp_tile_m"] = int(warp_tile_dims[0])
result["warp_tile_n"] = int(warp_tile_dims[1])
result["warp_tile_k"] = int(warp_tile_dims[2])
except (IndexError, ValueError):
pass
return result
def _layout_from_problem(problem: dict) -> str:
"""Derive layout shorthand (rcr/rrr/etc.) from problem JSON fields."""
la = problem.get("layout_a", "")
lb = problem.get("layout_b", "")
lc = problem.get("layout_c", "")
def _tag(s):
s = s.lower()
if "row" in s:
return "r"
if "col" in s:
return "c"
return "?"
return _tag(la) + _tag(lb) + _tag(lc)
def parse_streaming_log(
path: str | Path,
arch: str = "unknown",
run_id: Optional[str] = None,
op_type: str = "gemm_universal",
) -> pd.DataFrame:
"""Parse a CK Tile streaming benchmark log into a canonical DataFrame.
The log alternates between shape headers and JSON result blocks:
Shape N: M=16 N=1536 K=7168 dtype=fp8 layout=rcr
{
"name": "gemm_universal_...",
"problem": { ... },
"perf_result": { "latency(ms)": ..., "tflops(TFlops)": ..., "bandwidth(GB/s)": ... }
}
"""
path = Path(path)
if run_id is None:
run_id = hashlib.md5(path.name.encode()).hexdigest()[:12]
shape_re = re.compile(
r"Shape\s+\d+:\s+M=(\d+)\s+N=(\d+)\s+K=(\d+)\s+dtype=(\w+)\s+layout=(\w+)"
)
rows = []
current_m, current_n, current_k = 0, 0, 0
current_dtype, current_layout = "", ""
json_buf = []
brace_depth = 0
with open(path, "r") as f:
for line in f:
stripped = line.strip()
shape_match = shape_re.search(stripped)
if shape_match:
current_m = int(shape_match.group(1))
current_n = int(shape_match.group(2))
current_k = int(shape_match.group(3))
current_dtype = shape_match.group(4)
current_layout = shape_match.group(5)
continue
if brace_depth == 0 and stripped.startswith("{"):
json_buf = [stripped]
brace_depth = stripped.count("{") - stripped.count("}")
if brace_depth == 0:
raw = "\n".join(json_buf)
try:
obj = json.loads(raw)
except json.JSONDecodeError:
continue
else:
continue
elif brace_depth > 0:
json_buf.append(stripped)
brace_depth += stripped.count("{") - stripped.count("}")
if brace_depth <= 0:
brace_depth = 0
raw = "\n".join(json_buf)
try:
obj = json.loads(raw)
except json.JSONDecodeError:
continue
else:
continue
else:
continue
# If we get here, obj was successfully parsed
kernel_name = obj.get("name", "")
problem = obj.get("problem", {})
perf = obj.get("perf_result", {})
m = problem.get("m", current_m)
n = problem.get("n", current_n)
k = problem.get("k", current_k)
split_k = problem.get("split_k", 1)
dtype = problem.get("dtype_a", current_dtype)
layout = (
_layout_from_problem(problem)
if problem.get("layout_a")
else current_layout
)
tflops = perf.get("tflops(TFlops)", 0.0)
latency = perf.get("latency(ms)", 0.0)
bandwidth = perf.get("bandwidth(GB/s)", 0.0)
kp = parse_kernel_name(kernel_name)
row = {
"op_type": op_type,
"dtype": dtype,
"layout": layout,
"arch": arch,
"kernel_name": kernel_name,
"m": m,
"n": n,
"k": k,
"split_k": split_k,
"measured_tflops": tflops,
"latency_ms": latency,
"bandwidth_gb_s": bandwidth,
"is_valid": tflops > 0 and latency > 0,
"run_id": run_id,
}
row.update(kp)
rows.append(row)
df = pd.DataFrame(rows)
for col in CANONICAL_COLUMNS:
if col not in df.columns:
df[col] = None
return df
def get_hardware_profile() -> dict:
"""Capture GPU hardware profile from rocminfo."""
profile = {}
try:
result = subprocess.run(
["rocminfo"], capture_output=True, text=True, timeout=30
)
output = result.stdout
gpu_section = False
for line in output.split("\n"):
line = line.strip()
if "Device Type:" in line and "GPU" in line:
gpu_section = True
continue
if gpu_section and "Device Type:" in line and "GPU" not in line:
break
if not gpu_section:
continue
if ":" not in line:
continue
key, _, val = line.partition(":")
key = key.strip()
val = val.strip()
mapping = {
"Name": "gfx_name",
"Marketing Name": "marketing_name",
"Compute Unit": "num_cus",
"SIMDs per CU": "simds_per_cu",
"Shader Engines": "shader_engines",
"Shader Arrs. per Eng.": "shader_arrays_per_engine",
"Max Clock Freq. (MHz)": "max_clock_mhz",
"Wavefront Size": "wavefront_size",
"Max Waves Per CU": "max_waves_per_cu",
"Chip ID": "chip_id",
}
if key in mapping:
raw = val.split("(")[0].strip()
try:
profile[mapping[key]] = int(raw)
except ValueError:
profile[mapping[key]] = raw
for line in output.split("\n"):
line = line.strip()
if line.startswith("L1:") and "num_cus" in profile:
raw = line.split(":")[1].strip().split("(")[0].strip()
try:
profile["l1_cache_kb"] = int(raw)
except ValueError:
pass
elif line.startswith("L2:"):
raw = line.split(":")[1].strip().split("(")[0].strip()
try:
profile["l2_cache_kb"] = int(raw)
except ValueError:
pass
elif line.startswith("L3:"):
raw = line.split(":")[1].strip().split("(")[0].strip()
try:
profile["l3_cache_kb"] = int(raw)
except ValueError:
pass
except (subprocess.TimeoutExpired, FileNotFoundError):
pass
return profile
def load_parquet(path: str | Path) -> pd.DataFrame:
"""Load a canonical parquet dataset."""
return pd.read_parquet(path)
def save_parquet(df: pd.DataFrame, path: str | Path):
"""Save a DataFrame in canonical parquet format."""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
df.to_parquet(path, index=False, engine="pyarrow")
def build_training_dataset(
data_dir: str | Path,
op_type: str = "gemm_universal",
dtype: str = "fp8",
) -> pd.DataFrame:
"""Load and merge all parquet files matching the given op/dtype from a directory."""
data_dir = Path(data_dir)
frames = []
for f in sorted(data_dir.glob("*.parquet")):
df = pd.read_parquet(f)
if "op_type" in df.columns:
df = df[df["op_type"] == op_type]
if "dtype" in df.columns:
df = df[df["dtype"] == dtype]
if len(df) > 0:
frames.append(df)
if not frames:
raise FileNotFoundError(
f"No parquet files with op_type={op_type}, dtype={dtype} in {data_dir}"
)
return pd.concat(frames, ignore_index=True)
if __name__ == "__main__":
import argparse
import time
parser = argparse.ArgumentParser(description="Parse CK Tile benchmark data")
parser.add_argument("input", help="Input file (log or parquet)")
parser.add_argument("--output", "-o", required=True, help="Output parquet path")
parser.add_argument("--arch", default="gfx950", help="GPU architecture")
parser.add_argument("--op_type", default="gemm_universal", help="Operation type")
parser.add_argument(
"--capture_hw",
action="store_true",
help="Capture hardware profile from rocminfo",
)
args = parser.parse_args()
input_path = Path(args.input)
print(f"Parsing {input_path}...")
t0 = time.time()
if input_path.suffix == ".parquet":
df = load_parquet(input_path)
else:
df = parse_streaming_log(input_path, arch=args.arch, op_type=args.op_type)
elapsed = time.time() - t0
print(f"Parsed {len(df)} rows in {elapsed:.1f}s")
print(f" Unique shapes: {df.groupby(['m', 'n', 'k']).ngroups}")
print(f" Unique kernels: {df['kernel_name'].nunique()}")
print(f" Valid rows: {df['is_valid'].sum()} / {len(df)}")
if df["measured_tflops"].max() > 0:
print(
f" TFLOPS range: {df['measured_tflops'].min():.2f} - {df['measured_tflops'].max():.2f}"
)
if args.capture_hw:
hw = get_hardware_profile()
print(f" Hardware profile: {hw}")
for k, v in hw.items():
df[f"hw_{k}"] = v
save_parquet(df, args.output)
print(f"Saved to {args.output}")

324
dispatcher/heuristics/dispatcher_integration.py Normal file
View File

@@ -0,0 +1,324 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Dispatcher integration for ML-based kernel selection.
Bridges the trained LightGBM Predictor with the CK Tile dispatcher's
kernel selection flow. Provides heuristic functions compatible with
both the Python pre-selection pattern (08_heuristics.py style) and
the C++ HeuristicFunction signature.
Name mapping between feature engine and dispatcher KernelConfig:
Feature engine Dispatcher KernelConfig
--------------------- ----------------------
warp_m (warps/block) wave_m
warp_n wave_n
warp_k wave_k
warp_tile_m warp_m
warp_tile_n warp_n
warp_tile_k warp_k
Usage:
from dispatcher_integration import create_ml_heuristic
heuristic = create_ml_heuristic("models/gemm_universal_fp8_gfx950")
best_spec = heuristic(M=1024, N=1024, K=1024, kernel_pool=KERNEL_POOL)
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from data_pipeline import parse_kernel_name
from predict import Predictor
LAYOUT_TO_DISPATCHER = {
"rcr": ("row", "col", "row"),
"rrr": ("row", "row", "row"),
"crr": ("col", "row", "row"),
"ccr": ("col", "col", "row"),
}
DTYPE_TO_C_DTYPE = {
"fp8": "fp16",
"fp16": "fp16",
"bf16": "bf16",
"fp32": "fp32",
}
@dataclass
class MLKernelSpec:
"""Kernel spec returned by the ML heuristic, compatible with the dispatcher
example pattern. Carries both the feature-engine-space config and the
dispatcher-space KernelConfig fields."""
kernel_name: str
predicted_tflops: float
tile_m: int
tile_n: int
tile_k: int
wave_m: int
wave_n: int
wave_k: int
warp_m: int
warp_n: int
warp_k: int
pipeline: str
scheduler: str
epilogue: str
pad_m: bool
pad_n: bool
pad_k: bool
persistent: bool
def kernel_config_to_feature_dict(kernel_name: str) -> dict:
"""Parse a tile-engine kernel name into a feature-engine-compatible dict.
Returns a dict with fields matching what GemmUniversalFeatureEngine.extract()
expects for the kernel parameter: tile_m/n/k, warp_m/n/k (warps per block),
warp_tile_m/n/k, pipeline, scheduler, epilogue, pad_m/n/k, persistent.
"""
parsed = parse_kernel_name(kernel_name)
if not parsed:
return {}
parsed["kernel_name"] = kernel_name
return parsed
def feature_dict_to_dispatcher_config(
feat: dict, dtype: str = "fp8", arch: str = "gfx950"
) -> dict:
"""Convert a feature-engine kernel dict to dispatcher KernelConfig fields.
Handles the naming inversion:
feature engine warp_m -> KernelConfig wave_m (warps per block)
feature engine warp_tile_m -> KernelConfig warp_m (elements per warp)
"""
layout = feat.get("layout", "rcr")
la, lb, lc = LAYOUT_TO_DISPATCHER.get(layout, ("row", "col", "row"))
c_dtype = DTYPE_TO_C_DTYPE.get(dtype, dtype)
return {
"dtype_a": dtype,
"dtype_b": dtype,
"dtype_c": c_dtype,
"dtype_acc": "fp32",
"layout_a": la,
"layout_b": lb,
"layout_c": lc,
"tile_m": feat.get("tile_m", 128),
"tile_n": feat.get("tile_n", 128),
"tile_k": feat.get("tile_k", 64),
"wave_m": feat.get("warp_m", 2),
"wave_n": feat.get("warp_n", 2),
"wave_k": feat.get("warp_k", 1),
"warp_m": feat.get("warp_tile_m", 32),
"warp_n": feat.get("warp_tile_n", 32),
"warp_k": feat.get("warp_tile_k", 16),
"pipeline": feat.get("pipeline", "compv3"),
"scheduler": feat.get("scheduler", "intrawave"),
"epilogue": feat.get("epilogue", "cshuffle"),
"pad_m": feat.get("pad_m", True),
"pad_n": feat.get("pad_n", True),
"pad_k": feat.get("pad_k", True),
"gfx_arch": arch,
}
def feature_dict_to_ml_spec(feat: dict, predicted_tflops: float = 0.0) -> MLKernelSpec:
"""Convert a feature-engine kernel dict + prediction to an MLKernelSpec."""
return MLKernelSpec(
kernel_name=feat.get("kernel_name", "unknown"),
predicted_tflops=predicted_tflops,
tile_m=feat.get("tile_m", 128),
tile_n=feat.get("tile_n", 128),
tile_k=feat.get("tile_k", 64),
wave_m=feat.get("warp_m", 2),
wave_n=feat.get("warp_n", 2),
wave_k=feat.get("warp_k", 1),
warp_m=feat.get("warp_tile_m", 32),
warp_n=feat.get("warp_tile_n", 32),
warp_k=feat.get("warp_tile_k", 16),
pipeline=feat.get("pipeline", "compv3"),
scheduler=feat.get("scheduler", "intrawave"),
epilogue=feat.get("epilogue", "cshuffle"),
pad_m=feat.get("pad_m", False),
pad_n=feat.get("pad_n", False),
pad_k=feat.get("pad_k", False),
persistent=feat.get("persistent", False),
)
def load_kernel_pool_from_binaries(bin_dir: str | Path) -> list[dict]:
"""Discover benchmark executables and parse their names into feature dicts.
Each executable name encodes the full kernel config. This creates the
candidate pool for the ML heuristic without needing a registry JSON export.
"""
bin_dir = Path(bin_dir)
configs = []
for exe in sorted(bin_dir.glob("benchmark_gemm_universal_*")):
name = exe.stem.replace("benchmark_", "")
feat = kernel_config_to_feature_dict(name)
if feat and "tile_m" in feat:
configs.append(feat)
return configs
def create_ml_heuristic(
model_dir: str | Path,
dtype: str = "fp8",
arch: str = "gfx950",
layout: str = "rcr",
kernel_pool: Optional[list[dict]] = None,
bin_dir: Optional[str | Path] = None,
):
"""Create an ML heuristic function for kernel selection.
Returns a callable with signature:
(M: int, N: int, K: int) -> MLKernelSpec
The returned function scores all candidate kernels using the trained
LightGBM regressor and returns the best one as an MLKernelSpec.
Parameters
----------
model_dir : str or Path
Path to trained model directory (must contain model_tflops.lgbm or
model_tflops_log_big.lgbm and feature_spec.json).
dtype : str
Data type for the problem (fp8, fp16, bf16).
arch : str
GPU architecture (gfx942, gfx950).
layout : str
Matrix layout (rcr, rrr, crr, ccr).
kernel_pool : list of dict, optional
Pre-parsed kernel configs. If None, loads from bin_dir.
bin_dir : str or Path, optional
Directory with benchmark executables. Used to build kernel_pool if
kernel_pool is not provided. Defaults to /workspace/ck_tile/bin.
"""
model_dir = Path(model_dir)
predictor = Predictor(model_dir)
if kernel_pool is None:
if bin_dir is None:
bin_dir = Path("/workspace/ck_tile/bin")
kernel_pool = load_kernel_pool_from_binaries(bin_dir)
if not kernel_pool:
raise ValueError(
"No kernel configs found. Check bin_dir or provide kernel_pool."
)
def heuristic(M: int, N: int, K: int) -> MLKernelSpec:
problem = {
"m": M,
"n": N,
"k": K,
"dtype": dtype,
"layout": layout,
"split_k": 1,
}
ranked = predictor.rank_kernels(problem, kernel_pool)
if not ranked:
feat = kernel_pool[0]
return feature_dict_to_ml_spec(feat, 0.0)
best_name, best_tflops = ranked[0]
best_feat = next(
(kp for kp in kernel_pool if kp.get("kernel_name") == best_name),
kernel_pool[0],
)
return feature_dict_to_ml_spec(best_feat, best_tflops)
return heuristic
def create_ranked_heuristic(
model_dir: str | Path,
dtype: str = "fp8",
arch: str = "gfx950",
layout: str = "rcr",
kernel_pool: Optional[list[dict]] = None,
bin_dir: Optional[str | Path] = None,
top_k: int = 5,
):
"""Create an ML heuristic that returns the top-K ranked kernel specs.
Returns a callable with signature:
(M: int, N: int, K: int) -> list[MLKernelSpec]
Useful when you want fallback options if the top-1 kernel fails to build.
"""
model_dir = Path(model_dir)
predictor = Predictor(model_dir)
if kernel_pool is None:
if bin_dir is None:
bin_dir = Path("/workspace/ck_tile/bin")
kernel_pool = load_kernel_pool_from_binaries(bin_dir)
name_to_feat = {kp.get("kernel_name", ""): kp for kp in kernel_pool}
def heuristic(M: int, N: int, K: int) -> list[MLKernelSpec]:
problem = {
"m": M,
"n": N,
"k": K,
"dtype": dtype,
"layout": layout,
"split_k": 1,
}
ranked = predictor.rank_kernels(problem, kernel_pool)
results = []
for name, tflops in ranked[:top_k]:
feat = name_to_feat.get(name, kernel_pool[0])
results.append(feature_dict_to_ml_spec(feat, tflops))
return results
return heuristic
def ml_spec_to_dispatcher_config(
spec: MLKernelSpec, dtype: str = "fp8", arch: str = "gfx950"
) -> dict:
"""Convert an MLKernelSpec to a dict compatible with ctypes_utils.KernelConfig."""
layout_a, layout_b, layout_c = "row", "col", "row"
c_dtype = DTYPE_TO_C_DTYPE.get(dtype, dtype)
return {
"dtype_a": dtype,
"dtype_b": dtype,
"dtype_c": c_dtype,
"dtype_acc": "fp32",
"layout_a": layout_a,
"layout_b": layout_b,
"layout_c": layout_c,
"tile_m": spec.tile_m,
"tile_n": spec.tile_n,
"tile_k": spec.tile_k,
"wave_m": spec.wave_m,
"wave_n": spec.wave_n,
"wave_k": spec.wave_k,
"warp_m": spec.warp_m,
"warp_n": spec.warp_n,
"warp_k": spec.warp_k,
"pipeline": spec.pipeline,
"scheduler": spec.scheduler,
"epilogue": spec.epilogue,
"pad_m": spec.pad_m,
"pad_n": spec.pad_n,
"pad_k": spec.pad_k,
"gfx_arch": arch,
}

254
dispatcher/heuristics/evaluate.py Normal file
View File

@@ -0,0 +1,254 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Evaluation and reporting for CK Tile kernel performance models.
Computes:
- Global metrics: TFLOPS efficiency (mean, p10, p50, min), R2, NDCG@1, Top-K hit rate
- Per-slice breakdowns: by layout, shape family, K-depth regime, pipeline
- Cross-target consistency checks
- Feature importance analysis
Usage:
python evaluate.py --model_dir models/gemm_universal_fp8_gfx950 --data_dir data/
"""
import argparse
import json
import numpy as np
import pandas as pd
from data_pipeline import build_training_dataset
from feature_engine import GemmUniversalFeatureEngine
from predict import Predictor
from train import compute_tflops_efficiency
def classify_shape_family(m: int, n: int, k: int) -> str:
"""Classify a GEMM shape into a family for sliced evaluation.
Families:
- tiny_m: M < 32 (single-token / very small batch inference)
- small_m: 32 <= M < 256
- medium_m: 256 <= M < 4096
- large_m: M >= 4096
- square: 0.5 <= M/N <= 2.0 and 0.5 <= M/K <= 2.0
- tall: M/N > 2.0
- wide: M/N < 0.5
"""
if m < 32:
return "tiny_m"
elif m < 256:
return "small_m"
elif m < 4096:
return "medium_m"
elif m >= 4096:
return "large_m"
return "other"
def classify_k_regime(k: int) -> str:
"""Classify K dimension into depth regime."""
if k < 512:
return "shallow_k"
elif k < 4096:
return "medium_k"
else:
return "deep_k"
def evaluate_model(
predictor: Predictor,
df: pd.DataFrame,
feature_engine: GemmUniversalFeatureEngine,
) -> dict:
"""Run full evaluation on a dataset. Returns a metrics dictionary.
Parameters
----------
predictor : Predictor
Trained predictor with at least a TFLOPS model loaded.
df : pd.DataFrame
Benchmark data in canonical schema.
feature_engine : GemmUniversalFeatureEngine
Feature engine matching the trained model.
Returns
-------
dict with keys: global_metrics, shape_family_metrics, k_regime_metrics,
pipeline_metrics, per_shape_efficiency.
"""
valid = df[df["is_valid"].fillna(False) & (df["measured_tflops"] > 0)].copy()
valid = valid.reset_index(drop=True)
X = feature_engine.extract_batch(valid)
model = predictor._load_model("tflops")
if model is None:
raise FileNotFoundError("No TFLOPS model found")
# Predict and apply inverse log transform if model was trained in log-space
raw_pred = model.predict(X)
if "tflops" in predictor._log_targets:
valid["pred_tflops"] = np.expm1(raw_pred)
else:
# Clamp to non-negative even for non-log models
valid["pred_tflops"] = np.maximum(0.0, raw_pred)
y_true = valid["measured_tflops"].values
y_pred = valid["pred_tflops"].values
ss_res = np.sum((y_true - y_pred) ** 2)
ss_tot = np.sum((y_true - y_true.mean()) ** 2)
r2 = 1 - ss_res / max(ss_tot, 1e-10)
rmse = np.sqrt(np.mean((y_true - y_pred) ** 2))
mae = np.mean(np.abs(y_true - y_pred))
eff_df = compute_tflops_efficiency(valid, "pred_tflops")
ndcg1_count = 0
total_shapes = 0
topk_hits = {3: 0, 5: 0, 10: 0}
for (m, n, k), group in valid.groupby(["m", "n", "k"]):
if group["measured_tflops"].max() <= 0:
continue
total_shapes += 1
oracle_idx = group["measured_tflops"].idxmax()
pred_ranking = group.sort_values("pred_tflops", ascending=False).index.tolist()
if pred_ranking[0] == oracle_idx:
ndcg1_count += 1
oracle_rank = pred_ranking.index(oracle_idx)
for topk in topk_hits:
if oracle_rank < topk:
topk_hits[topk] += 1
global_metrics = {
"r2": r2,
"rmse": rmse,
"mae": mae,
"num_valid_rows": len(valid),
"num_shapes": total_shapes,
"efficiency_mean": float(eff_df["efficiency"].mean()) if len(eff_df) > 0 else 0,
"efficiency_p10": float(eff_df["efficiency"].quantile(0.1))
if len(eff_df) > 0
else 0,
"efficiency_p50": float(eff_df["efficiency"].quantile(0.5))
if len(eff_df) > 0
else 0,
"efficiency_min": float(eff_df["efficiency"].min()) if len(eff_df) > 0 else 0,
"ndcg_at_1": ndcg1_count / max(total_shapes, 1),
"top3_hit_rate": topk_hits[3] / max(total_shapes, 1),
"top5_hit_rate": topk_hits[5] / max(total_shapes, 1),
"top10_hit_rate": topk_hits[10] / max(total_shapes, 1),
}
def _slice_efficiency(slice_df):
if len(slice_df) == 0:
return {"count": 0}
eff = compute_tflops_efficiency(slice_df, "pred_tflops")
if len(eff) == 0:
return {"count": 0}
return {
"count": len(eff),
"mean": float(eff["efficiency"].mean()),
"p10": float(eff["efficiency"].quantile(0.1)),
"min": float(eff["efficiency"].min()),
}
valid["shape_family"] = valid.apply(
lambda r: classify_shape_family(r["m"], r["n"], r["k"]), axis=1
)
valid["k_regime"] = valid["k"].apply(classify_k_regime)
shape_family_metrics = {}
for family, group in valid.groupby("shape_family"):
shape_family_metrics[family] = _slice_efficiency(group)
k_regime_metrics = {}
for regime, group in valid.groupby("k_regime"):
k_regime_metrics[regime] = _slice_efficiency(group)
pipeline_metrics = {}
if "pipeline" in valid.columns:
for pipeline, group in valid.groupby("pipeline"):
pipeline_metrics[str(pipeline)] = _slice_efficiency(group)
return {
"global_metrics": global_metrics,
"shape_family_metrics": shape_family_metrics,
"k_regime_metrics": k_regime_metrics,
"pipeline_metrics": pipeline_metrics,
"per_shape_efficiency": eff_df.to_dict(orient="records")
if len(eff_df) > 0
else [],
}
def main():
parser = argparse.ArgumentParser(description="Evaluate CK Tile performance model")
parser.add_argument(
"--model_dir", required=True, help="Directory with trained models"
)
parser.add_argument("--data_dir", required=True, help="Directory with parquet data")
parser.add_argument("--op", default="gemm_universal")
parser.add_argument("--dtype", default="fp8")
parser.add_argument("--output", "-o", help="Output JSON path for metrics")
args = parser.parse_args()
print(f"Loading data from {args.data_dir}...")
df = build_training_dataset(args.data_dir, op_type=args.op, dtype=args.dtype)
print(f" {len(df)} rows, {df.groupby(['m', 'n', 'k']).ngroups} shapes")
fe = GemmUniversalFeatureEngine()
predictor = Predictor(args.model_dir, feature_engine=fe)
print("Evaluating...")
results = evaluate_model(predictor, df, fe)
gm = results["global_metrics"]
print("\nGlobal Metrics:")
print(f" R2: {gm['r2']:.4f}")
print(f" RMSE: {gm['rmse']:.2f}")
print(f" Efficiency Mean: {gm['efficiency_mean']:.4f}")
print(f" Efficiency P10: {gm['efficiency_p10']:.4f}")
print(f" Efficiency P50: {gm['efficiency_p50']:.4f}")
print(f" Efficiency Min: {gm['efficiency_min']:.4f}")
print(f" NDCG@1: {gm['ndcg_at_1']:.4f}")
print(f" Top-3 Hit Rate: {gm['top3_hit_rate']:.4f}")
print(f" Top-5 Hit Rate: {gm['top5_hit_rate']:.4f}")
print(f" Top-10 Hit Rate: {gm['top10_hit_rate']:.4f}")
print("\nShape Family Breakdown:")
for family, metrics in sorted(results["shape_family_metrics"].items()):
if metrics.get("count", 0) > 0:
print(
f" {family:12s}: mean={metrics['mean']:.4f} p10={metrics['p10']:.4f} min={metrics['min']:.4f} (n={metrics['count']})"
)
print("\nK-Depth Regime Breakdown:")
for regime, metrics in sorted(results["k_regime_metrics"].items()):
if metrics.get("count", 0) > 0:
print(
f" {regime:12s}: mean={metrics['mean']:.4f} p10={metrics['p10']:.4f} min={metrics['min']:.4f} (n={metrics['count']})"
)
print("\nPipeline Breakdown:")
for pipeline, metrics in sorted(results["pipeline_metrics"].items()):
if metrics.get("count", 0) > 0:
print(
f" {pipeline:15s}: mean={metrics['mean']:.4f} p10={metrics['p10']:.4f} (n={metrics['count']})"
)
if args.output:
with open(args.output, "w") as f:
json.dump(results, f, indent=2, default=str)
print(f"\nFull results saved to {args.output}")
if __name__ == "__main__":
main()

577
dispatcher/heuristics/feature_engine.py Normal file
View File

@@ -0,0 +1,577 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Feature engineering for CK Tile kernel performance prediction.
Provides a strict FeatureEngine interface with per-op subclasses.
All feature engines produce a consistent numpy array for LightGBM.
"""
import math
from abc import ABC, abstractmethod
import numpy as np
import pandas as pd
DTYPE_BYTES = {
"fp32": 4.0,
"fp16": 2.0,
"bf16": 2.0,
"fp8": 1.0,
"bf8": 1.0,
"int8": 1.0,
"int4": 0.5,
}
LAYOUT_MAP = {"rcr": 0, "rrr": 1, "crr": 2, "ccr": 3}
PIPELINE_MAP = {"compv3": 0, "compv4": 1, "compv5": 2, "mem": 3, "preshufflev2": 4}
SCHEDULER_MAP = {"intrawave": 0, "interwave": 1}
EPILOGUE_MAP = {"default": 0, "cshuffle": 1}
class FeatureEngine(ABC):
"""Abstract base for per-op feature extraction."""
@abstractmethod
def get_feature_names(self) -> list[str]:
"""Ordered list of feature names matching the output array columns."""
...
@abstractmethod
def get_categorical_features(self) -> list[str]:
"""Feature names that should be treated as categorical by LightGBM."""
...
@abstractmethod
def extract(self, problem: dict, kernel: dict) -> np.ndarray:
"""Extract a single feature vector from a (problem, kernel) pair."""
...
def extract_batch(self, df: pd.DataFrame) -> np.ndarray:
"""Vectorized batch extraction from a DataFrame. Override for speed."""
names = self.get_feature_names()
result = np.zeros((len(df), len(names)), dtype=np.float64)
for i in range(len(df)):
row = df.iloc[i]
prob = row.to_dict()
kern = row.to_dict()
result[i] = self.extract(prob, kern)
return result
def get_parameter_space(self) -> dict[str, list]:
"""Valid discrete values for each kernel parameter (for surrogate search)."""
return {}
def get_constraints(self) -> list:
"""Multi-param constraint functions returning True if config is valid."""
return []
def validate_config(self, config: dict) -> bool:
"""Check all constraints. Returns True if the config is valid."""
ps = self.get_parameter_space()
for k, valid_vals in ps.items():
if k in config and config[k] not in valid_vals:
return False
for constraint in self.get_constraints():
if not constraint(config):
return False
return True
def project_to_valid(self, config: dict) -> dict:
"""Snap a config to the nearest valid discrete point."""
ps = self.get_parameter_space()
result = dict(config)
for k, valid_vals in ps.items():
if k not in result:
continue
v = result[k]
if isinstance(valid_vals[0], (int, float)):
result[k] = min(valid_vals, key=lambda x: abs(x - v))
elif v not in valid_vals:
result[k] = valid_vals[0]
return result
class GemmUniversalFeatureEngine(FeatureEngine):
"""Feature engine for gemm_universal kernels."""
def __init__(
self,
num_cus: int = 256,
lds_capacity: int = 65536,
max_clock_mhz: int = 2400,
simds_per_cu: int = 4,
shader_engines: int = 32,
max_waves_per_cu: int = 32,
wavefront_size: int = 64,
l1_cache_kb: int = 32,
l2_cache_kb: int = 4096,
l3_cache_kb: int = 262144,
num_xcd: int = 8,
):
self._hw = {
"num_cus": num_cus,
"lds_capacity": lds_capacity,
"max_clock_mhz": max_clock_mhz,
"simds_per_cu": simds_per_cu,
"shader_engines": shader_engines,
"max_waves_per_cu": max_waves_per_cu,
"wavefront_size": wavefront_size,
"l1_cache_kb": l1_cache_kb,
"l2_cache_kb": l2_cache_kb,
"l3_cache_kb": l3_cache_kb,
"num_xcd": num_xcd,
"total_simds": num_cus * simds_per_cu,
}
def get_feature_names(self) -> list[str]:
return [
# Problem features
"M",
"N",
"K",
"split_k",
"log2_M",
"log2_N",
"log2_K",
"log2_MNK",
"arithmetic_intensity",
"aspect_ratio_mn",
"aspect_ratio_mk",
"aspect_ratio_nk",
"layout",
# Kernel features
"tile_m",
"tile_n",
"tile_k",
"warp_m",
"warp_n",
"warp_k",
"warp_tile_m",
"warp_tile_n",
"warp_tile_k",
"pipeline",
"scheduler",
"epilogue",
"pad_m",
"pad_n",
"pad_k",
"persistent",
"num_warps",
"tile_volume",
"tile_mn",
"lds_usage_estimate",
"lds_usage_ratio",
# Interaction features
"num_tiles_m",
"num_tiles_n",
"num_tiles_k",
"total_output_tiles",
"tile_eff_m",
"tile_eff_n",
"tile_eff_k",
"overall_tile_efficiency",
"cu_utilization",
# P0 FIX: Problem-to-tile ratio features
"ratio_M_to_tile_m",
"ratio_N_to_tile_n",
"ratio_K_to_tile_k",
"problem_smaller_than_tile_m",
"problem_smaller_than_tile_n",
"problem_smaller_than_tile_k",
"any_dim_too_small",
# P1 FIX: Padding requirement interaction features
"needs_padding_m",
"needs_padding_n",
"needs_padding_k",
"has_padding_when_needed_m",
"has_padding_when_needed_n",
"has_padding_when_needed_k",
"missing_required_padding_m",
"missing_required_padding_n",
"missing_required_padding_k",
"missing_any_required_padding",
# Hardware features
"hw_num_cus",
"hw_simds_per_cu",
"hw_total_simds",
"hw_shader_engines",
"hw_max_clock_mhz",
"hw_max_waves_per_cu",
"hw_wavefront_size",
"hw_lds_capacity",
"hw_l1_cache_kb",
"hw_l2_cache_kb",
"hw_l3_cache_kb",
"hw_num_xcd",
]
def get_categorical_features(self) -> list[str]:
return ["layout", "pipeline", "scheduler", "epilogue"]
def extract(self, problem: dict, kernel: dict) -> np.ndarray:
M = int(problem.get("m", problem.get("M", 0)))
N = int(problem.get("n", problem.get("N", 0)))
K = int(problem.get("k", problem.get("K", 0)))
split_k = int(problem.get("split_k", 1))
dtype = str(problem.get("dtype", "fp8"))
bpe = DTYPE_BYTES.get(dtype, 1.0)
log2_M = math.log2(max(M, 1))
log2_N = math.log2(max(N, 1))
log2_K = math.log2(max(K, 1))
log2_MNK = math.log2(max(M * N * K, 1))
mem_bytes = (M * K + K * N + M * N) * bpe
ai = (2.0 * M * N * K) / max(mem_bytes, 1)
ar_mn = M / max(N, 1)
ar_mk = M / max(K, 1)
ar_nk = N / max(K, 1)
layout_code = LAYOUT_MAP.get(str(problem.get("layout", "rcr")), 0)
tile_m = int(kernel.get("tile_m", 128))
tile_n = int(kernel.get("tile_n", 128))
tile_k = int(kernel.get("tile_k", 64))
warp_m = int(kernel.get("warp_m", 2))
warp_n = int(kernel.get("warp_n", 2))
warp_k = int(kernel.get("warp_k", 1))
warp_tile_m = int(kernel.get("warp_tile_m", 32))
warp_tile_n = int(kernel.get("warp_tile_n", 32))
warp_tile_k = int(kernel.get("warp_tile_k", 16))
pipeline_code = PIPELINE_MAP.get(str(kernel.get("pipeline", "compv4")), 0)
scheduler_code = SCHEDULER_MAP.get(str(kernel.get("scheduler", "intrawave")), 0)
epilogue_code = EPILOGUE_MAP.get(str(kernel.get("epilogue", "cshuffle")), 0)
pad_m = float(kernel.get("pad_m", False))
pad_n = float(kernel.get("pad_n", False))
pad_k = float(kernel.get("pad_k", False))
persistent = float(kernel.get("persistent", False))
num_warps = warp_m * warp_n * warp_k
tile_volume = tile_m * tile_n * tile_k
tile_mn = tile_m * tile_n
lds_est = (tile_m * tile_k + tile_n * tile_k) * bpe
lds_cap = self._hw["lds_capacity"]
if str(kernel.get("pipeline", "")).startswith("compv4"):
lds_cap = 32768
lds_ratio = lds_est / max(lds_cap, 1)
num_tiles_m = math.ceil(M / max(tile_m, 1))
num_tiles_n = math.ceil(N / max(tile_n, 1))
num_tiles_k = math.ceil(K / max(tile_k, 1))
total_output_tiles = num_tiles_m * num_tiles_n
rem_m = M % tile_m if tile_m > 0 else 0
tile_eff_m = rem_m / tile_m if rem_m > 0 else 1.0
rem_n = N % tile_n if tile_n > 0 else 0
tile_eff_n = rem_n / tile_n if rem_n > 0 else 1.0
rem_k = K % tile_k if tile_k > 0 else 0
tile_eff_k = rem_k / tile_k if rem_k > 0 else 1.0
overall_eff = tile_eff_m * tile_eff_n * tile_eff_k
cu_util = total_output_tiles / max(self._hw["num_cus"], 1)
# P0 FIX: Problem-to-tile ratio features (avoid oversized tiles for tiny problems)
ratio_M_to_tile_m = M / max(tile_m, 1)
ratio_N_to_tile_n = N / max(tile_n, 1)
ratio_K_to_tile_k = K / max(tile_k, 1)
# Binary features: is problem dimension smaller than tile?
problem_smaller_than_tile_m = float(M < tile_m)
problem_smaller_than_tile_n = float(N < tile_n)
problem_smaller_than_tile_k = float(K < tile_k)
any_dim_too_small = float((M < tile_m) or (N < tile_n) or (K < tile_k))
# P1 FIX: Padding requirement features (does this kernel have padding when needed?)
needs_padding_m = float(M % tile_m != 0) if tile_m > 0 else 0.0
needs_padding_n = float(N % tile_n != 0) if tile_n > 0 else 0.0
needs_padding_k = float(K % tile_k != 0) if tile_k > 0 else 0.0
# Interaction features: kernel has padding capability when problem needs it
has_padding_when_needed_m = float(needs_padding_m and pad_m)
has_padding_when_needed_n = float(needs_padding_n and pad_n)
has_padding_when_needed_k = float(needs_padding_k and pad_k)
# Critical feature: missing required padding (kernel will likely fail)
missing_required_padding_m = float(needs_padding_m and not pad_m)
missing_required_padding_n = float(needs_padding_n and not pad_n)
missing_required_padding_k = float(needs_padding_k and not pad_k)
missing_any_required_padding = float(
missing_required_padding_m
or missing_required_padding_n
or missing_required_padding_k
)
hw = self._hw
return np.array(
[
M,
N,
K,
split_k,
log2_M,
log2_N,
log2_K,
log2_MNK,
ai,
ar_mn,
ar_mk,
ar_nk,
layout_code,
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
pipeline_code,
scheduler_code,
epilogue_code,
pad_m,
pad_n,
pad_k,
persistent,
num_warps,
tile_volume,
tile_mn,
lds_est,
lds_ratio,
num_tiles_m,
num_tiles_n,
num_tiles_k,
total_output_tiles,
tile_eff_m,
tile_eff_n,
tile_eff_k,
overall_eff,
cu_util,
# P0 FIX: New ratio and binary features
ratio_M_to_tile_m,
ratio_N_to_tile_n,
ratio_K_to_tile_k,
problem_smaller_than_tile_m,
problem_smaller_than_tile_n,
problem_smaller_than_tile_k,
any_dim_too_small,
# P1 FIX: Padding requirement interaction features
needs_padding_m,
needs_padding_n,
needs_padding_k,
has_padding_when_needed_m,
has_padding_when_needed_n,
has_padding_when_needed_k,
missing_required_padding_m,
missing_required_padding_n,
missing_required_padding_k,
missing_any_required_padding,
hw["num_cus"],
hw["simds_per_cu"],
hw["total_simds"],
hw["shader_engines"],
hw["max_clock_mhz"],
hw["max_waves_per_cu"],
hw["wavefront_size"],
hw["lds_capacity"],
hw["l1_cache_kb"],
hw["l2_cache_kb"],
hw["l3_cache_kb"],
hw["num_xcd"],
],
dtype=np.float64,
)
def extract_batch(self, df: pd.DataFrame) -> np.ndarray:
"""Vectorized batch extraction -- much faster than row-by-row."""
n = len(df)
names = self.get_feature_names()
result = np.zeros((n, len(names)), dtype=np.float64)
M = df["m"].values.astype(np.float64)
N = df["n"].values.astype(np.float64)
K = df["k"].values.astype(np.float64)
split_k = df["split_k"].fillna(1).values.astype(np.float64)
dtype_col = df["dtype"].fillna("fp8")
bpe = dtype_col.map(DTYPE_BYTES).fillna(1.0).values
result[:, 0] = M
result[:, 1] = N
result[:, 2] = K
result[:, 3] = split_k
result[:, 4] = np.log2(np.maximum(M, 1))
result[:, 5] = np.log2(np.maximum(N, 1))
result[:, 6] = np.log2(np.maximum(K, 1))
result[:, 7] = np.log2(np.maximum(M * N * K, 1))
mem = (M * K + K * N + M * N) * bpe
result[:, 8] = (2.0 * M * N * K) / np.maximum(mem, 1)
result[:, 9] = M / np.maximum(N, 1)
result[:, 10] = M / np.maximum(K, 1)
result[:, 11] = N / np.maximum(K, 1)
result[:, 12] = df["layout"].map(LAYOUT_MAP).fillna(0).values
tile_m = df["tile_m"].fillna(128).values.astype(np.float64)
tile_n = df["tile_n"].fillna(128).values.astype(np.float64)
tile_k = df["tile_k"].fillna(64).values.astype(np.float64)
warp_m = df["warp_m"].fillna(2).values.astype(np.float64)
warp_n = df["warp_n"].fillna(2).values.astype(np.float64)
warp_k = df["warp_k"].fillna(1).values.astype(np.float64)
warp_tile_m = df["warp_tile_m"].fillna(32).values.astype(np.float64)
warp_tile_n = df["warp_tile_n"].fillna(32).values.astype(np.float64)
warp_tile_k = df["warp_tile_k"].fillna(16).values.astype(np.float64)
result[:, 13] = tile_m
result[:, 14] = tile_n
result[:, 15] = tile_k
result[:, 16] = warp_m
result[:, 17] = warp_n
result[:, 18] = warp_k
result[:, 19] = warp_tile_m
result[:, 20] = warp_tile_n
result[:, 21] = warp_tile_k
result[:, 22] = df["pipeline"].map(PIPELINE_MAP).fillna(0).values
result[:, 23] = df["scheduler"].map(SCHEDULER_MAP).fillna(0).values
result[:, 24] = df["epilogue"].map(EPILOGUE_MAP).fillna(0).values
result[:, 25] = df["pad_m"].fillna(False).astype(float).values
result[:, 26] = df["pad_n"].fillna(False).astype(float).values
result[:, 27] = df["pad_k"].fillna(False).astype(float).values
result[:, 28] = df["persistent"].fillna(False).astype(float).values
num_warps = warp_m * warp_n * warp_k
result[:, 29] = num_warps
result[:, 30] = tile_m * tile_n * tile_k
result[:, 31] = tile_m * tile_n
lds_est = (tile_m * tile_k + tile_n * tile_k) * bpe
result[:, 32] = lds_est
lds_cap = np.full(n, self._hw["lds_capacity"], dtype=np.float64)
is_compv4 = df["pipeline"].fillna("").str.startswith("compv4")
lds_cap[is_compv4] = 32768
result[:, 33] = lds_est / np.maximum(lds_cap, 1)
ntm = np.ceil(M / np.maximum(tile_m, 1))
ntn = np.ceil(N / np.maximum(tile_n, 1))
ntk = np.ceil(K / np.maximum(tile_k, 1))
result[:, 34] = ntm
result[:, 35] = ntn
result[:, 36] = ntk
result[:, 37] = ntm * ntn
rem_m = np.mod(M, np.maximum(tile_m, 1))
result[:, 38] = np.where(rem_m > 0, rem_m / tile_m, 1.0)
rem_n = np.mod(N, np.maximum(tile_n, 1))
result[:, 39] = np.where(rem_n > 0, rem_n / tile_n, 1.0)
rem_k = np.mod(K, np.maximum(tile_k, 1))
result[:, 40] = np.where(rem_k > 0, rem_k / tile_k, 1.0)
result[:, 41] = result[:, 38] * result[:, 39] * result[:, 40]
result[:, 42] = (ntm * ntn) / max(self._hw["num_cus"], 1)
# P0 FIX: Problem-to-tile ratio features
result[:, 43] = M / np.maximum(tile_m, 1) # ratio_M_to_tile_m
result[:, 44] = N / np.maximum(tile_n, 1) # ratio_N_to_tile_n
result[:, 45] = K / np.maximum(tile_k, 1) # ratio_K_to_tile_k
# Binary features: is problem smaller than tile?
result[:, 46] = (M < tile_m).astype(float) # problem_smaller_than_tile_m
result[:, 47] = (N < tile_n).astype(float) # problem_smaller_than_tile_n
result[:, 48] = (K < tile_k).astype(float) # problem_smaller_than_tile_k
result[:, 49] = ((M < tile_m) | (N < tile_n) | (K < tile_k)).astype(
float
) # any_dim_too_small
# P1 FIX: Padding requirement features
pad_m_bool = df["pad_m"].fillna(False).astype(bool).values
pad_n_bool = df["pad_n"].fillna(False).astype(bool).values
pad_k_bool = df["pad_k"].fillna(False).astype(bool).values
needs_padding_m = (np.mod(M, np.maximum(tile_m, 1)) != 0)
needs_padding_n = (np.mod(N, np.maximum(tile_n, 1)) != 0)
needs_padding_k = (np.mod(K, np.maximum(tile_k, 1)) != 0)
result[:, 50] = needs_padding_m.astype(float)
result[:, 51] = needs_padding_n.astype(float)
result[:, 52] = needs_padding_k.astype(float)
# Interaction features: kernel has padding when problem needs it
result[:, 53] = (needs_padding_m & pad_m_bool).astype(float) # has_padding_when_needed_m
result[:, 54] = (needs_padding_n & pad_n_bool).astype(float) # has_padding_when_needed_n
result[:, 55] = (needs_padding_k & pad_k_bool).astype(float) # has_padding_when_needed_k
# Critical feature: missing required padding
result[:, 56] = (needs_padding_m & ~pad_m_bool).astype(float) # missing_required_padding_m
result[:, 57] = (needs_padding_n & ~pad_n_bool).astype(float) # missing_required_padding_n
result[:, 58] = (needs_padding_k & ~pad_k_bool).astype(float) # missing_required_padding_k
result[:, 59] = ((needs_padding_m & ~pad_m_bool) | (needs_padding_n & ~pad_n_bool) | (needs_padding_k & ~pad_k_bool)).astype(float) # missing_any_required_padding
# Hardware profile features
hw = self._hw
result[:, 60] = hw["num_cus"]
result[:, 61] = hw["simds_per_cu"]
result[:, 62] = hw["total_simds"]
result[:, 63] = hw["shader_engines"]
result[:, 64] = hw["max_clock_mhz"]
result[:, 65] = hw["max_waves_per_cu"]
result[:, 66] = hw["wavefront_size"]
result[:, 67] = hw["lds_capacity"]
result[:, 68] = hw["l1_cache_kb"]
result[:, 69] = hw["l2_cache_kb"]
result[:, 70] = hw["l3_cache_kb"]
result[:, 71] = hw["num_xcd"]
return result
def get_parameter_space(self) -> dict[str, list]:
return {
"tile_m": [32, 64, 128, 192, 256],
"tile_n": [32, 64, 128, 192, 256],
"tile_k": [32, 64, 128, 256],
"warp_m": [1, 2, 4],
"warp_n": [1, 2, 4],
"warp_k": [1],
"warp_tile_m": [4, 16, 32, 64],
"warp_tile_n": [4, 16, 32, 64],
"warp_tile_k": [8, 16, 32, 64, 128],
"pipeline": list(PIPELINE_MAP.keys()),
"scheduler": list(SCHEDULER_MAP.keys()),
"epilogue": list(EPILOGUE_MAP.keys()),
"pad_m": [True, False],
"pad_n": [True, False],
"pad_k": [True, False],
"persistent": [True, False],
}
def get_constraints(self) -> list:
lds_cap = self._hw["lds_capacity"]
def _lds_constraint(cfg):
tm = cfg.get("tile_m", 128)
tn = cfg.get("tile_n", 128)
tk = cfg.get("tile_k", 64)
bpe = 1.0 # fp8 default
est = (tm * tk + tn * tk) * bpe
cap = (
32768 if str(cfg.get("pipeline", "")).startswith("compv4") else lds_cap
)
return est <= cap
def _warp_constraint(cfg):
wm = cfg.get("warp_m", 2)
wn = cfg.get("warp_n", 2)
wk = cfg.get("warp_k", 1)
return (wm * wn * wk) in [2, 4, 8]
return [_lds_constraint, _warp_constraint]

553
dispatcher/heuristics/generate_benchmark_data.py Normal file
View File

@@ -0,0 +1,553 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
GEMM Universal Benchmark Data Generation Script
This script generates training data for ML-based kernel selection heuristics by:
1. Reading kernel configurations from the tile engine
2. Building benchmark executables (in parallel)
3. Running benchmarks across multiple problem sizes
4. Outputting performance data in JSON format
Usage:
python generate_benchmark_data.py \
--build_dir /tmp/build \
--output_dir /tmp/benchmark_data \
--dtype fp16 \
--layout rcr \
--num_build_jobs 4 \
--num_benchmark_jobs 1
Requirements:
- ROCm-capable GPU
- CK tile engine built with CMake
"""
import argparse
import json
import subprocess
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import List, Optional, Tuple
import re
@dataclass
class KernelConfig:
"""Represents a single kernel configuration."""
name: str
dtype: str
layout: str
pipeline: str
epilogue: str
scheduler: str
pad_m: bool
pad_n: bool
pad_k: bool
persistent: bool
tile_m: int
tile_n: int
tile_k: int
warp_m: int
warp_n: int
warp_k: int
warp_tile_m: int
warp_tile_n: int
warp_tile_k: int
@classmethod
def from_kernel_name(cls, name: str, dtype: str, layout: str) -> "KernelConfig":
"""Parse kernel name to extract configuration."""
# Format: gemm_universal_{dtype}_{layout}_{pipeline}_{epilogue}_{scheduler}_{padM}_{padN}_{padK}_{persistent}_{tile_config}
# tile_config: {tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}
parts = name.split("_")
prefix = f"gemm_universal_{dtype}_{layout}_"
trait_and_tile = name[len(prefix) :]
trait_parts = trait_and_tile.split("_")
pipeline = trait_parts[0]
epilogue = trait_parts[1]
scheduler = trait_parts[2]
pad_m = trait_parts[3] == "True"
pad_n = trait_parts[4] == "True"
pad_k = trait_parts[5] == "True"
persistent = trait_parts[6] == "True"
# Parse tile config
tile_dims = trait_parts[7].split("x")
warp_dims = trait_parts[8].split("x")
warp_tile_dims = trait_parts[9].split("x")
return cls(
name=name,
dtype=dtype,
layout=layout,
pipeline=pipeline,
epilogue=epilogue,
scheduler=scheduler,
pad_m=pad_m,
pad_n=pad_n,
pad_k=pad_k,
persistent=persistent,
tile_m=int(tile_dims[0]),
tile_n=int(tile_dims[1]),
tile_k=int(tile_dims[2]),
warp_m=int(warp_dims[0]),
warp_n=int(warp_dims[1]),
warp_k=int(warp_dims[2]),
warp_tile_m=int(warp_tile_dims[0]),
warp_tile_n=int(warp_tile_dims[1]),
warp_tile_k=int(warp_tile_dims[2]),
)
@dataclass
class BenchmarkResult:
"""Result of a single benchmark run."""
kernel_name: str
m: int
n: int
k: int
avg_time_ms: float
tflops: float
is_valid: bool
error: Optional[str] = None
@dataclass
class ProblemSize:
"""GEMM problem dimensions."""
m: int
n: int
k: int
def get_problem_sizes() -> List[ProblemSize]:
"""
Generate diverse problem sizes for benchmarking.
Includes:
- Square matrices (powers of 2)
- Rectangular matrices (common in ML)
- LLM-specific sizes (attention, MLP)
- Edge cases (small, very large)
"""
sizes = []
# Powers of 2 (square)
for p in [6, 7, 8, 9, 10, 11, 12, 13]: # 64 to 8192
dim = 2**p
sizes.append(ProblemSize(dim, dim, dim))
# Common ML sizes (batch x hidden)
ml_sizes = [
(1, 4096, 4096), # Single token inference
(8, 4096, 4096), # Small batch
(32, 4096, 4096), # Medium batch
(128, 4096, 4096), # Large batch
(1, 4096, 11008), # LLaMA MLP up-projection
(1, 11008, 4096), # LLaMA MLP down-projection
(32, 4096, 11008),
(32, 11008, 4096),
(1, 8192, 8192), # Large model
(32, 8192, 8192),
(1, 8192, 28672), # LLaMA-70B MLP
(32, 8192, 28672),
]
for m, n, k in ml_sizes:
sizes.append(ProblemSize(m, n, k))
# Rectangular matrices
rect_sizes = [
(1024, 4096, 1024),
(4096, 1024, 4096),
(2048, 8192, 2048),
(256, 256, 8192), # Tall K
(8192, 8192, 256), # Short K
]
for m, n, k in rect_sizes:
sizes.append(ProblemSize(m, n, k))
# Remove duplicates while preserving order
seen = set()
unique_sizes = []
for s in sizes:
key = (s.m, s.n, s.k)
if key not in seen:
seen.add(key)
unique_sizes.append(s)
return unique_sizes
def load_kernel_list(build_dir: Path, dtype: str, layout: str) -> List[KernelConfig]:
"""Load kernel configurations from the tile engine build."""
kernel_list_path = (
build_dir
/ "tile_engine"
/ "ops"
/ "gemm"
/ "gemm_universal"
/ dtype
/ layout
/ "gemm_universal_kernel_list.txt"
)
if not kernel_list_path.exists():
raise FileNotFoundError(f"Kernel list not found: {kernel_list_path}")
kernels = []
with open(kernel_list_path, "r") as f:
for line in f:
line = line.strip()
if not line:
continue
# Format: kernel_name|tile_config|trait_combo
parts = line.split("|")
kernel_name = parts[0]
kernels.append(KernelConfig.from_kernel_name(kernel_name, dtype, layout))
return kernels
def build_kernel(build_dir: Path, kernel: KernelConfig) -> Tuple[str, bool, str]:
"""
Build a single kernel benchmark executable.
Returns: (kernel_name, success, error_message)
"""
target_name = f"benchmark_{kernel.name}"
try:
result = subprocess.run(
["ninja", "-j1", target_name],
cwd=build_dir,
capture_output=True,
text=True,
timeout=300, # 5 minute timeout
)
if result.returncode != 0:
return (kernel.name, False, result.stderr[:500])
return (kernel.name, True, "")
except subprocess.TimeoutExpired:
return (kernel.name, False, "Build timeout")
except Exception as e:
return (kernel.name, False, str(e))
def run_benchmark(
build_dir: Path,
kernel: KernelConfig,
problem: ProblemSize,
warmup: int = 10,
repeat: int = 50,
) -> BenchmarkResult:
"""
Run benchmark for a single kernel and problem size.
"""
exe_path = build_dir / "bin" / f"benchmark_{kernel.name}"
if not exe_path.exists():
return BenchmarkResult(
kernel_name=kernel.name,
m=problem.m,
n=problem.n,
k=problem.k,
avg_time_ms=0,
tflops=0,
is_valid=False,
error="Executable not found",
)
try:
result = subprocess.run(
[
str(exe_path),
f"-m={problem.m}",
f"-n={problem.n}",
f"-k={problem.k}",
f"-warmup={warmup}",
f"-repeat={repeat}",
"-verify=0",
"-json_output=true",
],
capture_output=True,
text=True,
timeout=120,
)
if result.returncode != 0:
# Try to parse error
error = result.stderr[:200] if result.stderr else result.stdout[:200]
return BenchmarkResult(
kernel_name=kernel.name,
m=problem.m,
n=problem.n,
k=problem.k,
avg_time_ms=0,
tflops=0,
is_valid=False,
error=error,
)
# Parse JSON output
output = result.stdout.strip()
# Try to find JSON in output
json_match = re.search(r"\{.*\}", output, re.DOTALL)
if json_match:
data = json.loads(json_match.group())
# Extract from nested perf_result object
perf = data.get("perf_result", {})
avg_time_ms = perf.get("latency(ms)", 0)
tflops = perf.get("tflops(TFlops)", 0)
return BenchmarkResult(
kernel_name=kernel.name,
m=problem.m,
n=problem.n,
k=problem.k,
avg_time_ms=avg_time_ms,
tflops=tflops,
is_valid=True,
)
else:
# Parse from text output
# Look for patterns like "avg_time: X ms" or "tflops: Y"
avg_time = 0.0
tflops = 0.0
time_match = re.search(
r"(?:avg[_\s]?time|latency)[:\s]+(\d+\.?\d*)\s*(?:ms)?", output, re.I
)
if time_match:
avg_time = float(time_match.group(1))
tflops_match = re.search(r"tflops[:\s]+(\d+\.?\d*)", output, re.I)
if tflops_match:
tflops = float(tflops_match.group(1))
# Calculate TFLOPs if not provided
if tflops == 0 and avg_time > 0:
flops = 2.0 * problem.m * problem.n * problem.k
tflops = flops / (avg_time * 1e-3) / 1e12
return BenchmarkResult(
kernel_name=kernel.name,
m=problem.m,
n=problem.n,
k=problem.k,
avg_time_ms=avg_time,
tflops=tflops,
is_valid=avg_time > 0,
error=None if avg_time > 0 else "Could not parse output",
)
except subprocess.TimeoutExpired:
return BenchmarkResult(
kernel_name=kernel.name,
m=problem.m,
n=problem.n,
k=problem.k,
avg_time_ms=0,
tflops=0,
is_valid=False,
error="Benchmark timeout",
)
except Exception as e:
return BenchmarkResult(
kernel_name=kernel.name,
m=problem.m,
n=problem.n,
k=problem.k,
avg_time_ms=0,
tflops=0,
is_valid=False,
error=str(e),
)
def main():
parser = argparse.ArgumentParser(
description="Generate GEMM benchmark data for ML training"
)
parser.add_argument(
"--build_dir", type=str, default="/tmp/build", help="CK build directory"
)
parser.add_argument(
"--output_dir",
type=str,
default="/tmp/benchmark_data",
help="Output directory for benchmark results",
)
parser.add_argument(
"--dtype",
type=str,
default="fp16",
choices=["fp16", "fp8", "bf16", "bf8"],
help="Data type to benchmark",
)
parser.add_argument(
"--layout",
type=str,
default="rcr",
choices=["rcr", "rrr", "crr", "ccr"],
help="Matrix layout to benchmark",
)
parser.add_argument(
"--num_build_jobs", type=int, default=4, help="Number of parallel build jobs"
)
parser.add_argument(
"--num_benchmark_jobs",
type=int,
default=1,
help="Number of parallel benchmark jobs (use 1 for accurate timing)",
)
parser.add_argument(
"--max_kernels",
type=int,
default=None,
help="Maximum number of kernels to benchmark (for testing)",
)
parser.add_argument(
"--skip_build",
action="store_true",
help="Skip building and only run benchmarks",
)
parser.add_argument(
"--warmup", type=int, default=10, help="Number of warmup iterations"
)
parser.add_argument(
"--repeat", type=int, default=50, help="Number of benchmark iterations"
)
args = parser.parse_args()
build_dir = Path(args.build_dir)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# Load kernel configurations
print(f"Loading kernel list for {args.dtype}/{args.layout}...")
kernels = load_kernel_list(build_dir, args.dtype, args.layout)
print(f"Found {len(kernels)} kernel configurations")
if args.max_kernels:
kernels = kernels[: args.max_kernels]
print(f"Limiting to {len(kernels)} kernels")
# Build kernels
if not args.skip_build:
print(
f"\nBuilding {len(kernels)} kernels with {args.num_build_jobs} parallel jobs..."
)
build_results = {"success": 0, "failed": 0, "failed_kernels": []}
with ProcessPoolExecutor(max_workers=args.num_build_jobs) as executor:
futures = {executor.submit(build_kernel, build_dir, k): k for k in kernels}
for i, future in enumerate(as_completed(futures)):
kernel_name, success, error = future.result()
if success:
build_results["success"] += 1
else:
build_results["failed"] += 1
build_results["failed_kernels"].append(
{"name": kernel_name, "error": error}
)
if (i + 1) % 10 == 0:
print(
f" Built {i + 1}/{len(kernels)} ({build_results['success']} success, {build_results['failed']} failed)"
)
print(
f"\nBuild complete: {build_results['success']} success, {build_results['failed']} failed"
)
# Save build results
with open(output_dir / "build_results.json", "w") as f:
json.dump(build_results, f, indent=2)
# Get problem sizes
problem_sizes = get_problem_sizes()
print(f"\nBenchmarking {len(problem_sizes)} problem sizes...")
# Run benchmarks
all_results = []
total_benchmarks = len(kernels) * len(problem_sizes)
completed = 0
print(f"Total benchmarks to run: {total_benchmarks}")
for kernel in kernels:
kernel_results = {
"kernel_config": asdict(kernel),
"benchmarks": [],
}
for problem in problem_sizes:
result = run_benchmark(
build_dir,
kernel,
problem,
warmup=args.warmup,
repeat=args.repeat,
)
kernel_results["benchmarks"].append(asdict(result))
completed += 1
if completed % 100 == 0:
print(f" Progress: {completed}/{total_benchmarks} benchmarks complete")
all_results.append(kernel_results)
# Save intermediate results
intermediate_file = (
output_dir / f"benchmark_results_{args.dtype}_{args.layout}_partial.json"
)
with open(intermediate_file, "w") as f:
json.dump(all_results, f, indent=2)
# Save final results
final_file = output_dir / f"benchmark_results_{args.dtype}_{args.layout}.json"
with open(final_file, "w") as f:
json.dump(
{
"metadata": {
"dtype": args.dtype,
"layout": args.layout,
"num_kernels": len(kernels),
"num_problems": len(problem_sizes),
"warmup": args.warmup,
"repeat": args.repeat,
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
},
"problem_sizes": [asdict(p) for p in problem_sizes],
"results": all_results,
},
f,
indent=2,
)
print(f"\nResults saved to {final_file}")
# Print summary
valid_count = sum(
1 for kr in all_results for br in kr["benchmarks"] if br["is_valid"]
)
print(f"Valid benchmarks: {valid_count}/{total_benchmarks}")
if __name__ == "__main__":
main()

166
dispatcher/heuristics/generate_edge_dims.py Normal file
View File

@@ -0,0 +1,166 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Supplementary edge-case benchmark generator for N=1 and K=1 dimensions.
These shapes represent vector-matrix multiply (N=1), rank-1 updates (K=1),
and other degenerate GEMM cases that stress tile efficiency and padding logic.
"""
import json
import subprocess
import sys
from pathlib import Path
def generate_edge_shapes():
"""Generate shapes with N=1, K=1, and other single-dimension edge cases."""
shapes = set()
# --- N=1: vector-matrix multiply / single output column ---
for m in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]:
for k in [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 7168, 8192]:
shapes.add((m, 1, k))
# --- K=1: rank-1 update / outer product ---
for m in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]:
for n in [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 7168, 8192]:
shapes.add((m, n, 1))
# --- M=1, N=1: dot product ---
for k in [1, 16, 64, 256, 1024, 4096, 8192]:
shapes.add((1, 1, k))
# --- M=1, K=1: scalar-vector ---
for n in [1, 16, 64, 256, 1024, 4096, 8192]:
shapes.add((1, n, 1))
# --- N=1, K=1: scalar-vector ---
for m in [1, 16, 64, 256, 1024, 4096, 8192]:
shapes.add((m, 1, 1))
# --- All ones: 1x1x1 ---
shapes.add((1, 1, 1))
# --- Small N (2-16) ---
for m in [64, 256, 1024, 4096]:
for n in [2, 3, 4, 7, 8, 15, 16]:
for k in [64, 256, 1024, 4096]:
shapes.add((m, n, k))
# --- Small K (2-16) ---
for m in [64, 256, 1024, 4096]:
for n in [64, 256, 1024, 4096]:
for k in [2, 3, 4, 7, 8, 15, 16]:
shapes.add((m, n, k))
return sorted(shapes)
def run_shapes(bin_dir, shapes, out_file, warmup=3, repeat=10):
"""Run all kernels against shapes, writing streaming log."""
executables = sorted(Path(bin_dir).glob("benchmark_gemm_universal_fp8_rcr_*"))
if not executables:
print(f"ERROR: No executables found in {bin_dir}", file=sys.stderr)
return 0
total = 0
for idx, (m, n, k) in enumerate(shapes):
out_file.write("\n========================================\n")
out_file.write(f"Shape {idx + 1}: M={m} N={n} K={k} dtype=fp8 layout=rcr\n")
out_file.write("========================================\n")
out_file.write(f"Found {len(executables)} kernels\n")
out_file.flush()
for exe in executables:
try:
result = subprocess.run(
[
str(exe),
f"-m={m}",
f"-n={n}",
f"-k={k}",
f"-warmup={warmup}",
f"-repeat={repeat}",
"-verify=0",
],
capture_output=True,
text=True,
timeout=60,
)
output = result.stdout
json_start = output.find("{")
json_end = output.rfind("}") + 1
if json_start >= 0 and json_end > json_start:
json_block = output[json_start:json_end]
try:
json.loads(json_block)
out_file.write(json_block + "\n")
total += 1
except json.JSONDecodeError:
pass
except (subprocess.TimeoutExpired, Exception):
pass
out_file.flush()
print(
f" Shape {idx + 1}/{len(shapes)}: M={m} N={n} K={k}",
file=sys.stderr,
flush=True,
)
return total
if __name__ == "__main__":
bin_dir = "/workspace/ck_tile/bin"
out_dir = Path("data/edge_dims")
out_dir.mkdir(parents=True, exist_ok=True)
shapes = generate_edge_shapes()
print(f"Generated {len(shapes)} edge-case shapes", file=sys.stderr, flush=True)
n1_count = sum(1 for m, n, k in shapes if n == 1)
k1_count = sum(1 for m, n, k in shapes if k == 1)
both1 = sum(1 for m, n, k in shapes if n == 1 and k == 1)
small_n = sum(1 for m, n, k in shapes if 2 <= n <= 16)
small_k = sum(1 for m, n, k in shapes if 2 <= k <= 16)
print(
f" N=1: {n1_count}, K=1: {k1_count}, both=1: {both1}",
file=sys.stderr,
flush=True,
)
print(
f" Small N(2-16): {small_n}, Small K(2-16): {small_k}",
file=sys.stderr,
flush=True,
)
batch_size = 25
total = 0
batch_idx = 0
for i in range(0, len(shapes), batch_size):
batch = shapes[i : i + batch_size]
batch_idx += 1
out_path = out_dir / f"edge_dims_batch_{batch_idx:03d}.log"
print(
f"\nBatch {batch_idx}: shapes {i + 1}-{i + len(batch)} -> {out_path}",
file=sys.stderr,
flush=True,
)
with open(out_path, "w") as f:
f.write(f"CK Tile Edge Dims Benchmark Batch {batch_idx}\n")
f.write("GPU ID: 0\nImplementation: gemm_universal\n\n")
count = run_shapes(bin_dir, batch, f, warmup=3, repeat=10)
total += count
print(f" Batch {batch_idx} done: {count} results", file=sys.stderr, flush=True)
print(
f"\nTotal: {total} benchmarks across {len(shapes)} shapes",
file=sys.stderr,
flush=True,
)

289
dispatcher/heuristics/generate_wide_coverage.py Normal file
View File

@@ -0,0 +1,289 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Wide-coverage benchmark data generator.
Generates benchmark results for hundreds of diverse GEMM shapes across all
corner cases: skinny M, tall N, deep K, M=1, prime dimensions, power-of-2,
LLM inference shapes, training shapes, and edge cases.
Runs all 4608 kernels in /workspace/ck_tile/bin/ against each shape and
writes streaming log output parseable by data_pipeline.py.
Usage:
python3 generate_wide_coverage.py --bin_dir /workspace/ck_tile/bin \
--out_dir data/ --batch_size 20 --warmup 3 --repeat 10
"""
import argparse
import json
import subprocess
import sys
from pathlib import Path
def generate_shape_list():
"""Generate a comprehensive list of (M, N, K) shapes covering all corner cases.
Categories:
1. M=1 (single token inference) -- the hardest case
2. Tiny M (2-16) -- small batch inference
3. Small M (32-128) -- medium batch
4. Medium M (256-2048) -- large batch / training
5. Large M (4096-20480) -- very large batch
6. Square shapes (powers of 2)
7. Skinny M, tall N (M << N)
8. Tall M, skinny N (M >> N)
9. Deep K (K >> M, N) -- compute-bound
10. Shallow K (K << M, N) -- memory-bound
11. Prime dimensions -- worst-case for tiling
12. LLM-specific shapes (DeepSeek, LLaMA, etc.)
13. Non-power-of-2 common sizes
"""
shapes = set()
# --- 1. M=1 (single token) across various N, K ---
for n in [512, 1024, 1536, 2048, 3072, 4096, 4608, 7168, 8192, 11008, 14336, 28672]:
for k in [256, 512, 1024, 1536, 2048, 2304, 4096, 7168, 8192]:
shapes.add((1, n, k))
# --- 2. Tiny M (2-16) ---
for m in [2, 4, 8, 16]:
for n in [512, 1536, 4096, 7168]:
for k in [256, 1024, 4096, 7168]:
shapes.add((m, n, k))
# --- 3. Small M (32-128) ---
for m in [32, 48, 64, 96, 128]:
for n in [512, 1536, 4096, 7168, 8192]:
for k in [256, 512, 2048, 4096, 7168]:
shapes.add((m, n, k))
# --- 4. Medium M (256-2048) ---
for m in [256, 384, 512, 768, 1024, 1536, 2048]:
for n in [512, 1536, 4096, 7168]:
for k in [256, 1024, 2048, 4096, 7168]:
shapes.add((m, n, k))
# --- 5. Large M (4096-20480) ---
for m in [4096, 8192, 12288, 16384, 20480]:
for n in [512, 1536, 4096, 7168]:
for k in [256, 1024, 2048, 7168]:
shapes.add((m, n, k))
# --- 6. Square shapes (powers of 2) ---
for p in range(5, 14): # 32 to 8192
d = 2**p
shapes.add((d, d, d))
# --- 7. Skinny M, tall N ---
for m in [1, 4, 16, 64]:
for n in [8192, 16384, 28672]:
for k in [1024, 4096, 8192]:
shapes.add((m, n, k))
# --- 8. Tall M, skinny N ---
for m in [4096, 8192, 16384]:
for n in [32, 64, 128, 256]:
for k in [1024, 4096]:
shapes.add((m, n, k))
# --- 9. Deep K (K >> M, N) ---
for m in [16, 64, 256]:
for n in [16, 64, 256]:
for k in [4096, 8192, 16384, 32768]:
shapes.add((m, n, k))
# --- 10. Shallow K (K << M, N) ---
for m in [1024, 4096, 8192]:
for n in [1024, 4096, 8192]:
for k in [16, 32, 64, 128]:
shapes.add((m, n, k))
# --- 11. Prime dimensions ---
primes = [17, 31, 37, 127, 251, 509, 1021, 2039, 4093]
for p in primes:
shapes.add((p, p, p))
for p in primes[:5]:
shapes.add((p, 4096, 4096))
shapes.add((4096, p, 4096))
shapes.add((4096, 4096, p))
# --- 12. LLM-specific shapes ---
llm_shapes = [
# DeepSeek MoE
(1, 1536, 7168),
(1, 4608, 7168),
(1, 7168, 2048),
(1, 7168, 2304),
(1, 7168, 256),
(1, 576, 7168),
(1, 512, 7168),
(1, 3072, 1536),
# LLaMA-7B
(1, 4096, 4096),
(32, 4096, 4096),
(128, 4096, 4096),
(1, 4096, 11008),
(32, 4096, 11008),
(1, 11008, 4096),
(32, 11008, 4096),
# LLaMA-70B
(1, 8192, 8192),
(32, 8192, 8192),
(128, 8192, 8192),
(1, 8192, 28672),
(32, 8192, 28672),
(1, 28672, 8192),
# GPT-style attention
(128, 128, 64),
(128, 128, 128),
(256, 256, 64),
(512, 512, 64),
(1024, 1024, 64),
(2048, 2048, 64),
]
for s in llm_shapes:
shapes.add(s)
# --- 13. Non-power-of-2 common sizes ---
for m in [48, 96, 192, 384, 576, 768, 1152, 1536, 2304, 3072, 4608, 6144]:
shapes.add((m, m, m))
shapes.add((m, 4096, 4096))
sorted_shapes = sorted(shapes)
return sorted_shapes
def run_shape_batch(bin_dir, shapes, out_file, warmup=3, repeat=10):
"""Run all kernels against a batch of shapes, writing streaming log output."""
executables = sorted(Path(bin_dir).glob("benchmark_gemm_universal_fp8_rcr_*"))
if not executables:
print(f"ERROR: No executables found in {bin_dir}", file=sys.stderr)
return 0
total_benchmarks = 0
for shape_idx, (m, n, k) in enumerate(shapes):
out_file.write("\n========================================\n")
out_file.write(
f"Shape {shape_idx + 1}: M={m} N={n} K={k} dtype=fp8 layout=rcr\n"
)
out_file.write("========================================\n")
out_file.write(f"Found {len(executables)} kernels\n")
out_file.flush()
for exe in executables:
try:
result = subprocess.run(
[
str(exe),
f"-m={m}",
f"-n={n}",
f"-k={k}",
f"-warmup={warmup}",
f"-repeat={repeat}",
"-verify=0",
],
capture_output=True,
text=True,
timeout=60,
)
output = result.stdout
# Extract JSON block from output
json_start = output.find("{")
json_end = output.rfind("}") + 1
if json_start >= 0 and json_end > json_start:
json_block = output[json_start:json_end]
try:
json.loads(json_block)
out_file.write(json_block + "\n")
total_benchmarks += 1
except json.JSONDecodeError:
pass
except (subprocess.TimeoutExpired, Exception):
pass
out_file.flush()
elapsed_kernels = len(executables)
print(
f" Shape {shape_idx + 1}/{len(shapes)}: M={m} N={n} K={k} "
f"({elapsed_kernels} kernels)",
file=sys.stderr,
flush=True,
)
return total_benchmarks
def main():
parser = argparse.ArgumentParser(
description="Generate wide-coverage benchmark data"
)
parser.add_argument(
"--bin_dir",
default="/workspace/ck_tile/bin",
help="Directory with benchmark executables",
)
parser.add_argument("--out_dir", default="data", help="Output directory")
parser.add_argument(
"--batch_size", type=int, default=25, help="Shapes per output file"
)
parser.add_argument("--warmup", type=int, default=3)
parser.add_argument("--repeat", type=int, default=10)
parser.add_argument(
"--max_shapes", type=int, default=None, help="Limit total shapes (for testing)"
)
args = parser.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
shapes = generate_shape_list()
if args.max_shapes:
shapes = shapes[: args.max_shapes]
print(f"Generated {len(shapes)} unique shapes", file=sys.stderr, flush=True)
print(f"Bin dir: {args.bin_dir}", file=sys.stderr, flush=True)
print(f"Output dir: {args.out_dir}", file=sys.stderr, flush=True)
print(f"Batch size: {args.batch_size}", file=sys.stderr, flush=True)
total = 0
batch_idx = 0
for i in range(0, len(shapes), args.batch_size):
batch = shapes[i : i + args.batch_size]
batch_idx += 1
out_path = out_dir / f"wide_coverage_batch_{batch_idx:03d}.log"
print(
f"\nBatch {batch_idx}: shapes {i + 1}-{i + len(batch)} -> {out_path}",
file=sys.stderr,
flush=True,
)
with open(out_path, "w") as f:
f.write(f"CK Tile Wide Coverage Benchmark Batch {batch_idx}\n")
f.write("GPU ID: 0\n")
f.write("Implementation: gemm_universal\n\n")
count = run_shape_batch(
args.bin_dir, batch, f, warmup=args.warmup, repeat=args.repeat
)
total += count
print(
f" Batch {batch_idx} complete: {count} benchmarks",
file=sys.stderr,
flush=True,
)
print(
f"\nTotal: {total} benchmarks across {len(shapes)} shapes",
file=sys.stderr,
flush=True,
)
if __name__ == "__main__":
main()

867
dispatcher/heuristics/ml_heuristic_sweep.py Normal file
View File

@@ -0,0 +1,867 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
ML Heuristic Sweep: Comprehensive GEMM Performance Evaluation
Sweeps across diverse problem shapes with ML-based kernel selection to measure
TFLOPS performance. Supports multiple dtypes (fp16, bf16, fp8) and validates
ML model predictions by executing kernels on GPU.
Shape Constraints (fp16/bf16 on gfx950):
- M >= 1 (any M is valid)
- N % 8 == 0 AND N >= 64
- K % 2 == 0 AND K >= 32
Usage:
python ml_heuristic_sweep.py --dtype fp16 --num_shapes 256
python ml_heuristic_sweep.py --dtypes fp16 bf16 --output sweep_results.csv
python ml_heuristic_sweep.py --dtype fp16 --dry_run # Prediction only, no GPU execution
"""
import sys
import argparse
import time
import csv
from pathlib import Path
from dataclasses import dataclass
from typing import List, Tuple
# Add parent directories to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent / "python"))
import numpy as np
from ctypes_utils import (
KernelConfig,
setup_gemm_dispatcher,
cleanup_gemm,
)
try:
from predict import Predictor
# from feature_engine import GemmUniversalFeatureEngine
HAS_ML = True
except ImportError:
HAS_ML = False
print("WARNING: ML heuristic modules not available. Will use first-fit selection.")
@dataclass
class KernelSpec:
"""Kernel specification for ML heuristic"""
name: str
tile_m: int
tile_n: int
tile_k: int
pipeline: str = "compv3"
scheduler: str = "intrawave"
wave_m: int = 2
wave_n: int = 2
wave_k: int = 1
warp_m: int = 32
warp_n: int = 32
warp_k: int = 16
# Comprehensive kernel pool covering diverse tile sizes and configurations
KERNEL_POOL = [
# Small tiles (64x64)
KernelSpec(
"s_64x64_k32_v3", 64, 64, 32, "compv3", "intrawave", 2, 2, 1, 16, 16, 16
),
KernelSpec(
"s_64x64_k64_v3", 64, 64, 64, "compv3", "intrawave", 2, 2, 1, 16, 16, 16
),
KernelSpec(
"s_64x64_k128_v3", 64, 64, 128, "compv3", "intrawave", 2, 2, 1, 16, 16, 16
),
KernelSpec(
"s_64x64_k64_v4", 64, 64, 64, "compv4", "intrawave", 2, 2, 1, 16, 16, 16
),
KernelSpec("s_64x64_k64_mem", 64, 64, 64, "mem", "intrawave", 2, 2, 1, 16, 16, 16),
KernelSpec(
"s_64x64_k128_mem", 64, 64, 128, "mem", "intrawave", 2, 2, 1, 16, 16, 16
),
# Medium tiles (128x128)
KernelSpec("m_128x128_k32_v3", 128, 128, 32, "compv3", "intrawave"),
KernelSpec("m_128x128_k64_v3", 128, 128, 64, "compv3", "intrawave"),
KernelSpec("m_128x128_k128_v3", 128, 128, 128, "compv3", "intrawave"),
KernelSpec("m_128x128_k64_v4", 128, 128, 64, "compv4", "intrawave"),
KernelSpec("m_128x128_k128_v4", 128, 128, 128, "compv4", "intrawave"),
KernelSpec("m_128x128_k64_mem", 128, 128, 64, "mem", "intrawave"),
KernelSpec("m_128x128_k128_mem", 128, 128, 128, "mem", "intrawave"),
# Rectangular medium (M != N)
KernelSpec(
"r_64x128_k32_v3", 64, 128, 32, "compv3", "intrawave", 2, 2, 1, 16, 32, 16
),
KernelSpec(
"r_128x64_k32_v3", 128, 64, 32, "compv3", "intrawave", 2, 2, 1, 32, 16, 16
),
KernelSpec(
"r_64x128_k64_v3", 64, 128, 64, "compv3", "intrawave", 2, 2, 1, 16, 32, 16
),
KernelSpec(
"r_128x64_k64_v3", 128, 64, 64, "compv3", "intrawave", 2, 2, 1, 32, 16, 16
),
KernelSpec(
"r_64x256_k32_v3", 64, 256, 32, "compv3", "intrawave", 2, 2, 1, 16, 32, 16
),
KernelSpec(
"r_256x64_k32_v3", 256, 64, 32, "compv3", "intrawave", 2, 2, 1, 32, 16, 16
),
# Large tiles (256x256)
KernelSpec("l_256x128_k32_v3", 256, 128, 32, "compv3", "intrawave"),
KernelSpec("l_128x256_k32_v3", 128, 256, 32, "compv3", "intrawave"),
KernelSpec("l_256x256_k32_v3", 256, 256, 32, "compv3", "intrawave"),
KernelSpec("l_256x256_k64_v3", 256, 256, 64, "compv3", "intrawave"),
KernelSpec("l_256x256_k64_v4", 256, 256, 64, "compv4", "intrawave"),
# Interwave variants
KernelSpec("m_128x128_k64_iw_v3", 128, 128, 64, "compv3", "interwave"),
KernelSpec("m_128x128_k128_iw_v3", 128, 128, 128, "compv3", "interwave"),
KernelSpec("l_256x256_k32_iw_v3", 256, 256, 32, "compv3", "interwave"),
]
def generate_problem_shapes(num_shapes: int = 1024) -> List[Tuple[int, int, int]]:
"""
Generate diverse problem shapes with hardware constraints:
- M >= 1 (any M is valid, including tiny M for inference)
- N % 8 == 0 AND N >= 64 (hardware alignment requirement)
- K % 2 == 0 AND K >= 32 (fp16 requirement)
Covers:
- Powers of 2 (square and rectangular)
- ML workloads (LLM attention, MLP, batch inference)
- Non-power-of-2 dimensions (aligned to constraints)
- Edge cases (tiny M, very large matrices, extreme aspect ratios)
"""
shapes = []
# 1. Powers of 2 - Square (64 to 8192) with K variations
for p in range(6, 14): # 2^6=64 to 2^13=8192
dim = 2**p
shapes.append((dim, dim, dim))
if dim >= 128:
# K variations (must be even and >= 32)
shapes.append((dim, dim, dim // 2))
shapes.append((dim, dim, dim * 2))
shapes.append((dim, dim, max(32, dim // 4)))
# 2. Small batch inference (1-256 batch, common hidden dims)
# N must be multiple of 8 and >= 64
hidden_dims = [768, 1024, 2048, 3072, 4096, 5120, 8192, 11008, 12288, 16384]
batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256]
for hidden in hidden_dims:
for batch in batch_sizes[:8]:
shapes.append((batch, hidden, hidden))
if hidden >= 4096:
# LLM MLP projections (ensure K is even)
k_mlp = hidden * 3 // 4
if k_mlp % 2 == 1:
k_mlp += 1 # Make even
if k_mlp >= 32:
shapes.append((batch, hidden, k_mlp))
shapes.append((batch, k_mlp, hidden))
# 3. Attention patterns (seq_len x head_dim)
# seq_len can be any value >= 1, total_dim must be multiple of 8
seq_lens = [128, 256, 512, 1024, 2048, 4096, 8192]
head_dims = [64, 80, 96, 128, 256]
num_heads = [8, 12, 16, 32, 40, 64]
for seq in seq_lens:
for head_dim in head_dims:
for nh in num_heads[:4]:
total_dim = nh * head_dim
# total_dim should be multiple of 8 (naturally satisfied for most cases)
if total_dim % 8 == 0 and total_dim >= 64:
# head_dim must be even for K
if head_dim % 2 == 0 and head_dim >= 32:
shapes.append((seq, total_dim, head_dim))
shapes.append((seq, head_dim, total_dim))
# 4. Rectangular matrices (extreme aspect ratios)
# All dims must satisfy constraints
dims_m = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
dims_n = [64, 128, 256, 512, 1024, 2048, 4096, 8192] # N >= 64, N % 8 == 0
dims_k = [
32,
64,
128,
256,
512,
1024,
2048,
4096,
8192,
16384,
] # K >= 32, K % 2 == 0
# Sample to avoid explosion
for i, m in enumerate(dims_m):
for j, n in enumerate(dims_n):
for _l, k in enumerate(dims_k):
if (i + j + _l) % 3 == 0: # Stratified sampling
shapes.append((m, n, k))
# 5. Non-power-of-2 dimensions (aligned to constraints)
# N values: multiples of 8, >= 64
non_pow2_n = [
72,
80,
88,
96,
104,
112,
120,
136,
144,
152,
160,
176,
184,
192,
200,
224,
240,
272,
288,
304,
320,
336,
352,
368,
384,
400,
416,
448,
480,
544,
576,
640,
672,
704,
736,
768,
800,
832,
896,
960,
1088,
1152,
1216,
1280,
1344,
1408,
1472,
1536,
1600,
1664,
1728,
1792,
1856,
1920,
2176,
2304,
2432,
2560,
2688,
2816,
2944,
3072,
3200,
3328,
3456,
3584,
3712,
3840,
3968,
4224,
4352,
4480,
4608,
4736,
4864,
4992,
]
# K values: even numbers >= 32
non_pow2_k = [
34,
36,
38,
40,
42,
44,
48,
50,
52,
56,
60,
66,
68,
72,
76,
80,
88,
96,
100,
112,
120,
136,
144,
160,
176,
192,
224,
240,
272,
288,
320,
352,
384,
416,
448,
480,
544,
576,
640,
672,
704,
768,
800,
832,
896,
960,
1088,
1152,
1280,
1344,
1408,
1536,
1600,
1664,
1792,
1920,
]
# M values: any value >= 1
non_pow2_m = [
1,
3,
5,
7,
9,
11,
13,
15,
17,
19,
23,
27,
31,
33,
37,
41,
47,
51,
57,
63,
65,
71,
79,
87,
95,
97,
111,
119,
127,
129,
143,
159,
175,
191,
193,
223,
239,
255,
257,
287,
319,
351,
383,
385,
447,
479,
511,
513,
575,
639,
703,
767,
769,
895,
959,
1023,
1025,
]
# Sample non-power-of-2 shapes
for i, m in enumerate(non_pow2_m[:30]):
for j, n in enumerate(non_pow2_n[:20]):
for _l, k in enumerate(non_pow2_k[:15]):
if (i + j + _l) % 4 == 0: # Stratified sampling
shapes.append((m, n, k))
# 6. Very tall K (memory-bound) - ensure N % 8 == 0, K % 2 == 0
for mn in [64, 128, 256, 512, 1024]:
for k in [4096, 8192, 16384]:
shapes.append((mn, mn, k))
# 7. Very short K (compute-bound) - ensure K >= 32, K % 2 == 0
for mn in [512, 1024, 2048, 4096]:
for k in [32, 64, 128]:
shapes.append((mn, mn, k))
# 8. Tiny M (edge cases for batch-1 inference)
for m in [1, 2, 4, 8, 16, 32]:
for n in [64, 128, 256, 512, 1024, 2048]: # N >= 64, N % 8 == 0
for k in [32, 64, 128, 256, 512]: # K >= 32, K % 2 == 0
shapes.append((m, n, k))
# 9. Stress test sizes (aligned to constraints)
stress_sizes = [
(10000, 10000, 10000),
(1000, 10000, 1000),
(1000, 1000, 10000),
(5000, 5000, 5000),
(7168, 7168, 7168), # Common LLM hidden dim
(8192, 11008, 8192), # LLaMA MLP dimensions
]
shapes.extend(stress_sizes)
# Remove duplicates while preserving order
seen = set()
unique_shapes = []
for s in shapes:
if s not in seen:
seen.add(s)
unique_shapes.append(s)
# Filter to ensure all shapes meet constraints
valid_shapes = []
for m, n, k in unique_shapes:
if m >= 1 and n >= 64 and n % 8 == 0 and k >= 32 and k % 2 == 0:
valid_shapes.append((m, n, k))
# Sample down to target number if we have too many
if len(valid_shapes) > num_shapes:
# Stratified sampling to preserve diversity
step = len(valid_shapes) / num_shapes
valid_shapes = [valid_shapes[int(i * step)] for i in range(num_shapes)]
return valid_shapes
def spec_to_feature_dict(spec: KernelSpec, dtype: str, layout: str) -> dict:
"""Convert KernelSpec to feature dict for ML predictor"""
return {
"kernel_name": spec.name,
"tile_m": spec.tile_m,
"tile_n": spec.tile_n,
"tile_k": spec.tile_k,
"warp_m": spec.wave_m,
"warp_n": spec.wave_n,
"warp_k": spec.wave_k,
"warp_tile_m": spec.warp_m,
"warp_tile_n": spec.warp_n,
"warp_tile_k": spec.warp_k,
"pipeline": spec.pipeline,
"scheduler": spec.scheduler,
"epilogue": "cshuffle",
"pad_m": True, # Enable padding to support arbitrary M dimensions
"pad_n": True, # Enable padding to support arbitrary N dimensions
"pad_k": True, # Enable padding to support arbitrary K dimensions
"persistent": False,
"dtype": dtype,
"layout": layout,
}
def spec_to_kernel_config(
spec: KernelSpec, dtype: str, arch: str, dtype_acc: str = "fp32"
) -> KernelConfig:
"""Convert KernelSpec to KernelConfig for dispatcher"""
return KernelConfig(
dtype_a=dtype,
dtype_b=dtype,
dtype_c=dtype,
dtype_acc=dtype_acc,
layout_a="row",
layout_b="col",
layout_c="row",
tile_m=spec.tile_m,
tile_n=spec.tile_n,
tile_k=spec.tile_k,
wave_m=spec.wave_m,
wave_n=spec.wave_n,
wave_k=spec.wave_k,
warp_m=spec.warp_m,
warp_n=spec.warp_n,
warp_k=spec.warp_k,
pipeline=spec.pipeline,
scheduler=spec.scheduler,
epilogue="cshuffle",
gfx_arch=arch,
)
def ml_select_kernel(
predictor, pool: List[KernelSpec], M: int, N: int, K: int, dtype: str, layout: str
) -> Tuple[KernelSpec, float]:
"""Use ML model to select best kernel"""
if not HAS_ML or predictor is None:
# Fallback: select first kernel
return pool[0], 0.0
problem = {"m": M, "n": N, "k": K, "dtype": dtype, "layout": layout, "split_k": 1}
kernel_dicts = [spec_to_feature_dict(s, dtype, layout) for s in pool]
ranked = predictor.rank_kernels(problem, kernel_dicts)
if not ranked:
return pool[0], 0.0
best_name, best_tflops = ranked[0]
best_spec = next((s for s in pool if s.name == best_name), pool[0])
return best_spec, best_tflops
def run_single_gemm(
M: int,
N: int,
K: int,
dtype: str,
arch: str,
predictor,
dry_run: bool = False,
dtype_acc: str = "fp32",
) -> dict:
"""Run a single GEMM with ML heuristic selection"""
# Select kernel via ML heuristic
t0 = time.time()
best_spec, pred_tflops = ml_select_kernel(
predictor, KERNEL_POOL, M, N, K, dtype, "rcr"
)
select_time_ms = (time.time() - t0) * 1000
result = {
"M": M,
"N": N,
"K": K,
"dtype": dtype,
"selected_kernel": best_spec.name,
"predicted_tflops": pred_tflops,
"selection_time_ms": select_time_ms,
"actual_time_ms": 0,
"actual_tflops": 0,
"status": "SKIP" if dry_run else "PENDING",
"error": None,
}
if dry_run:
return result
# Build and run kernel
config = spec_to_kernel_config(best_spec, dtype, arch, dtype_acc)
try:
setup = setup_gemm_dispatcher(
config=config,
registry_name=f"sweep_{dtype}_{best_spec.name}",
verbose=False,
auto_rebuild=True,
)
if not setup.success:
result["status"] = "BUILD_FAIL"
result["error"] = "Failed to build kernel"
cleanup_gemm()
return result
dispatcher = setup.dispatcher
if not dispatcher.is_supported(M, N, K):
result["status"] = "UNSUPPORTED"
result["error"] = "Problem size not supported by kernel"
cleanup_gemm()
return result
# Create input data
np_dtype = {"fp16": np.float16, "bf16": np.float16, "fp8": np.float16}[dtype]
np.random.seed(42)
A = (np.random.randn(M, K) * 0.1).astype(np_dtype)
B = (np.random.randn(K, N) * 0.1).astype(np_dtype)
# Run GEMM
exec_result = dispatcher.run(A, B, M, N, K)
if exec_result.success:
result["actual_time_ms"] = exec_result.time_ms
result["actual_tflops"] = exec_result.tflops
result["status"] = "SUCCESS"
else:
# Decode status code for better error message
status_messages = {
0: "Success",
-1: "GPU/HIP error (check permissions, memory, or kernel validity)",
-2: "No suitable kernel found for this problem size",
}
error_msg = status_messages.get(exec_result.status, f"Unknown error (status={exec_result.status})")
result["status"] = "RUN_FAIL"
result["error"] = f"{error_msg} (status_code={exec_result.status})"
# Print detailed error for debugging
print(f" ERROR: {error_msg}")
print(f" Status code: {exec_result.status}")
print(f" Time returned: {exec_result.time_ms}")
print(f" Kernel: {exec_result.kernel_name}")
cleanup_gemm()
except Exception as e:
result["status"] = "ERROR"
result["error"] = str(e)[:200]
cleanup_gemm()
return result
def main():
parser = argparse.ArgumentParser(
description="ML Heuristic Sweep: Test GEMM across many shapes and dtypes"
)
parser.add_argument(
"--dtypes",
nargs="+",
default=["fp16"],
choices=["fp16", "bf16", "fp8"],
help="Data types to test (default: fp16)",
)
parser.add_argument(
"--arch", default="gfx950", help="GPU architecture (default: gfx950)"
)
parser.add_argument(
"--dtype_acc",
default="fp32",
choices=["fp16", "fp32"],
help="Accumulator data type (default: fp32)",
)
parser.add_argument(
"--model_dir",
default=None,
help="Path to model directory (auto-detect if not specified)",
)
parser.add_argument(
"--num_shapes",
type=int,
default=256,
help="Number of problem shapes to test (default: 256)",
)
parser.add_argument(
"--output",
default="ml_heuristic_sweep_results.csv",
help="Output CSV file path",
)
parser.add_argument(
"--dry_run",
action="store_true",
help="Only predict, do not run kernels (fast validation)",
)
args = parser.parse_args()
# Setup ML predictor
predictor = None
if HAS_ML:
if args.model_dir is None:
# Auto-detect model directory based on first dtype
first_dtype = args.dtypes[0]
heuristics_dir = Path(__file__).parent
model_candidates = [
heuristics_dir / "models" / f"gemm_universal_{first_dtype}_{args.arch}",
]
for model_dir in model_candidates:
if model_dir.exists():
args.model_dir = str(model_dir)
break
if args.model_dir and Path(args.model_dir).exists():
try:
predictor = Predictor(args.model_dir)
print(f"✓ Loaded ML model from: {args.model_dir}")
except Exception as e:
print(f"⚠ Failed to load ML model: {e}")
print(" Will use first-fit selection instead")
else:
print(f"⚠ Model directory not found: {args.model_dir}")
print(" Will use first-fit selection instead")
# Generate problem shapes
print(f"\nGenerating {args.num_shapes} problem shapes...")
shapes = generate_problem_shapes(args.num_shapes)
print(
f"✓ Generated {len(shapes)} valid shapes (M>=1, N%8==0, N>=64, K%2==0, K>=32)"
)
# Validate all shapes meet constraints
invalid = [
(m, n, k)
for m, n, k in shapes
if not (m >= 1 and n >= 64 and n % 8 == 0 and k >= 32 and k % 2 == 0)
]
if invalid:
print(f"⚠ WARNING: {len(invalid)} shapes violate constraints!")
print(f" First few: {invalid[:5]}")
# Print configuration
print("\n" + "=" * 80)
print(" ML Heuristic Sweep Configuration")
print("=" * 80)
print(
f" Model: {args.model_dir if args.model_dir else 'first-fit (no ML)'}"
)
print(f" Data types: {', '.join(args.dtypes)}")
print(f" Accumulator: {args.dtype_acc}")
print(f" Architecture: {args.arch}")
print(f" Kernel pool: {len(KERNEL_POOL)} kernels")
print(f" Problem shapes: {len(shapes)}")
print(f" Total tests: {len(shapes) * len(args.dtypes)}")
print(
f" Mode: {'DRY RUN (prediction only)' if args.dry_run else 'FULL RUN (execute kernels)'}"
)
print(f" Output: {args.output}")
print("=" * 80)
# Open output CSV
csv_file = open(args.output, "w", newline="")
csv_writer = csv.DictWriter(
csv_file,
fieldnames=[
"dtype",
"M",
"N",
"K",
"selected_kernel",
"predicted_tflops",
"selection_time_ms",
"actual_time_ms",
"actual_tflops",
"status",
"error",
],
)
csv_writer.writeheader()
# Run sweep
total_tests = len(shapes) * len(args.dtypes)
completed = 0
start_time = time.time()
print("\nStarting sweep... (Ctrl+C to stop and save partial results)\n")
try:
for dtype in args.dtypes:
print(f"\n{'=' * 80}")
print(f" Testing dtype: {dtype.upper()}")
print(f"{'=' * 80}\n")
for i, (M, N, K) in enumerate(shapes):
result = run_single_gemm(
M, N, K, dtype, args.arch, predictor, args.dry_run, args.dtype_acc
)
# Write to CSV
csv_writer.writerow(result)
csv_file.flush()
completed += 1
# Progress update
if completed % 10 == 0 or result["status"] != "SUCCESS":
elapsed = time.time() - start_time
rate = completed / elapsed if elapsed > 0 else 0
eta = (total_tests - completed) / rate if rate > 0 else 0
status_emoji = {
"SUCCESS": "✓",
"SKIP": "→",
"BUILD_FAIL": "✗",
"UNSUPPORTED": "○",
"RUN_FAIL": "✗",
"ERROR": "✗",
}.get(result["status"], "?")
print(
f" [{completed:4d}/{total_tests}] {status_emoji} "
f"{dtype:4s} {M:5d}x{N:5d}x{K:5d} → "
f"{result['selected_kernel']:20s} "
f"pred={result['predicted_tflops']:6.1f} "
f"actual={result['actual_tflops']:6.1f} TFLOPS "
f"[{rate:.1f} tests/s, ETA {eta / 60:.1f}m]"
)
except KeyboardInterrupt:
print(f"\n\n⚠ Interrupted! Saving partial results to {args.output}...")
finally:
csv_file.close()
# Summary
print("\n" + "=" * 80)
print(" SWEEP COMPLETE")
print("=" * 80)
# Read back results and compute statistics
results = []
with open(args.output, "r") as f:
reader = csv.DictReader(f)
results = list(reader)
print(f"\n Total tests: {len(results)}")
print(f" Output file: {args.output}")
if not args.dry_run:
success = [r for r in results if r["status"] == "SUCCESS"]
print(
f" Successful: {len(success)} ({100 * len(success) / len(results):.1f}%)"
)
if success:
avg_tflops = np.mean([float(r["actual_tflops"]) for r in success])
max_tflops = max([float(r["actual_tflops"]) for r in success])
print(f" Avg TFLOPS: {avg_tflops:.2f}")
print(f" Max TFLOPS: {max_tflops:.2f}")
# Per-dtype breakdown
for dtype in args.dtypes:
dtype_results = [r for r in success if r["dtype"] == dtype]
if dtype_results:
avg = np.mean([float(r["actual_tflops"]) for r in dtype_results])
print(
f" {dtype:4s}: {avg:.2f} TFLOPS (n={len(dtype_results)})"
)
print("=" * 80)
print()
return 0
if __name__ == "__main__":
sys.exit(main())

113
dispatcher/heuristics/models/gemm_universal_fp16_gfx950/feature_spec.json Normal file
View File

@@ -0,0 +1,113 @@
{
"op_type": "gemm_universal",
"dtype": "fp16",
"arch": "gfx950",
"feature_names": [
"M",
"N",
"K",
"split_k",
"log2_M",
"log2_N",
"log2_K",
"log2_MNK",
"arithmetic_intensity",
"aspect_ratio_mn",
"aspect_ratio_mk",
"aspect_ratio_nk",
"layout",
"tile_m",
"tile_n",
"tile_k",
"warp_m",
"warp_n",
"warp_k",
"warp_tile_m",
"warp_tile_n",
"warp_tile_k",
"pipeline",
"scheduler",
"epilogue",
"pad_m",
"pad_n",
"pad_k",
"persistent",
"num_warps",
"tile_volume",
"tile_mn",
"lds_usage_estimate",
"lds_usage_ratio",
"num_tiles_m",
"num_tiles_n",
"num_tiles_k",
"total_output_tiles",
"tile_eff_m",
"tile_eff_n",
"tile_eff_k",
"overall_tile_efficiency",
"cu_utilization",
"ratio_M_to_tile_m",
"ratio_N_to_tile_n",
"ratio_K_to_tile_k",
"problem_smaller_than_tile_m",
"problem_smaller_than_tile_n",
"problem_smaller_than_tile_k",
"any_dim_too_small",
"needs_padding_m",
"needs_padding_n",
"needs_padding_k",
"has_padding_when_needed_m",
"has_padding_when_needed_n",
"has_padding_when_needed_k",
"missing_required_padding_m",
"missing_required_padding_n",
"missing_required_padding_k",
"missing_any_required_padding",
"hw_num_cus",
"hw_simds_per_cu",
"hw_total_simds",
"hw_shader_engines",
"hw_max_clock_mhz",
"hw_max_waves_per_cu",
"hw_wavefront_size",
"hw_lds_capacity",
"hw_l1_cache_kb",
"hw_l2_cache_kb",
"hw_l3_cache_kb",
"hw_num_xcd"
],
"categorical_features": [
"layout",
"pipeline",
"scheduler",
"epilogue"
],
"targets": [
"tflops",
"latency",
"bandwidth"
],
"log_targets": [
"bandwidth",
"tflops"
],
"params": {
"objective": "regression",
"metric": [
"rmse",
"mae"
],
"num_leaves": 255,
"max_depth": 15,
"n_estimators": 2000,
"learning_rate": 0.02,
"min_child_samples": 10,
"subsample": 0.85,
"colsample_bytree": 0.85,
"reg_alpha": 0.05,
"reg_lambda": 0.5,
"verbose": -1,
"n_jobs": 8,
"seed": 42
}
}

BIN
dispatcher/heuristics/models/gemm_universal_fp16_gfx950/model_tflops.lgbm.gz Normal file
View File

Binary file not shown.

10
dispatcher/heuristics/models/gemm_universal_fp16_gfx950/train_manifest.json Normal file
View File

@@ -0,0 +1,10 @@
{
"warm_start_from": null,
"prev_n_estimators": 0,
"new_n_estimators": 2000,
"total_n_estimators": 2000,
"data_rows": 25600,
"valid_rows": 21920,
"unique_shapes": 25,
"timestamp": "2026-03-20T05:00:55"
}

113
dispatcher/heuristics/models/gemm_universal_fp8_gfx950/feature_spec.json Normal file
View File

@@ -0,0 +1,113 @@
{
"op_type": "gemm_universal",
"dtype": "fp8",
"arch": "gfx950",
"feature_names": [
"M",
"N",
"K",
"split_k",
"log2_M",
"log2_N",
"log2_K",
"log2_MNK",
"arithmetic_intensity",
"aspect_ratio_mn",
"aspect_ratio_mk",
"aspect_ratio_nk",
"layout",
"tile_m",
"tile_n",
"tile_k",
"warp_m",
"warp_n",
"warp_k",
"warp_tile_m",
"warp_tile_n",
"warp_tile_k",
"pipeline",
"scheduler",
"epilogue",
"pad_m",
"pad_n",
"pad_k",
"persistent",
"num_warps",
"tile_volume",
"tile_mn",
"lds_usage_estimate",
"lds_usage_ratio",
"num_tiles_m",
"num_tiles_n",
"num_tiles_k",
"total_output_tiles",
"tile_eff_m",
"tile_eff_n",
"tile_eff_k",
"overall_tile_efficiency",
"cu_utilization",
"ratio_M_to_tile_m",
"ratio_N_to_tile_n",
"ratio_K_to_tile_k",
"problem_smaller_than_tile_m",
"problem_smaller_than_tile_n",
"problem_smaller_than_tile_k",
"any_dim_too_small",
"needs_padding_m",
"needs_padding_n",
"needs_padding_k",
"has_padding_when_needed_m",
"has_padding_when_needed_n",
"has_padding_when_needed_k",
"missing_required_padding_m",
"missing_required_padding_n",
"missing_required_padding_k",
"missing_any_required_padding",
"hw_num_cus",
"hw_simds_per_cu",
"hw_total_simds",
"hw_shader_engines",
"hw_max_clock_mhz",
"hw_max_waves_per_cu",
"hw_wavefront_size",
"hw_lds_capacity",
"hw_l1_cache_kb",
"hw_l2_cache_kb",
"hw_l3_cache_kb",
"hw_num_xcd"
],
"categorical_features": [
"layout",
"pipeline",
"scheduler",
"epilogue"
],
"targets": [
"tflops",
"latency",
"bandwidth"
],
"log_targets": [
"bandwidth",
"tflops"
],
"params": {
"objective": "regression",
"metric": [
"rmse",
"mae"
],
"num_leaves": 255,
"max_depth": 15,
"n_estimators": 2000,
"learning_rate": 0.02,
"min_child_samples": 10,
"subsample": 0.85,
"colsample_bytree": 0.85,
"reg_alpha": 0.05,
"reg_lambda": 0.5,
"verbose": -1,
"n_jobs": 8,
"seed": 42
}
}

BIN
dispatcher/heuristics/models/gemm_universal_fp8_gfx950/model_tflops.lgbm.gz Normal file
View File

Binary file not shown.

10
dispatcher/heuristics/models/gemm_universal_fp8_gfx950/train_manifest.json Normal file
View File

@@ -0,0 +1,10 @@
{
"warm_start_from": null,
"prev_n_estimators": 0,
"new_n_estimators": 2000,
"total_n_estimators": 2000,
"data_rows": 1296528,
"valid_rows": 1253076,
"unique_shapes": 168,
"timestamp": "2026-03-19T06:10:29"
}

243
dispatcher/heuristics/predict.py Normal file
View File

@@ -0,0 +1,243 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Predictor for CK Tile kernel performance.
Loads trained LightGBM models and provides:
- predict_tflops(): predicted TFLOPS for a single (problem, kernel) pair
- predict_latency(): predicted latency in ms
- predict_bandwidth(): predicted bandwidth in GB/s
- predict_all(): all three predictions at once
- rank_kernels(): rank all candidate kernels by predicted TFLOPS
- select_best(): return the best kernel ID
Usage:
predictor = Predictor("models/gemm_universal_fp8_gfx950")
best_kernel = predictor.select_best(
problem={"m": 128, "n": 1536, "k": 7168, "dtype": "fp8", "layout": "rcr"},
kernel_configs=[...],
)
"""
import gzip
import json
from pathlib import Path
from typing import Optional
import lightgbm as lgb
import numpy as np
import pandas as pd
from feature_engine import GemmUniversalFeatureEngine
class Predictor:
"""Loads trained models and feature spec for kernel performance prediction.
Parameters
----------
model_dir : str or Path
Directory containing model artifacts:
- model_tflops.lgbm (required)
- model_latency.lgbm (optional)
- model_bandwidth.lgbm (optional)
- feature_spec.json (required)
feature_engine : FeatureEngine, optional
Override the feature engine. If None, constructs one from feature_spec.json.
"""
def __init__(self, model_dir: str | Path, feature_engine=None):
self._model_dir = Path(model_dir)
self._models: dict[str, lgb.Booster] = {}
spec_path = self._model_dir / "feature_spec.json"
if spec_path.exists():
with open(spec_path) as f:
self._spec = json.load(f)
else:
self._spec = {}
self._log_targets = set(self._spec.get("log_targets", []))
if feature_engine is not None:
self._feature_engine = feature_engine
else:
self._feature_engine = GemmUniversalFeatureEngine()
def _load_model(self, target: str) -> Optional[lgb.Booster]:
"""Lazy-load a model for the given target.
Automatically decompresses .lgbm.gz files if the .lgbm file doesn't exist.
The decompressed file is cached to disk for subsequent loads.
"""
if target in self._models:
return self._models[target]
path = self._model_dir / f"model_{target}.lgbm"
gz_path = self._model_dir / f"model_{target}.lgbm.gz"
# Auto-decompress if needed
if not path.exists() and gz_path.exists():
with gzip.open(gz_path, 'rb') as f_in:
with open(path, 'wb') as f_out:
f_out.write(f_in.read())
if not path.exists():
return None
model = lgb.Booster(model_file=str(path))
self._models[target] = model
return model
def _predict_single(self, target: str, problem: dict, kernel_config: dict) -> float:
"""Predict a single target value, applying inverse log transform if needed."""
model = self._load_model(target)
if model is None:
raise FileNotFoundError(f"No model_{target}.lgbm in {self._model_dir}")
features = self._feature_engine.extract(problem, kernel_config)
raw = float(model.predict(features.reshape(1, -1))[0])
if target in self._log_targets:
return float(np.expm1(raw))
# Clamp to non-negative even for non-log models
return float(max(0.0, raw))
def predict_tflops(self, problem: dict, kernel_config: dict) -> float:
"""Predict TFLOPS for a single (problem, kernel) pair.
Returns a real TFLOPS estimate (interpretable, usable as DE surrogate).
If the model was trained in log-space, the inverse transform is applied
automatically.
"""
return self._predict_single("tflops", problem, kernel_config)
def predict_latency(self, problem: dict, kernel_config: dict) -> float:
"""Predict latency in milliseconds for a single (problem, kernel) pair."""
return self._predict_single("latency", problem, kernel_config)
def predict_bandwidth(self, problem: dict, kernel_config: dict) -> float:
"""Predict bandwidth in GB/s for a single (problem, kernel) pair."""
return self._predict_single("bandwidth", problem, kernel_config)
def predict_all(self, problem: dict, kernel_config: dict) -> dict[str, float]:
"""Predict all available targets for a single (problem, kernel) pair.
Returns dict with keys 'tflops', 'latency_ms', 'bandwidth_gb_s' (if models exist).
Note: Applies inverse log transform for targets in log_targets and clamps
negatives to 0.0, consistent with _predict_single().
"""
features = self._feature_engine.extract(problem, kernel_config).reshape(1, -1)
result = {}
for target, key in [
("tflops", "tflops"),
("latency", "latency_ms"),
("bandwidth", "bandwidth_gb_s"),
]:
model = self._load_model(target)
if model is not None:
raw = float(model.predict(features)[0])
# Apply inverse log transform if model was trained in log-space
if target in self._log_targets:
result[key] = float(np.expm1(raw))
else:
# Clamp to non-negative even for non-log models
result[key] = float(max(0.0, raw))
return result
def rank_kernels(
self, problem: dict, kernel_configs: list[dict]
) -> list[tuple[str, float]]:
"""Rank candidate kernels by predicted TFLOPS (descending).
Parameters
----------
problem : dict
Problem specification with keys: m, n, k, dtype, layout, split_k.
kernel_configs : list of dict
Each dict must have a 'kernel_name' key plus kernel parameters.
Returns
-------
list of (kernel_name, predicted_tflops) tuples, sorted descending.
"""
if not kernel_configs:
return []
model = self._load_model("tflops")
if model is None:
raise FileNotFoundError(f"No model_tflops.lgbm in {self._model_dir}")
rows = []
for kc in kernel_configs:
merged = {**problem, **kc}
rows.append(merged)
df = pd.DataFrame(rows)
X = self._feature_engine.extract_batch(df)
preds = model.predict(X)
if "tflops" in self._log_targets:
preds = np.expm1(preds)
results = []
for i, kc in enumerate(kernel_configs):
name = kc.get("kernel_name", f"kernel_{i}")
results.append((name, float(preds[i])))
results.sort(key=lambda x: -x[1])
return results
def select_best(self, problem: dict, kernel_configs: list[dict]) -> str:
"""Return the kernel_name of the best predicted kernel."""
ranked = self.rank_kernels(problem, kernel_configs)
if not ranked:
raise ValueError("No kernel configs provided")
return ranked[0][0]
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Predict kernel performance")
parser.add_argument(
"--model_dir", required=True, help="Directory with trained models"
)
parser.add_argument("--m", type=int, required=True)
parser.add_argument("--n", type=int, required=True)
parser.add_argument("--k", type=int, required=True)
parser.add_argument("--layout", default="rcr")
parser.add_argument("--dtype", default="fp8")
args = parser.parse_args()
predictor = Predictor(args.model_dir)
problem = {
"m": args.m,
"n": args.n,
"k": args.k,
"dtype": args.dtype,
"layout": args.layout,
"split_k": 1,
}
print(f"Loading models from {args.model_dir}...")
print(
f"Problem: M={args.m} N={args.n} K={args.k} dtype={args.dtype} layout={args.layout}"
)
data_dir = Path(args.model_dir).parent.parent / "data"
if data_dir.exists():
for pq in data_dir.glob("*.parquet"):
df = pd.read_parquet(pq)
kernel_names = df["kernel_name"].unique()
configs = []
for kn in kernel_names[:10]:
row = df[df["kernel_name"] == kn].iloc[0]
configs.append(row.to_dict())
if configs:
ranked = predictor.rank_kernels(problem, configs)
print(f"\nTop 5 kernels (from {len(configs)} candidates):")
for name, tflops in ranked[:5]:
print(f" {tflops:8.2f} TFLOPS {name}")
break

272
dispatcher/heuristics/search.py Normal file
View File

@@ -0,0 +1,272 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Surrogate search for CK Tile kernel configuration optimization.
Uses a trained LGBMRegressor as a cheap surrogate function to search the
discrete kernel parameter space (tile sizes, warp config, pipeline, etc.)
without running actual GPU benchmarks.
Strategies:
- 'random': Sample N random valid configs, score all, return top-K.
- 'de': Discrete Differential Evolution with mutation over valid parameter choices.
Usage:
from search import SurrogateSearch
from predict import Predictor
predictor = Predictor("models/gemm_universal_fp8_gfx950")
searcher = SurrogateSearch(predictor, strategy='random')
results = searcher.search(
problem={"m": 128, "n": 1536, "k": 7168, "dtype": "fp8", "layout": "rcr"},
budget=500,
)
# results: [(config_dict, predicted_tflops), ...] sorted descending
"""
import random
from typing import Optional
import numpy as np
from feature_engine import GemmUniversalFeatureEngine
from predict import Predictor
class SurrogateSearch:
"""Search kernel parameter space using ML regressor as surrogate objective.
Parameters
----------
predictor : Predictor
Trained predictor with a TFLOPS model.
feature_engine : GemmUniversalFeatureEngine, optional
Feature engine for parameter space and validation. If None, uses default.
strategy : str
Search strategy: 'random' or 'de' (Discrete Differential Evolution).
seed : int
Random seed for reproducibility.
"""
def __init__(
self,
predictor: Predictor,
feature_engine: Optional[GemmUniversalFeatureEngine] = None,
strategy: str = "random",
seed: int = 42,
):
self._predictor = predictor
self._fe = feature_engine or GemmUniversalFeatureEngine()
self._strategy = strategy
self._rng = random.Random(seed)
self._np_rng = np.random.RandomState(seed)
self._param_space = self._fe.get_parameter_space()
def _sample_random_config(self) -> dict:
"""Sample a single random config from the parameter space."""
config = {}
for param, values in self._param_space.items():
config[param] = self._rng.choice(values)
return config
def _sample_valid_config(self, max_attempts: int = 50) -> Optional[dict]:
"""Sample a random config that passes all validation constraints."""
for _ in range(max_attempts):
config = self._sample_random_config()
if self._fe.validate_config(config):
return config
return None
def _score_config(self, problem: dict, config: dict) -> float:
"""Score a config using the predictor."""
return self._predictor.predict_tflops(problem, config)
def _search_random(
self, problem: dict, budget: int, top_k: int
) -> list[tuple[dict, float]]:
"""Random search: sample valid configs, score all, return top-K."""
configs = []
for _ in range(budget):
cfg = self._sample_valid_config()
if cfg is not None:
configs.append(cfg)
if not configs:
return []
scored = []
for cfg in configs:
try:
score = self._score_config(problem, cfg)
scored.append((cfg, score))
except Exception:
continue
scored.sort(key=lambda x: -x[1])
return scored[:top_k]
def _search_de(
self,
problem: dict,
budget: int,
top_k: int,
pop_size: int = 20,
mutation_rate: float = 0.3,
crossover_rate: float = 0.7,
) -> list[tuple[dict, float]]:
"""Discrete Differential Evolution.
Uses discrete mutation: randomly swap parameters to other valid values
from the parameter space (no continuous relaxation + snap).
Each generation:
1. For each member of the population, create a trial vector by:
- Selecting 3 random other members (a, b, c)
- For each parameter, with probability mutation_rate, take the value
from a, b, or c (uniform choice among the three donors)
- With probability crossover_rate, take the trial value; otherwise keep original
2. Validate the trial; if invalid, resample that parameter from the space
3. Score the trial; if better than parent, replace
"""
param_names = list(self._param_space.keys())
population = []
for _ in range(pop_size):
cfg = self._sample_valid_config()
if cfg is not None:
score = self._score_config(problem, cfg)
population.append((cfg, score))
if len(population) < 4:
return self._search_random(problem, budget, top_k)
evals_used = len(population)
max_gens = (budget - evals_used) // pop_size
for gen in range(max_gens):
new_pop = []
for i, (parent, parent_score) in enumerate(population):
candidates = [j for j in range(len(population)) if j != i]
if len(candidates) < 3:
new_pop.append((parent, parent_score))
continue
a_idx, b_idx, c_idx = self._rng.sample(candidates, 3)
a, b, c = (
population[a_idx][0],
population[b_idx][0],
population[c_idx][0],
)
trial = dict(parent)
for param in param_names:
if self._rng.random() < mutation_rate:
donor = self._rng.choice([a, b, c])
trial[param] = donor.get(param, parent.get(param))
if self._rng.random() > crossover_rate:
trial[param] = parent.get(param)
if not self._fe.validate_config(trial):
for param in param_names:
if param in trial and trial[param] not in self._param_space.get(
param, [trial[param]]
):
trial[param] = self._rng.choice(self._param_space[param])
if not self._fe.validate_config(trial):
new_pop.append((parent, parent_score))
continue
try:
trial_score = self._score_config(problem, trial)
evals_used += 1
except Exception:
new_pop.append((parent, parent_score))
continue
if trial_score > parent_score:
new_pop.append((trial, trial_score))
else:
new_pop.append((parent, parent_score))
population = new_pop
population.sort(key=lambda x: -x[1])
return population[:top_k]
def search(
self,
problem: dict,
budget: int = 500,
top_k: int = 10,
**kwargs,
) -> list[tuple[dict, float]]:
"""Search the kernel parameter space for the best configuration.
Parameters
----------
problem : dict
Problem specification: m, n, k, dtype, layout, split_k.
budget : int
Maximum number of surrogate evaluations.
top_k : int
Number of top configurations to return.
**kwargs
Strategy-specific parameters (pop_size, mutation_rate, etc.).
Returns
-------
list of (config_dict, predicted_tflops), sorted descending by TFLOPS.
"""
if self._strategy == "random":
return self._search_random(problem, budget, top_k)
elif self._strategy == "de":
return self._search_de(problem, budget, top_k, **kwargs)
else:
raise ValueError(f"Unknown strategy: {self._strategy}")
if __name__ == "__main__":
import argparse
import time
parser = argparse.ArgumentParser(
description="Surrogate search for optimal kernel config"
)
parser.add_argument("--model_dir", required=True)
parser.add_argument("--m", type=int, required=True)
parser.add_argument("--n", type=int, required=True)
parser.add_argument("--k", type=int, required=True)
parser.add_argument("--dtype", default="fp8")
parser.add_argument("--layout", default="rcr")
parser.add_argument("--strategy", default="random", choices=["random", "de"])
parser.add_argument("--budget", type=int, default=500)
parser.add_argument("--top_k", type=int, default=10)
args = parser.parse_args()
predictor = Predictor(args.model_dir)
searcher = SurrogateSearch(predictor, strategy=args.strategy)
problem = {
"m": args.m,
"n": args.n,
"k": args.k,
"dtype": args.dtype,
"layout": args.layout,
"split_k": 1,
}
print(f"Searching with strategy={args.strategy}, budget={args.budget}...")
t0 = time.time()
results = searcher.search(problem, budget=args.budget, top_k=args.top_k)
elapsed = time.time() - t0
print(f"\nTop {len(results)} configs found in {elapsed * 1000:.1f}ms:")
for i, (cfg, tflops) in enumerate(results):
tile_str = f"{cfg.get('tile_m', '?')}x{cfg.get('tile_n', '?')}x{cfg.get('tile_k', '?')}"
warp_str = f"{cfg.get('warp_m', '?')}x{cfg.get('warp_n', '?')}x{cfg.get('warp_k', '?')}"
print(
f" #{i + 1}: {tflops:8.2f} TFLOPS tile={tile_str} warp={warp_str} "
f"pipeline={cfg.get('pipeline', '?')} scheduler={cfg.get('scheduler', '?')}"
)

2
dispatcher/heuristics/tests/__init__.py Normal file
View File

@@ -0,0 +1,2 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

368
dispatcher/heuristics/tests/test_data_pipeline.py Normal file
View File

@@ -0,0 +1,368 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for data_pipeline.py.
Covers: kernel name parsing, layout derivation, streaming log parsing,
schema validation, and corner cases (empty logs, malformed JSON, single-shape).
"""
import tempfile
from pathlib import Path
import pandas as pd
import pytest
import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from data_pipeline import (
parse_kernel_name,
_layout_from_problem,
parse_streaming_log,
save_parquet,
load_parquet,
CANONICAL_COLUMNS,
)
# ---------------------------------------------------------------------------
# parse_kernel_name
# ---------------------------------------------------------------------------
class TestParseKernelName:
def test_standard_name(self):
name = "gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128"
result = parse_kernel_name(name)
assert result["dtype"] == "fp8"
assert result["layout"] == "rcr"
assert result["pipeline"] == "compv3"
assert result["epilogue"] == "cshuffle"
assert result["scheduler"] == "intrawave"
assert result["pad_m"] is False
assert result["pad_n"] is False
assert result["pad_k"] is False
assert result["persistent"] is False
assert result["tile_m"] == 128
assert result["tile_n"] == 128
assert result["tile_k"] == 128
assert result["warp_m"] == 1
assert result["warp_n"] == 4
assert result["warp_k"] == 1
assert result["warp_tile_m"] == 16
assert result["warp_tile_n"] == 16
assert result["warp_tile_k"] == 128
def test_with_padding_and_persistent(self):
name = "gemm_universal_fp16_rrr_compv4_default_interwave_True_True_True_True_256x256x64_2x2x1_32x32x16"
result = parse_kernel_name(name)
assert result["dtype"] == "fp16"
assert result["layout"] == "rrr"
assert result["pad_m"] is True
assert result["pad_n"] is True
assert result["pad_k"] is True
assert result["persistent"] is True
assert result["tile_m"] == 256
def test_empty_name(self):
assert parse_kernel_name("") == {}
def test_malformed_name(self):
assert parse_kernel_name("not_a_kernel_name") == {}
def test_partial_name(self):
result = parse_kernel_name("gemm_universal_fp8_rcr_compv3")
assert result.get("dtype") == "fp8"
assert result.get("layout") == "rcr"
assert "tile_m" not in result # not enough parts
def test_all_layouts(self):
for layout in ["rcr", "rrr", "crr", "ccr"]:
name = f"gemm_universal_fp8_{layout}_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128"
result = parse_kernel_name(name)
assert result["layout"] == layout
# ---------------------------------------------------------------------------
# _layout_from_problem
# ---------------------------------------------------------------------------
class TestLayoutFromProblem:
def test_rcr(self):
assert (
_layout_from_problem(
{
"layout_a": "RowMajor",
"layout_b": "ColumnMajor",
"layout_c": "RowMajor",
}
)
== "rcr"
)
def test_rrr(self):
assert (
_layout_from_problem(
{"layout_a": "RowMajor", "layout_b": "RowMajor", "layout_c": "RowMajor"}
)
== "rrr"
)
def test_empty(self):
assert _layout_from_problem({}) == "???"
def test_case_insensitive(self):
assert (
_layout_from_problem(
{
"layout_a": "rowmajor",
"layout_b": "COLUMNMAJOR",
"layout_c": "RowMajor",
}
)
== "rcr"
)
# ---------------------------------------------------------------------------
# parse_streaming_log
# ---------------------------------------------------------------------------
SAMPLE_LOG = """\
================================================================================
LOG FILE: test.log
================================================================================
CK Tile Profiling Run
GPU ID: 0
--- Running CK Tile benchmarks on GPU 0 ---
========================================
Shape 1: M=16 N=1536 K=7168 dtype=fp8 layout=rcr
========================================
Found 2 kernels
{
"name": "gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128",
"problem": {
"split_k":1,
"m":16,
"n":1536,
"k":7168,
"stride_a":7168,
"stride_b":7168,
"stride_c":1536,
"dtype_a":"fp8",
"dtype_b":"fp8",
"dtype_acc":"fp32",
"dtype_c":"fp16",
"layout_a":"RowMajor",
"layout_b":"ColumnMajor",
"layout_c":"RowMajor",
"structured_sparsity":false
},
"perf_result": {
"latency(ms)": 0.04,
"tflops(TFlops)": 8.81,
"bandwidth(GB/s)": 279.51
}
}
{
"name": "gemm_universal_fp8_rcr_compv4_default_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16",
"problem": {
"split_k":1,
"m":16,
"n":1536,
"k":7168,
"stride_a":7168,
"stride_b":7168,
"stride_c":1536,
"dtype_a":"fp8",
"dtype_b":"fp8",
"dtype_acc":"fp32",
"dtype_c":"fp16",
"layout_a":"RowMajor",
"layout_b":"ColumnMajor",
"layout_c":"RowMajor",
"structured_sparsity":false
},
"perf_result": {
"latency(ms)": 0.05,
"tflops(TFlops)": 7.22,
"bandwidth(GB/s)": 228.85
}
}
========================================
Shape 2: M=20480 N=7168 K=256 dtype=fp8 layout=rcr
========================================
Found 1 kernels
{
"name": "gemm_universal_fp8_rcr_mem_cshuffle_intrawave_False_False_False_True_64x64x128_1x4x1_16x16x32",
"problem": {
"split_k":1,
"m":20480,
"n":7168,
"k":256,
"stride_a":256,
"stride_b":256,
"stride_c":7168,
"dtype_a":"fp8",
"dtype_b":"fp8",
"dtype_acc":"fp32",
"dtype_c":"fp16",
"layout_a":"RowMajor",
"layout_b":"ColumnMajor",
"layout_c":"RowMajor",
"structured_sparsity":false
},
"perf_result": {
"latency(ms)": 0.15,
"tflops(TFlops)": 505.00,
"bandwidth(GB/s)": 1200.50
}
}
"""
class TestParseStreamingLog:
def _write_log(self, content: str) -> Path:
f = tempfile.NamedTemporaryFile(mode="w", suffix=".log", delete=False)
f.write(content)
f.close()
return Path(f.name)
def test_basic_parse(self):
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path, arch="gfx950")
assert len(df) == 3
assert df["arch"].iloc[0] == "gfx950"
assert df["m"].tolist() == [16, 16, 20480]
assert df["n"].tolist() == [1536, 1536, 7168]
assert df["k"].tolist() == [7168, 7168, 256]
def test_tflops_values(self):
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path)
assert df["measured_tflops"].tolist() == pytest.approx([8.81, 7.22, 505.0])
def test_kernel_config_parsed(self):
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path)
assert df["tile_m"].iloc[0] == 128
assert df["pipeline"].iloc[0] == "compv3"
assert df["pipeline"].iloc[1] == "compv4"
def test_layout_derived_from_json(self):
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path)
assert all(df["layout"] == "rcr")
def test_empty_log(self):
path = self._write_log("No shapes here\nJust noise\n")
df = parse_streaming_log(path)
assert len(df) == 0
for col in CANONICAL_COLUMNS:
assert col in df.columns
def test_single_kernel(self):
log = """\
Shape 1: M=1 N=1 K=1 dtype=fp8 layout=rcr
{
"name": "gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128",
"problem": {"split_k":1, "m":1, "n":1, "k":1, "dtype_a":"fp8", "dtype_b":"fp8", "layout_a":"RowMajor", "layout_b":"ColumnMajor", "layout_c":"RowMajor"},
"perf_result": {"latency(ms)": 0.001, "tflops(TFlops)": 0.002, "bandwidth(GB/s)": 0.01}
}
"""
path = self._write_log(log)
df = parse_streaming_log(path)
assert len(df) == 1
assert df["m"].iloc[0] == 1
assert bool(df["is_valid"].iloc[0]) is True
def test_zero_tflops_marked_invalid(self):
log = """\
Shape 1: M=16 N=16 K=16 dtype=fp8 layout=rcr
{
"name": "test_kernel",
"problem": {"split_k":1, "m":16, "n":16, "k":16, "dtype_a":"fp8"},
"perf_result": {"latency(ms)": 0.0, "tflops(TFlops)": 0.0, "bandwidth(GB/s)": 0.0}
}
"""
path = self._write_log(log)
df = parse_streaming_log(path)
assert len(df) == 1
assert bool(df["is_valid"].iloc[0]) is False
def test_malformed_json_skipped(self):
log = """\
Shape 1: M=16 N=16 K=16 dtype=fp8 layout=rcr
{
"name": "good_kernel",
"problem": {"split_k":1, "m":16, "n":16, "k":16, "dtype_a":"fp8"},
"perf_result": {"latency(ms)": 0.01, "tflops(TFlops)": 1.0, "bandwidth(GB/s)": 10.0}
}
{ this is not valid json }
{
"name": "another_good",
"problem": {"split_k":1, "m":16, "n":16, "k":16, "dtype_a":"fp8"},
"perf_result": {"latency(ms)": 0.02, "tflops(TFlops)": 2.0, "bandwidth(GB/s)": 20.0}
}
"""
path = self._write_log(log)
df = parse_streaming_log(path)
assert len(df) == 2
def test_extreme_shapes(self):
"""Tiny M=1 (single token) and very large M=20480."""
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path)
assert 1 not in df["m"].values # sample has M=16, M=20480
assert 16 in df["m"].values
assert 20480 in df["m"].values
def test_run_id_assigned(self):
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path, run_id="test_run_123")
assert all(df["run_id"] == "test_run_123")
def test_op_type_assigned(self):
path = self._write_log(SAMPLE_LOG)
df = parse_streaming_log(path, op_type="gemm_streamk")
assert all(df["op_type"] == "gemm_streamk")
# ---------------------------------------------------------------------------
# Parquet round-trip
# ---------------------------------------------------------------------------
class TestParquetIO:
def test_round_trip(self, tmp_path):
df = pd.DataFrame(
{
"m": [16, 32],
"n": [1536, 1536],
"k": [7168, 7168],
"measured_tflops": [8.81, 15.0],
}
)
path = tmp_path / "test.parquet"
save_parquet(df, path)
loaded = load_parquet(path)
assert len(loaded) == 2
assert loaded["m"].tolist() == [16, 32]
def test_creates_parent_dirs(self, tmp_path):
path = tmp_path / "sub" / "dir" / "test.parquet"
df = pd.DataFrame({"x": [1]})
save_parquet(df, path)
assert path.exists()
if __name__ == "__main__":
pytest.main([__file__, "-v"])

264
dispatcher/heuristics/tests/test_dispatcher_integration.py Normal file
View File

@@ -0,0 +1,264 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for dispatcher_integration.py.
Covers: kernel name parsing to feature dict, feature dict to dispatcher config
(name mapping inversion), MLKernelSpec creation, binary pool loading, and
the ML heuristic function.
"""
import json
import sys
from pathlib import Path
import lightgbm as lgb
import numpy as np
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from dispatcher_integration import (
kernel_config_to_feature_dict,
feature_dict_to_dispatcher_config,
feature_dict_to_ml_spec,
ml_spec_to_dispatcher_config,
create_ml_heuristic,
load_kernel_pool_from_binaries,
MLKernelSpec,
LAYOUT_TO_DISPATCHER,
)
from feature_engine import GemmUniversalFeatureEngine
SAMPLE_KERNEL_NAME = (
"gemm_universal_fp8_rcr_compv3_cshuffle_intrawave"
"_False_False_False_False_128x128x128_1x4x1_16x16x128"
)
# ---------------------------------------------------------------------------
# kernel_config_to_feature_dict
# ---------------------------------------------------------------------------
class TestKernelConfigToFeatureDict:
def test_parses_standard_name(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
assert feat["tile_m"] == 128
assert feat["tile_n"] == 128
assert feat["tile_k"] == 128
assert feat["warp_m"] == 1 # warps per block
assert feat["warp_n"] == 4
assert feat["warp_k"] == 1
assert feat["warp_tile_m"] == 16
assert feat["warp_tile_n"] == 16
assert feat["warp_tile_k"] == 128
assert feat["pipeline"] == "compv3"
assert feat["scheduler"] == "intrawave"
assert feat["epilogue"] == "cshuffle"
assert feat["kernel_name"] == SAMPLE_KERNEL_NAME
def test_empty_name_returns_empty(self):
assert kernel_config_to_feature_dict("") == {}
def test_invalid_name_returns_empty(self):
assert kernel_config_to_feature_dict("not_a_kernel") == {}
# ---------------------------------------------------------------------------
# Name mapping: feature dict <-> dispatcher config
# ---------------------------------------------------------------------------
class TestNameMapping:
"""The critical inversion: feature engine warp_m/n/k (warps per block)
maps to dispatcher wave_m/n/k, and feature engine warp_tile_m/n/k
maps to dispatcher warp_m/n/k."""
def test_warp_to_wave_mapping(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
disp = feature_dict_to_dispatcher_config(feat)
assert disp["wave_m"] == feat["warp_m"] # 1
assert disp["wave_n"] == feat["warp_n"] # 4
assert disp["wave_k"] == feat["warp_k"] # 1
def test_warp_tile_to_warp_mapping(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
disp = feature_dict_to_dispatcher_config(feat)
assert disp["warp_m"] == feat["warp_tile_m"] # 16
assert disp["warp_n"] == feat["warp_tile_n"] # 16
assert disp["warp_k"] == feat["warp_tile_k"] # 128
def test_tile_dims_pass_through(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
disp = feature_dict_to_dispatcher_config(feat)
assert disp["tile_m"] == 128
assert disp["tile_n"] == 128
assert disp["tile_k"] == 128
def test_pipeline_passes_through(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
disp = feature_dict_to_dispatcher_config(feat)
assert disp["pipeline"] == "compv3"
assert disp["scheduler"] == "intrawave"
assert disp["epilogue"] == "cshuffle"
def test_rcr_layout_mapping(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
disp = feature_dict_to_dispatcher_config(feat, dtype="fp8")
assert disp["layout_a"] == "row"
assert disp["layout_b"] == "col"
assert disp["layout_c"] == "row"
def test_all_layouts(self):
for layout, (la, lb, lc) in LAYOUT_TO_DISPATCHER.items():
feat = {"layout": layout, "tile_m": 128}
disp = feature_dict_to_dispatcher_config(feat)
assert disp["layout_a"] == la
assert disp["layout_b"] == lb
assert disp["layout_c"] == lc
# ---------------------------------------------------------------------------
# MLKernelSpec
# ---------------------------------------------------------------------------
class TestMLKernelSpec:
def test_from_feature_dict(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
spec = feature_dict_to_ml_spec(feat, predicted_tflops=123.4)
assert spec.kernel_name == SAMPLE_KERNEL_NAME
assert spec.predicted_tflops == 123.4
assert spec.tile_m == 128
assert spec.wave_m == 1 # was warp_m in feature space
assert spec.warp_m == 16 # was warp_tile_m in feature space
def test_spec_to_dispatcher_config(self):
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
spec = feature_dict_to_ml_spec(feat, 100.0)
disp = ml_spec_to_dispatcher_config(spec, dtype="fp8", arch="gfx950")
assert disp["tile_m"] == 128
assert disp["wave_m"] == 1
assert disp["warp_m"] == 16
assert disp["gfx_arch"] == "gfx950"
assert disp["dtype_a"] == "fp8"
def test_roundtrip_preserves_values(self):
"""feature_dict -> MLKernelSpec -> dispatcher_config should be consistent."""
feat = kernel_config_to_feature_dict(SAMPLE_KERNEL_NAME)
spec = feature_dict_to_ml_spec(feat, 0.0)
disp_from_spec = ml_spec_to_dispatcher_config(spec)
disp_from_feat = feature_dict_to_dispatcher_config(feat)
for key in [
"tile_m",
"tile_n",
"tile_k",
"wave_m",
"wave_n",
"wave_k",
"warp_m",
"warp_n",
"warp_k",
"pipeline",
"scheduler",
"epilogue",
]:
assert disp_from_spec[key] == disp_from_feat[key], f"Mismatch on {key}"
# ---------------------------------------------------------------------------
# Binary pool loading
# ---------------------------------------------------------------------------
class TestLoadKernelPool:
def test_loads_from_real_bin_dir(self):
bin_dir = Path("/workspace/ck_tile/bin")
if not bin_dir.exists():
pytest.skip("No /workspace/ck_tile/bin")
pool = load_kernel_pool_from_binaries(bin_dir)
assert len(pool) > 0
assert "tile_m" in pool[0]
assert "kernel_name" in pool[0]
def test_empty_dir_returns_empty(self, tmp_path):
pool = load_kernel_pool_from_binaries(tmp_path)
assert pool == []
# ---------------------------------------------------------------------------
# ML heuristic function
# ---------------------------------------------------------------------------
class TestCreateMLHeuristic:
@pytest.fixture
def mock_model_dir(self, tmp_path):
"""Create a minimal model for testing the heuristic flow."""
fe = GemmUniversalFeatureEngine()
n_features = len(fe.get_feature_names())
np.random.seed(42)
X = np.random.rand(100, n_features)
y = np.random.rand(100) * 500
model = lgb.LGBMRegressor(n_estimators=5, verbose=-1)
model.fit(X, y)
model.booster_.save_model(str(tmp_path / "model_tflops.lgbm"))
spec = {
"feature_names": fe.get_feature_names(),
"categorical_features": fe.get_categorical_features(),
}
with open(tmp_path / "feature_spec.json", "w") as f:
json.dump(spec, f)
return tmp_path
def _make_pool(self):
"""Create a small synthetic kernel pool."""
names = [
"gemm_universal_fp8_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x128_1x4x1_16x16x128",
"gemm_universal_fp8_rcr_compv4_default_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16",
"gemm_universal_fp8_rcr_mem_cshuffle_interwave_False_False_False_False_64x64x128_1x4x1_16x16x32",
]
return [kernel_config_to_feature_dict(n) for n in names]
def test_returns_ml_kernel_spec(self, mock_model_dir):
pool = self._make_pool()
heuristic = create_ml_heuristic(mock_model_dir, kernel_pool=pool)
result = heuristic(1024, 1024, 1024)
assert isinstance(result, MLKernelSpec)
assert result.tile_m > 0
assert isinstance(result.predicted_tflops, float)
def test_returns_valid_kernel_from_pool(self, mock_model_dir):
pool = self._make_pool()
pool_names = {p["kernel_name"] for p in pool}
heuristic = create_ml_heuristic(mock_model_dir, kernel_pool=pool)
result = heuristic(1024, 1024, 1024)
assert result.kernel_name in pool_names
def test_different_shapes_may_select_different_kernels(self, mock_model_dir):
pool = self._make_pool()
heuristic = create_ml_heuristic(mock_model_dir, kernel_pool=pool)
r1 = heuristic(16, 1536, 7168)
r2 = heuristic(8192, 8192, 256)
# At minimum both should return valid specs
assert r1.tile_m > 0
assert r2.tile_m > 0
def test_m1_corner_case(self, mock_model_dir):
pool = self._make_pool()
heuristic = create_ml_heuristic(mock_model_dir, kernel_pool=pool)
result = heuristic(1, 4096, 4096)
assert isinstance(result, MLKernelSpec)
assert np.isfinite(result.predicted_tflops)
def test_empty_pool_raises(self, mock_model_dir):
with pytest.raises(ValueError, match="No kernel configs"):
create_ml_heuristic(mock_model_dir, kernel_pool=[])
if __name__ == "__main__":
pytest.main([__file__, "-v"])

55
dispatcher/heuristics/tests/test_evaluate.py Normal file
View File

@@ -0,0 +1,55 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for evaluate.py.
Covers: shape family classification, K-depth regime classification,
and basic evaluation metric checks.
"""
import sys
from pathlib import Path
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from evaluate import classify_shape_family, classify_k_regime
class TestClassifyShapeFamily:
def test_tiny_m(self):
assert classify_shape_family(1, 4096, 4096) == "tiny_m"
assert classify_shape_family(16, 1536, 7168) == "tiny_m"
def test_small_m(self):
assert classify_shape_family(32, 1536, 7168) == "small_m"
assert classify_shape_family(128, 4096, 4096) == "small_m"
def test_medium_m(self):
assert classify_shape_family(256, 1024, 1024) == "medium_m"
assert classify_shape_family(2048, 2048, 2048) == "medium_m"
def test_large_m(self):
assert classify_shape_family(4096, 4096, 4096) == "large_m"
assert classify_shape_family(20480, 7168, 256) == "large_m"
class TestClassifyKRegime:
def test_shallow(self):
assert classify_k_regime(256) == "shallow_k"
assert classify_k_regime(32) == "shallow_k"
def test_medium(self):
assert classify_k_regime(1024) == "medium_k"
assert classify_k_regime(2048) == "medium_k"
def test_deep(self):
assert classify_k_regime(4096) == "deep_k"
assert classify_k_regime(7168) == "deep_k"
if __name__ == "__main__":
pytest.main([__file__, "-v"])

409
dispatcher/heuristics/tests/test_feature_engine.py Normal file
View File

@@ -0,0 +1,409 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for feature_engine.py.
Covers: feature count consistency, formula correctness (tile efficiency, LDS,
arithmetic intensity), corner-case shapes (M=1, huge M, square, skinny-K),
parameter space validity, config validation, and batch vs single extraction parity.
"""
import sys
from pathlib import Path
import numpy as np
import pandas as pd
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from feature_engine import (
GemmUniversalFeatureEngine,
)
@pytest.fixture
def fe():
"""Default feature engine with MI355X-like hardware."""
return GemmUniversalFeatureEngine(
num_cus=256,
lds_capacity=65536,
max_clock_mhz=2400,
simds_per_cu=4,
shader_engines=32,
max_waves_per_cu=32,
wavefront_size=64,
l1_cache_kb=32,
l2_cache_kb=4096,
l3_cache_kb=262144,
num_xcd=8,
)
def _make_problem(m=1024, n=1024, k=1024, dtype="fp8", layout="rcr", split_k=1):
return {
"m": m,
"n": n,
"k": k,
"dtype": dtype,
"layout": layout,
"split_k": split_k,
}
def _make_kernel(
tile_m=128,
tile_n=128,
tile_k=64,
warp_m=2,
warp_n=2,
warp_k=1,
warp_tile_m=32,
warp_tile_n=32,
warp_tile_k=16,
pipeline="compv3",
scheduler="intrawave",
epilogue="cshuffle",
pad_m=False,
pad_n=False,
pad_k=False,
persistent=False,
):
return {
"tile_m": tile_m,
"tile_n": tile_n,
"tile_k": tile_k,
"warp_m": warp_m,
"warp_n": warp_n,
"warp_k": warp_k,
"warp_tile_m": warp_tile_m,
"warp_tile_n": warp_tile_n,
"warp_tile_k": warp_tile_k,
"pipeline": pipeline,
"scheduler": scheduler,
"epilogue": epilogue,
"pad_m": pad_m,
"pad_n": pad_n,
"pad_k": pad_k,
"persistent": persistent,
}
# ---------------------------------------------------------------------------
# Basic consistency
# ---------------------------------------------------------------------------
class TestFeatureConsistency:
def test_feature_count_matches_names(self, fe):
prob = _make_problem()
kern = _make_kernel()
vec = fe.extract(prob, kern)
assert len(vec) == len(fe.get_feature_names())
def test_feature_count_is_72(self, fe):
assert len(fe.get_feature_names()) == 72
def test_no_nan_in_output(self, fe):
prob = _make_problem()
kern = _make_kernel()
vec = fe.extract(prob, kern)
assert not np.any(np.isnan(vec))
def test_no_inf_in_output(self, fe):
prob = _make_problem()
kern = _make_kernel()
vec = fe.extract(prob, kern)
assert not np.any(np.isinf(vec))
def test_categorical_features_in_names(self, fe):
names = fe.get_feature_names()
for cat in fe.get_categorical_features():
assert cat in names
# ---------------------------------------------------------------------------
# Formula correctness
# ---------------------------------------------------------------------------
class TestTileEfficiency:
"""Tile efficiency: fraction of the last tile that is useful work."""
def test_perfectly_divisible(self, fe):
prob = _make_problem(m=256, n=256, k=128)
kern = _make_kernel(tile_m=128, tile_n=128, tile_k=64)
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
assert vec[names.index("tile_eff_m")] == 1.0
assert vec[names.index("tile_eff_n")] == 1.0
assert vec[names.index("tile_eff_k")] == 1.0
assert vec[names.index("overall_tile_efficiency")] == 1.0
def test_not_divisible(self, fe):
prob = _make_problem(m=100, n=100, k=100)
kern = _make_kernel(tile_m=128, tile_n=128, tile_k=64)
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
assert vec[names.index("tile_eff_m")] == pytest.approx(100 / 128)
assert vec[names.index("tile_eff_n")] == pytest.approx(100 / 128)
assert vec[names.index("tile_eff_k")] == pytest.approx(36 / 64)
def test_m_equals_1(self, fe):
"""Single-token inference: M=1, tile_m=128 => eff = 1/128."""
prob = _make_problem(m=1)
kern = _make_kernel(tile_m=128)
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
assert vec[names.index("tile_eff_m")] == pytest.approx(1.0 / 128.0)
class TestLDSUsage:
def test_lds_formula(self, fe):
prob = _make_problem(dtype="fp8")
kern = _make_kernel(tile_m=128, tile_n=128, tile_k=64)
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
expected = (128 * 64 + 128 * 64) * 1.0 # fp8 = 1 byte
assert vec[names.index("lds_usage_estimate")] == expected
def test_lds_ratio_compv4(self, fe):
"""compv4 has 32KB LDS limit, not 64KB."""
prob = _make_problem(dtype="fp8")
kern = _make_kernel(tile_m=128, tile_n=128, tile_k=64, pipeline="compv4")
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
lds_est = (128 * 64 + 128 * 64) * 1.0
assert vec[names.index("lds_usage_ratio")] == pytest.approx(lds_est / 32768)
def test_lds_fp16_doubles(self, fe):
prob = _make_problem(dtype="fp16")
kern = _make_kernel(tile_m=128, tile_n=128, tile_k=64)
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
expected = (128 * 64 + 128 * 64) * 2.0 # fp16 = 2 bytes
assert vec[names.index("lds_usage_estimate")] == expected
class TestArithmeticIntensity:
def test_square_shape(self, fe):
M, N, K = 1024, 1024, 1024
prob = _make_problem(m=M, n=N, k=K, dtype="fp8")
kern = _make_kernel()
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
mem = (M * K + K * N + M * N) * 1.0
expected = (2.0 * M * N * K) / mem
assert vec[names.index("arithmetic_intensity")] == pytest.approx(expected)
def test_skinny_k(self, fe):
"""Small K => low arithmetic intensity (memory-bound)."""
prob = _make_problem(m=8192, n=8192, k=32, dtype="fp8")
kern = _make_kernel()
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
assert vec[names.index("arithmetic_intensity")] < 100
def test_deep_k(self, fe):
"""Large K => high arithmetic intensity (compute-bound)."""
prob = _make_problem(m=256, n=256, k=8192, dtype="fp8")
kern = _make_kernel()
vec = fe.extract(prob, kern)
names = fe.get_feature_names()
assert vec[names.index("arithmetic_intensity")] > 100
# ---------------------------------------------------------------------------
# Corner-case shapes
# ---------------------------------------------------------------------------
class TestCornerCaseShapes:
def test_m1_single_token(self, fe):
vec = fe.extract(_make_problem(m=1, n=4096, k=4096), _make_kernel())
assert not np.any(np.isnan(vec))
def test_m1_n1_k1_minimum(self, fe):
vec = fe.extract(_make_problem(m=1, n=1, k=1), _make_kernel())
assert not np.any(np.isnan(vec))
assert not np.any(np.isinf(vec))
def test_very_large_m(self, fe):
vec = fe.extract(_make_problem(m=20480, n=7168, k=7168), _make_kernel())
assert not np.any(np.isnan(vec))
def test_non_power_of_2(self, fe):
vec = fe.extract(_make_problem(m=1536, n=7168, k=2304), _make_kernel())
assert not np.any(np.isnan(vec))
def test_prime_dimensions(self, fe):
vec = fe.extract(_make_problem(m=17, n=31, k=127), _make_kernel())
assert not np.any(np.isnan(vec))
def test_tall_matrix(self, fe):
"""M >> N (tall matrix)."""
prob = _make_problem(m=16384, n=64, k=1024)
vec = fe.extract(prob, _make_kernel())
names = fe.get_feature_names()
assert vec[names.index("aspect_ratio_mn")] > 100
def test_wide_matrix(self, fe):
"""N >> M (wide matrix)."""
prob = _make_problem(m=64, n=16384, k=1024)
vec = fe.extract(prob, _make_kernel())
names = fe.get_feature_names()
assert vec[names.index("aspect_ratio_mn")] < 0.01
# ---------------------------------------------------------------------------
# Batch vs single extraction parity
# ---------------------------------------------------------------------------
class TestBatchParity:
def test_batch_matches_single(self, fe):
"""Vectorized batch should produce identical results to row-by-row."""
rows = [
{
"m": 16,
"n": 1536,
"k": 7168,
"split_k": 1,
"dtype": "fp8",
"layout": "rcr",
"tile_m": 128,
"tile_n": 128,
"tile_k": 128,
"warp_m": 1,
"warp_n": 4,
"warp_k": 1,
"warp_tile_m": 16,
"warp_tile_n": 16,
"warp_tile_k": 128,
"pipeline": "compv3",
"scheduler": "intrawave",
"epilogue": "cshuffle",
"pad_m": False,
"pad_n": False,
"pad_k": False,
"persistent": False,
},
{
"m": 20480,
"n": 7168,
"k": 256,
"split_k": 1,
"dtype": "fp8",
"layout": "rcr",
"tile_m": 64,
"tile_n": 64,
"tile_k": 128,
"warp_m": 2,
"warp_n": 2,
"warp_k": 1,
"warp_tile_m": 32,
"warp_tile_n": 32,
"warp_tile_k": 16,
"pipeline": "mem",
"scheduler": "interwave",
"epilogue": "default",
"pad_m": True,
"pad_n": True,
"pad_k": True,
"persistent": True,
},
]
df = pd.DataFrame(rows)
batch_result = fe.extract_batch(df)
for i, row_dict in enumerate(rows):
single_result = fe.extract(row_dict, row_dict)
np.testing.assert_allclose(
batch_result[i],
single_result,
rtol=1e-5,
atol=1e-5,
err_msg=f"Mismatch at row {i}",
)
# ---------------------------------------------------------------------------
# Parameter space and validation
# ---------------------------------------------------------------------------
class TestParameterSpace:
def test_parameter_space_non_empty(self, fe):
ps = fe.get_parameter_space()
assert len(ps) > 0
assert "tile_m" in ps
assert "pipeline" in ps
def test_valid_config_passes(self, fe):
config = {
"tile_m": 128,
"tile_n": 128,
"tile_k": 64,
"warp_m": 2,
"warp_n": 2,
"warp_k": 1,
"pipeline": "compv3",
"scheduler": "intrawave",
"epilogue": "cshuffle",
"pad_m": False,
"pad_n": False,
"pad_k": False,
"persistent": False,
}
assert fe.validate_config(config) is True
def test_invalid_tile_rejected(self, fe):
config = {"tile_m": 999}
assert fe.validate_config(config) is False
def test_lds_constraint_rejects_huge_tile(self, fe):
config = {
"tile_m": 256,
"tile_n": 256,
"tile_k": 256,
"warp_m": 2,
"warp_n": 2,
"warp_k": 1,
"pipeline": "compv4",
}
assert fe.validate_config(config) is False
def test_project_to_valid_snaps(self, fe):
config = {"tile_m": 100, "tile_n": 200, "pipeline": "compv3"}
projected = fe.project_to_valid(config)
assert projected["tile_m"] == 128
assert projected["tile_n"] == 192
assert projected["pipeline"] == "compv3"
# ---------------------------------------------------------------------------
# Hardware features
# ---------------------------------------------------------------------------
class TestHardwareFeatures:
def test_hardware_values_propagated(self, fe):
vec = fe.extract(_make_problem(), _make_kernel())
names = fe.get_feature_names()
assert vec[names.index("hw_num_cus")] == 256
assert vec[names.index("hw_max_clock_mhz")] == 2400
assert vec[names.index("hw_total_simds")] == 256 * 4
assert vec[names.index("hw_num_xcd")] == 8
def test_different_hardware(self):
fe_small = GemmUniversalFeatureEngine(num_cus=120, max_clock_mhz=1800)
vec = fe_small.extract(_make_problem(), _make_kernel())
names = fe_small.get_feature_names()
assert vec[names.index("hw_num_cus")] == 120
assert vec[names.index("hw_max_clock_mhz")] == 1800
if __name__ == "__main__":
pytest.main([__file__, "-v"])

357
dispatcher/heuristics/tests/test_feature_parity.py Normal file
View File

@@ -0,0 +1,357 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Test that the C++ extract_features() in ml_heuristic.hpp produces identical
values to the Python GemmUniversalFeatureEngine.extract().
This test uses ctypes to call the C++ feature extraction compiled into a
small shared library, then compares against Python output. If compilation
fails (no HIP/ROCm), it falls back to verifying the Python feature engine
against manually computed expected values for specific test cases.
"""
import math
import sys
from pathlib import Path
import numpy as np
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from feature_engine import (
GemmUniversalFeatureEngine,
PIPELINE_MAP,
SCHEDULER_MAP,
EPILOGUE_MAP,
LAYOUT_MAP,
)
def _compute_features_manually(
M,
N,
K,
split_k,
dtype,
layout,
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
pipeline,
scheduler,
epilogue,
pad_m,
pad_n,
pad_k,
persistent,
hw,
):
"""Recompute features independently to verify the Python engine."""
bpe_map = {"fp8": 1.0, "fp16": 2.0, "bf16": 2.0, "fp32": 4.0}
bpe = bpe_map.get(dtype, 1.0)
log2_M = math.log2(max(M, 1))
log2_N = math.log2(max(N, 1))
log2_K = math.log2(max(K, 1))
log2_MNK = math.log2(max(M * N * K, 1))
mem = (M * K + K * N + M * N) * bpe
ai = (2.0 * M * N * K) / max(mem, 1)
lds_est = (tile_m * tile_k + tile_n * tile_k) * bpe
lds_cap = 32768 if pipeline == "compv4" else hw["lds_capacity"]
ntm = math.ceil(M / max(tile_m, 1))
ntn = math.ceil(N / max(tile_n, 1))
ntk = math.ceil(K / max(tile_k, 1))
def eff(d, t):
if t <= 0:
return 1.0
r = d % t
return r / t if r > 0 else 1.0
# Problem-to-tile ratios
ratio_M_to_tile_m = M / max(tile_m, 1)
ratio_N_to_tile_n = N / max(tile_n, 1)
ratio_K_to_tile_k = K / max(tile_k, 1)
# Binary features: problem smaller than tile
problem_smaller_than_tile_m = float(M < tile_m)
problem_smaller_than_tile_n = float(N < tile_n)
problem_smaller_than_tile_k = float(K < tile_k)
any_dim_too_small = float((M < tile_m) or (N < tile_n) or (K < tile_k))
# Padding requirement features
needs_padding_m = float(tile_m > 0 and M % tile_m != 0)
needs_padding_n = float(tile_n > 0 and N % tile_n != 0)
needs_padding_k = float(tile_k > 0 and K % tile_k != 0)
# Interaction features
has_padding_when_needed_m = float(needs_padding_m and pad_m)
has_padding_when_needed_n = float(needs_padding_n and pad_n)
has_padding_when_needed_k = float(needs_padding_k and pad_k)
# Missing padding features
missing_required_padding_m = float(needs_padding_m and not pad_m)
missing_required_padding_n = float(needs_padding_n and not pad_n)
missing_required_padding_k = float(needs_padding_k and not pad_k)
missing_any_required_padding = float(
missing_required_padding_m or missing_required_padding_n or missing_required_padding_k
)
return [
M, # 0
N, # 1
K, # 2
split_k, # 3
log2_M, # 4
log2_N, # 5
log2_K, # 6
log2_MNK, # 7
ai, # 8
M / max(N, 1), # 9 (aspect_ratio_mn)
M / max(K, 1), # 10 (aspect_ratio_mk)
N / max(K, 1), # 11 (aspect_ratio_nk)
LAYOUT_MAP.get(layout, 0), # 12
tile_m, # 13
tile_n, # 14
tile_k, # 15
warp_m, # 16
warp_n, # 17
warp_k, # 18
warp_tile_m, # 19
warp_tile_n, # 20
warp_tile_k, # 21
PIPELINE_MAP.get(pipeline, 0), # 22
SCHEDULER_MAP.get(scheduler, 0), # 23
EPILOGUE_MAP.get(epilogue, 0), # 24
float(pad_m), # 25
float(pad_n), # 26
float(pad_k), # 27
float(persistent), # 28
warp_m * warp_n * warp_k, # 29 (num_warps)
tile_m * tile_n * tile_k, # 30 (tile_volume)
tile_m * tile_n, # 31 (tile_mn)
lds_est, # 32 (lds_usage_estimate)
lds_est / max(lds_cap, 1), # 33 (lds_usage_ratio)
ntm, # 34 (num_tiles_m)
ntn, # 35 (num_tiles_n)
ntk, # 36 (num_tiles_k)
ntm * ntn, # 37 (total_output_tiles)
eff(M, tile_m), # 38 (tile_eff_m)
eff(N, tile_n), # 39 (tile_eff_n)
eff(K, tile_k), # 40 (tile_eff_k)
eff(M, tile_m) * eff(N, tile_n) * eff(K, tile_k), # 41 (overall_tile_efficiency)
ntm * ntn / max(hw["num_cus"], 1), # 42 (cu_utilization)
ratio_M_to_tile_m, # 43
ratio_N_to_tile_n, # 44
ratio_K_to_tile_k, # 45
problem_smaller_than_tile_m, # 46
problem_smaller_than_tile_n, # 47
problem_smaller_than_tile_k, # 48
any_dim_too_small, # 49
needs_padding_m, # 50
needs_padding_n, # 51
needs_padding_k, # 52
has_padding_when_needed_m, # 53
has_padding_when_needed_n, # 54
has_padding_when_needed_k, # 55
missing_required_padding_m, # 56
missing_required_padding_n, # 57
missing_required_padding_k, # 58
missing_any_required_padding, # 59
hw["num_cus"], # 60
hw["simds_per_cu"], # 61
hw["num_cus"] * hw["simds_per_cu"], # 62 (total_simds)
hw["shader_engines"], # 63
hw["max_clock_mhz"], # 64
hw["max_waves_per_cu"], # 65
hw["wavefront_size"], # 66
hw["lds_capacity"], # 67
hw["l1_cache_kb"], # 68
hw["l2_cache_kb"], # 69
hw["l3_cache_kb"], # 70
hw["num_xcd"], # 71
]
TEST_CASES = [
{
"problem": {
"m": 1024,
"n": 1024,
"k": 1024,
"split_k": 1,
"dtype": "fp8",
"layout": "rcr",
},
"kernel": {
"tile_m": 128,
"tile_n": 128,
"tile_k": 64,
"warp_m": 2,
"warp_n": 2,
"warp_k": 1,
"warp_tile_m": 32,
"warp_tile_n": 32,
"warp_tile_k": 16,
"pipeline": "compv3",
"scheduler": "intrawave",
"epilogue": "cshuffle",
"pad_m": False,
"pad_n": False,
"pad_k": False,
"persistent": False,
},
},
{
"problem": {
"m": 1,
"n": 4096,
"k": 4096,
"split_k": 1,
"dtype": "fp8",
"layout": "rcr",
},
"kernel": {
"tile_m": 64,
"tile_n": 64,
"tile_k": 128,
"warp_m": 1,
"warp_n": 4,
"warp_k": 1,
"warp_tile_m": 16,
"warp_tile_n": 16,
"warp_tile_k": 128,
"pipeline": "compv4",
"scheduler": "interwave",
"epilogue": "default",
"pad_m": True,
"pad_n": True,
"pad_k": True,
"persistent": True,
},
},
{
"problem": {
"m": 20480,
"n": 7168,
"k": 256,
"split_k": 1,
"dtype": "fp16",
"layout": "rrr",
},
"kernel": {
"tile_m": 256,
"tile_n": 256,
"tile_k": 32,
"warp_m": 4,
"warp_n": 1,
"warp_k": 1,
"warp_tile_m": 32,
"warp_tile_n": 32,
"warp_tile_k": 16,
"pipeline": "mem",
"scheduler": "interwave",
"epilogue": "cshuffle",
"pad_m": False,
"pad_n": False,
"pad_k": False,
"persistent": False,
},
},
]
HW = {
"num_cus": 256,
"simds_per_cu": 4,
"shader_engines": 32,
"max_clock_mhz": 2400,
"max_waves_per_cu": 32,
"wavefront_size": 64,
"lds_capacity": 65536,
"l1_cache_kb": 32,
"l2_cache_kb": 4096,
"l3_cache_kb": 262144,
"num_xcd": 8,
}
class TestFeatureParity:
"""Verify Python feature engine matches manual computation (C++ uses same logic)."""
@pytest.fixture
def fe(self):
return GemmUniversalFeatureEngine(**HW)
@pytest.mark.parametrize("case_idx", range(len(TEST_CASES)))
def test_python_matches_manual(self, fe, case_idx):
case = TEST_CASES[case_idx]
prob = case["problem"]
kern = case["kernel"]
py_features = fe.extract(prob, kern)
manual = _compute_features_manually(
prob["m"],
prob["n"],
prob["k"],
prob["split_k"],
prob["dtype"],
prob["layout"],
kern["tile_m"],
kern["tile_n"],
kern["tile_k"],
kern["warp_m"],
kern["warp_n"],
kern["warp_k"],
kern["warp_tile_m"],
kern["warp_tile_n"],
kern["warp_tile_k"],
kern["pipeline"],
kern["scheduler"],
kern["epilogue"],
kern["pad_m"],
kern["pad_n"],
kern["pad_k"],
kern["persistent"],
HW,
)
manual_arr = np.array(manual, dtype=np.float64)
assert len(py_features) == len(manual_arr) == 72
for i in range(72):
assert py_features[i] == pytest.approx(
manual_arr[i], rel=1e-10, abs=1e-15
), (
f"Feature {i} ({fe.get_feature_names()[i]}): Python={py_features[i]}, Manual={manual_arr[i]}"
)
def test_feature_count(self, fe):
assert len(fe.get_feature_names()) == 72
def test_encoding_maps_match_cpp(self):
"""The C++ encode_* functions must use the same mapping as Python."""
assert PIPELINE_MAP == {
"compv3": 0,
"compv4": 1,
"compv5": 2,
"mem": 3,
"preshufflev2": 4,
}
assert SCHEDULER_MAP == {"intrawave": 0, "interwave": 1}
assert EPILOGUE_MAP == {"default": 0, "cshuffle": 1}
assert LAYOUT_MAP == {"rcr": 0, "rrr": 1, "crr": 2, "ccr": 3}
if __name__ == "__main__":
pytest.main([__file__, "-v"])

90
dispatcher/heuristics/tests/test_model_compression.py Normal file
View File

@@ -0,0 +1,90 @@
#!/usr/bin/env python3
"""Test that compressed models can be loaded and used."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from predict import Predictor
def test_fp16_model_decompression():
"""Test that fp16 model is auto-decompressed and usable."""
model_dir = Path(__file__).parent.parent / "models" / "gemm_universal_fp16_gfx950"
# Ensure .lgbm.gz exists
gz_file = model_dir / "model_tflops.lgbm.gz"
assert gz_file.exists(), f"Compressed model not found: {gz_file}"
# Load predictor - should auto-decompress
predictor = Predictor(model_dir)
# Test prediction
problem = {"m": 128, "n": 1536, "k": 7168, "dtype": "fp16", "layout": "rcr"}
kernel_config = {
"tile_shape": {"m0": 128, "n0": 128, "k0": 16},
"wave_shape": {"m1": 2, "n1": 2, "k1": 1},
"warp_tile": {"m2": 32, "n2": 32, "k2": 8},
}
tflops = predictor.predict_tflops(problem, kernel_config)
assert isinstance(tflops, float), f"Expected float, got {type(tflops)}"
assert tflops > 0, f"Expected positive TFLOPS, got {tflops}"
# Verify decompressed file was created
lgbm_file = model_dir / "model_tflops.lgbm"
assert lgbm_file.exists(), "Model should have been decompressed"
print(f"✅ FP16 model decompression test passed")
print(f" Predicted TFLOPS: {tflops:.2f}")
print(f" Decompressed to: {lgbm_file}")
return True
def test_fp8_model_decompression():
"""Test that fp8 model is auto-decompressed and usable."""
model_dir = Path(__file__).parent.parent / "models" / "gemm_universal_fp8_gfx950"
# Ensure .lgbm.gz exists
gz_file = model_dir / "model_tflops.lgbm.gz"
assert gz_file.exists(), f"Compressed model not found: {gz_file}"
# Load predictor - should auto-decompress
predictor = Predictor(model_dir)
# Test prediction
problem = {"m": 2048, "n": 2048, "k": 2048, "dtype": "fp8", "layout": "rcr"}
kernel_config = {
"tile_shape": {"m0": 256, "n0": 256, "k0": 64},
"wave_shape": {"m1": 2, "n1": 2, "k1": 1},
"warp_tile": {"m2": 32, "n2": 32, "k2": 16},
}
tflops = predictor.predict_tflops(problem, kernel_config)
assert isinstance(tflops, float), f"Expected float, got {type(tflops)}"
assert tflops > 0, f"Expected positive TFLOPS, got {tflops}"
# Verify decompressed file was created
lgbm_file = model_dir / "model_tflops.lgbm"
assert lgbm_file.exists(), "Model should have been decompressed"
print(f"✅ FP8 model decompression test passed")
print(f" Predicted TFLOPS: {tflops:.2f}")
print(f" Decompressed to: {lgbm_file}")
return True
if __name__ == "__main__":
print("Testing compressed model auto-decompression...")
print()
test_fp16_model_decompression()
print()
test_fp8_model_decompression()
print()
print("✅ All model compression tests passed!")

181
dispatcher/heuristics/tests/test_predict.py Normal file
View File

@@ -0,0 +1,181 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for predict.py.
Covers: Predictor initialization, single prediction, ranking, select_best,
missing model handling, and edge cases (single kernel, empty list).
"""
import json
import sys
from pathlib import Path
import lightgbm as lgb
import numpy as np
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from feature_engine import GemmUniversalFeatureEngine
from predict import Predictor
@pytest.fixture
def model_dir(tmp_path):
"""Create a minimal trained model for testing."""
fe = GemmUniversalFeatureEngine()
n_features = len(fe.get_feature_names())
np.random.seed(42)
X = np.random.rand(200, n_features)
y = np.random.rand(200) * 100
model = lgb.LGBMRegressor(n_estimators=10, verbose=-1)
model.fit(X, y)
model.booster_.save_model(str(tmp_path / "model_tflops.lgbm"))
y_lat = np.random.rand(200) * 0.1
model_lat = lgb.LGBMRegressor(n_estimators=10, verbose=-1)
model_lat.fit(X, y_lat)
model_lat.booster_.save_model(str(tmp_path / "model_latency.lgbm"))
spec = {
"feature_names": fe.get_feature_names(),
"categorical_features": fe.get_categorical_features(),
}
with open(tmp_path / "feature_spec.json", "w") as f:
json.dump(spec, f)
return tmp_path
@pytest.fixture
def predictor(model_dir):
return Predictor(model_dir)
def _problem():
return {
"m": 1024,
"n": 1024,
"k": 1024,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
}
def _kernel(tile_m=128, pipeline="compv3"):
return {
"kernel_name": f"test_kernel_{tile_m}_{pipeline}",
"tile_m": tile_m,
"tile_n": 128,
"tile_k": 64,
"warp_m": 2,
"warp_n": 2,
"warp_k": 1,
"warp_tile_m": 32,
"warp_tile_n": 32,
"warp_tile_k": 16,
"pipeline": pipeline,
"scheduler": "intrawave",
"epilogue": "cshuffle",
"pad_m": False,
"pad_n": False,
"pad_k": False,
"persistent": False,
}
class TestPredictor:
def test_predict_tflops_returns_float(self, predictor):
result = predictor.predict_tflops(_problem(), _kernel())
assert isinstance(result, float)
def test_predict_latency_returns_float(self, predictor):
result = predictor.predict_latency(_problem(), _kernel())
assert isinstance(result, float)
def test_predict_all_returns_dict(self, predictor):
result = predictor.predict_all(_problem(), _kernel())
assert "tflops" in result
assert "latency_ms" in result
def test_rank_kernels_sorted_descending(self, predictor):
kernels = [_kernel(64, "compv3"), _kernel(128, "compv4"), _kernel(256, "mem")]
ranked = predictor.rank_kernels(_problem(), kernels)
assert len(ranked) == 3
scores = [s for _, s in ranked]
assert scores == sorted(scores, reverse=True)
def test_select_best_returns_name(self, predictor):
kernels = [_kernel(64), _kernel(128)]
best = predictor.select_best(_problem(), kernels)
assert isinstance(best, str)
assert best in [k["kernel_name"] for k in kernels]
def test_single_kernel(self, predictor):
kernels = [_kernel(128)]
ranked = predictor.rank_kernels(_problem(), kernels)
assert len(ranked) == 1
def test_missing_bandwidth_model(self, model_dir):
pred = Predictor(model_dir)
with pytest.raises(FileNotFoundError):
pred.predict_bandwidth(_problem(), _kernel())
def test_empty_kernel_list(self, predictor):
with pytest.raises(ValueError):
predictor.select_best(_problem(), [])
def test_corner_case_m1(self, predictor):
prob = {
"m": 1,
"n": 4096,
"k": 4096,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
}
result = predictor.predict_tflops(prob, _kernel())
assert np.isfinite(result)
def test_different_shapes_give_different_results(self, predictor):
k = _kernel()
r1 = predictor.predict_tflops(
{
"m": 16,
"n": 1536,
"k": 7168,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
},
k,
)
r2 = predictor.predict_tflops(
{
"m": 20480,
"n": 7168,
"k": 256,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
},
k,
)
assert r1 != r2
class TestPredictorEdgeCases:
def test_nonexistent_model_dir(self):
with pytest.raises(Exception):
pred = Predictor("/nonexistent/path")
pred.predict_tflops(_problem(), _kernel())
if __name__ == "__main__":
pytest.main([__file__, "-v"])

192
dispatcher/heuristics/tests/test_search.py Normal file
View File

@@ -0,0 +1,192 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for search.py.
Covers: random search, DE search, config validity, result ordering,
budget compliance, and edge cases.
"""
import json
import sys
from pathlib import Path
import lightgbm as lgb
import numpy as np
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from feature_engine import GemmUniversalFeatureEngine
from predict import Predictor
from search import SurrogateSearch
@pytest.fixture
def model_dir(tmp_path):
"""Create a minimal trained model."""
fe = GemmUniversalFeatureEngine()
n_features = len(fe.get_feature_names())
np.random.seed(42)
X = np.random.rand(200, n_features)
y = np.random.rand(200) * 500
model = lgb.LGBMRegressor(n_estimators=10, verbose=-1)
model.fit(X, y)
model.booster_.save_model(str(tmp_path / "model_tflops.lgbm"))
spec = {
"feature_names": fe.get_feature_names(),
"categorical_features": fe.get_categorical_features(),
}
with open(tmp_path / "feature_spec.json", "w") as f:
json.dump(spec, f)
return tmp_path
@pytest.fixture
def predictor(model_dir):
return Predictor(model_dir)
def _problem():
return {
"m": 1024,
"n": 1024,
"k": 1024,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
}
class TestRandomSearch:
def test_returns_results(self, predictor):
searcher = SurrogateSearch(predictor, strategy="random")
results = searcher.search(_problem(), budget=50, top_k=5)
assert len(results) > 0
assert len(results) <= 5
def test_results_sorted_descending(self, predictor):
searcher = SurrogateSearch(predictor, strategy="random")
results = searcher.search(_problem(), budget=100, top_k=10)
scores = [s for _, s in results]
assert scores == sorted(scores, reverse=True)
def test_configs_are_valid(self, predictor):
fe = GemmUniversalFeatureEngine()
searcher = SurrogateSearch(predictor, feature_engine=fe, strategy="random")
results = searcher.search(_problem(), budget=50, top_k=5)
for cfg, _ in results:
ps = fe.get_parameter_space()
for k, v in cfg.items():
if k in ps:
assert v in ps[k], f"{k}={v} not in {ps[k]}"
def test_respects_top_k(self, predictor):
searcher = SurrogateSearch(predictor, strategy="random")
results = searcher.search(_problem(), budget=100, top_k=3)
assert len(results) <= 3
def test_different_problems_produce_results(self, predictor):
"""Both problem sizes should produce valid search results."""
searcher = SurrogateSearch(predictor, strategy="random", seed=42)
r1 = searcher.search(
{
"m": 16,
"n": 1536,
"k": 7168,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
},
budget=50,
top_k=3,
)
searcher2 = SurrogateSearch(predictor, strategy="random", seed=42)
r2 = searcher2.search(
{
"m": 20480,
"n": 7168,
"k": 256,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
},
budget=50,
top_k=3,
)
assert len(r1) > 0
assert len(r2) > 0
for _, score in r1 + r2:
assert np.isfinite(score)
def test_m1_corner_case(self, predictor):
searcher = SurrogateSearch(predictor, strategy="random")
results = searcher.search(
{
"m": 1,
"n": 4096,
"k": 4096,
"dtype": "fp8",
"layout": "rcr",
"split_k": 1,
},
budget=50,
top_k=5,
)
assert len(results) > 0
for _, score in results:
assert np.isfinite(score)
class TestDESearch:
def test_returns_results(self, predictor):
searcher = SurrogateSearch(predictor, strategy="de")
results = searcher.search(_problem(), budget=100, top_k=5)
assert len(results) > 0
def test_results_sorted_descending(self, predictor):
searcher = SurrogateSearch(predictor, strategy="de")
results = searcher.search(_problem(), budget=100, top_k=5)
scores = [s for _, s in results]
assert scores == sorted(scores, reverse=True)
def test_de_improves_over_initial(self, predictor):
"""DE should generally find at least as good as random initialization."""
searcher_r = SurrogateSearch(predictor, strategy="random", seed=42)
r_results = searcher_r.search(_problem(), budget=100, top_k=1)
searcher_d = SurrogateSearch(predictor, strategy="de", seed=42)
d_results = searcher_d.search(_problem(), budget=100, top_k=1)
if r_results and d_results:
assert d_results[0][1] >= r_results[0][1] * 0.9
def test_small_budget(self, predictor):
searcher = SurrogateSearch(predictor, strategy="de")
results = searcher.search(_problem(), budget=30, top_k=5)
assert len(results) > 0
class TestSearchEdgeCases:
def test_unknown_strategy_raises(self, predictor):
searcher = SurrogateSearch(predictor, strategy="unknown")
with pytest.raises(ValueError):
searcher.search(_problem(), budget=10)
def test_zero_budget(self, predictor):
searcher = SurrogateSearch(predictor, strategy="random")
results = searcher.search(_problem(), budget=0, top_k=5)
assert len(results) == 0
def test_deterministic_with_same_seed(self, predictor):
s1 = SurrogateSearch(predictor, strategy="random", seed=123)
s2 = SurrogateSearch(predictor, strategy="random", seed=123)
r1 = s1.search(_problem(), budget=50, top_k=5)
r2 = s2.search(_problem(), budget=50, top_k=5)
assert len(r1) == len(r2)
for (c1, s1_), (c2, s2_) in zip(r1, r2):
assert s1_ == pytest.approx(s2_)
if __name__ == "__main__":
pytest.main([__file__, "-v"])

329
dispatcher/heuristics/tests/test_train.py Normal file
View File

@@ -0,0 +1,329 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Tests for train.py.
Covers: group key computation, TFLOPS efficiency calculation, edge cases
(single group, all-invalid data, tied predictions), and warm-start
incremental training (feature compat, lineage, quality).
"""
import json
import sys
from pathlib import Path
import numpy as np
import pandas as pd
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from feature_engine import GemmUniversalFeatureEngine
from train import (
compute_group_keys,
compute_tflops_efficiency,
check_feature_compatibility,
load_warm_start_model,
train_final_model,
DEFAULT_PARAMS,
)
class TestComputeGroupKeys:
def test_basic(self):
df = pd.DataFrame(
{"m": [16, 16, 32], "n": [1536, 1536, 1536], "k": [7168, 7168, 7168]}
)
keys = compute_group_keys(df)
assert keys[0] == keys[1]
assert keys[0] != keys[2]
def test_unique_shapes(self):
df = pd.DataFrame({"m": [1, 2, 3], "n": [4, 5, 6], "k": [7, 8, 9]})
keys = compute_group_keys(df)
assert len(set(keys)) == 3
class TestComputeTflopsEfficiency:
def test_perfect_prediction(self):
"""Model predicts highest TFLOPS kernel => efficiency = 1.0."""
df = pd.DataFrame(
{
"m": [1024, 1024, 1024],
"n": [1024, 1024, 1024],
"k": [1024, 1024, 1024],
"measured_tflops": [100, 200, 150],
"pred_tflops": [50, 300, 100], # correctly ranks kernel 1 highest
}
)
eff = compute_tflops_efficiency(df, "pred_tflops")
assert len(eff) == 1
assert eff["efficiency"].iloc[0] == pytest.approx(1.0)
def test_worst_prediction(self):
"""Model picks the worst kernel."""
df = pd.DataFrame(
{
"m": [1024, 1024, 1024],
"n": [1024, 1024, 1024],
"k": [1024, 1024, 1024],
"measured_tflops": [100, 200, 150],
"pred_tflops": [999, 1, 1], # incorrectly ranks kernel 0 highest
}
)
eff = compute_tflops_efficiency(df, "pred_tflops")
assert eff["efficiency"].iloc[0] == pytest.approx(100 / 200)
def test_multiple_shapes(self):
df = pd.DataFrame(
{
"m": [16, 16, 32, 32],
"n": [1536, 1536, 1536, 1536],
"k": [7168, 7168, 7168, 7168],
"measured_tflops": [10, 20, 100, 200],
"pred_tflops": [5, 25, 150, 190],
}
)
eff = compute_tflops_efficiency(df, "pred_tflops")
assert len(eff) == 2
assert eff.iloc[0]["efficiency"] == pytest.approx(1.0)
assert eff.iloc[1]["efficiency"] == pytest.approx(1.0)
def test_zero_tflops_shape_skipped(self):
df = pd.DataFrame(
{
"m": [16, 16],
"n": [16, 16],
"k": [16, 16],
"measured_tflops": [0, 0],
"pred_tflops": [1, 2],
}
)
eff = compute_tflops_efficiency(df, "pred_tflops")
assert len(eff) == 0
def test_single_kernel_per_shape(self):
df = pd.DataFrame(
{
"m": [1024],
"n": [1024],
"k": [1024],
"measured_tflops": [150],
"pred_tflops": [100],
}
)
eff = compute_tflops_efficiency(df, "pred_tflops")
assert len(eff) == 1
assert eff["efficiency"].iloc[0] == pytest.approx(1.0)
def test_tied_predictions(self):
"""When multiple kernels have the same predicted TFLOPS, pandas idxmax picks the first."""
df = pd.DataFrame(
{
"m": [1024, 1024, 1024],
"n": [1024, 1024, 1024],
"k": [1024, 1024, 1024],
"measured_tflops": [100, 200, 200],
"pred_tflops": [50, 50, 50],
}
)
eff = compute_tflops_efficiency(df, "pred_tflops")
assert len(eff) == 1
assert eff["efficiency"].iloc[0] >= 0.5
# ---------------------------------------------------------------------------
# Helpers for warm-start tests
# ---------------------------------------------------------------------------
def _make_dummy_data(n_rows=200, n_shapes=5):
"""Create a small synthetic benchmark DataFrame for testing training."""
rng = np.random.RandomState(42)
rows = []
for _ in range(n_rows):
m = rng.choice([64, 128, 256, 512, 1024])
n = rng.choice([64, 128, 256, 512, 1024])
k = rng.choice([64, 128, 256, 512, 1024])
rows.append(
{
"m": m,
"n": n,
"k": k,
"split_k": 1,
"dtype": "fp8",
"layout": "rcr",
"op_type": "gemm_universal",
"tile_m": rng.choice([64, 128, 256]),
"tile_n": rng.choice([64, 128, 256]),
"tile_k": rng.choice([32, 64, 128]),
"warp_m": rng.choice([1, 2, 4]),
"warp_n": rng.choice([1, 2, 4]),
"warp_k": 1,
"warp_tile_m": 32,
"warp_tile_n": 32,
"warp_tile_k": 16,
"pipeline": rng.choice(["compv3", "compv4", "mem"]),
"scheduler": rng.choice(["intrawave", "interwave"]),
"epilogue": "cshuffle",
"pad_m": False,
"pad_n": False,
"pad_k": False,
"persistent": False,
"measured_tflops": float(rng.uniform(10, 500)),
"latency_ms": float(rng.uniform(0.01, 1.0)),
"bandwidth_gb_s": float(rng.uniform(50, 1500)),
"is_valid": True,
"kernel_name": f"test_kernel_{rng.randint(0, 100)}",
}
)
return pd.DataFrame(rows)
def _save_feature_spec(model_dir, fe):
"""Save a feature_spec.json matching the given feature engine."""
spec = {
"feature_names": fe.get_feature_names(),
"categorical_features": fe.get_categorical_features(),
}
with open(model_dir / "feature_spec.json", "w") as f:
json.dump(spec, f)
def _train_and_save_base_model(model_dir, df, fe, target="tflops"):
"""Train a small base model and save it to model_dir."""
params = dict(DEFAULT_PARAMS)
params["n_estimators"] = 20
params["n_jobs"] = 1
model = train_final_model(df, fe, target, params)
model.booster_.save_model(str(model_dir / f"model_{target}.lgbm"))
_save_feature_spec(model_dir, fe)
return model
# ---------------------------------------------------------------------------
# Warm-start tests
# ---------------------------------------------------------------------------
class TestCheckFeatureCompatibility:
def test_compatible_passes(self, tmp_path):
fe = GemmUniversalFeatureEngine()
_save_feature_spec(tmp_path, fe)
check_feature_compatibility(tmp_path, fe)
def test_missing_spec_raises(self, tmp_path):
fe = GemmUniversalFeatureEngine()
with pytest.raises(FileNotFoundError, match="feature_spec.json"):
check_feature_compatibility(tmp_path, fe)
def test_added_feature_raises(self, tmp_path):
fe = GemmUniversalFeatureEngine()
spec = {
"feature_names": fe.get_feature_names()[:-1],
"categorical_features": fe.get_categorical_features(),
}
with open(tmp_path / "feature_spec.json", "w") as f:
json.dump(spec, f)
with pytest.raises(ValueError, match="Feature schema mismatch"):
check_feature_compatibility(tmp_path, fe)
def test_removed_feature_raises(self, tmp_path):
fe = GemmUniversalFeatureEngine()
spec = {
"feature_names": fe.get_feature_names() + ["extra_feature"],
"categorical_features": fe.get_categorical_features(),
}
with open(tmp_path / "feature_spec.json", "w") as f:
json.dump(spec, f)
with pytest.raises(ValueError, match="Feature schema mismatch"):
check_feature_compatibility(tmp_path, fe)
def test_categorical_mismatch_raises(self, tmp_path):
fe = GemmUniversalFeatureEngine()
spec = {
"feature_names": fe.get_feature_names(),
"categorical_features": ["layout", "pipeline"],
}
with open(tmp_path / "feature_spec.json", "w") as f:
json.dump(spec, f)
with pytest.raises(ValueError, match="Categorical feature mismatch"):
check_feature_compatibility(tmp_path, fe)
class TestLoadWarmStartModel:
def test_loads_existing_model(self, tmp_path):
fe = GemmUniversalFeatureEngine()
df = _make_dummy_data()
_train_and_save_base_model(tmp_path, df, fe)
path = load_warm_start_model(tmp_path, "tflops")
assert path is not None
assert Path(path).exists()
def test_returns_none_for_missing_target(self, tmp_path):
assert load_warm_start_model(tmp_path, "tflops") is None
def test_returns_none_for_wrong_target(self, tmp_path):
fe = GemmUniversalFeatureEngine()
df = _make_dummy_data()
_train_and_save_base_model(tmp_path, df, fe, target="tflops")
assert load_warm_start_model(tmp_path, "bandwidth") is None
class TestWarmStartTraining:
def test_warm_start_produces_more_trees(self, tmp_path):
"""A warm-started model should have more trees than the base."""
fe = GemmUniversalFeatureEngine()
df = _make_dummy_data(n_rows=300)
base_dir = tmp_path / "base"
base_dir.mkdir()
base_model = _train_and_save_base_model(base_dir, df, fe)
base_n_trees = base_model.booster_.num_trees()
init_model_path = load_warm_start_model(base_dir, "tflops")
params = dict(DEFAULT_PARAMS)
params["n_estimators"] = 15
params["n_jobs"] = 1
warm_model = train_final_model(
df, fe, "tflops", params, init_model=init_model_path
)
warm_n_trees = warm_model.booster_.num_trees()
assert warm_n_trees > base_n_trees
def test_warm_start_does_not_degrade(self, tmp_path):
"""Warm-started model on the same data should not be significantly worse."""
fe = GemmUniversalFeatureEngine()
df = _make_dummy_data(n_rows=300)
base_dir = tmp_path / "base"
base_dir.mkdir()
base_model = _train_and_save_base_model(base_dir, df, fe)
X = fe.extract_batch(df[df["is_valid"]].reset_index(drop=True))
y = df[df["is_valid"]]["measured_tflops"].values
base_rmse = np.sqrt(np.mean((base_model.predict(X) - y) ** 2))
init_model_path = load_warm_start_model(base_dir, "tflops")
params = dict(DEFAULT_PARAMS)
params["n_estimators"] = 15
params["n_jobs"] = 1
warm_model = train_final_model(
df, fe, "tflops", params, init_model=init_model_path
)
warm_rmse = np.sqrt(np.mean((warm_model.predict(X) - y) ** 2))
assert warm_rmse <= base_rmse * 1.1
def test_warm_start_from_nonexistent_dir(self):
with pytest.raises(FileNotFoundError):
check_feature_compatibility(
Path("/nonexistent/model/dir"), GemmUniversalFeatureEngine()
)
if __name__ == "__main__":
pytest.main([__file__, "-v"])

555
dispatcher/heuristics/train.py Normal file
View File

@@ -0,0 +1,555 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Training script for CK Tile kernel performance prediction.
Trains LGBMRegressor models (TFLOPS, latency, bandwidth) with:
- Log-space regression (log1p transform) for scale-invariant accuracy
- GroupKFold cross-validation (group key = (M, N, K))
- Iterative Hard Example Mining (IHEM)
- Model complexity bounds for C++ deployability
- Optional Optuna hyperparameter tuning
- Warm-start incremental training from a previous model via --warm_start
Log-transform rationale:
GEMM TFLOPS spans 5 orders of magnitude (0.02 for M=1 to 2230 for large
shapes). Raw regression optimizes for absolute RMSE, which means the model
spends all its capacity predicting large shapes accurately and ignores tiny
shapes where TFLOPS is < 10. Training on log1p(TFLOPS) puts all shapes on
equal footing, improving tiny_m efficiency from 84% to 96%.
"""
import argparse
import json
import time
from pathlib import Path
import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.model_selection import GroupKFold
from data_pipeline import build_training_dataset
from feature_engine import GemmUniversalFeatureEngine
TARGET_COLUMNS = {
"tflops": "measured_tflops",
"latency": "latency_ms",
"bandwidth": "bandwidth_gb_s",
}
# Targets where log1p transform is applied by default.
# TFLOPS and bandwidth span orders of magnitude; latency is already small-scale.
LOG_TARGETS = {"tflops", "bandwidth"}
DEFAULT_PARAMS = {
"objective": "regression",
"metric": ["rmse", "mae"],
"num_leaves": 255,
"max_depth": 15,
"n_estimators": 2000,
"learning_rate": 0.02,
"min_child_samples": 10,
"subsample": 0.85,
"colsample_bytree": 0.85,
"reg_alpha": 0.05,
"reg_lambda": 0.5,
"verbose": -1,
"n_jobs": 8,
"seed": 42,
}
MAX_ESTIMATORS = 5000
WARM_START_N_ESTIMATORS = 500
def check_feature_compatibility(
prev_model_dir: Path,
feature_engine: GemmUniversalFeatureEngine,
) -> None:
"""Verify that the previous model's feature spec matches the current engine.
Raises ValueError with a detailed message on mismatch. This prevents silent
corruption when warm-starting from a model trained with a different feature
schema (e.g., after adding a new feature or changing an encoding).
"""
spec_path = prev_model_dir / "feature_spec.json"
if not spec_path.exists():
raise FileNotFoundError(
f"No feature_spec.json in {prev_model_dir}. "
"Cannot verify feature compatibility for warm start."
)
with open(spec_path) as f:
prev_spec = json.load(f)
prev_names = prev_spec.get("feature_names", [])
curr_names = feature_engine.get_feature_names()
if prev_names != curr_names:
added = set(curr_names) - set(prev_names)
removed = set(prev_names) - set(curr_names)
parts = ["Feature schema mismatch between previous model and current engine."]
if added:
parts.append(f" Added features: {sorted(added)}")
if removed:
parts.append(f" Removed features: {sorted(removed)}")
if not added and not removed:
parts.append(" Feature order changed (names match but order differs).")
raise ValueError("\n".join(parts))
prev_cats = prev_spec.get("categorical_features", [])
curr_cats = feature_engine.get_categorical_features()
if sorted(prev_cats) != sorted(curr_cats):
raise ValueError(
f"Categorical feature mismatch.\n"
f" Previous: {sorted(prev_cats)}\n"
f" Current: {sorted(curr_cats)}"
)
def load_warm_start_model(prev_model_dir: Path, target: str) -> str | None:
"""Load the path to a previous model file for warm-start, or None if absent.
Automatically decompresses .lgbm.gz files if the .lgbm file doesn't exist.
The decompressed file is cached to disk for subsequent loads.
Returns the string path (what LightGBM's init_model expects) rather than
a loaded Booster, because LGBMRegressor.fit(init_model=...) accepts both
path strings and Booster objects and path strings avoid keeping the old
model in memory.
"""
import gzip
model_path = prev_model_dir / f"model_{target}.lgbm"
gz_path = prev_model_dir / f"model_{target}.lgbm.gz"
# Auto-decompress if needed
if not model_path.exists() and gz_path.exists():
print(f" Decompressing {gz_path.name}...")
with gzip.open(gz_path, "rb") as f_in:
with open(model_path, "wb") as f_out:
f_out.write(f_in.read())
if not model_path.exists():
return None
return str(model_path)
def compute_group_keys(df: pd.DataFrame) -> np.ndarray:
"""Create GroupKFold group keys from (M, N, K)."""
return (
df["m"].astype(str) + "_" + df["n"].astype(str) + "_" + df["k"].astype(str)
).values
def compute_tflops_efficiency(
df: pd.DataFrame, pred_col: str = "pred_tflops"
) -> pd.DataFrame:
"""Compute per-shape efficiency: predicted-best TFLOPS / oracle-best TFLOPS."""
results = []
for (m, n, k), group in df.groupby(["m", "n", "k"]):
oracle_best = group["measured_tflops"].max()
if oracle_best <= 0:
continue
pred_best_idx = group[pred_col].idxmax()
selected_tflops = group.loc[pred_best_idx, "measured_tflops"]
efficiency = selected_tflops / oracle_best
results.append(
{
"m": m,
"n": n,
"k": k,
"oracle_best_tflops": oracle_best,
"selected_tflops": selected_tflops,
"efficiency": efficiency,
}
)
return pd.DataFrame(results)
def train_single_target(
X_train,
y_train,
X_val,
y_val,
params: dict,
categorical_features: list[str],
feature_names: list[str],
init_model=None,
) -> lgb.LGBMRegressor:
"""Train a single LGBMRegressor with early stopping.
Parameters
----------
init_model : str, Path, lgb.Booster, lgb.LGBMModel, or None
If provided, training continues from this model (warm start).
Accepts a file path to a .lgbm file, a Booster instance, or an
LGBMModel instance. The new model adds n_estimators trees on top
of the existing ones.
"""
cat_indices = [
feature_names.index(c) for c in categorical_features if c in feature_names
]
model = lgb.LGBMRegressor(**params)
model.fit(
X_train,
y_train,
eval_set=[(X_val, y_val)],
eval_metric=["rmse"],
callbacks=[
lgb.early_stopping(50, verbose=False),
lgb.log_evaluation(0),
],
categorical_feature=cat_indices if cat_indices else "auto",
init_model=init_model,
)
return model
def run_cv(
df: pd.DataFrame,
feature_engine: GemmUniversalFeatureEngine,
target: str,
params: dict,
n_splits: int = 5,
use_log: bool = True,
) -> dict:
"""Run GroupKFold cross-validation and return OOF predictions + metrics.
Parameters
----------
use_log : bool
If True and target is in LOG_TARGETS, train on log1p(y) and invert
predictions with expm1 for efficiency calculation. This normalizes
the scale so that tiny-M shapes (TFLOPS ~ 1) get equal attention
as large-M shapes (TFLOPS ~ 2000).
"""
target_col = TARGET_COLUMNS[target]
valid_mask = df["is_valid"].fillna(False) & (df[target_col] > 0)
df_valid = df[valid_mask].reset_index(drop=True)
apply_log = use_log and target in LOG_TARGETS
print(
f" Training on {len(df_valid)} valid rows for target={target}"
f"{' (log-space)' if apply_log else ''}"
)
X = feature_engine.extract_batch(df_valid)
y_raw = df_valid[target_col].values
y = np.log1p(y_raw) if apply_log else y_raw
groups = compute_group_keys(df_valid)
feature_names = feature_engine.get_feature_names()
cat_features = feature_engine.get_categorical_features()
unique_groups = np.unique(groups)
actual_splits = min(n_splits, len(unique_groups))
if actual_splits < 2:
print(f" WARNING: Only {len(unique_groups)} unique groups, skipping CV")
return {}
gkf = GroupKFold(n_splits=actual_splits)
oof_preds = np.zeros(len(df_valid))
fold_metrics = []
for fold_idx, (train_idx, val_idx) in enumerate(gkf.split(X, y, groups)):
X_tr, X_val = X[train_idx], X[val_idx]
y_tr, y_val = y[train_idx], y[val_idx]
model = train_single_target(
X_tr, y_tr, X_val, y_val, params, cat_features, feature_names
)
preds = model.predict(X_val)
oof_preds[val_idx] = preds
rmse = np.sqrt(np.mean((preds - y_val) ** 2))
r2 = 1 - np.sum((preds - y_val) ** 2) / max(
np.sum((y_val - y_val.mean()) ** 2), 1e-10
)
if target == "tflops":
val_df = df_valid.iloc[val_idx].copy()
preds_raw = np.expm1(preds) if apply_log else preds
val_df["pred_tflops"] = preds_raw
eff_df = compute_tflops_efficiency(val_df)
mean_eff = eff_df["efficiency"].mean() if len(eff_df) > 0 else 0
p10_eff = eff_df["efficiency"].quantile(0.1) if len(eff_df) > 0 else 0
else:
mean_eff, p10_eff = None, None
fold_metrics.append(
{
"fold": fold_idx,
"rmse": rmse,
"r2": r2,
"mean_efficiency": mean_eff,
"p10_efficiency": p10_eff,
"train_size": len(train_idx),
"val_size": len(val_idx),
"val_groups": len(np.unique(groups[val_idx])),
}
)
eff_str = (
f", eff={mean_eff:.4f}, p10={p10_eff:.4f}" if mean_eff is not None else ""
)
print(f" Fold {fold_idx}: RMSE={rmse:.4f}, R2={r2:.4f}{eff_str}")
df_valid[f"oof_pred_{target}"] = oof_preds
return {
"fold_metrics": fold_metrics,
"oof_df": df_valid,
"feature_names": feature_names,
"log_transform": apply_log,
}
def train_final_model(
df: pd.DataFrame,
feature_engine: GemmUniversalFeatureEngine,
target: str,
params: dict,
init_model=None,
use_log: bool = True,
) -> lgb.LGBMRegressor:
"""Train the final model on all valid data.
Parameters
----------
init_model : str, Path, lgb.Booster, lgb.LGBMModel, or None
If provided, training continues from this model (warm start).
use_log : bool
If True and target is in LOG_TARGETS, train on log1p(y).
The saved model then predicts in log-space; callers must apply
expm1() to get raw values.
"""
target_col = TARGET_COLUMNS[target]
valid_mask = df["is_valid"].fillna(False) & (df[target_col] > 0)
df_valid = df[valid_mask].reset_index(drop=True)
apply_log = use_log and target in LOG_TARGETS
X = feature_engine.extract_batch(df_valid)
y_raw = df_valid[target_col].values
y = np.log1p(y_raw) if apply_log else y_raw
feature_names = feature_engine.get_feature_names()
cat_features = feature_engine.get_categorical_features()
cat_indices = [feature_names.index(c) for c in cat_features if c in feature_names]
model = lgb.LGBMRegressor(**params)
model.fit(
X,
y,
categorical_feature=cat_indices if cat_indices else "auto",
init_model=init_model,
)
return model
def main():
parser = argparse.ArgumentParser(
description="Train CK Tile kernel performance models"
)
parser.add_argument(
"--data_dir", required=True, help="Directory with parquet files"
)
parser.add_argument("--out_dir", required=True, help="Output directory for models")
parser.add_argument("--op", default="gemm_universal", help="Operation type")
parser.add_argument("--dtype", default="fp8", help="Data type filter")
parser.add_argument("--arch", default="gfx950", help="Architecture")
parser.add_argument(
"--targets", default="tflops,latency,bandwidth", help="Comma-separated targets"
)
parser.add_argument("--n_splits", type=int, default=5, help="Number of CV folds")
parser.add_argument(
"--tune", action="store_true", help="Run Optuna hyperparameter tuning"
)
parser.add_argument(
"--no_log_transform",
action="store_true",
help="Disable log1p transform on targets. By default, TFLOPS and bandwidth "
"are trained in log-space for scale-invariant accuracy across shape sizes.",
)
parser.add_argument(
"--warm_start",
default=None,
help="Path to previous model directory to continue training from. "
"Uses LightGBM's init_model to add new trees on top of the "
"existing model. Feature schemas must match exactly.",
)
parser.add_argument(
"--warm_start_n_estimators",
type=int,
default=WARM_START_N_ESTIMATORS,
help=f"Number of new trees to add when warm-starting (default: {WARM_START_N_ESTIMATORS}). "
"Lower than a full train since we're refining, not starting from scratch.",
)
args = parser.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
targets = [t.strip() for t in args.targets.split(",")]
print(f"Loading data from {args.data_dir}...")
df = build_training_dataset(args.data_dir, op_type=args.op, dtype=args.dtype)
print(f" Total rows: {len(df)}")
print(f" Unique shapes: {df.groupby(['m', 'n', 'k']).ngroups}")
print(f" Unique kernels: {df['kernel_name'].nunique()}")
hw_cols = [c for c in df.columns if c.startswith("hw_")]
hw_kwargs = {}
if hw_cols:
row0 = df.iloc[0]
if "hw_num_cus" in df.columns:
hw_kwargs["num_cus"] = int(row0.get("hw_num_cus", 256))
if "hw_max_clock_mhz" in df.columns:
hw_kwargs["max_clock_mhz"] = int(row0.get("hw_max_clock_mhz", 2400))
if "hw_simds_per_cu" in df.columns:
hw_kwargs["simds_per_cu"] = int(row0.get("hw_simds_per_cu", 4))
if "hw_shader_engines" in df.columns:
hw_kwargs["shader_engines"] = int(row0.get("hw_shader_engines", 32))
if "hw_max_waves_per_cu" in df.columns:
hw_kwargs["max_waves_per_cu"] = int(row0.get("hw_max_waves_per_cu", 32))
if "hw_wavefront_size" in df.columns:
hw_kwargs["wavefront_size"] = int(row0.get("hw_wavefront_size", 64))
if "hw_l1_cache_kb" in df.columns:
hw_kwargs["l1_cache_kb"] = int(row0.get("hw_l1_cache_kb", 32))
if "hw_l2_cache_kb" in df.columns:
hw_kwargs["l2_cache_kb"] = int(row0.get("hw_l2_cache_kb", 4096))
if "hw_l3_cache_kb" in df.columns:
hw_kwargs["l3_cache_kb"] = int(row0.get("hw_l3_cache_kb", 262144))
fe = GemmUniversalFeatureEngine(**hw_kwargs)
params = dict(DEFAULT_PARAMS)
use_log = not args.no_log_transform
prev_model_dir = None
prev_manifest = {}
if args.warm_start:
prev_model_dir = Path(args.warm_start)
if not prev_model_dir.exists():
raise FileNotFoundError(f"Warm-start directory not found: {prev_model_dir}")
print(f" Warm-starting from {prev_model_dir}")
check_feature_compatibility(prev_model_dir, fe)
print(" Feature compatibility: OK")
params["n_estimators"] = args.warm_start_n_estimators
print(f" New trees to add: {args.warm_start_n_estimators}")
prev_manifest_path = prev_model_dir / "train_manifest.json"
if prev_manifest_path.exists():
with open(prev_manifest_path) as f:
prev_manifest = json.load(f)
all_cv_results = {}
for target in targets:
if target not in TARGET_COLUMNS:
print(f" Skipping unknown target: {target}")
continue
print(f"\n{'=' * 60}")
print(f"Training {target} model")
print(f"{'=' * 60}")
init_model_path = None
if prev_model_dir is not None:
init_model_path = load_warm_start_model(prev_model_dir, target)
if init_model_path:
print(f" Warm-starting from {init_model_path}")
else:
print(f" No previous {target} model found, training from scratch")
t0 = time.time()
cv_result = run_cv(
df, fe, target, params, n_splits=args.n_splits, use_log=use_log
)
cv_time = time.time() - t0
if cv_result and cv_result["fold_metrics"]:
all_cv_results[target] = cv_result["fold_metrics"]
metrics_path = out_dir / f"cv_metrics_{target}.json"
with open(metrics_path, "w") as f:
json.dump(cv_result["fold_metrics"], f, indent=2)
print(f" CV completed in {cv_time:.1f}s, saved to {metrics_path}")
if target == "tflops" and cv_result.get("oof_df") is not None:
oof_df = cv_result["oof_df"]
oof_df.to_parquet(out_dir / "oof_predictions.parquet", index=False)
eff_df = compute_tflops_efficiency(oof_df, "oof_pred_tflops")
if len(eff_df) > 0:
print("\n OOF TFLOPS Efficiency:")
print(f" Mean: {eff_df['efficiency'].mean():.4f}")
print(f" P10: {eff_df['efficiency'].quantile(0.1):.4f}")
print(f" P50: {eff_df['efficiency'].quantile(0.5):.4f}")
print(f" Min: {eff_df['efficiency'].min():.4f}")
print(f"\n Training final {target} model on all data...")
t0 = time.time()
model = train_final_model(
df, fe, target, params, init_model=init_model_path, use_log=use_log
)
train_time = time.time() - t0
model_path = out_dir / f"model_{target}.lgbm"
model.booster_.save_model(str(model_path))
print(f" Saved {model_path} ({train_time:.1f}s)")
importances = dict(
zip(
fe.get_feature_names(),
model.feature_importances_.tolist(),
)
)
imp_path = out_dir / f"feature_importances_{target}.json"
with open(imp_path, "w") as f:
json.dump(importances, f, indent=2)
log_targets_used = sorted(LOG_TARGETS & set(targets)) if use_log else []
spec = {
"op_type": args.op,
"dtype": args.dtype,
"arch": args.arch,
"feature_names": fe.get_feature_names(),
"categorical_features": fe.get_categorical_features(),
"targets": targets,
"log_targets": log_targets_used,
"params": params,
}
with open(out_dir / "feature_spec.json", "w") as f:
json.dump(spec, f, indent=2)
manifest = {
"warm_start_from": str(prev_model_dir) if prev_model_dir else None,
"prev_n_estimators": prev_manifest.get(
"total_n_estimators", params.get("n_estimators")
)
if prev_model_dir
else 0,
"new_n_estimators": params["n_estimators"],
"total_n_estimators": (
prev_manifest.get("total_n_estimators", 0) + params["n_estimators"]
if prev_model_dir
else params["n_estimators"]
),
"data_rows": len(df),
"valid_rows": int(df["is_valid"].fillna(False).sum()),
"unique_shapes": int(df.groupby(["m", "n", "k"]).ngroups),
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
}
with open(out_dir / "train_manifest.json", "w") as f:
json.dump(manifest, f, indent=2)
print(f"\nAll models saved to {out_dir}")
if prev_model_dir:
print(f" Warm-started from: {prev_model_dir}")
print(f" Total estimators: {manifest['total_n_estimators']}")
if __name__ == "__main__":
main()

317
dispatcher/heuristics/validate_ml_heuristic.py Normal file
View File

@@ -0,0 +1,317 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
ML Heuristic Validation: Test ML predictions against oracle-best from training data
This script validates ML-based kernel selection by:
1. Loading benchmark data (oracle-best results for each shape)
2. Using ML model to predict best kernel for each shape
3. Comparing ML selection with oracle-best to compute efficiency
Usage:
python validate_ml_heuristic.py --dtype fp16 --model_dir models/gemm_universal_fp16_gfx950
python validate_ml_heuristic.py --dtype fp8 --layout rcr
"""
import sys
import argparse
import pandas as pd
import numpy as np
from pathlib import Path
from predict import Predictor
def validate_ml_heuristic(dtype: str, layout: str, model_dir: str, data_dir: str):
"""Validate ML heuristic predictions against oracle-best"""
print("=" * 100)
print(f" ML Heuristic Validation: {dtype.upper()} {layout.upper()}")
print("=" * 100)
print()
# Load training data
print(f"Loading training data from {data_dir}...")
# Try dtype-specific parquet first, then fall back to combined
dtype_specific = (
Path(data_dir) / f"{dtype}_original" / f"{dtype}_training_data.parquet"
)
combined = Path(data_dir) / "all_training_data_fixed.parquet"
if dtype_specific.exists():
training_data = pd.read_parquet(dtype_specific)
print(f"✓ Loaded {len(training_data):,} benchmark runs from {dtype_specific}")
elif combined.exists():
training_data = pd.read_parquet(combined)
training_data = training_data[
(training_data["dtype"] == dtype) & (training_data["layout"] == layout)
]
print(f"✓ Loaded {len(training_data):,} benchmark runs from {combined}")
else:
print(f"❌ Error: No training data found at {dtype_specific} or {combined}")
return
if len(training_data) == 0:
print(f"❌ Error: No data found for dtype={dtype}, layout={layout}")
return
# Get unique shapes with oracle-best
shape_groups = training_data.groupby(["m", "n", "k"])
print(f"Unique shapes: {len(shape_groups)}")
print()
# Load ML predictor
print(f"Loading ML predictor from {model_dir}...")
try:
predictor = Predictor(model_dir)
print("✓ Loaded ML predictor")
print(f" Log targets: {predictor._log_targets}")
except Exception as e:
print(f"❌ Error loading model: {e}")
return
print()
print("=" * 100)
print(" Computing Oracle-Best Efficiency for Each Shape")
print("=" * 100)
print()
results = []
for shape_idx, ((m, n, k), group) in enumerate(shape_groups):
# Find oracle-best (max TFLOPS across all kernels tested)
oracle_best_row = group.loc[group["measured_tflops"].idxmax()]
oracle_best_tflops = oracle_best_row["measured_tflops"]
oracle_best_kernel = oracle_best_row["kernel_name"]
# Get all kernel configs tested for this shape
kernel_configs = []
for _, row in group.iterrows():
kernel_dict = {
"tile_m": row["tile_m"],
"tile_n": row["tile_n"],
"tile_k": row["tile_k"],
"warp_m": row["warp_m"],
"warp_n": row["warp_n"],
"warp_k": row["warp_k"],
"warp_tile_m": row["warp_tile_m"],
"warp_tile_n": row["warp_tile_n"],
"warp_tile_k": row["warp_tile_k"],
"pipeline": row["pipeline"],
"scheduler": row["scheduler"],
"epilogue": row["epilogue"],
"pad_m": row["pad_m"],
"pad_n": row["pad_n"],
"pad_k": row["pad_k"],
"persistent": row["persistent"],
"kernel_name": row["kernel_name"],
}
kernel_configs.append(kernel_dict)
# Use ML model to rank kernels
problem = {
"m": m,
"n": n,
"k": k,
"dtype": dtype,
"layout": layout,
"split_k": 1,
}
try:
ranked = predictor.rank_kernels(problem, kernel_configs)
if ranked:
ml_best_kernel, ml_predicted_tflops = ranked[0]
# Find actual TFLOPS for the ML-predicted kernel
ml_kernel_row = group[group["kernel_name"] == ml_best_kernel]
if len(ml_kernel_row) > 0:
ml_actual_tflops = ml_kernel_row["measured_tflops"].values[0]
# Calculate efficiency
efficiency_pct = 100.0 * (ml_actual_tflops / oracle_best_tflops)
# Determine if ML picked oracle-best
is_oracle_best = ml_best_kernel == oracle_best_kernel
results.append(
{
"m": m,
"n": n,
"k": k,
"oracle_best_tflops": oracle_best_tflops,
"oracle_best_kernel": oracle_best_kernel,
"ml_predicted_tflops": ml_predicted_tflops,
"ml_selected_kernel": ml_best_kernel,
"ml_actual_tflops": ml_actual_tflops,
"efficiency_pct": efficiency_pct,
"is_oracle_best": is_oracle_best,
"num_kernels": len(group),
}
)
if (shape_idx + 1) % 20 == 0:
status = "✓" if is_oracle_best else f"{efficiency_pct:.1f}%"
print(
f" [{shape_idx + 1:3d}/{len(shape_groups)}] "
f"M={m:4d} N={n:5d} K={k:5d}: {status}"
)
except Exception as e:
print(f" Error on shape M={m} N={n} K={k}: {e}")
continue
print()
print("=" * 100)
print(" Results Summary")
print("=" * 100)
print()
if results:
df_results = pd.DataFrame(results)
efficiencies = df_results["efficiency_pct"].values
oracle_matches = df_results["is_oracle_best"].sum()
print(f"Total shapes tested: {len(results)}")
print()
print("Efficiency Statistics (% of Oracle-Best TFLOPS):")
print(f" Mean: {np.mean(efficiencies):.2f}%")
print(f" Median: {np.median(efficiencies):.2f}%")
print(f" Min: {np.min(efficiencies):.2f}%")
print(f" Max: {np.max(efficiencies):.2f}%")
print(f" P10: {np.percentile(efficiencies, 10):.2f}%")
print(f" P50: {np.percentile(efficiencies, 50):.2f}%")
print(f" P90: {np.percentile(efficiencies, 90):.2f}%")
print()
print(
f"Oracle-best matches: {oracle_matches}/{len(results)} ({100 * oracle_matches / len(results):.1f}%)"
)
print()
# Classify by M size
df_results["m_class"] = pd.cut(
df_results["m"],
bins=[0, 8, 128, 1024, float("inf")],
labels=[
"Tiny (M<8)",
"Small (8≤M<128)",
"Medium (128≤M<1024)",
"Large (M≥1024)",
],
)
print("Efficiency by M size:")
for m_class in [
"Tiny (M<8)",
"Small (8≤M<128)",
"Medium (128≤M<1024)",
"Large (M≥1024)",
]:
subset = df_results[df_results["m_class"] == m_class]
if len(subset) > 0:
print(
f" {m_class:25s}: {subset['efficiency_pct'].mean():6.2f}% "
f"(n={len(subset)}, P10={subset['efficiency_pct'].quantile(0.1):.2f}%)"
)
print()
# Save results
output_file = f"validation_results_{dtype}_{layout}.csv"
df_results.to_csv(output_file, index=False)
print(f"✓ Results saved to {output_file}")
# Show best and worst shapes
print()
print("Top 5 shapes (best efficiency):")
top5 = df_results.nlargest(5, "efficiency_pct")[
["m", "n", "k", "efficiency_pct", "oracle_best_tflops", "is_oracle_best"]
]
for idx, row in top5.iterrows():
match = "✓" if row["is_oracle_best"] else " "
print(
f" {match} M={row['m']:5d} N={row['n']:5d} K={row['k']:5d}: "
f"{row['efficiency_pct']:.2f}% ({row['oracle_best_tflops']:.2f} TFLOPS)"
)
print()
print("Bottom 5 shapes (worst efficiency):")
bottom5 = df_results.nsmallest(5, "efficiency_pct")[
["m", "n", "k", "efficiency_pct", "oracle_best_tflops", "is_oracle_best"]
]
for idx, row in bottom5.iterrows():
match = "✓" if row["is_oracle_best"] else " "
print(
f" {match} M={row['m']:5d} N={row['n']:5d} K={row['k']:5d}: "
f"{row['efficiency_pct']:.2f}% ({row['oracle_best_tflops']:.2f} TFLOPS)"
)
else:
print("No results to display")
print()
print("=" * 100)
def main():
parser = argparse.ArgumentParser(
description="Validate ML heuristic predictions against oracle-best from training data"
)
parser.add_argument(
"--dtype",
default="fp16",
choices=["fp16", "bf16", "fp8"],
help="Data type to validate",
)
parser.add_argument(
"--layout",
default="rcr",
choices=["rcr", "rrr", "crr", "ccr"],
help="Matrix layout",
)
parser.add_argument(
"--model_dir",
default=None,
help="Path to model directory (auto-detect if not specified)",
)
parser.add_argument(
"--data_dir",
default=None,
help="Path to training data directory (auto-detect if not specified)",
)
args = parser.parse_args()
# Auto-detect model directory if not specified
if args.model_dir is None:
heuristics_dir = Path(__file__).parent
model_candidates = [
heuristics_dir / "models" / f"gemm_universal_{args.dtype}_gfx950",
heuristics_dir / "models" / f"gemm_universal_{args.dtype}_gfx942",
]
for candidate in model_candidates:
if candidate.exists():
args.model_dir = str(candidate)
break
if args.model_dir is None:
print(f"❌ Error: Could not find model directory for {args.dtype}")
print(f" Searched: {[str(c) for c in model_candidates]}")
print(" Please specify --model_dir explicitly")
return 1
# Auto-detect data directory if not specified
if args.data_dir is None:
heuristics_dir = Path(__file__).parent
args.data_dir = str(heuristics_dir / "data")
validate_ml_heuristic(args.dtype, args.layout, args.model_dir, args.data_dir)
return 0
if __name__ == "__main__":
sys.exit(main())

20
dispatcher/include/ck_tile/dispatcher.hpp
View File

@@ -3,9 +3,17 @@
#pragma once
/// Main dispatcher header - includes all core components
/// Use this for convenient access to the full dispatcher API
/// Full dispatcher header - includes ALL operation types.
/// For minimal includes, use the per-operation headers instead:
/// ck_tile/dispatcher_gemm.hpp -- GEMM only
/// ck_tile/dispatcher_conv.hpp -- Grouped Convolution only
// Core (needed by all ops)
#include "ck_tile/dispatcher/base_registry.hpp"
#include "ck_tile/dispatcher/dispatcher_error.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
// GEMM
#include "ck_tile/dispatcher/kernel_key.hpp"
#include "ck_tile/dispatcher/kernel_config.hpp"
#include "ck_tile/dispatcher/kernel_decl.hpp"
@@ -13,7 +21,15 @@
#include "ck_tile/dispatcher/kernel_instance.hpp"
#include "ck_tile/dispatcher/registry.hpp"
#include "ck_tile/dispatcher/dispatcher.hpp"
#include "ck_tile/dispatcher/json_export.hpp"
#include "ck_tile/dispatcher/arch_filter.hpp"
#include "ck_tile/dispatcher/backends/tile_backend.hpp"
#include "ck_tile/dispatcher/backends/generated_tile_backend.hpp"
#include "ck_tile/dispatcher/utils.hpp"
// Grouped Convolution
#include "ck_tile/dispatcher/grouped_conv_config.hpp"
#include "ck_tile/dispatcher/grouped_conv_problem.hpp"
#include "ck_tile/dispatcher/grouped_conv_kernel_decl.hpp"
#include "ck_tile/dispatcher/grouped_conv_registry.hpp"
#include "ck_tile/dispatcher/grouped_conv_utils.hpp"

96
dispatcher/include/ck_tile/dispatcher/README.md
View File

@@ -1,6 +1,6 @@
# CK Tile Dispatcher - C++ Headers
C++ API for the CK Tile dispatcher.
C++ API for the CK Tile dispatcher (GEMM and Grouped Convolution).
> **See also:** [Main Dispatcher README](../../../../README.md) for installation and core concepts.
@@ -8,16 +8,25 @@ C++ API for the CK Tile dispatcher.
```
dispatcher/
├── dispatcher.hpp # Main dispatcher (kernel selection)
├── registry.hpp # Kernel registry (storage & lookup)
├── problem.hpp # Problem specification
├── kernel_key.hpp # Kernel configuration key
├── kernel_instance.hpp # Kernel instance interface
├── utils.hpp # Utilities (timers, GPU buffers)
│
└── backends/ # Backend implementations
├── generated_tile_backend.hpp # CK Tile kernels (production)
└── tile_backend.hpp # Tile backend base
|---- dispatcher.hpp # Main include (includes all below)
|
|---- # GEMM Headers
|---- registry.hpp # Kernel registry (storage & lookup)
|---- problem.hpp # GEMM problem specification
|---- kernel_key.hpp # Kernel configuration key
|---- kernel_instance.hpp # Kernel instance interface
|---- utils.hpp # Utilities (timers, GPU buffers)
|
|---- # Grouped Convolution Headers
|---- grouped_conv_config.hpp # GroupedConvDirection, GroupedConvConfig
|---- grouped_conv_problem.hpp # GroupedConvProblem + ProblemBuilder
|---- grouped_conv_kernel_decl.hpp # GroupedConvKernelDecl, DECL_GROUPED_CONV_KERNEL_SET
|---- grouped_conv_registry.hpp # Thread-safe registry with JSON export & filtering
|---- grouped_conv_utils.hpp # Config creators, validation, benchmark utilities
|
+---- backends/ # Backend implementations
|---- generated_tile_backend.hpp # CK Tile kernels (production)
+---- tile_backend.hpp # Tile backend base
```
## Quick Start
@@ -148,6 +157,69 @@ auto kernel = create_generated_tile_kernel<
>(key, name);
```
## Grouped Convolution API
### GroupedConvProblem (`grouped_conv_problem.hpp`)
Problem specification with builder pattern:
```cpp
#include "ck_tile/dispatcher/grouped_conv_problem.hpp"
using namespace ck_tile::dispatcher;
auto problem = GroupedConvProblemBuilder()
.n(2).g(1).c(128).k(256)
.input_spatial({28, 28})
.filter_spatial({3, 3})
.strides({1, 1})
.dilations({1, 1})
.left_pads({1, 1})
.right_pads({1, 1})
.build();
bool ok = problem.is_valid();
```
### GroupedConvRegistry (`grouped_conv_registry.hpp`)
Thread-safe registry with JSON export and filtering:
```cpp
#include "ck_tile/dispatcher/grouped_conv_registry.hpp"
auto& registry = GroupedConvRegistry::instance();
// Thread-safe registration
registry.register_kernel(kernel);
// JSON export
std::string json = registry.export_json();
registry.export_json_to_file("kernels.json");
// Filtering
auto gfx942_kernels = registry.filter_by_arch("gfx942");
auto matched = registry.filter([](const auto& k) { return k.is_fwd(); });
```
### DECL_GROUPED_CONV_KERNEL_SET (`grouped_conv_kernel_decl.hpp`)
Declarative kernel definition:
```cpp
DECL_GROUPED_CONV_KERNEL_SET(my_conv_kernels,
.add(
GroupedConvSignature().dtype("fp16").layout("nhwgc"),
GroupedConvAlgorithm().tile(128, 128, 32).wave(2, 2, 1)
.warp(32, 32, 16).pipeline("compv4"),
"gfx942"
)
);
// Register all matching current arch
DECL_GROUPED_CONV_KERNEL_ALL(all_conv_kernels, "gfx942");
```
## Best Practices
1. Use `Release` build for performance
@@ -155,6 +227,8 @@ auto kernel = create_generated_tile_kernel<
3. Use `Priority::High` for hand-tuned kernels
4. Reuse dispatcher instances
5. Clear registry between test runs
6. Use `GroupedConvProblemBuilder` for validated problem construction
7. Leverage `export_json()` for kernel inventory and debugging
---

152
dispatcher/include/ck_tile/dispatcher/backends/generated_conv_backend.hpp Normal file
View File

@@ -0,0 +1,152 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
// Generated Convolution Kernel Backend
//
// Wraps CK Tile grouped convolution launchers for use through the
// GroupedConvDispatcher. Each generated kernel launcher is wrapped in
// a ConvKernelRunFn that builds the correct host-args type (forward,
// bwd-data, or bwd-weight) and calls Launcher::launch().
#pragma once
#include "ck_tile/dispatcher/grouped_conv_problem.hpp"
#include "ck_tile/dispatcher/grouped_conv_registry.hpp"
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/convolution_parameter.hpp"
#include "ck_tile/ops/grouped_convolution.hpp"
#include <hip/hip_runtime.h>
#include <functional>
namespace ck_tile {
namespace dispatcher {
namespace backends {
// Buffer context is defined in grouped_conv_registry.hpp (g_conv_dispatch_buffers)
// so there's no circular dependency.
// Helper: build ck_tile::conv::ConvParam from GroupedConvProblem
inline ck_tile::conv::ConvParam make_conv_param_2d(const GroupedConvProblem& p)
{
return ck_tile::conv::ConvParam{
2,
static_cast<ck_tile::index_t>(p.G),
static_cast<ck_tile::index_t>(p.N),
static_cast<ck_tile::index_t>(p.K),
static_cast<ck_tile::index_t>(p.C),
{static_cast<ck_tile::index_t>(p.filter_spatial[1]),
static_cast<ck_tile::index_t>(p.filter_spatial[2])},
{static_cast<ck_tile::index_t>(p.input_spatial[1]),
static_cast<ck_tile::index_t>(p.input_spatial[2])},
{static_cast<ck_tile::index_t>(p.stride[1]), static_cast<ck_tile::index_t>(p.stride[2])},
{static_cast<ck_tile::index_t>(p.dilation[1]),
static_cast<ck_tile::index_t>(p.dilation[2])},
{static_cast<ck_tile::index_t>(p.padding[1]), static_cast<ck_tile::index_t>(p.padding[2])},
{static_cast<ck_tile::index_t>(p.padding[1]), static_cast<ck_tile::index_t>(p.padding[2])}};
}
inline ck_tile::conv::ConvParam make_conv_param_3d(const GroupedConvProblem& p)
{
return ck_tile::conv::ConvParam{3,
static_cast<ck_tile::index_t>(p.G),
static_cast<ck_tile::index_t>(p.N),
static_cast<ck_tile::index_t>(p.K),
static_cast<ck_tile::index_t>(p.C),
{static_cast<ck_tile::index_t>(p.filter_spatial[0]),
static_cast<ck_tile::index_t>(p.filter_spatial[1]),
static_cast<ck_tile::index_t>(p.filter_spatial[2])},
{static_cast<ck_tile::index_t>(p.input_spatial[0]),
static_cast<ck_tile::index_t>(p.input_spatial[1]),
static_cast<ck_tile::index_t>(p.input_spatial[2])},
{static_cast<ck_tile::index_t>(p.stride[0]),
static_cast<ck_tile::index_t>(p.stride[1]),
static_cast<ck_tile::index_t>(p.stride[2])},
{static_cast<ck_tile::index_t>(p.dilation[0]),
static_cast<ck_tile::index_t>(p.dilation[1]),
static_cast<ck_tile::index_t>(p.dilation[2])},
{static_cast<ck_tile::index_t>(p.padding[0]),
static_cast<ck_tile::index_t>(p.padding[1]),
static_cast<ck_tile::index_t>(p.padding[2])},
{static_cast<ck_tile::index_t>(p.padding[0]),
static_cast<ck_tile::index_t>(p.padding[1]),
static_cast<ck_tile::index_t>(p.padding[2])}};
}
// Create a RunFn for a forward convolution launcher (2D or 3D)
template <typename LauncherType, int NDim>
inline GroupedConvKernelInstance::RunFn make_conv_fwd_run_fn()
{
return [](const GroupedConvProblem& problem, void* stream) -> float {
auto& ctx = g_conv_dispatch_buffers;
auto param = (NDim == 2) ? make_conv_param_2d(problem) : make_conv_param_3d(problem);
ck_tile::GroupedConvFwdHostArgs<> args(
param, ctx.input_ptr, ctx.weight_ptr, {}, ctx.output_ptr, 1);
ck_tile::stream_config sc;
sc.stream_id_ = reinterpret_cast<hipStream_t>(stream);
sc.time_kernel_ = ctx.benchmarking;
sc.log_level_ = 0;
sc.cold_niters_ = ctx.benchmarking ? ctx.warmup : 0;
sc.nrepeat_ = ctx.benchmarking ? ctx.repeat : 1;
sc.is_gpu_timer_ = ctx.benchmarking;
return LauncherType::launch(args, sc);
};
}
// Create a RunFn for a backward-data convolution launcher.
// Dispatcher convention: run(dY, W, dX, problem) where dX is computed.
// BwdDataHostArgs(param, in_ptr=dX, wei_ptr=W, {}, out_ptr=dY, k_batch)
template <typename LauncherType, int NDim>
inline GroupedConvKernelInstance::RunFn make_conv_bwd_data_run_fn()
{
return [](const GroupedConvProblem& problem, void* stream) -> float {
auto& ctx = g_conv_dispatch_buffers;
auto param = (NDim == 2) ? make_conv_param_2d(problem) : make_conv_param_3d(problem);
ck_tile::GroupedConvBwdDataHostArgs args(
param,
ctx.output_ptr, // in_ptr = dX (being computed)
ctx.weight_ptr, // wei_ptr = W
{},
ctx.input_ptr, // out_ptr = dY (gradient from next layer)
1);
ck_tile::stream_config sc;
sc.stream_id_ = reinterpret_cast<hipStream_t>(stream);
sc.time_kernel_ = ctx.benchmarking;
sc.log_level_ = 0;
sc.cold_niters_ = ctx.benchmarking ? ctx.warmup : 0;
sc.nrepeat_ = ctx.benchmarking ? ctx.repeat : 1;
sc.is_gpu_timer_ = ctx.benchmarking;
return LauncherType::launch(args, sc);
};
}
// Create a RunFn for a backward-weight convolution launcher.
// Dispatcher convention: run(X, dY, dW, problem) where dW is computed.
// BwdWeightHostArgs(param, in_ptr=X, wei_ptr=dW, {}, out_ptr=dY, k_batch)
template <typename LauncherType, int NDim>
inline GroupedConvKernelInstance::RunFn make_conv_bwd_weight_run_fn()
{
return [](const GroupedConvProblem& problem, void* stream) -> float {
auto& ctx = g_conv_dispatch_buffers;
auto param = (NDim == 2) ? make_conv_param_2d(problem) : make_conv_param_3d(problem);
const int k_batch = (ctx.split_k > 1) ? ctx.split_k : 1;
ck_tile::GroupedConvBwdWeightHostArgs args(param,
ctx.input_ptr, // in_ptr = X
ctx.output_ptr, // wei_ptr = dW (being computed)
{},
ctx.weight_ptr, // out_ptr = dY
k_batch);
ck_tile::stream_config sc;
sc.stream_id_ = reinterpret_cast<hipStream_t>(stream);
sc.time_kernel_ = ctx.benchmarking;
sc.log_level_ = 0;
sc.cold_niters_ = ctx.benchmarking ? ctx.warmup : 0;
sc.nrepeat_ = ctx.benchmarking ? ctx.repeat : 1;
sc.is_gpu_timer_ = ctx.benchmarking;
return LauncherType::launch(args, sc);
};
}
} // namespace backends
} // namespace dispatcher
} // namespace ck_tile

199
dispatcher/include/ck_tile/dispatcher/base_registry.hpp Normal file
View File

@@ -0,0 +1,199 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <atomic>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
namespace ck_tile {
namespace dispatcher {
/// Shared priority enum used by all registry types
enum class Priority
{
Low = 0,
Normal = 1,
High = 2
};
/// BaseRegistry: Thread-safe, priority-aware kernel storage shared by GEMM and Conv registries.
///
/// Template Parameters:
/// Derived - CRTP derived class (e.g., Registry, ConvRegistry)
/// KeyType - primary key type (std::string for GEMM, ConvKernelKey for Conv)
/// InstanceType - kernel instance type (KernelInstance, ConvKernelInstance)
/// KeyHash - hash functor for KeyType (defaults to std::hash<KeyType>)
template <typename Derived,
typename KeyType,
typename InstanceType,
typename KeyHash = std::hash<KeyType>>
class BaseRegistry
{
public:
using InstancePtr = std::shared_ptr<InstanceType>;
struct Entry
{
InstancePtr instance;
Priority priority;
};
BaseRegistry() = default;
virtual ~BaseRegistry() = default;
BaseRegistry(BaseRegistry&& other) noexcept
{
std::lock_guard<std::mutex> lock(other.mutex_);
entries_ = std::move(other.entries_);
name_ = std::move(other.name_);
}
BaseRegistry& operator=(BaseRegistry&& other) noexcept
{
if(this != &other)
{
std::scoped_lock lock(mutex_, other.mutex_);
entries_ = std::move(other.entries_);
name_ = std::move(other.name_);
}
return *this;
}
BaseRegistry(const BaseRegistry&) = delete;
BaseRegistry& operator=(const BaseRegistry&) = delete;
/// Register a kernel. If the key already exists, the new entry replaces it
/// unless the existing entry has strictly higher priority.
/// Same-priority registration overwrites (last-writer-wins at equal priority).
bool
register_kernel(const KeyType& key, InstancePtr instance, Priority priority = Priority::Normal)
{
std::lock_guard<std::mutex> lock(mutex_);
auto it = entries_.find(key);
if(it != entries_.end() && it->second.priority > priority)
{
return false;
}
entries_[key] = Entry{std::move(instance), priority};
return true;
}
[[nodiscard]] std::size_t size() const
{
std::lock_guard<std::mutex> lock(mutex_);
return entries_.size();
}
[[nodiscard]] bool empty() const
{
std::lock_guard<std::mutex> lock(mutex_);
return entries_.empty();
}
void clear()
{
std::lock_guard<std::mutex> lock(mutex_);
entries_.clear();
}
[[nodiscard]] std::string get_name() const
{
std::lock_guard<std::mutex> lock(mutex_);
return name_; // return by value to avoid dangling reference
}
void set_name(const std::string& name)
{
std::lock_guard<std::mutex> lock(mutex_);
name_ = name;
}
[[nodiscard]] std::vector<InstancePtr> get_all_instances() const
{
std::lock_guard<std::mutex> lock(mutex_);
std::vector<InstancePtr> result;
result.reserve(entries_.size());
for(const auto& [key, entry] : entries_)
{
result.push_back(entry.instance);
}
return result;
}
std::size_t merge_from(const BaseRegistry& other, Priority priority = Priority::Normal)
{
std::scoped_lock lock(mutex_, other.mutex_);
std::size_t merged = 0;
for(const auto& [key, entry] : other.entries_)
{
auto it = entries_.find(key);
if(it == entries_.end() || it->second.priority <= priority)
{
entries_[key] = Entry{entry.instance, priority};
++merged;
}
}
return merged;
}
/// Enable automatic JSON export after every kernel registration.
/// Requires the derived class to implement export_json_to_file(path, stats).
void enable_auto_export(const std::string& path,
bool include_statistics = true,
bool export_on_every_registration = true)
{
std::lock_guard<std::mutex> lock(mutex_);
auto_export_path_ = path;
auto_export_stats_ = include_statistics;
auto_export_on_register_ = export_on_every_registration;
auto_export_enabled_.store(true, std::memory_order_release);
}
void disable_auto_export() { auto_export_enabled_.store(false, std::memory_order_release); }
[[nodiscard]] bool is_auto_export_enabled() const
{
return auto_export_enabled_.load(std::memory_order_acquire);
}
/// Call after registration to trigger auto-export if enabled.
void perform_auto_export()
{
if(!auto_export_enabled_.load(std::memory_order_acquire))
return;
std::lock_guard<std::mutex> lock(mutex_);
if(auto_export_on_register_)
{
static_cast<Derived*>(this)->export_json_to_file(auto_export_path_, auto_export_stats_);
}
}
protected:
[[nodiscard]] const std::unordered_map<KeyType, Entry, KeyHash>& entries() const
{
return entries_;
}
[[nodiscard]] std::unordered_map<KeyType, Entry, KeyHash>& entries_mut() { return entries_; }
std::mutex& mutex() const { return mutex_; }
private:
mutable std::mutex mutex_;
std::unordered_map<KeyType, Entry, KeyHash> entries_;
std::string name_ = "default";
std::atomic<bool> auto_export_enabled_{false};
bool auto_export_on_register_ = true;
bool auto_export_stats_ = true;
std::string auto_export_path_;
};
} // namespace dispatcher
} // namespace ck_tile

22
dispatcher/include/ck_tile/dispatcher/dispatcher.hpp
View File

@@ -23,6 +23,7 @@
#pragma once
#include "ck_tile/dispatcher/dispatcher_error.hpp"
#include "ck_tile/dispatcher/kernel_instance.hpp"
#include "ck_tile/dispatcher/problem.hpp"
#include "ck_tile/dispatcher/registry.hpp"
@@ -52,7 +53,11 @@ class Dispatcher
/// Constructor
/// @param registry Registry instance to use (default: global singleton)
explicit Dispatcher(Registry* registry = nullptr);
/// @param gfx_arch Target GPU architecture (e.g. "gfx950")
explicit Dispatcher(Registry* registry = nullptr, const std::string& gfx_arch = "");
void set_arch(const std::string& arch) { gfx_arch_ = arch; }
[[nodiscard]] const std::string& arch() const { return gfx_arch_; }
/// Register a heuristic function for kernel selection
/// @param heuristic Function that maps problems to ranked kernel identifiers
@@ -74,7 +79,7 @@ class Dispatcher
/// @param problem Problem configuration
/// @param stream HIP stream for kernel launch (nullptr = default stream)
/// @return Kernel execution time in milliseconds
/// @throws std::runtime_error if no suitable kernel found
/// @throws NoKernelFound if no suitable kernel found
[[nodiscard]] float run(const void* a_ptr,
const void* b_ptr,
void* c_ptr,
@@ -89,7 +94,7 @@ class Dispatcher
/// @param problem Problem configuration
/// @param stream HIP stream for kernel launch (nullptr = default stream)
/// @return Kernel execution time in milliseconds
/// @throws std::runtime_error if no suitable kernel found
/// @throws NoKernelFound if no suitable kernel found
[[nodiscard]] float run_fused(const void* a_ptr,
const void* b_ptr,
void* c_ptr,
@@ -106,7 +111,8 @@ class Dispatcher
/// @param problem Problem configuration
/// @param stream HIP stream for kernel launch (nullptr = default stream)
/// @return Kernel execution time in milliseconds
/// @throws std::runtime_error if kernel not found or doesn't support problem
/// @throws NoKernelFound if the kernel identifier is not registered
/// @throws UnsupportedProblem if the selected kernel does not support the problem
[[nodiscard]] float run_explicit(const std::string& kernel_id,
const void* a_ptr,
const void* b_ptr,
@@ -130,10 +136,18 @@ class Dispatcher
const Problem& problem,
float tolerance = 1e-3f) const;
/// Enable or disable GPU benchmarking (timing) on all kernels.
/// When disabled, kernels execute once with no timing overhead
/// (one-shot mode for production plugins).
void set_benchmarking(bool enable) { benchmarking_ = enable; }
[[nodiscard]] bool benchmarking_enabled() const { return benchmarking_; }
private:
Registry* registry_;
HeuristicFunction heuristic_;
SelectionStrategy strategy_;
std::string gfx_arch_;
bool benchmarking_ = true;
/// Select kernel using first-fit strategy
[[nodiscard]] KernelInstancePtr select_first_fit(const Problem& problem) const;

28
dispatcher/include/ck_tile/dispatcher/dispatcher_error.hpp Normal file
View File

@@ -0,0 +1,28 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <stdexcept>
#include <string>
namespace ck_tile {
namespace dispatcher {
struct DispatcherError : std::runtime_error
{
using std::runtime_error::runtime_error;
};
struct NoKernelFound : DispatcherError
{
using DispatcherError::DispatcherError;
};
struct UnsupportedProblem : DispatcherError
{
using DispatcherError::DispatcherError;
};
} // namespace dispatcher
} // namespace ck_tile

55
dispatcher/include/ck_tile/dispatcher/dispatcher_log.hpp Normal file
View File

@@ -0,0 +1,55 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <cstdlib>
#include <iostream>
#include <string>
namespace ck_tile {
namespace dispatcher {
/// Log levels for dispatcher transparency:
/// 0 = silent (default)
/// 1 = print selected kernel name
/// 2 = print all candidates considered and acceptance/rejection reasons
inline int get_log_level()
{
static int level = []() {
const char* env = std::getenv("CK_DISPATCHER_LOG_LEVEL");
return env ? std::atoi(env) : 0;
}();
return level;
}
inline void log_kernel_selected(const std::string& kernel_name, const std::string& problem_desc)
{
if(get_log_level() >= 1)
{
std::cerr << "[CK Dispatcher] Selected kernel: " << kernel_name << " for " << problem_desc
<< std::endl;
}
}
inline void
log_kernel_candidate(const std::string& kernel_name, bool accepted, const std::string& reason)
{
if(get_log_level() >= 2)
{
std::cerr << "[CK Dispatcher] Candidate: " << kernel_name << " -> "
<< (accepted ? "ACCEPTED" : "REJECTED")
<< (reason.empty() ? "" : " (" + reason + ")") << std::endl;
}
}
inline void log_no_kernel_found(const std::string& problem_desc)
{
if(get_log_level() >= 1)
{
std::cerr << "[CK Dispatcher] No kernel found for " << problem_desc << std::endl;
}
}
} // namespace dispatcher
} // namespace ck_tile

Some files were not shown because too many files have changed in this diff Show More