diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index daf3c258d9d..bfa77b84452 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,6 +26,11 @@ repos: verbose: false language: script types_or: [c++, python, shell, cmake] + - id: ascii-only-checker + name: Check for non-ASCII characters in C/C++ sources + entry: projects/composablekernel/script/check_ascii_only.sh + language: script + types_or: [c++, inc] - id: remove-exec-bit name: Remove executable bit from non-executable files entry: projects/composablekernel/script/remove_exec_bit.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ce054255cc..7dc7dfbae23 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -741,6 +741,10 @@ SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV") if(BUILD_DEV) add_compile_options(-Werror) add_compile_options(-Weverything) + add_compile_options(-Wno-lifetime-safety-intra-tu-suggestions) + add_compile_options(-Wno-lifetime-safety-cross-tu-suggestions) + add_compile_options(-Wno-lifetime-safety-lifetimebound-violation) + add_compile_options(-Wno-unknown-warning-option) endif() message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}") @@ -761,8 +765,6 @@ if(NOT MIOPEN_REQ_LIBS_ONLY AND NOT HIPTENSOR_REQ_LIBS_ONLY) endif() -option(MIOPEN_REQ_LIBS_ONLY "Build only the MIOpen required libraries" OFF) -option(HIPTENSOR_REQ_LIBS_ONLY "Build only the HipTensor required libraries" OFF) option(DISABLE_OFFLOAD_COMPRESS "Disable offload compress compiler flag when building instances" OFF) option(BUILD_MHA_LIB "Build the static library for flash attention" OFF) option(BUILD_CK_DEVICE_INSTANCES "Build device operation instances in library/" ON) diff --git a/Jenkinsfile b/Jenkinsfile index 03475924055..a57f3432a92 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -21,1229 +21,18 @@ // - Forces full build if dependency cache stale (>7 days) // - Manual override: set DISABLE_SMART_BUILD=true // -// Benefits: PR builds 5h → 30min (typical), nightly builds unchanged +// Benefits: PR builds 5h -> 30min (typical), nightly builds unchanged // See: script/dependency-parser/README.md for details // - -@NonCPS -String getGitHubCommitHash(def build) -{ - def scmAction = build?.actions.find { action -> - action instanceof jenkins.scm.api.SCMRevisionAction - } - if (scmAction?.revision instanceof org.jenkinsci.plugins.github_branch_source.PullRequestSCMRevision) - { - return scmAction.revision.pullHash - } - else if (scmAction?.revision instanceof jenkins.plugins.git.AbstractGitSCMSource$SCMRevisionImpl) - { - return scmAction.revision.hash - } - return null -} - def rocmnode(name) { return '(rocmtest || miopen) && (' + name + ')' } -def show_node_info() { - sh """ - echo "NODE_NAME = \$NODE_NAME" - hostname - lsb_release -sd - uname -r - cat /sys/module/amdgpu/version - ls /opt/ -la - """ -} - -def setGithubStatus(String context, String state, String description) { - def sha = env.GIT_COMMIT - def targetUrl = env.RUN_DISPLAY_URL ?: env.BUILD_URL - def statusUrl = "https://api.github.com/repos/ROCm/rocm-libraries/statuses/${sha}" - withCredentials([usernamePassword(credentialsId: 'github-app-miopen', usernameVariable: 'GITHUB_APP', passwordVariable: 'GITHUB_TOKEN')]) { - def code = '0' - try { - retry(3) { - code = sh(returnStdout: true, script: """ - curl -s -w "%{http_code}" -o /dev/null -X POST '${statusUrl}' \\ - -H "Authorization: token \$GITHUB_TOKEN" \\ - -H 'Content-Type: application/json' \\ - -d '{"state":"${state}","context":"${context}","description":"${description}","target_url":"${targetUrl}"}' - """).trim() - if (!code.startsWith('2')) { - error("GitHub status POST returned ${code}") - } - } - } catch (Exception e) { - echo "WARNING: GitHub status POST failed after retries (context=${context}, state=${state}, code=${code})" - } - } -} - -def cloneUpdateRefRepo() { - def refRepoPath = "/var/jenkins/ref-repo/rocm-libraries" - def lockLabel = "git ref repo lock - ${env.NODE_NAME}" - def folderExists = sh( - script: "test -d ${refRepoPath}/refs", - returnStatus: true - ) == 0 - - if (!folderExists) { - echo "rocm-libraries repo does not exist at ${refRepoPath}, creating mirror clone..." - echo "locking on label: ${lockLabel}" - lock(lockLabel) { - def cloneCommand = """ - set -ex - rm -rf ${refRepoPath} && mkdir -p ${refRepoPath} - git clone --mirror https://github.com/ROCm/rocm-libraries.git ${refRepoPath} - """ - sh(script: cloneCommand, label: "clone ref repo") - } - echo "Completed git clone, lock released" - } - echo "rocm-libraries repo exists at ${refRepoPath}, performing git remote update..." - echo "locking on label: ${lockLabel}" - lock(lockLabel) { - def fetchCommand = """ - set -ex - cd ${refRepoPath} - git remote prune origin - git remote update - """ - sh(script: fetchCommand, label: "update ref repo") - } - echo "Completed git ref repo fetch, lock released" -} - -def checkoutComposableKernel() -{ - //update ref repo - cloneUpdateRefRepo() - // checkout project - def scmVars = checkout scm - // getGitHubCommitHash reads SCMRevisionAction recorded before any local merge, - // giving the true PR branch tip (pullHash) or branch HEAD (hash). - // Falls back to ORIG_HEAD (pre-merge HEAD set by git merge) when SCMRevisionAction - // is unavailable, then to HEAD for branch builds where no merge occurred. - env.GIT_COMMIT = getGitHubCommitHash(currentBuild.rawBuild) ?: sh(returnStdout: true, script: ''' - git rev-parse ORIG_HEAD 2>/dev/null || git rev-parse HEAD - ''').trim() -} - -def generateAndArchiveBuildTraceVisualization(String buildTraceFileName) { - try { - checkoutComposableKernel() - - // Retrieve the build trace artifact - def traceFileExists = false - try { - copyArtifacts( - projectName: env.JOB_NAME, - selector: specific(env.BUILD_NUMBER), - filter: buildTraceFileName - ) - traceFileExists = fileExists(buildTraceFileName) - } catch (Exception e) { - echo "Could not copy build trace artifact: ${e.getMessage()}" - traceFileExists = false - return - } - - sh """ - echo "post artifact download:" - ls -la - """ - - // Pull image - def image = "ghcr.io/puppeteer/puppeteer:24.30.0" - echo "Pulling image: ${image}" - def retimage = docker.image("${image}") - retimage.pull() - - // Create a temporary workspace - sh """#!/bin/bash - ls -la - mkdir -p workspace - cp ./projects/composablekernel/script/infra_helper/capture_build_trace.js ./workspace - cp ${buildTraceFileName} ./workspace/${buildTraceFileName} - chmod 777 ./workspace - ls -la ./workspace - """ - - // Run container to get snapshot - def dockerOpts = "--cap-add=SYS_ADMIN -v \"\$(pwd)/workspace:/workspace\" -e NODE_PATH=/home/pptruser/node_modules -e BUILD_TRACE_FILE=${buildTraceFileName}" - // Create unique image name by sanitizing job name - def sanitizedJobName = env.JOB_NAME.replaceAll(/[\/\\:*?"<>| ]/, '_').replaceAll('%2F', '_') - def architectureName = (buildTraceFileName =~ /(gfx[0-9a-zA-Z]+)/)[0][1] - def imageName = "perfetto_snapshot_${sanitizedJobName}_build_${env.BUILD_NUMBER}_${architectureName}.png" - sh """ - docker run --rm ${dockerOpts} ${image} node /workspace/capture_build_trace.js - mv ./workspace/perfetto_snapshot_build.png ./workspace/${imageName} - """ - - // Archive the snapshot - sh """ - mv ./workspace/${imageName} ${imageName} - """ - archiveArtifacts "${imageName}" - - // Notify the channel - withCredentials([string(credentialsId: 'ck_ci_build_perf_webhook_url', variable: 'WEBHOOK_URL')]) { - sh ''' - # Create build trace filename with build number based on the original filename - BUILD_TRACE_WITH_NUMBER=$(echo "''' + buildTraceFileName + '''" | sed 's/.json/_''' + sanitizedJobName + '''_''' + env.BUILD_NUMBER + '''_''' + architectureName + '''.json/') - - # Convert image to base64 - echo "Converting image to base64..." - IMAGE_BASE64=$(base64 -w 0 ''' + imageName + ''') - echo "Image base64 length: ${#IMAGE_BASE64}" - - # Convert build trace to base64 - echo "Converting build trace to base64..." - BUILD_TRACE_BASE64=$(base64 -w 0 ''' + buildTraceFileName + ''') - echo "Build trace base64 length: ${#BUILD_TRACE_BASE64}" - - # Create JSON payload with base64 data - echo "Creating JSON payload..." - { - printf '{\n' - printf ' "jobName": "%s",\n' "''' + env.JOB_NAME + '''" - printf ' "buildNumber": "%s",\n' "''' + env.BUILD_NUMBER + '''" - printf ' "jobUrl": "%s",\n' "''' + env.RUN_DISPLAY_URL + '''" - printf ' "imageName": "%s",\n' "''' + imageName + '''" - printf ' "architecture": "%s",\n' "''' + architectureName + '''" - printf ' "imageData": "%s",\n' "$IMAGE_BASE64" - printf ' "buildTraceName": "%s",\n' "$BUILD_TRACE_WITH_NUMBER" - printf ' "buildTraceData": "%s"\n' "$BUILD_TRACE_BASE64" - printf '}\n' - } > webhook_payload.json - - echo "JSON payload created, size: $(wc -c < webhook_payload.json) bytes" - - curl -X POST "${WEBHOOK_URL}" \ - -H "Content-Type: application/json" \ - -d @webhook_payload.json - - # Clean up temporary file - rm -f webhook_payload.json - ''' - } - } catch (Exception e) { - echo "Throwing error exception while generating build trace visualization" - echo 'Exception occurred: ' + e.toString() - } -} - -class Version { - int major, minor, patch - @Override - String toString() { - return [major, minor, patch].findAll().join('.') - } -} -def parseVersion(String versionString) { - if (!versionString) return null - int[] tokens = versionString.split(/\./).collect { it as int } // Splits the string by '.' and converts each part to an integer. - return new Version( - major: tokens[0], - minor: tokens.length > 1 ? tokens[1] : null, - patch: tokens.length > 2 ? tokens[2] : null, - ) -} - -def nthreads() { - def nproc = sh(returnStdout: true, script: 'nproc') - echo "Number of cores: ${nproc}" - def n = nproc.toInteger() - if (n > 64){ - n = 64 - } - echo "Number of threads used for building: ${n}" - return n -} - -def runShell(String command){ - def responseCode = sh returnStatus: true, script: "${command} > tmp.txt" - def output = readFile(file: "tmp.txt") - return (output != "") -} - -def shouldRunCICheck() { - // File patterns that should not trigger CI - def skipFilePatterns = [ - /^projects\/composablekernel\/\.github\/.*/, // GitHub workflow files - /^projects\/composablekernel\/docs\/.*/, // Documentation files - /^projects\/composablekernel\/LICENSE$/, // License file - /^projects\/composablekernel\/.*\.gitignore$/, // Git ignore files - /^projects\/composablekernel\/.*\.md$/ // Markdown files - ] - - try { - // Always run if this is a base branch build - def baseBranch = "develop" - def isBaseBranchBuild = (env.CHANGE_ID == null && env.BRANCH_NAME == baseBranch) - - if (isBaseBranchBuild) { - echo "Base branch (${baseBranch}) build detected - always running CI for safety" - return true - } - - // Get the list of changed files (all files touched in any commit, even if reverted) - def changedFiles = sh( - returnStdout: true, - script: ''' - BASE_BRANCH="develop" - - if [ "$CHANGE_ID" != "" ]; then - # For PR builds, get all files touched in any commit - echo "PR build detected, checking all touched files against origin/$CHANGE_TARGET" >&2 - git log --name-only --pretty=format: origin/$CHANGE_TARGET..HEAD -- projects/composablekernel/ | sort -u | grep -v '^$' || true - else - # For feature branch builds, compare against merge-base with base branch - MERGE_BASE=$(git merge-base HEAD origin/$BASE_BRANCH 2>/dev/null || echo "HEAD~1") - echo "Branch build detected, checking all touched files since merge-base: $MERGE_BASE" >&2 - git log --name-only --pretty=format: $MERGE_BASE..HEAD -- projects/composablekernel/ | sort -u | grep -v '^$' || true - fi - ''' - ).trim().split('\n') - - if (changedFiles.size() == 1 && changedFiles[0] == '') { - echo "No changed files detected - this might be a manual trigger or merge commit, running CI for safety" - return true - } - - echo "Changed files: ${changedFiles.join(', ')}" - - // Separate files into those requiring CI and those that can be skipped - def filesRequiringCI = [] - def skippedFiles = [] - - changedFiles.each { file -> - def shouldSkip = skipFilePatterns.any { pattern -> - file ==~ pattern - } - - if (shouldSkip) { - skippedFiles.add(file) - } else { - filesRequiringCI.add(file) - } - } - - // Debug output - if (skippedFiles.size() > 0) { - echo "Files that don't require CI (${skippedFiles.size()}):" - skippedFiles.each { echo " - ${it}" } - } - - if (filesRequiringCI.size() > 0) { - echo "Files that require CI (${filesRequiringCI.size()}):" - filesRequiringCI.each { echo " - ${it}" } - return true - } else { - echo "Only non-relevant files changed, skipping CI" - return false - } - } catch (Exception e) { - echo "Error checking changed files: ${e.getMessage()}, running CI by default" - return true - } -} - -def getBaseDockerImageName(){ - def img - if (params.USE_CUSTOM_DOCKER != ""){ - img = "${params.USE_CUSTOM_DOCKER}" - } - else{ - img = "${env.CK_DOCKERHUB}:ck_ub24.04_rocm${params.ROCMVERSION}" - } - return img -} - -def getDockerImageName(){ - def img - def base_name = getBaseDockerImageName() - if (params.USE_CUSTOM_DOCKER != ""){ - img = "${params.USE_CUSTOM_DOCKER}" - } - else{ - if (params.COMPILER_VERSION == "") { - img = "${base_name}" - } - else{ - if (params.COMPILER_COMMIT == ""){ - img = "${base_name}_${params.COMPILER_VERSION}" - } - else{ - def commit = "${params.COMPILER_COMMIT}"[0..6] - img = "${base_name}_${params.COMPILER_VERSION}_${commit}" - } - } - } - return img -} - -def check_host() { - if ("${env.CK_SCCACHE}" != "null"){ - def SCCACHE_SERVER="${env.CK_SCCACHE.split(':')[0]}" - echo "sccache server: ${SCCACHE_SERVER}" - sh "chmod +w -R ${env.WORKSPACE}" - sh '''ping -c 1 -p 6379 "${SCCACHE_SERVER}" | echo $? > tmp.txt''' - def output = readFile(file: "tmp.txt") - echo "tmp.txt contents: \$output" - return (output != "0") - } - else{ - return 1 - } -} - -def check_arch_name(){ - sh 'rocminfo | tee rocminfo.log' - if ( runShell('grep -n "gfx90a" rocminfo.log') ){ - return "gfx90a" - } - else if ( runShell('grep -n "gfx942" rocminfo.log') ) { - return "gfx942" - } - else if ( runShell('grep -n "gfx101" rocminfo.log') ) { - return "gfx101" - } - else if ( runShell('grep -n "gfx103" rocminfo.log') ) { - return "gfx103" - } - else if ( runShell('grep -n "gfx11" rocminfo.log') ) { - return "gfx11" - } - else if ( runShell('grep -n "gfx120" rocminfo.log') ) { - return "gfx12" - } - else if ( runShell('grep -n "gfx908" rocminfo.log') ) { - return "gfx908" - } - else if ( runShell('grep -n "gfx950" rocminfo.log') ) { - return "gfx950" - } - else { - return "" - } -} - -def getDockerImage(Map conf=[:]){ - def image - if ( conf.get("docker_name", "") != "" ){ - image = conf.get("docker_name", "") - echo "Using special docker: ${image}" - } - else{ - image = getDockerImageName() - echo "Using default docker: ${image}" - } - //Check if image exists - def retimage - 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}" - } - return [retimage, image] -} - -// 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 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." - } - } - - 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) - } -} - -def get_docker_options(){ - def dockerOpts - if ( params.BUILD_INSTANCES_ONLY ){ - dockerOpts = "--network=host --group-add video --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" - } - 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-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 --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') - def render_id = sh(returnStdout: true, script: 'getent group render | cut -d: -f3') - dockerOpts = dockerOpts + " --group-add=${video_id} --group-add=${render_id} -v /var/jenkins/ref-repo/:/var/jenkins/ref-repo/ " - echo "Docker flags: ${dockerOpts}" - return dockerOpts -} - -def build_client_examples(String arch){ - def cmd = """ cd ../client_example && rm -rf build && mkdir build && cd build && \ - cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ - -DGPU_TARGETS="${arch}" \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j """ - return cmd -} - -def build_client_examples_and_codegen_tests(String arch){ - def cmd = """ cd ../codegen && rm -rf build && mkdir build && cd build && \ - cmake -DCMAKE_PREFIX_PATH=/opt/rocm -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" .. && \ - make -j64 check && \ - cd ../../client_example && rm -rf build && mkdir build && cd build && \ - cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ - -DGPU_TARGETS="${arch}" \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j """ - return cmd -} - -def build_and_run_fmha(String arch){ - def cmd = """ cmake -G Ninja -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ - -DGPU_TARGETS="${arch}" \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" .. && \ - ninja -j128 tile_example_fmha_fwd tile_example_fmha_bwd && \ - cd ../ && - example/ck_tile/01_fmha/script/run_full_test.sh "CI_${params.COMPILER_VERSION}" "${env.BRANCH_NAME}" "${NODE_NAME}" "${arch}" """ - return cmd -} - -def cmake_build(Map conf=[:]){ - - def config_targets = conf.get("config_targets","check") - def build_envs = "CTEST_PARALLEL_LEVEL=4 " + conf.get("build_env","") - def prefixpath = conf.get("prefixpath","/opt/rocm") - def setup_args = conf.get("setup_args","") - // make sure all unit tests always run on develop branch - def runAllUnitTests = (env.BRANCH_NAME == "develop") ? true : params.RUN_ALL_UNIT_TESTS - - if (prefixpath != "/usr/local"){ - setup_args = setup_args + " -DCMAKE_PREFIX_PATH=${prefixpath} " - } - - //cmake_env can overwrite default CXX variables. - def cmake_envs - if(!setup_args.contains("gfx1250")){ - cmake_envs = "CXX=${params.BUILD_COMPILER} CXXFLAGS='-Werror' " + conf.get("cmake_ex_env","") - } - else{ //use default compiler for gfx1250 - cmake_envs = "CXX=/opt/rocm/llvm/bin/clang++ CXXFLAGS='-Werror' " + conf.get("cmake_ex_env","") - } - - if(conf.get("build_install","") == "true") - { - config_targets = 'install ' + config_targets - setup_args = ' -DBUILD_DEV=On -DCMAKE_INSTALL_PREFIX=../install' + setup_args - } else{ - setup_args = ' -DBUILD_DEV=On' + setup_args - } - if (params.DISABLE_DL_KERNELS){ - setup_args = setup_args + " -DDISABLE_DL_KERNELS=ON " - } - - setup_args = " -DCMAKE_BUILD_TYPE=release " + setup_args - - def pre_setup_cmd = """ - #!/bin/bash - cd projects/composablekernel - ulimit -c unlimited - rm -rf build - mkdir build - rm -rf install - mkdir install - cd build - """ - def invocation_tag="" - if (setup_args.contains("gfx12")){ - invocation_tag="gfx12" - } - if (setup_args.contains("gfx11")){ - invocation_tag="gfx11" - } - if (setup_args.contains("gfx101")){ - invocation_tag="gfx101" - } - if (setup_args.contains("gfx103")){ - invocation_tag="gfx103" - } - if (setup_args.contains("gfx908")){ - invocation_tag="gfx908" - } - if (setup_args.contains("gfx90a")){ - invocation_tag="gfx90a" - } - if (setup_args.contains("gfx94")){ - invocation_tag="gfx94" - } - if (setup_args.contains("gfx95")){ - invocation_tag="gfx95" - } - echo "invocation tag: ${invocation_tag}" - def redis_pre_setup_cmd = pre_setup_cmd - if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { - redis_pre_setup_cmd = pre_setup_cmd + """ - #!/bin/bash - export ROCM_PATH=/opt/rocm - export SCCACHE_ENABLED=true - export SCCACHE_LOG_LEVEL=debug - export SCCACHE_IDLE_TIMEOUT=14400 - export COMPILERS_HASH_DIR=/tmp/.sccache - export SCCACHE_BIN=/usr/local/.cargo/bin/sccache - export SCCACHE_EXTRAFILES=/tmp/.sccache/rocm_compilers_hash_file - export SCCACHE_REDIS="redis://${env.CK_SCCACHE}" - echo "connect = ${env.CK_SCCACHE}" >> ../script/redis-cli.conf - export SCCACHE_C_CUSTOM_CACHE_BUSTER="${invocation_tag}" - echo \$SCCACHE_C_CUSTOM_CACHE_BUSTER - stunnel ../script/redis-cli.conf - ../script/sccache_wrapper.sh --enforce_redis - """ - try { - def cmd1 = conf.get("cmd1", """ - ${redis_pre_setup_cmd} - """) - sh cmd1 - setup_args = " -DCMAKE_HIP_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache -DCMAKE_C_COMPILER_LAUNCHER=sccache " + setup_args - } - catch(Exception err){ - echo "could not connect to redis server: ${err.getMessage()}. will not use sccache." - def cmd2 = conf.get("cmd2", """ - ${pre_setup_cmd} - """) - sh cmd2 - } - } - else{ - def cmd3 = conf.get("cmd3", """ - ${pre_setup_cmd} - """) - sh cmd3 - } - - // reduce parallelism when compiling, clang uses too much memory - def nt = nthreads() - def cmd - def setup_cmd - def build_cmd - def execute_cmd = conf.get("execute_cmd", "") - //check the node gpu architecture - def arch_name = check_arch_name() - if(!setup_args.contains("NO_CK_BUILD")){ - if (params.NINJA_BUILD_TRACE) { - echo "running ninja build trace" - } - if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { - setup_args = " -D CK_EXPERIMENTAL_BUILDER=ON " + setup_args - } - if (params.RUN_ROCM_CK_TESTS) { - setup_args = " -D CK_ENABLE_ROCM_CK=ON " + setup_args - } - setup_cmd = conf.get( - "setup_cmd", - """${cmake_envs} cmake -G Ninja ${setup_args} -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DCMAKE_CXX_FLAGS=" -O3 " .. """ - ) - - // Smart-build: Only build if running all tests or forced - // Otherwise, smart-build will determine what to build after cmake configure - if (runAllUnitTests) { - build_cmd = conf.get( - "build_cmd", - "${build_envs} ninja -j${nt} ${config_targets}" - ) - } else { - // Smart-build enabled: skip full build and execute_cmd (client examples) - build_cmd = "" - execute_cmd = "" - } - - cmd = conf.get("cmd", """ - ${setup_cmd} - ${build_cmd} - ${execute_cmd} - """) - } - else{ - cmd = conf.get("cmd", """ - ${execute_cmd} - """) - } - - echo cmd - - dir("projects/composablekernel/build"){ - // Start sccache monitoring - if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { - sh """ - chmod +x ../script/monitor_sccache_during_build.sh - mkdir -p logs - export SCCACHE_C_CUSTOM_CACHE_BUSTER="${invocation_tag}" - ../script/monitor_sccache_during_build.sh build_monitor & - MONITOR_PID=\$! - echo "Monitor PID: \$MONITOR_PID" - echo \$MONITOR_PID > monitor.pid - """ - } - try { - //build CK - sh cmd - if (runAllUnitTests){ - // Archive artifacts if they were generated - if (fileExists("ck_build_trace_${arch_name}.json")) { - archiveArtifacts "ck_build_trace_${arch_name}.json" - } - if (fileExists("clang_build_analysis_${arch_name}.log")) { - archiveArtifacts "clang_build_analysis_${arch_name}.log" - } - // Process ninja build trace after full build - if(fileExists(".ninja_log")) { - sh "python3 ../script/ninja_json_converter.py .ninja_log --legacy-format --output ck_build_trace_${arch_name}.json" - archiveArtifacts "ck_build_trace_${arch_name}.json" - sh "python3 ../script/parse_ninja_trace.py ck_build_trace_${arch_name}.json" - } - - if (params.NINJA_FTIME_TRACE) { - echo "running ClangBuildAnalyzer" - sh "/ClangBuildAnalyzer/build/ClangBuildAnalyzer --all . clang_build.log" - sh "/ClangBuildAnalyzer/build/ClangBuildAnalyzer --analyze clang_build.log > clang_build_analysis_${arch_name}.log" - archiveArtifacts "clang_build_analysis_${arch_name}.log" - } - } - } catch (Exception buildError) { - echo "Build failed: ${buildError.getMessage()}" - throw buildError - } finally { - // Stop sccache monitoring - if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { - sh """ - # Stop monitoring - if [ -f monitor.pid ]; then - MONITOR_PID=\$(cat monitor.pid) - kill \$MONITOR_PID 2>/dev/null || echo "Monitor already stopped" - rm -f monitor.pid - fi - """ - - // Archive the monitoring logs - try { - archiveArtifacts artifacts: "logs/*monitor*.log", allowEmptyArchive: true - } catch (Exception e) { - echo "Could not archive sccache monitoring logs: ${e.getMessage()}" - } - } - } - - //run tests except when NO_CK_BUILD is set and except on gfx1250 - if(!setup_args.contains("NO_CK_BUILD")){ - // run unit tests unless building library for all targets - // Note: This else block is when NINJA_BUILD_TRACE=false and BUILD_INSTANCES_ONLY=false - // So no ninja trace processing needed here - if (!params.BUILD_INSTANCES_ONLY){ - if (!runAllUnitTests && !setup_args.contains("gfx1250") ){ - // Smart Build: Run smart_build_and_test.sh - sh """ - export WORKSPACE_ROOT=${env.WORKSPACE} - export PARALLEL=32 - export NINJA_JOBS=${nt} - export ARCH_NAME=${arch_name} - export PROCESS_NINJA_TRACE=false - export NINJA_FTIME_TRACE=false - bash ../script/dependency-parser/smart_build_and_test.sh - """ - } - else{ //run all tests - if(!setup_args.contains("gfx1250")){ - echo "Full test suite requested (RUN_ALL_UNIT_TESTS=true or develop branch)" - sh "ninja -j${nt} check" - } - else{ //do not run tests on gfx1250, just build everything - echo "Building for gfx1250" - sh "ninja -j${nt}" - } - if (params.RUN_ROCM_CK_TESTS) { - sh 'ninja check-rocm-ck' - } - if(params.BUILD_PACKAGES || params.BUILD_INSTANCES_ONLY){ - echo "Build ckProfiler packages" - sh 'ninja -j64 package' - sh "mv composablekernel-ckprofiler_*.deb composablekernel-ckprofiler_1.2.0_amd64_${arch_name}.deb" - stash includes: "composablekernel-ckprofiler**.deb", name: "profiler_package_${arch_name}" - } - } - if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { - sh 'ninja check-builder' - } - } - } - } - - if (params.RUN_CK_TILE_FMHA_TESTS){ - try{ - dir("projects/composablekernel"){ - archiveArtifacts "perf_fmha_*.log" - stash includes: "perf_fmha_**.log", name: "perf_fmha_log_${arch_name}" - } - } - catch(Exception err){ - echo "could not locate the requested artifacts: ${err.getMessage()}. will skip the stashing." - } - } -} - -def buildHipClangJob(Map conf=[:]){ - show_node_info() - checkoutComposableKernel() - def prefixpath = conf.get("prefixpath", "/opt/rocm") - def dockerOpts = get_docker_options() - def image - def retimage - (retimage, image) = getDockerImage(conf) - - setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") - try { - withDockerContainer(image: image, args: dockerOpts) { - timeout(time: 20, unit: 'HOURS') - { - cmake_build(conf) - } - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - return retimage -} - -def buildHipClangJobAndReboot(Map conf=[:]){ - try{ - buildHipClangJob(conf) - } - catch(e){ - echo "throwing error exception for the stage" - echo 'Exception occurred: ' + e.toString() - throw e - } -} - -def Build_CK(Map conf=[:]){ - show_node_info() - checkoutComposableKernel() - def prefixpath = conf.get("prefixpath", "/opt/rocm") - def dockerOpts=get_docker_options() - def image - def retimage - - setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") - try { - try { - (retimage, image) = getDockerImage(conf) - withDockerContainer(image: image, args: dockerOpts) { - timeout(time: 2, unit: 'MINUTES'){ - sh 'rocminfo | tee rocminfo.log' - if ( !runShell('grep -n "gfx" rocminfo.log') ){ - throw new Exception ("GPU not found") - } - else{ - echo "GPU is OK" - } - } - } - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - echo "The job was cancelled or aborted" - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - withDockerContainer(image: image, args: dockerOpts) { - timeout(time: 20, unit: 'HOURS') - { - //check whether to run performance tests on this node - def arch = check_arch_name() - cmake_build(conf) - if ( params.RUN_INDUCTOR_TESTS && arch == "gfx90a" ){ - echo "Run inductor codegen tests" - sh "projects/composablekernel/script/run_inductor_tests.sh" - } - // run performance tests, stash the logs, results will be processed on the master node - dir("projects/composablekernel/script"){ - if (params.RUN_PERFORMANCE_TESTS){ - if (params.RUN_FULL_QA && (arch == "gfx90a" || arch == "gfx942")){ - // run full tests on gfx90a or gfx942 - echo "Run full performance tests" - sh "./run_full_performance_tests.sh 0 QA_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" - archiveArtifacts "perf_*.log" - stash includes: "perf_**.log", name: "perf_log_${arch}" - } - else if (!params.RUN_FULL_QA && (arch == "gfx90a" || arch == "gfx942")){ - // run standard tests on gfx90a or gfx942 - echo "Run performance tests" - sh "./run_performance_tests.sh 0 CI_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" - archiveArtifacts "perf_*.log" - stash includes: "perf_**.log", name: "perf_log_${arch}" - } - else if ( arch != "gfx10"){ - // run basic tests on gfx11/gfx12/gfx908/gfx950, but not on gfx10, it takes too long - echo "Run gemm performance tests" - sh "./run_gemm_performance_tests.sh 0 CI_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" - archiveArtifacts "perf_onnx_gemm_*.log" - stash includes: "perf_onnx_gemm_**.log", name: "perf_log_${arch}" - } - } - } - if (params.hipTensor_test && arch == "gfx90a" ){ - // build and test hipTensor on gfx90a node - sh """#!/bin/bash - rm -rf rocm-libraries - git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-libraries.git - cd rocm-libraries - git sparse-checkout init --cone - git sparse-checkout set projects/hiptensor - git checkout "${params.hipTensor_branch}" - """ - dir("rocm-libraries/projects/hiptensor"){ - sh """#!/bin/bash - mkdir -p build - ls -ltr - CC=hipcc CXX=hipcc cmake -Bbuild . -D CMAKE_PREFIX_PATH="${env.WORKSPACE}/install" - cmake --build build -- -j - ctest --test-dir build - """ - } - } - } - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - return retimage -} - -def Build_CK_and_Reboot(Map conf=[:]){ - try{ - Build_CK(conf) - } - catch(e){ - echo "throwing error exception while building CK" - echo 'Exception occurred: ' + e.toString() - throw e - } -} - -def process_results(Map conf=[:]){ - checkoutComposableKernel() - //use older image that has user jenkins - def image = "${env.CK_DOCKERHUB}:ck_ub22.04_rocm6.3" - - setGithubStatus("${env.STAGE_NAME}", 'pending', 'Processing results...') - try { - try - { - echo "Pulling image: ${image}" - def retimage = docker.image("${image}") - withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { - retimage.pull() - } - } - catch(Exception ex) - { - error "Unable to locate image: ${image}" - } - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - - withDockerContainer(image: image, args: '--cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v=/var/jenkins/:/var/jenkins') { - timeout(time: 15, unit: 'MINUTES'){ - try{ - dir("projects/composablekernel/script"){ - if (params.RUN_CK_TILE_FMHA_TESTS){ - try{ - unstash "perf_fmha_log_gfx942" - } - catch(Exception err){ - echo "could not locate the FMHA performance logs for gfx942: ${err.getMessage()}." - } - try{ - unstash "perf_fmha_log_gfx90a" - } - catch(Exception err){ - echo "could not locate the FMHA performance logs for gfx90a: ${err.getMessage()}." - } - try{ - unstash "perf_fmha_log_gfx950" - } - catch(Exception err){ - echo "could not locate the FMHA performance logs for gfx950: ${err.getMessage()}." - } - - } - if (params.BUILD_INSTANCES_ONLY){ - // unstash deb packages - try{ - unstash "lib_package" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate lib_package." - } - } - if (params.BUILD_PACKAGES){ - // unstash deb packages - try{ - unstash "profiler_package_gfx90a" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx90a." - } - try{ - unstash "profiler_package_gfx942" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx942." - } - try{ - unstash "profiler_package_gfx950" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx950." - } - try{ - unstash "profiler_package_gfx12" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx12." - } - } - else{ - // unstash perf files to master - try{ - unstash "perf_log_gfx90a" - } - catch(Exception err){ - echo "could not locate the gfx90a performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx942" - } - catch(Exception err){ - echo "could not locate the gfx942 performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx950" - } - catch(Exception err){ - echo "could not locate the gfx950 performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx908" - } - catch(Exception err){ - echo "could not locate the gfx908 performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx11" - } - catch(Exception err){ - echo "could not locate the gfx11 performance logs: ${err.getMessage()}." - } - try{ - - unstash "perf_log_gfx12" - } - catch(Exception err){ - echo "could not locate the gfx12 performance logs: ${err.getMessage()}." - } - } - // process the logs - sh "./process_perf_data.sh" - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - finally{ - echo "Finished processing performance test results" - } - } - } -} - -def run_downstream_tests(Map conf=[:]){ - show_node_info() - checkoutComposableKernel() - def dockerOpts = get_docker_options() + ' --group-add irc ' - - setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") - try { - try - { - 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: ${conf.image}" - } - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - - withDockerContainer(image: conf.image, args: dockerOpts) { - timeout(time: conf.get("timeoutHours", 2), unit: 'HOURS'){ - try{ - sh "rocminfo" - sh "python3 --version" - for (cmd in conf.execute_cmds) { - sh "${cmd}" - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch(e){ - echo "Throwing error exception while running ${env.STAGE_NAME}" - echo 'Exception occurred: ' + e.toString() - setGithubStatus("${env.STAGE_NAME}", 'error', "Stage ${env.STAGE_NAME} failed") - throw e - } - finally{ - echo "Finished running ${env.STAGE_NAME}" - } - } - } -} - -def getPytorchTestsCmds() { - return [ - "mkdir pytorch", - "cp -r /var/jenkins/workspace/pytorch/* pytorch/", - "ls -ltr pytorch", - "python3 pytorch/tools/amd_build/build_amd.py", - "cd pytorch && USE_ROCM_CK_SDPA=1 PYTORCH_ROCM_ARCH=gfx942 python3 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" - ] +def loadCk() { + def branch = (params.USE_CURRENT_BRANCH_FOR_CK_GROOVY + ? (env.CHANGE_BRANCH ?: env.BRANCH_NAME) + : 'develop') + library("ck@${branch}") } //launch develop branch daily jobs @@ -1463,6 +252,10 @@ pipeline { name: "FORCE_CI", defaultValue: false, description: "Force CI to run even when only non-relevant files are changed (default: OFF)") + booleanParam( + name: 'USE_CURRENT_BRANCH_FOR_CK_GROOVY', + defaultValue: false, + description: 'Load ck.groovy from the current branch instead of develop. Enable when testing pipeline changes (default: OFF).') } environment{ dbuser = "${dbuser}" @@ -1480,8 +273,9 @@ pipeline { agent{ label rocmnode("nogpu") } steps { script { - checkoutComposableKernel() - env.SHOULD_RUN_CI = String.valueOf(params.FORCE_CI.toBoolean() || shouldRunCICheck()) + loadCk() + ck.checkoutComposableKernel() + env.SHOULD_RUN_CI = String.valueOf(params.FORCE_CI.toBoolean() || ck.shouldRunCICheck()) echo "SHOULD_RUN_CI: ${env.SHOULD_RUN_CI}" } } @@ -1496,7 +290,10 @@ pipeline { agent{ label rocmnode("nogpu") } steps{ deleteDir() - buildDocker('/opt/rocm') + script { + loadCk() + ck.buildDocker('/opt/rocm') + } cleanWs() } } @@ -1514,21 +311,9 @@ pipeline { expression { params.RUN_CPPCHECK.toBoolean() } } agent{ label rocmnode("nogpu") } - environment{ - setup_args = "NO_CK_BUILD" - execute_cmd = """cd .. && \ - find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ - -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ - xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)' && \ - /cppcheck/build/bin/cppcheck ../* -v -j \$(nproc) -I ../include -I ../profiler/include -I ../library/include \ - -D CK_ENABLE_FP64 -D CK_ENABLE_FP32 -D CK_ENABLE_FP16 -D CK_ENABLE_FP8 -D CK_ENABLE_BF16 -D CK_ENABLE_BF8 -D CK_ENABLE_INT8 \ - -D __gfx908__ -D __gfx90a__ -D __gfx942__ -D __gfx1030__ -D __gfx1100__ -D __gfx1101__ -D __gfx1102__ \ - -U __gfx803__ -U __gfx900__ -U __gfx906__ -U CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4 \ - --file-filter=*.cpp --force --enable=all --output-file=ck_cppcheck.log""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) + script { loadCk(); ck.runClangFormatAndCppcheck() } archiveArtifacts "build/ck_cppcheck.log" cleanWs() } @@ -1538,17 +323,28 @@ pipeline { beforeAgent true expression { !params.RUN_CPPCHECK.toBoolean() } } + agent{ label rocmnode("nogpu") } + steps{ + deleteDir() + script { loadCk(); ck.runClangFormat() } + cleanWs() + } + } + stage('ASCII Only Check') { agent{ label rocmnode("nogpu") } environment{ setup_args = "NO_CK_BUILD" execute_cmd = """cd .. && \ - find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ - -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ - xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)'""" + find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.inc' -o -name '*.cl' \\) \ + -not -path '*/build/*' -not -path '*/include/rapidjson/*' \ + -print0 | xargs -0 -P 8 -n 64 script/check_ascii_only.sh""" } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) + script { + loadCk(); + ck.buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) + } cleanWs() } } @@ -1570,7 +366,10 @@ pipeline { } agent{ label rocmnode("gfx942")} steps{ - run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: getPytorchTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: ck.getPytorchTestsCmds()) + } cleanWs() } } @@ -1582,7 +381,10 @@ pipeline { } agent{ label rocmnode("gfx942")} steps{ - run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: getAiterTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) + } cleanWs() } } @@ -1594,7 +396,10 @@ pipeline { } agent{ label rocmnode("gfx950")} steps{ - run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: getAiterTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) + } cleanWs() } } @@ -1606,7 +411,10 @@ pipeline { } agent{ label rocmnode("gfx942")} steps{ - run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: getFaTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) + } cleanWs() } } @@ -1618,7 +426,10 @@ pipeline { } agent{ label rocmnode("gfx950")} steps{ - run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: getFaTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) + } cleanWs() } } @@ -1639,17 +450,9 @@ pipeline { expression { params.RUN_FULL_CONV_TILE_TESTS.toBoolean() } } agent{ label rocmnode("gfx90a")} - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ python3 ../experimental/grouped_convolution_tile_instances/generate_instances.py --mode=profiler && \ - cmake .. --preset dev-gfx90a -D CK_EXPERIMENTAL_BUILDER=ON && \ - make -j64 test_grouped_convnd_fwd_tile test_grouped_convnd_bwd_weight_tile && \ - ./bin/test_grouped_convnd_bwd_weight_tile && \ - ./bin/test_grouped_convnd_fwd_tile""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runFullGroupedConvTileTests() } cleanWs() } } @@ -1670,15 +473,9 @@ pipeline { expression { params.RUN_GROUPED_CONV_LARGE_CASES_TESTS.toBoolean() } } agent{ label rocmnode("gfx90a")} - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake .. --preset dev-gfx90a && \ - make -j64 test_grouped_convnd_fwd_large_cases test_grouped_convnd_bwd_data_large_cases test_grouped_convnd_fwd_bias_clamp_large_cases && \ - ./bin/test_grouped_convnd_fwd_large_cases && ./bin/test_grouped_convnd_bwd_data_large_cases && ./bin/test_grouped_convnd_fwd_bias_clamp_large_cases""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runGroupedConvLargeCaseTests() } cleanWs() } } @@ -1699,27 +496,9 @@ pipeline { expression { params.RUN_CONV_COMPREHENSIVE_DATASET.toBoolean() } } agent{ label rocmnode("gfx90a")} - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cd ../build && \ - cmake .. --preset dev-gfx90a && \ - make -j64 test_grouped_convnd_fwd_dataset_xdl && \ - test_grouped_convnd_bwd_data_dataset_xdl \ - test_grouped_convnd_bwd_weight_dataset_xdl && \ - cd ../test_data && \ - # Dataset generation modes: - # - small: ~60 test cases (minimal, quick testing - 3 models, 2 batch sizes, 2 image sizes) - # - half: ~300 test cases (moderate coverage - 16 models, 3 batch sizes, 5 image sizes), ~ 17 hours testing time - # - full: ~600 test cases (comprehensive - 16 models, 5 batch sizes, 9 image sizes), ~ 40 hours testing time - ./generate_test_dataset.sh small && \ - cd ../build && \ - ./bin/test_grouped_convnd_fwd_dataset_xdl && \ - ./bin/test_grouped_convnd_bwd_data_dataset_xdl && \ - ./bin/test_grouped_convnd_bwd_weight_dataset_xdl""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runComprehensiveConvDatasetTests() } cleanWs() } } @@ -1742,11 +521,14 @@ pipeline { agent{ label rocmnode("gfx90a") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx90a") + execute_args = ck.build_and_run_fmha("gfx90a") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1759,11 +541,14 @@ pipeline { agent{ label rocmnode("gfx942") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx942") + execute_args = ck.build_and_run_fmha("gfx942") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1776,11 +561,14 @@ pipeline { agent{ label rocmnode("gfx950") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx950") + execute_args = ck.build_and_run_fmha("gfx950") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1793,11 +581,14 @@ pipeline { agent{ label rocmnode("gfx1201") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx1201") + execute_args = ck.build_and_run_fmha("gfx1201") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1818,30 +609,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_BASIC_TESTS.toBoolean() } } agent{ label rocmnode("gfx942") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx942" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_UNIVERSAL_CONFIG_FILE="default_ci_config.json" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_MULTI_D_CONFIG_FILE="default_ci_config.json" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D GEMM_PRESHUFFLE_CONFIG_FILE="default_ci_config.json" .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineBasicTests(params.BUILD_COMPILER) } cleanWs() } } @@ -1862,33 +632,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } } agent{ label rocmnode("gfx942") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx942" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16;bf8;bf16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_STREAMK_DATATYPE="fp8;fp16" \ - -D GEMM_STREAMK_LAYOUT="rcr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D GROUPED_GEMM_DATATYPE="fp8;fp16" \ - -D GROUPED_GEMM_LAYOUT="rcr;rrr;crr;ccr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all benchmark_gemm_streamk_all benchmark_grouped_gemm_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py . --problem-sizes "1024,1024,1024" --group-counts 8 --warmup 5 --repeat 5 --verbose --json grouped_gemm_results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineGemmTests("gfx942", params.BUILD_COMPILER) } cleanWs() } } @@ -1899,28 +645,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } } agent{ label rocmnode("gfx950") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx950" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineGemmTests("gfx950", params.BUILD_COMPILER) } cleanWs() } } @@ -1931,22 +658,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } } agent{ label rocmnode("gfx1201") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx1201" \ - -D GEMM_UNIVERSAL_DATATYPE="fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineGemmTests("gfx1201", params.BUILD_COMPILER) } cleanWs() } } @@ -1968,13 +682,9 @@ pipeline { expression { (params.BUILD_GFX942.toBoolean() || params.RUN_FULL_QA.toBoolean()) && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx942") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx942" """ - execute_args = build_client_examples("gfx942") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx942") } cleanWs() } } @@ -1985,13 +695,9 @@ pipeline { expression { params.BUILD_GFX950.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx950") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx950" """ - execute_args = build_client_examples("gfx950") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx950") } cleanWs() } } @@ -2003,13 +709,9 @@ pipeline { expression { params.BUILD_GFX908.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx908") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx908" """ - execute_args = build_client_examples("gfx908") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx908") } cleanWs() } } @@ -2021,13 +723,9 @@ pipeline { expression { params.BUILD_GFX90A.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx90a") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx90a" -DCK_CXX_STANDARD="17" """ - execute_args = build_client_examples_and_codegen_tests("gfx90a") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx90a") } cleanWs() } } @@ -2048,7 +746,7 @@ pipeline { } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args: setup_args, build_cmd: "", build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runBuildInstancesOnly(params.BUILD_COMPILER) } cleanWs() } } @@ -2060,13 +758,9 @@ pipeline { expression { params.BUILD_GFX101.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx1010") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx10-1-generic" """ - execute_args = build_client_examples("gfx10-1-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx10-1-generic") } cleanWs() } } @@ -2078,13 +772,9 @@ pipeline { expression { params.BUILD_GFX103.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx1030") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx10-3-generic" """ - execute_args = build_client_examples("gfx10-3-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx10-3-generic") } cleanWs() } } @@ -2095,13 +785,9 @@ pipeline { expression { params.BUILD_GFX11.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label 'miopen && (gfx1101 || gfx1100)' } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx11-generic" """ - execute_args = build_client_examples("gfx11-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx11-generic") } cleanWs() } } @@ -2112,13 +798,9 @@ pipeline { expression { params.BUILD_GFX12.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx1201") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx12-generic" """ - execute_args = build_client_examples("gfx12-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx12-generic") } cleanWs() } } @@ -2129,12 +811,9 @@ pipeline { expression { params.BUILD_GFX1250.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx90a") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx1250" -DDISABLE_DL_KERNELS="ON" """ - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, docker_name: "${env.CK_DOCKERHUB_PRIVATE}:npi-mi450-latest", config_targets: "install", no_reboot:true, build_type: 'Release', prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx1250") } cleanWs() } } @@ -2143,12 +822,13 @@ pipeline { always { node(rocmnode("nogpu")) { script { + loadCk() // Simulate capture - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx11.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx12.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx90a.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx942.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx950.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx11.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx12.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx90a.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx942.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx950.json") } cleanWs() } @@ -2156,8 +836,9 @@ pipeline { success { script { node(rocmnode("nogpu")) { + loadCk() // Report the parent stage build ck and run tests status - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + ck.setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") echo "Reporting success status for build ck and run tests" } } @@ -2176,7 +857,10 @@ pipeline { agent { label 'mici' } steps{ deleteDir() - process_results() + script { + loadCk() + ck.process_results() + } cleanWs() } } @@ -2185,8 +869,9 @@ pipeline { success { script { node(rocmnode("nogpu")) { + loadCk() // Report the skipped parent's stage status - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + ck.setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") echo "Process Performance Test Results stage skipped." } } @@ -2198,17 +883,17 @@ pipeline { success { script { node(rocmnode("nogpu")) { - setGithubStatus('Math CI Summary', 'success', "Math CI passed") + loadCk() + ck.setGithubStatus('Math CI Summary', 'success', "Math CI passed") } } } failure { script { node(rocmnode("nogpu")) { - setGithubStatus('Math CI Summary', 'failure', "Math CI failed") - script { - checkoutComposableKernel() - } + loadCk() + ck.setGithubStatus('Math CI Summary', 'failure', "Math CI failed") + ck.checkoutComposableKernel() withCredentials([string(credentialsId: 'ck_ci_errors_webhook_url', variable: 'WEBHOOK_URL')]) { sh 'bash projects/composablekernel/script/infra_helper/send_failure_notifications.sh' } diff --git a/README.md b/README.md index d48f7ed6765..9d5affa13a9 100644 --- a/README.md +++ b/README.md @@ -230,7 +230,7 @@ Additional cmake flags can be used to significantly speed-up the build: These instances offer a slightly better performance of fp16 gemms on NAVI2x. But on other architectures faster alternatives are available. * `CK_USE_FP8_ON_UNSUPPORTED_ARCH` (default is OFF) must be set to ON in order to build instances, - such as `gemm_universal`, `gemm_universal_streamk` and `gemm_multiply_multiply` for fp8 data type for GPU targets which do not have native support for fp8 data type, such as gfx908 or gfx90a. These instances are useful on + such as `gemm_universal`, and `gemm_multiply_multiply` for fp8 data type for GPU targets which do not have native support for fp8 data type, such as gfx908 or gfx90a. These instances are useful on architectures like the MI100/MI200 for the functional support only. ## Using sccache for building diff --git a/cmake/EnableCompilerWarnings.cmake b/cmake/EnableCompilerWarnings.cmake index 9cc960cc234..2f9a04f4855 100644 --- a/cmake/EnableCompilerWarnings.cmake +++ b/cmake/EnableCompilerWarnings.cmake @@ -50,6 +50,10 @@ else() -Wsign-compare -Wno-extra-semi-stmt -Wno-unused-template + -Wno-lifetime-safety-intra-tu-suggestions + -Wno-lifetime-safety-cross-tu-suggestions + -Wno-lifetime-safety-lifetimebound-violation + -Wno-unknown-warning-option ) if (CMAKE_${COMPILER}_COMPILER_ID MATCHES "Clang") list(APPEND CMAKE_COMPILER_WARNINGS @@ -76,6 +80,10 @@ else() -Wno-unsafe-buffer-usage -Wno-unused-lambda-capture -Wno-nvcc-compat + -Wno-lifetime-safety-intra-tu-suggestions + -Wno-lifetime-safety-cross-tu-suggestions + -Wno-lifetime-safety-lifetimebound-violation + -Wno-unknown-warning-option ) if(CK_CXX_STANDARD GREATER_EQUAL 20) list(APPEND CMAKE_COMPILER_WARNINGS -Wno-c++20-compat) diff --git a/cmake/TestUtilities.cmake b/cmake/TestUtilities.cmake new file mode 100644 index 00000000000..f95570b3984 --- /dev/null +++ b/cmake/TestUtilities.cmake @@ -0,0 +1,30 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +# Helper function to conditionally link device_conv libraries only if they exist as targets. +# This is useful when device_conv libraries may be filtered out based on GPU targets, +# DTYPES, or build configuration flags. +# +# Usage: +# target_link_device_conv_libraries_if_exist(my_test PRIVATE utility device_conv2d_nhwgc_operations ...) +# +# Only device_conv* libraries are checked with if(TARGET). +# All other libraries (utility, gtest_main, etc.) are always linked. +function(target_link_device_conv_libraries_if_exist TARGET_NAME VISIBILITY) + set(_libs_to_link) + foreach(lib ${ARGN}) + if(lib MATCHES "^device_conv") + # Only add device_conv libraries if they exist + if(TARGET ${lib}) + list(APPEND _libs_to_link ${lib}) + endif() + else() + # Always add non-device_conv libraries + list(APPEND _libs_to_link ${lib}) + endif() + endforeach() + # Single target_link_libraries call with all libraries + if(_libs_to_link) + target_link_libraries(${TARGET_NAME} ${VISIBILITY} ${_libs_to_link}) + endif() +endfunction() diff --git a/dispatcher/CMakeLists.txt b/dispatcher/CMakeLists.txt index ed9b20d33c9..79bdde45e87 100644 --- a/dispatcher/CMakeLists.txt +++ b/dispatcher/CMakeLists.txt @@ -59,7 +59,7 @@ endif() # Compiler warnings if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") target_compile_options(ck_tile_dispatcher PRIVATE - -Wall -Wextra -Wpedantic + -Wall -Wextra -Wpedantic -Wno-lifetime-safety-intra-tu-suggestions -Wno-lifetime-safety-cross-tu-suggestions -Wno-lifetime-safety-lifetimebound-violation -Wno-unknown-warning-option ) elseif(CMAKE_CXX_COMPILER_ID MATCHES "MSVC") target_compile_options(ck_tile_dispatcher PRIVATE diff --git a/dispatcher/codegen/fmha/validation.py b/dispatcher/codegen/fmha/validation.py index 20b3a00540d..2fb791a8ba0 100644 --- a/dispatcher/codegen/fmha/validation.py +++ b/dispatcher/codegen/fmha/validation.py @@ -29,8 +29,10 @@ # Ensure this directory and parent codegen/ are on sys.path for sibling imports _THIS_DIR = Path(__file__).resolve().parent _CODEGEN_DIR = _THIS_DIR.parent +_DISPATCHER_PYTHON_DIR = _CODEGEN_DIR.parent / "python" sys.path.insert(0, str(_THIS_DIR)) sys.path.insert(0, str(_CODEGEN_DIR)) +sys.path.insert(0, str(_DISPATCHER_PYTHON_DIR)) from symbol_map import ( # noqa: E402 BWD_DTYPE_MAP, @@ -39,6 +41,10 @@ canonical_mask, canonical_qscale, ) +from fmha_dtype_contract import ( # noqa: E402 + FmhaDTypeContractKind, + dtype_contract_from_signature, +) # Import shared hardware data from parent arch_specs_generated (generated from # arch_specs.json by generate_arch_specs.py). Falls back to inline defaults if @@ -872,6 +878,15 @@ def validate_config( # --- Family-specific rules --- if family == "batch_prefill": + dtype_contract = dtype_contract_from_signature(sig) + if dtype_contract.kind == FmhaDTypeContractKind.MIXED_Q_FP8_KV: + result.add_error( + "batch_prefill mixed activation/FP8-KV dtype contract is not implemented " + f"(Q={dtype_contract.q_dtype}, K={dtype_contract.k_dtype}, " + f"V={dtype_contract.v_dtype}, O={dtype_contract.o_dtype}); " + "current generated kernels use one data_type token for Q/K/V, and fp8bf16 " + "means FP8 Q/K/V with BF16 output" + ) if sig.get("vlayout", "r") != "r": result.add_error("batch_prefill only supports row-major V layout") if not sig.get("paged_kv", False): diff --git a/tile_engine/ops/grouped_conv/compare_ml_vs_oracle.py b/dispatcher/heuristics/validation/grouped_conv/compare_ml_vs_oracle.py similarity index 94% rename from tile_engine/ops/grouped_conv/compare_ml_vs_oracle.py rename to dispatcher/heuristics/validation/grouped_conv/compare_ml_vs_oracle.py index 974b85e4f83..ce8dca980b7 100644 --- a/tile_engine/ops/grouped_conv/compare_ml_vs_oracle.py +++ b/dispatcher/heuristics/validation/grouped_conv/compare_ml_vs_oracle.py @@ -1,4 +1,7 @@ #!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + """ Compare ML heuristic predictions against oracle benchmark results. @@ -114,7 +117,15 @@ def run_end_to_end_workflow(args): elif args.problem_set: print(f"Problem set: {args.problem_set}") # Import problem set dynamically - sys.path.insert(0, str(Path(__file__).parent / "problems")) + # Problem sets live with the benchmarking harness in tile_engine. + _THIS_DIR = Path(__file__).parent + _TILE_ENGINE_GROUPED_CONV = ( + _THIS_DIR.parent.parent.parent.parent + / "tile_engine" + / "ops" + / "grouped_conv" + ) + sys.path.insert(0, str(_TILE_ENGINE_GROUPED_CONV / "problems")) try: problem_module = __import__(args.problem_set) problem_attr = ( @@ -165,15 +176,15 @@ def run_end_to_end_workflow(args): print() print("Please use the manual workflow documented in README.md:") print() - print(" 1. Create problem set file in problems/") + print(" 1. Create problem set file in tile_engine/ops/grouped_conv/problems/") print( - " 2. Run: python grouped_conv_full_benchmark.py --problems --csv oracle.csv" + " 2. Run: cd tile_engine/ops/grouped_conv && python grouped_conv_full_benchmark.py --problems --csv oracle.csv" ) print( - " 3. Run: cd ../../dispatcher/heuristics && python predict_cli.py --problem-module --output ml.csv" + " 3. Run: cd dispatcher/heuristics && python predict_cli.py --problem-module --output ml.csv" ) print( - " 4. Run: cd ../../tile_engine/ops/grouped_conv && python compare_ml_vs_oracle.py --oracle-csv oracle.csv --ml-csv ml.csv --plot result.png" + " 4. Run: cd dispatcher/heuristics/validation/grouped_conv && python compare_ml_vs_oracle.py --oracle-csv oracle.csv --ml-csv ml.csv --plot result.png" ) print() diff --git a/tile_engine/ops/grouped_conv/validate_ml_vs_oracle.py b/dispatcher/heuristics/validation/grouped_conv/validate_ml_vs_oracle.py similarity index 87% rename from tile_engine/ops/grouped_conv/validate_ml_vs_oracle.py rename to dispatcher/heuristics/validation/grouped_conv/validate_ml_vs_oracle.py index 9e5124caf8a..0da88839ae3 100755 --- a/tile_engine/ops/grouped_conv/validate_ml_vs_oracle.py +++ b/dispatcher/heuristics/validation/grouped_conv/validate_ml_vs_oracle.py @@ -12,18 +12,24 @@ 4. Reports efficiency metrics """ +import argparse import sys from pathlib import Path import pandas as pd import numpy as np _THIS_DIR = Path(__file__).parent -_DISPATCHER_ROOT = _THIS_DIR.parent.parent.parent / "dispatcher" +# This file lives at: /projects/composablekernel/dispatcher/heuristics/validation/grouped_conv/ +# Walk up three levels (validation -> heuristics -> dispatcher) to find the dispatcher root. +_DISPATCHER_ROOT = _THIS_DIR.parent.parent.parent +_CK_ROOT = _DISPATCHER_ROOT.parent +# Problem definitions still live with the benchmarking harness in tile_engine. +_TILE_ENGINE_GROUPED_CONV = _CK_ROOT / "tile_engine" / "ops" / "grouped_conv" sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) sys.path.insert(0, str(_DISPATCHER_ROOT / "heuristics")) sys.path.insert(0, str(_DISPATCHER_ROOT / "codegen")) -sys.path.insert(0, str(_THIS_DIR / "problems")) +sys.path.insert(0, str(_TILE_ENGINE_GROUPED_CONV / "problems")) from validation_holdout import VALIDATION_PROBLEMS # noqa: E402 from predict import Predictor # noqa: E402 @@ -81,11 +87,31 @@ def _build_kernel_name(kconf, ndim): ) -# Load model -model_dir = ( - _DISPATCHER_ROOT - / "heuristics/models/grouped_conv_forward_bf16_gfx950_2d_3d_no_compv5" +# Parse CLI args +_parser = argparse.ArgumentParser(description=__doc__) +_parser.add_argument( + "--oracle-csv", + type=Path, + default=_TILE_ENGINE_GROUPED_CONV / "validation_oracle_results.csv", + help="Oracle benchmark CSV (produced by tile_engine/ops/grouped_conv/grouped_conv_full_benchmark.py)", +) +_parser.add_argument( + "--model-dir", + type=Path, + default=_DISPATCHER_ROOT + / "heuristics/models/grouped_conv_forward_bf16_gfx950_2d_3d_no_compv5", + help="Trained LightGBM model directory.", +) +_parser.add_argument( + "--output", + type=Path, + default=_THIS_DIR / "validation_heuristic_vs_oracle.csv", + help="Where to write the per-problem comparison CSV.", ) +_args = _parser.parse_args() + +# Load model +model_dir = _args.model_dir feature_engine = GroupedConvFeatureEngine() predictor = Predictor(model_dir, feature_engine=feature_engine) @@ -98,7 +124,7 @@ def _build_kernel_name(kconf, ndim): print() # Load oracle benchmark results -oracle_df = pd.read_csv(_THIS_DIR / "validation_oracle_results.csv") +oracle_df = pd.read_csv(_args.oracle_csv) print(f"Oracle measurements: {len(oracle_df)}") print() @@ -281,7 +307,7 @@ def _build_kernel_name(kconf, ndim): print() # Save detailed results - results_df.to_csv(_THIS_DIR / "validation_heuristic_vs_oracle.csv", index=False) - print("Detailed results saved to: validation_heuristic_vs_oracle.csv") + results_df.to_csv(_args.output, index=False) + print(f"Detailed results saved to: {_args.output}") else: print("ERROR: No predictions could be compared with oracle data") diff --git a/dispatcher/include/ck_tile/dispatcher/backends/generated_conv_backend.hpp b/dispatcher/include/ck_tile/dispatcher/backends/generated_conv_backend.hpp index b8e4964b132..75a71777b46 100644 --- a/dispatcher/include/ck_tile/dispatcher/backends/generated_conv_backend.hpp +++ b/dispatcher/include/ck_tile/dispatcher/backends/generated_conv_backend.hpp @@ -148,7 +148,7 @@ inline GroupedConvKernelInstance::RunFn make_conv_bwd_weight_run_fn() } // ------------------------------------------------------------------------- -// IsSupportedFn factories — check kernel applicability without launching +// IsSupportedFn factories -- check kernel applicability without launching // ------------------------------------------------------------------------- template @@ -181,7 +181,7 @@ inline GroupedConvKernelInstance::IsSupportedFn make_conv_bwd_data_is_supported_ } // ------------------------------------------------------------------------- -// Instance string extraction — get CK Tile GetInstanceString() representation +// Instance string extraction -- get CK Tile GetInstanceString() representation // ------------------------------------------------------------------------- #ifdef CK_EXPERIMENTAL_BUILDER diff --git a/dispatcher/python/fmha_dtype_contract.py b/dispatcher/python/fmha_dtype_contract.py new file mode 100644 index 00000000000..21b6e932962 --- /dev/null +++ b/dispatcher/python/fmha_dtype_contract.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +from dataclasses import dataclass +from enum import Enum +from typing import Mapping, Optional + + +class FmhaDTypeContractKind(Enum): + HOMOGENEOUS = "homogeneous" + ALL_FP8_WITH_BF16_OUTPUT = "all_fp8_with_bf16_output" + ALL_FP8_WITH_FP32_OUTPUT = "all_fp8_with_fp32_output" + MIXED_Q_FP8_KV = "mixed_q_fp8_kv" + UNSUPPORTED = "unsupported" + + +@dataclass(frozen=True) +class FmhaDTypeContract: + data_type: str + q_dtype: str + k_dtype: str + v_dtype: str + o_dtype: str + kind: FmhaDTypeContractKind + + @property + def uses_fp8_kv(self) -> bool: + return _is_fp8(self.k_dtype) and _is_fp8(self.v_dtype) + + +_TOKEN_CONTRACTS = { + "fp16": ("fp16", "fp16", "fp16", "fp16"), + "bf16": ("bf16", "bf16", "bf16", "bf16"), + "fp32": ("fp32", "fp32", "fp32", "fp32"), + "fp8": ("fp8", "fp8", "fp8", "fp8"), + "bf8": ("bf8", "bf8", "bf8", "bf8"), + "fp8bf16": ("fp8", "fp8", "fp8", "bf16"), + "fp8fp32": ("fp8", "fp8", "fp8", "fp32"), + "fp8fp16": ("fp8", "fp8", "fp8", "fp16"), + "mxfp8": ("fp8", "fp8", "fp8", "fp32"), +} + + +def _normalize_dtype(dtype: Optional[str]) -> Optional[str]: + if dtype is None: + return None + + normalized = str(dtype).lower() + aliases = { + "float16": "fp16", + "half": "fp16", + "uint16": "bf16", + "bfloat16": "bf16", + "float32": "fp32", + "uint8": "fp8", + "fp8_e4m3": "fp8", + "fp8_e4m3fnuz": "fp8", + "float8_e4m3fnuz": "fp8", + } + return aliases.get(normalized, normalized) + + +def _is_fp8(dtype: str) -> bool: + return _normalize_dtype(dtype) in {"fp8", "bf8", "mxfp8"} + + +def _classify( + q_dtype: str, k_dtype: str, v_dtype: str, o_dtype: str +) -> FmhaDTypeContractKind: + q_dtype = _normalize_dtype(q_dtype) or "" + k_dtype = _normalize_dtype(k_dtype) or "" + v_dtype = _normalize_dtype(v_dtype) or "" + o_dtype = _normalize_dtype(o_dtype) or "" + + if q_dtype == k_dtype == v_dtype == o_dtype: + return FmhaDTypeContractKind.HOMOGENEOUS + if _is_fp8(q_dtype) and _is_fp8(k_dtype) and _is_fp8(v_dtype): + if o_dtype == "bf16": + return FmhaDTypeContractKind.ALL_FP8_WITH_BF16_OUTPUT + if o_dtype == "fp32": + return FmhaDTypeContractKind.ALL_FP8_WITH_FP32_OUTPUT + if ( + q_dtype in {"fp16", "bf16"} + and _is_fp8(k_dtype) + and _is_fp8(v_dtype) + and o_dtype in {"fp16", "bf16"} + ): + return FmhaDTypeContractKind.MIXED_Q_FP8_KV + return FmhaDTypeContractKind.UNSUPPORTED + + +def dtype_contract_from_components( + data_type: str, + q_dtype: str, + k_dtype: str, + v_dtype: str, + o_dtype: str, +) -> FmhaDTypeContract: + data_type = _normalize_dtype(data_type) or data_type + q_dtype = _normalize_dtype(q_dtype) or "" + k_dtype = _normalize_dtype(k_dtype) or "" + v_dtype = _normalize_dtype(v_dtype) or "" + o_dtype = _normalize_dtype(o_dtype) or "" + return FmhaDTypeContract( + data_type=data_type, + q_dtype=q_dtype, + k_dtype=k_dtype, + v_dtype=v_dtype, + o_dtype=o_dtype, + kind=_classify(q_dtype, k_dtype, v_dtype, o_dtype), + ) + + +def dtype_contract_from_data_type(data_type: str) -> FmhaDTypeContract: + data_type = _normalize_dtype(data_type) or data_type + q_dtype, k_dtype, v_dtype, o_dtype = _TOKEN_CONTRACTS.get( + data_type, (data_type, data_type, data_type, data_type) + ) + return dtype_contract_from_components(data_type, q_dtype, k_dtype, v_dtype, o_dtype) + + +def dtype_contract_from_signature(signature: Mapping[str, object]) -> FmhaDTypeContract: + data_type = str(signature.get("data_type", "fp16")) + inferred = dtype_contract_from_data_type(data_type) + kv_dtype = signature.get("kv_data_type", signature.get("kv_dtype")) + + q_dtype = signature.get("q_data_type", signature.get("q_dtype", inferred.q_dtype)) + k_dtype = signature.get( + "k_data_type", signature.get("k_dtype", kv_dtype or inferred.k_dtype) + ) + v_dtype = signature.get( + "v_data_type", signature.get("v_dtype", kv_dtype or inferred.v_dtype) + ) + o_dtype = signature.get("o_data_type", signature.get("o_dtype", inferred.o_dtype)) + + return dtype_contract_from_components( + data_type, + str(q_dtype), + str(k_dtype), + str(v_dtype), + str(o_dtype), + ) diff --git a/dispatcher/python/fmha_utils.py b/dispatcher/python/fmha_utils.py index 5d3d0854960..0c30d823b72 100644 --- a/dispatcher/python/fmha_utils.py +++ b/dispatcher/python/fmha_utils.py @@ -28,6 +28,13 @@ import numpy as np +from fmha_dtype_contract import ( + FmhaDTypeContract, + FmhaDTypeContractKind, + dtype_contract_from_components, + dtype_contract_from_data_type, +) + # ============================================================================= # Utility helpers @@ -350,6 +357,60 @@ def _bf16_to_float32(arr: np.ndarray) -> np.ndarray: return (arr.astype(np.uint32) << 16).view(np.float32) +def _array_contract_dtype(arr: np.ndarray, fallback: str) -> str: + if arr.dtype == np.uint8: + return "fp8" + if arr.dtype == np.uint16: + return "bf16" + if arr.dtype == np.float16: + return "fp16" + if arr.dtype == np.float32: + return "fp32" + return fallback + + +def get_batch_prefill_dtype_contract( + data_type: str, + Q: np.ndarray, + K: np.ndarray, + V: np.ndarray, +) -> FmhaDTypeContract: + """Classify the public batch_prefill dtype contract requested by the caller.""" + inferred = dtype_contract_from_data_type(data_type) + return dtype_contract_from_components( + data_type=data_type, + q_dtype=_array_contract_dtype(Q, inferred.q_dtype), + k_dtype=_array_contract_dtype(K, inferred.k_dtype), + v_dtype=_array_contract_dtype(V, inferred.v_dtype), + o_dtype=inferred.o_dtype, + ) + + +def _validate_batch_prefill_input_dtypes( + api_family: str, + data_type: str, + Q: np.ndarray, + K: np.ndarray, + V: np.ndarray, +) -> None: + """Reject dtype contracts that CK Tile batch_prefill cannot dispatch yet.""" + if api_family != "batch_prefill": + return + + contract = get_batch_prefill_dtype_contract(data_type, Q, K, V) + if contract.kind == FmhaDTypeContractKind.MIXED_Q_FP8_KV: + raise ValueError( + "CK Tile batch_prefill does not yet support the mixed activation/FP8-KV " + "dtype contract " + f"(data_type={data_type}, Q={contract.q_dtype}, K={contract.k_dtype}, " + f"V={contract.v_dtype}, O={contract.o_dtype}). " + "The current dispatcher and generated kernels select Q/K/V types from a single " + "data_type token; fp8bf16 means FP8 Q/K/V with BF16 output. Use AITER " + "paged_attention_ragged or another fallback for BF16/FP16 Q with FP8 KV until " + "CK Tile has mixed Q/KV kernel instances." + ) + + def cpu_attention_fwd( Q: np.ndarray, K: np.ndarray, @@ -811,6 +872,8 @@ def run( Returns: FmhaResult with output array, timing, TFLOPS """ + _validate_batch_prefill_input_dtypes(api_family, data_type, Q, K, V) + # Map CK dtype to numpy dtype for buffer allocation. # bf16 is stored as uint16 (upper 16 bits of float32). # fp8 uses uint8 (1 byte per element). diff --git a/dispatcher/tests/test_fmha_rules.py b/dispatcher/tests/test_fmha_rules.py index b2bcd99c092..29b316b9784 100644 --- a/dispatcher/tests/test_fmha_rules.py +++ b/dispatcher/tests/test_fmha_rules.py @@ -133,6 +133,44 @@ def test_batch_prefill_valid_group(self): r = validate_config(cfg, SPECS) self.assertTrue(r.valid, r.errors) + def test_batch_prefill_rejects_mixed_activation_fp8_kv_contract(self): + cfg = _base_config( + family="batch_prefill", + dtype="bf16", + pipeline="qr_async", + mode="group", + paged_kv=True, + page_size=16, + q_data_type="bf16", + kv_data_type="fp8", + o_data_type="bf16", + ) + + r = validate_config(cfg, SPECS) + + self.assertFalse(r.valid) + self.assertTrue( + any("mixed activation/FP8-KV dtype contract" in e for e in r.errors), + r.errors, + ) + + def test_batch_prefill_keeps_all_fp8_bf16_output_contract_valid(self): + cfg = _base_config( + family="batch_prefill", + dtype="fp8bf16", + pipeline="qr_async", + mode="group", + paged_kv=True, + page_size=16, + q_data_type="fp8", + kv_data_type="fp8", + o_data_type="bf16", + ) + + r = validate_config(cfg, SPECS) + + self.assertTrue(r.valid, r.errors) + def test_splitkv_combine_bn1_must_be_32(self): cfg = _base_config(family="fwd_splitkv_combine", pipeline="qr") cfg["algorithm"]["tile"][3] = 64 diff --git a/dispatcher/tests/test_fmha_utils.py b/dispatcher/tests/test_fmha_utils.py new file mode 100644 index 00000000000..622a487cc74 --- /dev/null +++ b/dispatcher/tests/test_fmha_utils.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +import sys +import unittest +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "python")) + +from fmha_dtype_contract import FmhaDTypeContractKind # noqa: E402 +from fmha_utils import ( # noqa: E402 + _validate_batch_prefill_input_dtypes, + get_batch_prefill_dtype_contract, +) + + +class TestBatchPrefillDtypeValidation(unittest.TestCase): + def test_mixed_bf16_q_fp8_kv_gqa_decode_reports_unsupported(self): + batch = 4 + q_len = 1 + ctx_len = 1024 + num_q_heads = 96 + num_kv_heads = 8 + head_dim = 128 + + q = np.zeros((batch, num_q_heads, q_len, head_dim), dtype=np.uint16) + k = np.zeros((batch, num_kv_heads, ctx_len, head_dim), dtype=np.uint8) + v = np.zeros((batch, num_kv_heads, ctx_len, head_dim), dtype=np.uint8) + + with self.assertRaisesRegex( + ValueError, + "mixed activation/FP8-KV dtype contract", + ): + _validate_batch_prefill_input_dtypes("batch_prefill", "bf16", q, k, v) + + def test_mixed_fp16_q_fp8_kv_reports_unsupported(self): + q = np.zeros((4, 96, 1, 128), dtype=np.float16) + k = np.zeros((4, 8, 1024, 128), dtype=np.uint8) + v = np.zeros((4, 8, 1024, 128), dtype=np.uint8) + + contract = get_batch_prefill_dtype_contract("fp16", q, k, v) + + self.assertEqual(contract.kind, FmhaDTypeContractKind.MIXED_Q_FP8_KV) + with self.assertRaisesRegex(ValueError, "AITER paged_attention_ragged"): + _validate_batch_prefill_input_dtypes("batch_prefill", "fp16", q, k, v) + + def test_all_fp8_bf16_output_path_remains_allowed(self): + q = np.zeros((1, 96, 1, 128), dtype=np.uint8) + k = np.zeros((1, 8, 128, 128), dtype=np.uint8) + v = np.zeros((1, 8, 128, 128), dtype=np.uint8) + + contract = get_batch_prefill_dtype_contract("fp8bf16", q, k, v) + + self.assertEqual(contract.kind, FmhaDTypeContractKind.ALL_FP8_WITH_BF16_OUTPUT) + _validate_batch_prefill_input_dtypes("batch_prefill", "fp8bf16", q, k, v) + + def test_all_bf16_batch_prefill_path_remains_allowed(self): + q = np.zeros((1, 96, 1, 128), dtype=np.uint16) + k = np.zeros((1, 8, 128, 128), dtype=np.uint16) + v = np.zeros((1, 8, 128, 128), dtype=np.uint16) + + contract = get_batch_prefill_dtype_contract("bf16", q, k, v) + + self.assertEqual(contract.kind, FmhaDTypeContractKind.HOMOGENEOUS) + _validate_batch_prefill_input_dtypes("batch_prefill", "bf16", q, k, v) + + def test_non_batch_prefill_paths_are_unchanged(self): + q = np.zeros((1, 96, 1, 128), dtype=np.uint16) + k = np.zeros((1, 8, 128, 128), dtype=np.uint8) + v = np.zeros((1, 8, 128, 128), dtype=np.uint8) + + _validate_batch_prefill_input_dtypes("fwd", "bf16", q, k, v) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/conceptual/ck_tile/adaptors.rst b/docs/conceptual/ck_tile/adaptors.rst index 8720199eab6..de71c405a71 100644 --- a/docs/conceptual/ck_tile/adaptors.rst +++ b/docs/conceptual/ck_tile/adaptors.rst @@ -15,50 +15,28 @@ TensorAdaptor Basics A TensorAdaptor encapsulates a sequence of :ref:`coordinate transformations `, managing the flow of coordinates through multiple transform stages: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Adaptor Composition" - subgraph "Single Transform" - direction TB - I1["Input Coords
[0,1,2]"] - T1["Transform
(e.g., Transpose)"] - O1["Output Coords
[2,0,1]"] - I1 --> T1 --> O1 - end - - subgraph "Chained Transforms" - direction TB - I2["Input
2D"] - T2A["Transform A
(e.g., Merge)"] - M2["Intermediate
1D"] - T2B["Transform B
(e.g., Pad)"] - O2["Output
1D Padded"] - I2 --> T2A --> M2 --> T2B --> O2 - end - end - - style T1 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style T2A fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style T2B fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - - - -.. image:: diagrams/adaptors_1.svg - :alt: Diagram - :align: center - -.. image:: diagrams/adaptors_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Adaptor Composition" + subgraph "Single Transform" + direction TB + I1["Input Coords
[0,1,2]"] + T1["Transform
(e.g., Transpose)"] + O1["Output Coords
[2,0,1]"] + I1 --> T1 --> O1 + end + + subgraph "Chained Transforms" + direction TB + I2["Input
2D"] + T2A["Transform A
(e.g., Merge)"] + M2["Intermediate
1D"] + T2B["Transform B
(e.g., Pad)"] + O2["Output
1D Padded"] + I2 --> T2A --> M2 --> T2B --> O2 + end + end Core Components ~~~~~~~~~~~~~~~ @@ -125,59 +103,36 @@ Chaining Adaptors: Building Complex Transformations The real power of adaptors comes from chaining multiple transformations together to create advanced data access patterns: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Adaptor Chaining Flow" - subgraph "Adaptor 1" - A1I["Bottom Dims
[0,1]"] - A1T["Transform:
Merge[2,3]"] - A1O["Top Dims
[0]"] - end - - subgraph "Adaptor 2" - A2I["Bottom Dims
[0]"] - A2T["Transform:
Unmerge[2,3]"] - A2O["Top Dims
[0,1]"] - end - - subgraph "Chained Result" - CI["Input 2D
Bottom[0,1]"] - CO["Output 2D
Top[0,1]"] - end - end - - A1I --> A1T - A1T --> A1O - A1O --> A2I - A2I --> A2T - A2T --> A2O - - CI --> A1I - A2O --> CO - - style A1T fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style A2T fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style CI fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style CO fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - -.. image:: diagrams/adaptors_2.svg - :alt: Diagram - :align: center - -.. image:: diagrams/adaptors_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Adaptor Chaining Flow" + subgraph "Adaptor 1" + A1I["Bottom Dims
[0,1]"] + A1T["Transform:
Merge[2,3]"] + A1O["Top Dims
[0]"] + end + + subgraph "Adaptor 2" + A2I["Bottom Dims
[0]"] + A2T["Transform:
Unmerge[2,3]"] + A2O["Top Dims
[0,1]"] + end + + subgraph "Chained Result" + CI["Input 2D
Bottom[0,1]"] + CO["Output 2D
Top[0,1]"] + end + end + + A1I --> A1T + A1T --> A1O + A1O --> A2I + A2I --> A2T + A2T --> A2O + + CI --> A1I + A2O --> CO .. code-block:: cpp diff --git a/docs/conceptual/ck_tile/buffer_views.rst b/docs/conceptual/ck_tile/buffer_views.rst index ca574724ab6..3ec7aeb88c3 100644 --- a/docs/conceptual/ck_tile/buffer_views.rst +++ b/docs/conceptual/ck_tile/buffer_views.rst @@ -24,53 +24,32 @@ Memory coherence and caching policies represent another layer of complexity that Address Space Usage Patterns ---------------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart TB - subgraph CF ["Compute Flow"] - direction LR - GM1["Global Memory
Input Data"] --> LDS["LDS
Tile Cache"] - LDS --> VGPR["VGPR
Working Set"] - VGPR --> Compute["Compute
Operations"] - Compute --> VGPR - VGPR --> LDS2["LDS
Reduction"] - LDS2 --> GM2["Global Memory
Output Data"] - end - - subgraph UP ["Usage Pattern"] - direction LR - P1["1. Load tile from Global → LDS"] - P2["2. Load working set LDS → VGPR"] - P3["3. Compute in VGPR"] - P4["4. Store results VGPR → LDS"] - P5["5. Reduce in LDS"] - P6["6. Write final LDS → Global"] - - P1 --> P2 --> P3 --> P4 --> P5 --> P6 - end - - CF ~~~ UP - - style GM1 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style LDS fill:#fed7aa,stroke:#f59e0b,stroke-width:2px - style VGPR fill:#d1fae5,stroke:#10b981,stroke-width:2px - style Compute fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - - - - - - -.. image:: diagrams/buffer_views_1.svg - :alt: Diagram - :align: center - +.. mermaid:: + + flowchart TB + subgraph CF ["Compute Flow"] + direction LR + GM1["Global Memory
Input Data"] --> LDS["LDS
Tile Cache"] + LDS --> VGPR["VGPR
Working Set"] + VGPR --> Compute["Compute
Operations"] + Compute --> VGPR + VGPR --> LDS2["LDS
Reduction"] + LDS2 --> GM2["Global Memory
Output Data"] + end + + subgraph UP ["Usage Pattern"] + direction LR + P1["1. Load tile from Global → LDS"] + P2["2. Load working set LDS → VGPR"] + P3["3. Compute in VGPR"] + P4["4. Store results VGPR → LDS"] + P5["5. Reduce in LDS"] + P6["6. Write final LDS → Global"] + + P1 --> P2 --> P3 --> P4 --> P5 --> P6 + end + + CF ~~~ UP C++ Implementation ------------------ @@ -190,101 +169,59 @@ The implementation of vector access maintains the same parameter structure as sc Scalar vs Vectorized Memory Access ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Scalar Access (4 instructions)" - S1["Load float[0]"] --> R1["Register 1"] - S2["Load float[1]"] --> R2["Register 2"] - S3["Load float[2]"] --> R3["Register 3"] - S4["Load float[3]"] --> R4["Register 4"] - end - - subgraph "Vectorized Access (1 instruction)" - V1["Load float4[0]"] --> VR["Vector Register
(4 floats)"] - end - - subgraph "Performance Impact" - Perf["4x fewer instructions
Better memory bandwidth
Reduced latency"] - end - - R1 & R2 & R3 & R4 --> Perf - VR --> Perf - - style S1 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style S2 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style S3 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style S4 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style V1 fill:#d1fae5,stroke:#10b981,stroke-width:2px - style Perf fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - - - - - +.. mermaid:: + + graph LR + subgraph "Scalar Access (4 instructions)" + S1["Load float[0]"] --> R1["Register 1"] + S2["Load float[1]"] --> R2["Register 2"] + S3["Load float[2]"] --> R3["Register 3"] + S4["Load float[3]"] --> R4["Register 4"] + end + + subgraph "Vectorized Access (1 instruction)" + V1["Load float4[0]"] --> VR["Vector Register
(4 floats)"] + end -.. image:: diagrams/buffer_views_2.svg - :alt: Diagram - :align: center + subgraph "Performance Impact" + Perf["4x fewer instructions
Better memory bandwidth
Reduced latency"] + end + + R1 & R2 & R3 & R4 --> Perf + VR --> Perf Understanding BufferView Indexing ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart LR - subgraph "Input Parameters" - Offset["Offset
(e.g., 5)"] - ValidFlag["Valid Flag
(optional)"] - end - - subgraph "Processing" - BoundsCheck{{"Bounds Check
offset < buffer_size?"}} - FlagCheck{{"Flag Check
valid_flag == True?"}} - Access["Access Memory
buffer[offset]"] - end - - subgraph "Output" - ValidResult["Valid Result
Return value"] - Invalid["Invalid Result
Return 0 or default"] - end - - Offset --> BoundsCheck - ValidFlag --> FlagCheck - - BoundsCheck -->|Yes| FlagCheck - BoundsCheck -->|No| Invalid - - FlagCheck -->|Yes| Access - FlagCheck -->|No| Invalid - - Access --> ValidResult - - style Offset fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - style ValidFlag fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - style ValidResult fill:#d1fae5,stroke:#10b981,stroke-width:2px - style Invalid fill:#fee2e2,stroke:#ef4444,stroke-width:2px - - - - - +.. mermaid:: + + flowchart LR + subgraph "Input Parameters" + Offset["Offset
(e.g., 5)"] + ValidFlag["Valid Flag
(optional)"] + end + + subgraph "Processing" + BoundsCheck{{"Bounds Check
offset < buffer_size?"}} + FlagCheck{{"Flag Check
valid_flag == True?"}} + Access["Access Memory
buffer[offset]"] + end -.. image:: diagrams/buffer_views_3.svg - :alt: Diagram - :align: center + subgraph "Output" + ValidResult["Valid Result
Return value"] + Invalid["Invalid Result
Return 0 or default"] + end + + Offset --> BoundsCheck + ValidFlag --> FlagCheck + + BoundsCheck -->|Yes| FlagCheck + BoundsCheck -->|No| Invalid + + FlagCheck -->|Yes| Access + FlagCheck -->|No| Invalid + + Access --> ValidResult C++ Get Operations ~~~~~~~~~~~~~~~~~~ @@ -381,40 +318,22 @@ Atomic Operations Atomic vs Non-Atomic Operations ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Non-Atomic Operation (Race Condition)" - NA1["Thread 1: Read value (10)"] --> NA2["Thread 1: Add 5 (15)"] - NA3["Thread 2: Read value (10)"] --> NA4["Thread 2: Add 3 (13)"] - NA2 --> NA5["Thread 1: Write 15"] - NA4 --> NA6["Thread 2: Write 13"] - NA5 & NA6 --> NA7["Final value: 13 ❌
(Lost update from Thread 1)"] - end - - subgraph "Atomic Operation (Thread-Safe)" - A1["Thread 1: atomic_add(5)"] --> A2["Hardware ensures
serialization"] - A3["Thread 2: atomic_add(3)"] --> A2 - A2 --> A4["Final value: 18 ✓
(Both updates applied)"] - end - - style NA7 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style A4 fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - - -.. image:: diagrams/buffer_views_4.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Non-Atomic Operation (Race Condition)" + NA1["Thread 1: Read value (10)"] --> NA2["Thread 1: Add 5 (15)"] + NA3["Thread 2: Read value (10)"] --> NA4["Thread 2: Add 3 (13)"] + NA2 --> NA5["Thread 1: Write 15"] + NA4 --> NA6["Thread 2: Write 13"] + NA5 & NA6 --> NA7["Final value: 13 ❌
(Lost update from Thread 1)"] + end + + subgraph "Atomic Operation (Thread-Safe)" + A1["Thread 1: atomic_add(5)"] --> A2["Hardware ensures
serialization"] + A3["Thread 2: atomic_add(3)"] --> A2 + A2 --> A4["Final value: 18 ✓
(Both updates applied)"] + end C++ Atomic Operations ~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/conceptual/ck_tile/convolution_example.rst b/docs/conceptual/ck_tile/convolution_example.rst index a857f9ae9e3..7dbd879bccf 100644 --- a/docs/conceptual/ck_tile/convolution_example.rst +++ b/docs/conceptual/ck_tile/convolution_example.rst @@ -15,53 +15,36 @@ This section covers how CK Tile's :ref:`tensor descriptor ` The key insight is that convolution can be transformed from a complex nested loop operation into a highly parallel matrix multiplication through the image to column (im2col) transformation. CK Tile's tensor descriptors provide the perfect abstraction for implementing this transformation efficiently without data duplication. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Convolution Process" - I["Input Image
6×6"] - K["Kernel
3×3"] - SW["Sliding Window
Extract 3×3 patches"] - DP["Dot Product
Element-wise multiply & sum"] - O["Output
4×4"] - end - - subgraph "Im2col Optimization" - W["Windows Matrix
16×9
(all patches)"] - KF["Kernel Flattened
9×1"] - MM["Matrix Multiply
W @ K"] - OF["Output Flattened
16×1"] - end - - I --> SW - K --> DP - SW --> DP - DP --> O - - SW --> W - K --> KF - W --> MM - KF --> MM - MM --> OF - OF --> O - - style I fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style O fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style MM fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - - - -.. image:: diagrams/convolution_example.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Convolution Process" + I["Input Image
6×6"] + K["Kernel
3×3"] + SW["Sliding Window
Extract 3×3 patches"] + DP["Dot Product
Element-wise multiply & sum"] + O["Output
4×4"] + end + + subgraph "Im2col Optimization" + W["Windows Matrix
16×9
(all patches)"] + KF["Kernel Flattened
9×1"] + MM["Matrix Multiply
W @ K"] + OF["Output Flattened
16×1"] + end + + I --> SW + K --> DP + SW --> DP + DP --> O + + SW --> W + K --> KF + W --> MM + KF --> MM + MM --> OF + OF --> O + Understanding Sliding Windows ============================= diff --git a/docs/conceptual/ck_tile/coordinate_movement.rst b/docs/conceptual/ck_tile/coordinate_movement.rst index 73633afa884..b8c917ec4c7 100644 --- a/docs/conceptual/ck_tile/coordinate_movement.rst +++ b/docs/conceptual/ck_tile/coordinate_movement.rst @@ -17,51 +17,28 @@ The coordinate movement system provides two key abstractions: TensorCoordinate f For the mathematical foundations of coordinate systems, see :ref:`ck_tile_coordinate_systems`. For simpler coordinate concepts, see :ref:`ck_tile_tensor_coordinates`. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Coordinate Movement System" - TC["TensorCoordinate
Position + Descriptor Context"] - TAC["TensorAdaptorCoordinate
Position + Transform Context"] - MC["move_coordinate()
Efficient Navigation"] - end - - subgraph "Movement Example" - S["Start: [1,1]
Offset: 5"] - M1["Move [0,1]
→ [1,2]
Offset: 6"] - M2["Move [1,0]
→ [2,2]
Offset: 10"] - M3["Move [1,1]
→ [3,3]
Offset: 15"] - end - - TC --> MC - TAC --> MC - - S --> M1 - M1 --> M2 - M2 --> M3 - - style TC fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style TAC fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style MC fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - - -.. image:: diagrams/coordinate_movement.svg - :alt: Diagram - :align: center - -.. image:: diagrams/coordinate_movement.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Coordinate Movement System" + TC["TensorCoordinate
Position + Descriptor Context"] + TAC["TensorAdaptorCoordinate
Position + Transform Context"] + MC["move_coordinate()
Efficient Navigation"] + end + + subgraph "Movement Example" + S["Start: [1,1]
Offset: 5"] + M1["Move [0,1]
→ [1,2]
Offset: 6"] + M2["Move [1,0]
→ [2,2]
Offset: 10"] + M3["Move [1,1]
→ [3,3]
Offset: 15"] + end + + TC --> MC + TAC --> MC + + S --> M1 + M1 --> M2 + M2 --> M3 TensorCoordinate: Descriptor-Aware Navigation ============================================= diff --git a/docs/conceptual/ck_tile/coordinate_systems.rst b/docs/conceptual/ck_tile/coordinate_systems.rst index 13a96190108..bca3a605e0e 100644 --- a/docs/conceptual/ck_tile/coordinate_systems.rst +++ b/docs/conceptual/ck_tile/coordinate_systems.rst @@ -15,51 +15,30 @@ The Five Coordinate Spaces The CK framework employs five interconnected coordinate spaces, each serving a specific purpose in the journey from thread identification to memory access. These spaces work together to solve the fundamental challenge of GPU programming: efficiently distributing work across thousands of parallel threads while maintaining optimal memory access patterns. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Coordinate Spaces Overview" - P["P-space
Thread Identification
Which thread am I?"] - Y["Y-space
Logical Tile
Which element in my tile?"] - X["X-space
Physical Tensor
Where in the tensor?"] - R["R-space
Replication
Data sharing pattern"] - D["D-space
Linear Storage
Memory address"] - end - - subgraph "Transformations" - T1["P + Y → X
Thread + Element → Position"] - T2["X → D
Position → Address"] - end - - P --> T1 - Y --> T1 - T1 --> X - X --> T2 - T2 --> D - - R -.-> P - R -.-> Y - - style P fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style Y fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style X fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style R fill:#fce4ec,stroke:#c2185b,stroke-width:2px - style D fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px - - - - - - -.. image:: diagrams/coordinate_systems_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Coordinate Spaces Overview" + P["P-space
Thread Identification
Which thread am I?"] + Y["Y-space
Logical Tile
Which element in my tile?"] + X["X-space
Physical Tensor
Where in the tensor?"] + R["R-space
Replication
Data sharing pattern"] + D["D-space
Linear Storage
Memory address"] + end + + subgraph "Transformations" + T1["P + Y → X
Thread + Element → Position"] + T2["X → D
Position → Address"] + end + + P --> T1 + Y --> T1 + T1 --> X + X --> T2 + T2 --> D + + R -.-> P + R -.-> Y The Challenge and Solution ~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -82,53 +61,35 @@ Partition Space (P-space) represents the foundation of the coordinate system hie GPU Thread Hierarchy ~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "GPU Thread Hierarchy" - subgraph "Block" - subgraph "Warp 0" - T0["Thread 0
P=[0,0]"] - T1["Thread 1
P=[0,1]"] - T2["Thread 2
P=[0,2]"] - T31["..."] - T3["Thread 31
P=[0,31]"] - end - subgraph "Warp 1" - T32["Thread 32
P=[1,0]"] - T33["Thread 33
P=[1,1]"] - T34["..."] - T63["Thread 63
P=[1,31]"] - end - W2["Warp 2..."] - W7["Warp 7"] - end - end - - subgraph "P-space Mapping" - PM["P-coordinates = [warp_id, lane_id]
or
P-coordinates = [block_x, block_y, thread_x, thread_y]"] - end - - T0 --> PM - T32 --> PM - - style T0 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style T32 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - - - - - - -.. image:: diagrams/coordinate_systems_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "GPU Thread Hierarchy" + subgraph "Block" + subgraph "Warp 0" + T0["Thread 0
P=[0,0]"] + T1["Thread 1
P=[0,1]"] + T2["Thread 2
P=[0,2]"] + T31["..."] + T3["Thread 31
P=[0,31]"] + end + subgraph "Warp 1" + T32["Thread 32
P=[1,0]"] + T33["Thread 33
P=[1,1]"] + T34["..."] + T63["Thread 63
P=[1,31]"] + end + W2["Warp 2..."] + W7["Warp 7"] + end + end + + subgraph "P-space Mapping" + PM["P-coordinates = [warp_id, lane_id]
or
P-coordinates = [block_x, block_y, thread_x, thread_y]"] + end + + T0 --> PM + T32 --> PM The structure of P-space directly reflects the :ref:`hardware organization ` of GPUs. Each thread receives a unique P-coordinate that encodes its position within the execution hierarchy. For simple distributions, P-space might be one-dimensional, containing only a thread ID. For complex hierarchical distributions, P-space can have multiple dimensions representing different levels of the GPU's thread organization. @@ -173,56 +134,36 @@ Yield Space (Y-space) represents the logical organization of work within each th Work Assignment Structure ~~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Thread's Tile (2x2 elements)" - Y00["Y=[0,0]
Element 0"] - Y01["Y=[0,1]
Element 1"] - Y10["Y=[1,0]
Element 2"] - Y11["Y=[1,1]
Element 3"] - end - - subgraph "Y-space Structure" - YS["Each thread processes
the same Y-space pattern
but at different X locations"] - end - - subgraph "Example: 4 Threads" - T0["Thread 0
P=[0,0]"] - T1["Thread 1
P=[0,1]"] - T2["Thread 2
P=[1,0]"] - T3["Thread 3
P=[1,1]"] - end - - Y00 --> YS - Y01 --> YS - Y10 --> YS - Y11 --> YS - - T0 --> YS - T1 --> YS - T2 --> YS - T3 --> YS - - style Y00 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style Y01 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style Y10 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style Y11 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - - - - -.. image:: diagrams/coordinate_systems_3.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Thread's Tile (2x2 elements)" + Y00["Y=[0,0]
Element 0"] + Y01["Y=[0,1]
Element 1"] + Y10["Y=[1,0]
Element 2"] + Y11["Y=[1,1]
Element 3"] + end + + subgraph "Y-space Structure" + YS["Each thread processes
the same Y-space pattern
but at different X locations"] + end + + subgraph "Example: 4 Threads" + T0["Thread 0
P=[0,0]"] + T1["Thread 1
P=[0,1]"] + T2["Thread 2
P=[1,0]"] + T3["Thread 3
P=[1,1]"] + end + + Y00 --> YS + Y01 --> YS + Y10 --> YS + Y11 --> YS + + T0 --> YS + T1 --> YS + T2 --> YS + T3 --> YS The power of Y-space lies in its ability to express different iteration patterns without changing the underlying distribution logic. A thread might traverse its Y-space in row-major order for one algorithm, column-major for another, or even use :ref:`space-filling curves ` for optimal cache utilization. This flexibility enables algorithm-specific optimizations while maintaining a consistent framework. @@ -304,49 +245,30 @@ The transformation from P and Y coordinates to X coordinates represents the hear Transformation Pipeline ~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Input" - P["P-coordinates
Thread identity
P=[1,0]"] - Y["Y-coordinates
Element in tile
Y=[0,1]"] - end - - subgraph "Transformation" - T["P + Y → X
Base position + Offset"] - end - - subgraph "Output" - X["X-coordinates
Tensor position
X=[2,1]"] - end - - subgraph "Example" - E["Thread P=[1,0] at base (2,0)
Element Y=[0,1] adds offset (0,1)
Result X=[2,1] in tensor"] - end - - P --> T - Y --> T - T --> X - X --> E - - style P fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style Y fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style X fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - - -.. image:: diagrams/coordinate_systems_4.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Input" + P["P-coordinates
Thread identity
P=[1,0]"] + Y["Y-coordinates
Element in tile
Y=[0,1]"] + end + + subgraph "Transformation" + T["P + Y → X
Base position + Offset"] + end + + subgraph "Output" + X["X-coordinates
Tensor position
X=[2,1]"] + end + + subgraph "Example" + E["Thread P=[1,0] at base (2,0)
Element Y=[0,1] adds offset (0,1)
Result X=[2,1] in tensor"] + end + + P --> T + Y --> T + T --> X + X --> E Mathematical Foundation ~~~~~~~~~~~~~~~~~~~~~~~ @@ -413,45 +335,27 @@ D-space represents the final transformation in the coordinate pipeline: converti Linearization Strategies ~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "X-coordinates" - X["X = [2, 3]
2D Position"] - end - - subgraph "Layout Options" - RM["Row-Major
D = 2×width + 3"] - CM["Column-Major
D = 3×height + 2"] - BL["Blocked
Complex pattern"] - end - - subgraph "D-coordinate" - D["D = 11
Linear Address"] - end - - X --> RM - X --> CM - X --> BL - RM --> D - - style X fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style D fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px - - - - - - -.. image:: diagrams/coordinate_systems_5.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "X-coordinates" + X["X = [2, 3]
2D Position"] + end + + subgraph "Layout Options" + RM["Row-Major
D = 2×width + 3"] + CM["Column-Major
D = 3×height + 2"] + BL["Blocked
Complex pattern"] + end + + subgraph "D-coordinate" + D["D = 11
Linear Address"] + end + + X --> RM + X --> CM + X --> BL + RM --> D The linearization process must consider multiple factors: @@ -482,58 +386,39 @@ Complete Pipeline Example The following is a complete example showing how all coordinate spaces work together: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Step 1: Thread Identification" - TID["Thread ID = 5"] - P["P-coordinates
P = [0, 5]
(warp 0, lane 5)"] - end - - subgraph "Step 2: Work Assignment" - Y["Y-coordinates
Y = [1, 0]
(element in tile)"] - end - - subgraph "Step 3: P+Y Transformation" - TRANS["P + Y → X
Thread position + Element offset"] - X["X-coordinates
X = [1, 5]
(tensor position)"] - end - - subgraph "Step 4: Linearization" - LIN["X → D
Row-major: D = x₀ × width + x₁"] - D["D-coordinate
D = 13
(memory address)"] - end - - subgraph "Step 5: Memory Access" - MEM["Hardware accesses
memory[13]"] - end - - TID --> P - P --> TRANS - Y --> TRANS - TRANS --> X - X --> LIN - LIN --> D - D --> MEM - - style P fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style Y fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style X fill:#e8f5e9,stroke:#388e3c,stroke-width:3px - style D fill:#f3e5f5,stroke:#7b1fa2,stroke-width:3px - style MEM fill:#ffebee,stroke:#c62828,stroke-width:3px - - - - -.. image:: diagrams/coordinate_systems_6.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Step 1: Thread Identification" + TID["Thread ID = 5"] + P["P-coordinates
P = [0, 5]
(warp 0, lane 5)"] + end + + subgraph "Step 2: Work Assignment" + Y["Y-coordinates
Y = [1, 0]
(element in tile)"] + end + + subgraph "Step 3: P+Y Transformation" + TRANS["P + Y → X
Thread position + Element offset"] + X["X-coordinates
X = [1, 5]
(tensor position)"] + end + + subgraph "Step 4: Linearization" + LIN["X → D
Row-major: D = x₀ × width + x₁"] + D["D-coordinate
D = 13
(memory address)"] + end + + subgraph "Step 5: Memory Access" + MEM["Hardware accesses
memory[13]"] + end + + TID --> P + P --> TRANS + Y --> TRANS + TRANS --> X + X --> LIN + LIN --> D + D --> MEM Real-World Example: Matrix Multiplication ----------------------------------------- diff --git a/docs/conceptual/ck_tile/descriptors.rst b/docs/conceptual/ck_tile/descriptors.rst index 449e7bc4b1a..6a028dffdd0 100644 --- a/docs/conceptual/ck_tile/descriptors.rst +++ b/docs/conceptual/ck_tile/descriptors.rst @@ -96,46 +96,28 @@ The Pipeline Concept Every TensorDescriptor in CK Tile can be thought of as a **transformation pipeline**. The functions above create the *first stage* of this pipeline, defining the initial :ref:`transformation ` that takes a simple, one-dimensional block of memory and presents it as a logical, multi-dimensional tensor view. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Pipeline Stages" - S1["Stage 1
Base Layout
[M, N]"] - S2["Stage 2
Transform
Unmerge"] - S3["Stage 3
New View
[M1, M2, N]"] - S4["Stage N
Final View
[...]"] - end - - subgraph "Same Data" - D["Physical Memory
No data movement"] - end - - S1 --> S2 - S2 --> S3 - S3 --> S4 - - S1 -.-> D - S2 -.-> D - S3 -.-> D - S4 -.-> D - - style D fill:#ffebee,stroke:#d32f2f,stroke-width:2px - style S1 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style S3 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - -.. image:: diagrams/descriptors_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Pipeline Stages" + S1["Stage 1
Base Layout
[M, N]"] + S2["Stage 2
Transform
Unmerge"] + S3["Stage 3
New View
[M1, M2, N]"] + S4["Stage N
Final View
[...]"] + end + + subgraph "Same Data" + D["Physical Memory
No data movement"] + end + + S1 --> S2 + S2 --> S3 + S3 --> S4 -.. image:: diagrams/descriptors_1.svg - :alt: Diagram - :align: center + S1 -.-> D + S2 -.-> D + S3 -.-> D + S4 -.-> D The Initial Pipeline Stage ~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -200,52 +182,32 @@ To get from [2, 6] to [2, 2, 3], we need: Analysis of the Final Pipeline ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Transform Pipeline" - T0["Transform 0
Base Unmerge
Input: [0]
Output: [1,2]"] - T1["Transform 1
PassThrough
Input: [1]
Output: [3]"] - T2["Transform 2
Unmerge
Input: [2]
Output: [4,5]"] - end - - subgraph "Hidden Dimensions" - H0["Hidden ID 0
Raw Buffer"] - H1["Hidden ID 1
Dim 0 (size 2)"] - H2["Hidden ID 2
Dim 1 (size 6)"] - H3["Hidden ID 3
Final Dim 0"] - H4["Hidden ID 4
Final Dim 1"] - H5["Hidden ID 5
Final Dim 2"] - end - - H0 --> T0 - T0 --> H1 - T0 --> H2 - H1 --> T1 - H2 --> T2 - T1 --> H3 - T2 --> H4 - T2 --> H5 - - style H0 fill:#ffebee,stroke:#d32f2f,stroke-width:2px - style H3 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style H4 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style H5 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - -.. image:: diagrams/descriptors_2.svg - :alt: Diagram - :align: center - -.. image:: diagrams/descriptors_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Transform Pipeline" + T0["Transform 0
Base Unmerge
Input: [0]
Output: [1,2]"] + T1["Transform 1
PassThrough
Input: [1]
Output: [3]"] + T2["Transform 2
Unmerge
Input: [2]
Output: [4,5]"] + end + + subgraph "Hidden Dimensions" + H0["Hidden ID 0
Raw Buffer"] + H1["Hidden ID 1
Dim 0 (size 2)"] + H2["Hidden ID 2
Dim 1 (size 6)"] + H3["Hidden ID 3
Final Dim 0"] + H4["Hidden ID 4
Final Dim 1"] + H5["Hidden ID 5
Final Dim 2"] + end + + H0 --> T0 + T0 --> H1 + T0 --> H2 + H1 --> T1 + H2 --> T2 + T1 --> H3 + T2 --> H4 + T2 --> H5 The pipeline now has three stages: diff --git a/docs/conceptual/ck_tile/encoding_internals.rst b/docs/conceptual/ck_tile/encoding_internals.rst index 499ec0bd4a0..1b64504e6f9 100644 --- a/docs/conceptual/ck_tile/encoding_internals.rst +++ b/docs/conceptual/ck_tile/encoding_internals.rst @@ -15,53 +15,39 @@ The tile distribution encoding system represents the core mathematical framework At its heart, the encoding system defines how multi-dimensional tensor data is distributed across GPU processing elements through a hierarchical decomposition scheme. By specifying relationships between different coordinate spaces of replication (R), hierarchical (H), partition (P), and yield (Y) dimension, the encoding provides a complete blueprint for data layout and access patterns that can be resolved entirely at compile time. This is the internal mechanism behind :ref:`ck_tile_tile_distribution`. See :ref:`ck_tile_coordinate_systems` for more information about coordinate spaces. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Encoding Components" - RS["R-space Lengths
Replication dimensions"] - HS["H-space Lengths
Hierarchical decomposition
[[2,2],[2,2]]"] - P2RH["P→RH Mappings
Thread to hierarchy
Major/Minor"] - Y2RH["Y→RH Mappings
Element to hierarchy
Major/Minor"] - end - - subgraph "Generated Components" - ADAPTOR["ps_ys_to_xs_adaptor
Coordinate transformer"] - DESC["ys_to_d_descriptor
Memory linearizer"] - ENC["Encoding
Original specification"] - end - - subgraph "Transformation Chain" - T1["Replicate
Transform"] - T2["Unmerge
Transform"] - T3["Merge
Transform"] - end - - RS --> T1 - HS --> T2 - P2RH --> ADAPTOR - Y2RH --> ADAPTOR - - T1 --> T2 - T2 --> T3 - T3 --> ADAPTOR - - HS --> DESC - Y2RH --> DESC - - style RS fill:#fce4ec,stroke:#c2185b,stroke-width:2px - style HS fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style ADAPTOR fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style DESC fill:#fff3e0,stroke:#f57c00,stroke-width:3px - - - -.. image:: diagrams/encoding_internals_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Encoding Components" + RS["R-space Lengths
Replication dimensions"] + HS["H-space Lengths
Hierarchical decomposition
[[2,2],[2,2]]"] + P2RH["P→RH Mappings
Thread to hierarchy
Major/Minor"] + Y2RH["Y→RH Mappings
Element to hierarchy
Major/Minor"] + end + + subgraph "Generated Components" + ADAPTOR["ps_ys_to_xs_adaptor
Coordinate transformer"] + DESC["ys_to_d_descriptor
Memory linearizer"] + ENC["Encoding
Original specification"] + end + + subgraph "Transformation Chain" + T1["Replicate
Transform"] + T2["Unmerge
Transform"] + T3["Merge
Transform"] + end + + RS --> T1 + HS --> T2 + P2RH --> ADAPTOR + Y2RH --> ADAPTOR + + T1 --> T2 + T2 --> T3 + T3 --> ADAPTOR + + HS --> DESC + Y2RH --> DESC Encoding Structure ================== @@ -203,44 +189,31 @@ Transformation Pipeline The encoding generates a transformation pipeline that converts coordinates using the concepts from :ref:`ck_tile_transforms` and :ref:`ck_tile_adaptors`: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart LR - subgraph "Input Coordinates" - P["P-coordinates
[warp_id, lane_id]"] - Y["Y-coordinates
[y0, y1, y2, y3]"] - end - - subgraph "Transformation Pipeline" - C1["Combine P+Y"] - T1["Replicate
Transform
(if R-dims exist)"] - T2["Unmerge
Transform
(break into H-dims)"] - T3["Merge
Transform
(combine to X-dims)"] - end - - subgraph "Output" - X["X-coordinates
[x0, x1]
Tensor position"] - end - - P --> C1 - Y --> C1 - C1 --> T1 - T1 --> T2 - T2 --> T3 - T3 --> X - - style P fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style Y fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style X fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - -.. image:: diagrams/encoding_internals_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + flowchart LR + subgraph "Input Coordinates" + P["P-coordinates
[warp_id, lane_id]"] + Y["Y-coordinates
[y0, y1, y2, y3]"] + end + + subgraph "Transformation Pipeline" + C1["Combine P+Y"] + T1["Replicate
Transform
(if R-dims exist)"] + T2["Unmerge
Transform
(break into H-dims)"] + T3["Merge
Transform
(combine to X-dims)"] + end + + subgraph "Output" + X["X-coordinates
[x0, x1]
Tensor position"] + end + + P --> C1 + Y --> C1 + C1 --> T1 + T1 --> T2 + T2 --> T3 + T3 --> X Building the Transformation Chain --------------------------------- diff --git a/docs/conceptual/ck_tile/introduction_motivation.rst b/docs/conceptual/ck_tile/introduction_motivation.rst index a939aef4c85..33550e49d09 100644 --- a/docs/conceptual/ck_tile/introduction_motivation.rst +++ b/docs/conceptual/ck_tile/introduction_motivation.rst @@ -17,81 +17,62 @@ In this introduction, we establish the fundamental problems that tile distributi The GPU Memory Problem ---------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Random Access Pattern (Inefficient)" - subgraph "Threads" - T0_R["Thread 0"] - T1_R["Thread 1"] - T2_R["Thread 2"] - T3_R["Thread 3"] - end - - subgraph "Memory" - M0["Mem[0]"] - M7["Mem[7]"] - M15["Mem[15]"] - M23["Mem[23]"] - M31["Mem[31]"] - M39["Mem[39]"] - M47["Mem[47]"] - M55["Mem[55]"] - end - - T0_R -.-> M23 - T1_R -.-> M7 - T2_R -.-> M47 - T3_R -.-> M15 - end - - subgraph "Tile Distribution Pattern (Efficient)" - subgraph "Threads_TD" - T0_TD["Thread 0"] - T1_TD["Thread 1"] - T2_TD["Thread 2"] - T3_TD["Thread 3"] - end - - subgraph "Memory_TD" - M0_TD["Mem[0]"] - M1_TD["Mem[1]"] - M2_TD["Mem[2]"] - M3_TD["Mem[3]"] - M4_TD["Mem[4]"] - M5_TD["Mem[5]"] - M6_TD["Mem[6]"] - M7_TD["Mem[7]"] - end - - T0_TD --> M0_TD - T0_TD --> M1_TD - T1_TD --> M2_TD - T1_TD --> M3_TD - T2_TD --> M4_TD - T2_TD --> M5_TD - T3_TD --> M6_TD - T3_TD --> M7_TD - end - - style T0_R fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style T1_R fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style T2_R fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style T3_R fill:#fee2e2,stroke:#ef4444,stroke-width:2px - - style T0_TD fill:#d1fae5,stroke:#10b981,stroke-width:2px - style T1_TD fill:#d1fae5,stroke:#10b981,stroke-width:2px - style T2_TD fill:#d1fae5,stroke:#10b981,stroke-width:2px - style T3_TD fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - -.. image:: diagrams/introduction_motivation_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Random Access Pattern (Inefficient)" + subgraph "Threads" + T0_R["Thread 0"] + T1_R["Thread 1"] + T2_R["Thread 2"] + T3_R["Thread 3"] + end + + subgraph "Memory" + M0["Mem[0]"] + M7["Mem[7]"] + M15["Mem[15]"] + M23["Mem[23]"] + M31["Mem[31]"] + M39["Mem[39]"] + M47["Mem[47]"] + M55["Mem[55]"] + end + + T0_R -.-> M23 + T1_R -.-> M7 + T2_R -.-> M47 + T3_R -.-> M15 + end + + subgraph "Tile Distribution Pattern (Efficient)" + subgraph "Threads_TD" + T0_TD["Thread 0"] + T1_TD["Thread 1"] + T2_TD["Thread 2"] + T3_TD["Thread 3"] + end + + subgraph "Memory_TD" + M0_TD["Mem[0]"] + M1_TD["Mem[1]"] + M2_TD["Mem[2]"] + M3_TD["Mem[3]"] + M4_TD["Mem[4]"] + M5_TD["Mem[5]"] + M6_TD["Mem[6]"] + M7_TD["Mem[7]"] + end + + T0_TD --> M0_TD + T0_TD --> M1_TD + T1_TD --> M2_TD + T1_TD --> M3_TD + T2_TD --> M4_TD + T2_TD --> M5_TD + T3_TD --> M6_TD + T3_TD --> M7_TD + end Why Random Memory Access is Slow ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -216,42 +197,26 @@ The Coordinate Mapping Insight At the heart of tile distribution lies a profound mathematical insight: efficient GPU computation requires a systematic framework for mapping between different coordinate spaces. This framework transforms the complex problem of thread-to-data assignment into a series of well-defined mathematical transformations, each serving a specific purpose in the journey from abstract algorithm to concrete hardware execution. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Coordinate Spaces" - P["P-space
Thread Position
(thread_x, thread_y,
warp_id, block_id)"] - Y["Y-space
Local Data
(y0, y1, y2, y3)"] - X["X-space
Global Position
(x0, x1)"] - D["D-space
Memory Address
(linearized)"] - end - - subgraph "Transformations" - T1["P + Y → X
Thread data mapping"] - T2["X → D
Memory linearization"] - end - - P --> T1 - Y --> T1 - T1 --> X - X --> T2 - T2 --> D - - style P fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style Y fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style X fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style D fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px - style T1 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - style T2 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - - - -.. image:: diagrams/introduction_motivation_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Coordinate Spaces" + P["P-space
Thread Position
(thread_x, thread_y,
warp_id, block_id)"] + Y["Y-space
Local Data
(y0, y1, y2, y3)"] + X["X-space
Global Position
(x0, x1)"] + D["D-space
Memory Address
(linearized)"] + end + + subgraph "Transformations" + T1["P + Y → X
Thread data mapping"] + T2["X → D
Memory linearization"] + end + + P --> T1 + Y --> T1 + T1 --> X + X --> T2 + T2 --> D The elegance of this approach emerges from its separation of concerns. Each coordinate space represents a distinct aspect of the computation, and the transformations between them encapsulate specific optimization strategies. This separation allows developers to reason about their algorithms in natural terms while the framework handles the complex mapping to efficient hardware execution patterns. diff --git a/docs/conceptual/ck_tile/lds_index_swapping.rst b/docs/conceptual/ck_tile/lds_index_swapping.rst index b0a2b320100..f0d47494a47 100644 --- a/docs/conceptual/ck_tile/lds_index_swapping.rst +++ b/docs/conceptual/ck_tile/lds_index_swapping.rst @@ -25,48 +25,31 @@ Step 1: XOR Transform The original K coordinate is split into K0 and K1, where K1 represents the thread vector size along the K dimension (KPack) and K0 is KPerBlock/KPack. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "3D LDS coordinate [K0, M, K1]" - K0["KPerBlock/KPack * MLdsLayer
K0"] - M["MPerBlock/MLdsLayer
M"] - K1["KPack
K1"] - end - - subgraph "XOR Transform" - XT["make_xor_transform"] - end - - subgraph "Update K0 with XOR transformation" - K01["KPerBlock/KPack * MLdsLayer
K0'"] - M1["MPerBlock/MLdsLayer
M"] - K11["KPack
K1"] - end - - K0 --> XT - M --> XT - K1 --> K11 - - XT --> K01 - XT --> M1 - - style K0 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style K01 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style M fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style M1 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - - style K1 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style K11 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - -.. image:: diagrams/lds_index_swapping_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "3D LDS coordinate [K0, M, K1]" + K0["KPerBlock/KPack * MLdsLayer
K0"] + M["MPerBlock/MLdsLayer
M"] + K1["KPack
K1"] + end + + subgraph "XOR Transform" + XT["make_xor_transform"] + end + + subgraph "Update K0 with XOR transformation" + K01["KPerBlock/KPack * MLdsLayer
K0'"] + M1["MPerBlock/MLdsLayer
M"] + K11["KPack
K1"] + end + + K0 --> XT + M --> XT + K1 --> K11 + + XT --> K01 + XT --> M1 The XOR transformation updates the K0 coordinate using the formula: @@ -81,54 +64,33 @@ Step 2: Unmerge Transform The transformed K0' is split into L and K0'' components, creating an intermediate 4D coordinate space. This is necessary when MLdsLayer > 1, allowing multiple rows to share the same set of memory banks for better utilization with smaller tile sizes. - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "3D LDS coordinate [K0', M, K1]" - K0["KPerBlock/KPack * MLdsLayer
K0'"] - M["MPerBlock/MLdsLayer
M"] - K1["KPack
K1"] - end - - subgraph "Unmerge into 2 components" - UM["make_unmerge_transform"] - end - - subgraph "4D intermediate transformation space" - L["MLdsLayer
L"] - M1["MPerBlock/MLdsLayer
M"] - K01["KPerBlock/KPack
K0''"] - K11["KPack
K1"] - end - - K0 --> UM - M --> M1 - K1 --> K11 - - UM --> L - UM --> K01 - - style K0 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style L fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style K01 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - - style M fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style M1 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - style K1 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style K11 fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - - - -.. image:: diagrams/lds_index_swapping_2.svg - :alt: Diagram - :align: center + +.. mermaid:: + + graph TB + subgraph "3D LDS coordinate [K0', M, K1]" + K0["KPerBlock/KPack * MLdsLayer
K0'"] + M["MPerBlock/MLdsLayer
M"] + K1["KPack
K1"] + end + + subgraph "Unmerge into 2 components" + UM["make_unmerge_transform"] + end + + subgraph "4D intermediate transformation space" + L["MLdsLayer
L"] + M1["MPerBlock/MLdsLayer
M"] + K01["KPerBlock/KPack
K0''"] + K11["KPack
K1"] + end + + K0 --> UM + M --> M1 + K1 --> K11 + + UM --> L + UM --> K01 The unmerge operation: @@ -144,56 +106,38 @@ Step 3: Merge Transform The final step merges the 4D coordinates back into 2D transformed coordinates (M', K'). - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "4D LDS Coordinates [L, M, K0'', K1]" - L["MLdsLayer
L"] - M1["MPerBlock/MLdsLayer
M"] - K0["KPerBlock/KPack
K0''"] - K1["KPack
K1"] - end - - subgraph "Merge into 1 component" - ME0["make_merge_transform"] - end - - subgraph "Merge into 1 component" - ME1["make_merge_transform"] - end - - subgraph "Transformed 2D coordinates [M', K']" - M11["MPerBlock
M'"] - K01["KPerBlock
K'"] - end - - L --> ME0 - M1 --> ME0 - - K0 --> ME1 - K1 --> ME1 - - ME0 --> M11 - ME1 --> K01 - - style K0 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style K1 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style K01 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - - style M1 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style L fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style M11 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - -.. image:: diagrams/lds_index_swapping_3.svg - :alt: Diagram - :align: center + +.. mermaid:: + + graph TB + subgraph "4D LDS Coordinates [L, M, K0'', K1]" + L["MLdsLayer
L"] + M1["MPerBlock/MLdsLayer
M"] + K0["KPerBlock/KPack
K0''"] + K1["KPack
K1"] + end + + subgraph "Merge into 1 component" + ME0["make_merge_transform"] + end + + subgraph "Merge into 1 component" + ME1["make_merge_transform"] + end + + subgraph "Transformed 2D coordinates [M', K']" + M11["MPerBlock
M'"] + K01["KPerBlock
K'"] + end + + L --> ME0 + M1 --> ME0 + + K0 --> ME1 + K1 --> ME1 + + ME0 --> M11 + ME1 --> K01 C++ Implementation diff --git a/docs/conceptual/ck_tile/load_store_traits.rst b/docs/conceptual/ck_tile/load_store_traits.rst index bf2decc37e7..7be70c770f3 100644 --- a/docs/conceptual/ck_tile/load_store_traits.rst +++ b/docs/conceptual/ck_tile/load_store_traits.rst @@ -102,34 +102,20 @@ Vectorization Selection Algorithm LoadStoreTraits employs an advanced algorithm to select the best dimension for vectorization: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TD - A[Analyze Distribution] --> B{Check Each Dimension} - B --> C[Calculate Stride] - C --> D{Stride == 1?} - D -->|Yes| E[Candidate for Vectorization] - D -->|No| F[Skip Dimension] - E --> G[Check Alignment] - G --> H[Check Vector Size] - H --> I[Score Dimension] - F --> B - I --> J[Select Best Dimension] - J --> K[Configure Vector Access] - - style A fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style J fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style K fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - - -.. image:: diagrams/load_store_traits_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TD + A[Analyze Distribution] --> B{Check Each Dimension} + B --> C[Calculate Stride] + C --> D{Stride == 1?} + D -->|Yes| E[Candidate for Vectorization] + D -->|No| F[Skip Dimension] + E --> G[Check Alignment] + G --> H[Check Vector Size] + H --> I[Score Dimension] + F --> B + I --> J[Select Best Dimension] + J --> K[Configure Vector Access] **Example: Comparing Different Memory Layouts** @@ -172,42 +158,30 @@ Memory Access Patterns LoadStoreTraits creates efficient access patterns using space-filling curves: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Linear Traversal" - L1["0→1→2→3"] - L2["4→5→6→7"] - L3["Cache miss"] - L4["8→9→10→11"] - end - - subgraph "Snake Pattern" - S1["0→1→2→3"] - S2["7←6←5←4"] - S3["Cache hit!"] - S4["8→9→10→11"] - end - - L1 --> L2 - L2 --> L3 - L3 --> L4 - - S1 --> S2 - S2 --> S3 - S3 --> S4 - - style L3 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style S3 fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - -.. image:: diagrams/load_store_traits_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Linear Traversal" + L1["0→1→2→3"] + L2["4→5→6→7"] + L3["Cache miss"] + L4["8→9→10→11"] + end + + subgraph "Snake Pattern" + S1["0→1→2→3"] + S2["7←6←5←4"] + S3["Cache hit!"] + S4["8→9→10→11"] + end + + L1 --> L2 + L2 --> L3 + L3 --> L4 + + S1 --> S2 + S2 --> S3 + S3 --> S4 **C++ Access Pattern Example:** diff --git a/docs/conceptual/ck_tile/space_filling_curve.rst b/docs/conceptual/ck_tile/space_filling_curve.rst index 869285b4621..43b4ce0b78c 100644 --- a/docs/conceptual/ck_tile/space_filling_curve.rst +++ b/docs/conceptual/ck_tile/space_filling_curve.rst @@ -190,44 +190,30 @@ Snake Pattern for Cache Optimization The snake pattern reverses traversal direction on alternate rows, minimizing the distance between consecutive accesses: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Linear Pattern" - L1["Row 0: →"] - L2["Row 1: →"] - L3["Jump back"] - L4["Row 2: →"] - end - - subgraph "Snake Pattern" - S1["Row 0: →"] - S2["Row 1: ←"] - S3["Continue"] - S4["Row 2: →"] - end - - L1 --> L3 - L3 --> L2 - L2 --> L3 - L3 --> L4 - - S1 --> S2 - S2 --> S4 - - style L3 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style S3 fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - -.. image:: diagrams/space_filling_curve.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Linear Pattern" + L1["Row 0: →"] + L2["Row 1: →"] + L3["Jump back"] + L4["Row 2: →"] + end + + subgraph "Snake Pattern" + S1["Row 0: →"] + S2["Row 1: ←"] + S3["Continue"] + S4["Row 2: →"] + end + + L1 --> L3 + L3 --> L2 + L2 --> L3 + L3 --> L4 + + S1 --> S2 + S2 --> S4 .. code-block:: cpp diff --git a/docs/conceptual/ck_tile/static_distributed_tensor.rst b/docs/conceptual/ck_tile/static_distributed_tensor.rst index 1f7a93657f0..075a2150b28 100644 --- a/docs/conceptual/ck_tile/static_distributed_tensor.rst +++ b/docs/conceptual/ck_tile/static_distributed_tensor.rst @@ -89,29 +89,18 @@ Understanding how static distributed tensors organize memory is important for pe The memory layout follows a hierarchical pattern: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TD - A[Global Tensor 64x64] --> B[Thread Block 16x16] - B --> C[Thread 0,0
Elements 0:3,0:3] - B --> D[Thread 0,1
Elements 0:3,4:7] - B --> E[Thread 1,0
Elements 4:7,0:3] - B --> F[...] - - C --> G[Local Array
16 elements] - D --> H[Local Array
16 elements] - E --> I[Local Array
16 elements] - - - - - -.. image:: diagrams/static_distributed_tensor.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TD + A[Global Tensor 64x64] --> B[Thread Block 16x16] + B --> C[Thread 0,0
Elements 0:3,0:3] + B --> D[Thread 0,1
Elements 0:3,4:7] + B --> E[Thread 1,0
Elements 4:7,0:3] + B --> F[...] + + C --> G[Local Array
16 elements] + D --> H[Local Array
16 elements] + E --> I[Local Array
16 elements] Element Access and Indexing =========================== diff --git a/docs/conceptual/ck_tile/sweep_tile.rst b/docs/conceptual/ck_tile/sweep_tile.rst index 4dfb6a2ad10..c8aace2de82 100644 --- a/docs/conceptual/ck_tile/sweep_tile.rst +++ b/docs/conceptual/ck_tile/sweep_tile.rst @@ -15,46 +15,29 @@ Sweep operations are the clean way to iterate over distributed data in CK Tile. Sweep operations use the "load once, use many times" pattern. Load X data once into registers, then sweep through Y positions while keeping X in fast memory. This maximizes data reuse and minimizes memory bandwidth requirements. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart LR - subgraph "X-Tile (Reused)" - XT["X data loaded once
Stays in registers"] - end - - subgraph "Y-Sweep" - Y1["Y position 0"] - Y2["Y position 1"] - Y3["Y position 2"] - YN["Y position N"] - end - - subgraph "Computation" - C["Process(X, Y)"] - end - - XT --> C - Y1 --> C - Y2 --> C - Y3 --> C - YN --> C - - style XT fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style C fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - - - - - -.. image:: diagrams/sweep_tile_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + flowchart LR + subgraph "X-Tile (Reused)" + XT["X data loaded once
Stays in registers"] + end + + subgraph "Y-Sweep" + Y1["Y position 0"] + Y2["Y position 1"] + Y3["Y position 2"] + YN["Y position N"] + end + + subgraph "Computation" + C["Process(X, Y)"] + end + + XT --> C + Y1 --> C + Y2 --> C + Y3 --> C + YN --> C The Complete GPU Workflow ========================= @@ -123,38 +106,24 @@ Memory Efficiency Pattern The sweep pattern provides significant memory efficiency benefits. This is particularly important for GPU architectures (see :ref:`ck_tile_gpu_basics`) where memory bandwidth is often the limiting factor: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Traditional Approach" - T1["Load X[0]"] --> P1["Process"] - T2["Load Y[0]"] --> P1 - T3["Load X[0]"] --> P2["Process"] - T4["Load Y[1]"] --> P2 - T5["Load X[0]"] --> P3["Process"] - T6["Load Y[2]"] --> P3 - Note1["X loaded 3 times!"] - end - - subgraph "Sweep Approach" - S1["Load X[0]"] --> SP["Process with
Y[0], Y[1], Y[2]"] - S2["Load Y[0,1,2]"] --> SP - Note2["X loaded once!"] - end - - style Note1 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style Note2 fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - -.. image:: diagrams/sweep_tile_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Traditional Approach" + T1["Load X[0]"] --> P1["Process"] + T2["Load Y[0]"] --> P1 + T3["Load X[0]"] --> P2["Process"] + T4["Load Y[1]"] --> P2 + T5["Load X[0]"] --> P3["Process"] + T6["Load Y[2]"] --> P3 + Note1["X loaded 3 times!"] + end + + subgraph "Sweep Approach" + S1["Load X[0]"] --> SP["Process with
Y[0], Y[1], Y[2]"] + S2["Load Y[0,1,2]"] --> SP + Note2["X loaded once!"] + end Practical Sweep Patterns ======================== @@ -382,45 +351,32 @@ Performance Characteristics Sweep operations provide several performance benefits: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Sweep Performance Benefits" - B1["Zero runtime overhead
Compile-time unrolling"] - B2["Perfect memory coalescing
Sequential access patterns"] - B3["Automatic vectorization
Compiler optimizations"] - B4["Register reuse
X data stays in VGPR"] - end - - subgraph "Use Cases" - U1["Matrix Multiplication
Reuse A columns"] - U2["Convolution
Reuse filter weights"] - U3["Reduction
Accumulate over Y"] - U4["Broadcast
Apply X to all Y"] - end - - B1 --> Performance["High Performance"] - B2 --> Performance - B3 --> Performance - B4 --> Performance - - Performance --> U1 - Performance --> U2 - Performance --> U3 - Performance --> U4 - - style Performance fill:#d1fae5,stroke:#10b981,stroke-width:3px - - - - - -.. image:: diagrams/sweep_tile_3.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Sweep Performance Benefits" + B1["Zero runtime overhead
Compile-time unrolling"] + B2["Perfect memory coalescing
Sequential access patterns"] + B3["Automatic vectorization
Compiler optimizations"] + B4["Register reuse
X data stays in VGPR"] + end + + subgraph "Use Cases" + U1["Matrix Multiplication
Reuse A columns"] + U2["Convolution
Reuse filter weights"] + U3["Reduction
Accumulate over Y"] + U4["Broadcast
Apply X to all Y"] + end + + B1 --> Performance["High Performance"] + B2 --> Performance + B3 --> Performance + B4 --> Performance + + Performance --> U1 + Performance --> U2 + Performance --> U3 + Performance --> U4 Compiler Optimizations ---------------------- @@ -457,39 +413,21 @@ Integration with CK Tile Components Complete workflow example: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart TB - subgraph "Complete Workflow" - TD["TileDistribution
Define data layout"] - TW["TileWindow
Create view"] - DT["DistributedTensor
Load X data"] - ST["SweepTile
Iterate Y positions"] - R["Results
Store outputs"] - end - - TD --> TW - TW --> DT - DT --> ST - ST --> R - - style TD fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style ST fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style R fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - -.. image:: diagrams/sweep_tile_4.svg - :alt: Diagram - :align: center +.. mermaid:: + + flowchart TB + subgraph "Complete Workflow" + TD["TileDistribution
Define data layout"] + TW["TileWindow
Create view"] + DT["DistributedTensor
Load X data"] + ST["SweepTile
Iterate Y positions"] + R["Results
Store outputs"] + end + + TD --> TW + TW --> DT + DT --> ST + ST --> R .. code-block:: cpp diff --git a/docs/conceptual/ck_tile/tensor_coordinates.rst b/docs/conceptual/ck_tile/tensor_coordinates.rst index 4e9240b83c4..ef047776dbf 100644 --- a/docs/conceptual/ck_tile/tensor_coordinates.rst +++ b/docs/conceptual/ck_tile/tensor_coordinates.rst @@ -15,49 +15,31 @@ Before diving into transforms and adaptors (see :ref:`ck_tile_transforms` and :r MultiIndex serves as the common currency between different coordinate spaces (see :ref:`ck_tile_coordinate_systems`), enabling seamless transformation and navigation through complex tensor layouts. Every transform, adaptor, and descriptor in CK Tile operates on these coordinate containers. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "MultiIndex Structure" - MI["MultiIndex
Container for N integers"] - D0["Dimension 0"] - D1["Dimension 1"] - D2["Dimension 2"] - DN["Dimension N-1"] - end - - subgraph "Usage Context" - T["Transforms
"] - A["Adaptors
"] - TV["Tensors
"] - end - - MI --> D0 - MI --> D1 - MI --> D2 - MI --> DN - - T --> MI - A --> MI - TV --> MI - - style MI fill:#f3e5f5,stroke:#7b1fa2,stroke-width:3px - style D0 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style D1 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style D2 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style DN fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style T fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style A fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style TV fill:#ffebee,stroke:#d32f2f,stroke-width:2px - - - -.. image:: diagrams/tensor_coordinates_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "MultiIndex Structure" + MI["MultiIndex
Container for N integers"] + D0["Dimension 0"] + D1["Dimension 1"] + D2["Dimension 2"] + DN["Dimension N-1"] + end + + subgraph "Usage Context" + T["Transforms
"] + A["Adaptors
"] + TV["Tensors
"] + end + + MI --> D0 + MI --> D1 + MI --> D2 + MI --> DN + + T --> MI + A --> MI + TV --> MI MultiIndex Implementation ========================= @@ -176,36 +158,22 @@ MultiIndex in Coordinate Flow MultiIndex serves as the interface between user code and the transformation pipeline: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart TB - subgraph CF ["Coordinate Flow"] - direction LR - UI["User Input
[1, 2, 3]"] --> MI["MultiIndex
Storage"] - MI --> TR["Transform
Processing"] - TR --> MO["MultiIndex
Output"] - MO --> TA["Tensor Access
element(coord)"] - end - - subgraph EX ["Example: 3D Tensor Access"] - direction LR - T3D["3D Tensor
shape=[4,5,6]"] --> COORD["MultiIndex(3, [1,2,3])"] - COORD --> ELEM["Element at
position [1,2,3]"] - end - - style UI fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - style MI fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px - style MO fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px - style COORD fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - -.. image:: diagrams/tensor_coordinates_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + flowchart TB + subgraph CF ["Coordinate Flow"] + direction LR + UI["User Input
[1, 2, 3]"] --> MI["MultiIndex
Storage"] + MI --> TR["Transform
Processing"] + TR --> MO["MultiIndex
Output"] + MO --> TA["Tensor Access
element(coord)"] + end + + subgraph EX ["Example: 3D Tensor Access"] + direction LR + T3D["3D Tensor
shape=[4,5,6]"] --> COORD["MultiIndex(3, [1,2,3])"] + COORD --> ELEM["Element at
position [1,2,3]"] + end Common Usage Patterns ===================== diff --git a/docs/conceptual/ck_tile/tensor_views.rst b/docs/conceptual/ck_tile/tensor_views.rst index 0c46e1e5930..22e6a61c3d3 100644 --- a/docs/conceptual/ck_tile/tensor_views.rst +++ b/docs/conceptual/ck_tile/tensor_views.rst @@ -13,52 +13,31 @@ TensorView presents different logical views of the same underlying memory withou TensorView Architecture ----------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Memory Foundation" - Memory["Flat Memory Array
0 1 2 3 4 5 6 7 8 9 10 11"] - end - - subgraph "Access Layer" - BufferView["BufferView
Linear Memory Access"] - Descriptor["TensorDescriptor
Shape & Stride Info"] - end - - subgraph "Tensor Layer" - TensorView["TensorView
Multi-dimensional Access"] - end - - subgraph "Logical View" - Matrix["2D Matrix View
[3×4]
[[0,1,2,3]
[4,5,6,7]
[8,9,10,11]]"] - end - - Memory --> BufferView - Memory --> Descriptor - BufferView --> TensorView - Descriptor --> TensorView - TensorView --> Matrix - - style Memory fill:#d1fae5,stroke:#10b981,stroke-width:2px - style BufferView fill:#dbeafe,stroke:#3b82f6,stroke-width:2px - style Descriptor fill:#fed7aa,stroke:#f59e0b,stroke-width:2px - style TensorView fill:#fce7f3,stroke:#ec4899,stroke-width:2px - style Matrix fill:#e9d5ff,stroke:#9333ea,stroke-width:2px - - - - - +.. mermaid:: + + graph TB + subgraph "Memory Foundation" + Memory["Flat Memory Array
0 1 2 3 4 5 6 7 8 9 10 11"] + end + + subgraph "Access Layer" + BufferView["BufferView
Linear Memory Access"] + Descriptor["TensorDescriptor
Shape & Stride Info"] + end + + subgraph "Tensor Layer" + TensorView["TensorView
Multi-dimensional Access"] + end + + subgraph "Logical View" + Matrix["2D Matrix View
[3×4]
[[0,1,2,3]
[4,5,6,7]
[8,9,10,11]]"] + end -.. image:: diagrams/tensor_views_1.svg - :alt: Diagram - :align: center + Memory --> BufferView + Memory --> Descriptor + BufferView --> TensorView + Descriptor --> TensorView + TensorView --> Matrix The Foundation: BufferView and TensorDescriptor ------------------------------------------------ @@ -122,93 +101,53 @@ Coordinate-Based Access The fundamental operation of TensorView is translating multi-dimensional coordinates into memory accesses. This translation happens through an advanced pipeline that maintains efficiency while providing flexibility: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart LR - subgraph "User Input" - Coord["Coordinate
(1, 2)"] - end - - subgraph "TensorView Processing" - Shape["Shape Check
row < 3?
col < 4?"] - Stride["Apply Strides
offset = 1×4 + 2×1"] - Buffer["BufferView Access
buffer[6]"] - end - - subgraph "Result" - Value["Value: 6"] - end - - Coord --> Shape - Shape -->|Valid| Stride - Stride --> Buffer - Buffer --> Value - - style Coord fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - style Shape fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - style Stride fill:#dcfce7,stroke:#10b981,stroke-width:2px - style Buffer fill:#dbeafe,stroke:#3b82f6,stroke-width:2px - style Value fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - +.. mermaid:: + + flowchart LR + subgraph "User Input" + Coord["Coordinate
(1, 2)"] + end + + subgraph "TensorView Processing" + Shape["Shape Check
row < 3?
col < 4?"] + Stride["Apply Strides
offset = 1×4 + 2×1"] + Buffer["BufferView Access
buffer[6]"] + end -.. image:: diagrams/tensor_views_2.svg - :alt: Diagram - :align: center + subgraph "Result" + Value["Value: 6"] + end + + Coord --> Shape + Shape -->|Valid| Stride + Stride --> Buffer + Buffer --> Value Memory Layouts and Strides -------------------------- A key feature of TensorView is its ability to represent different memory layouts through stride manipulation. This capability enables zero-copy transformations that would otherwise require expensive memory operations: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Row-Major Layout (C-style)" - RM["Memory: [0,1,2,3,4,5,6,7,8,9,10,11]
Shape: (3,4)
Strides: (4,1)"] - RMMatrix["[[0, 1, 2, 3]
[4, 5, 6, 7]
[8, 9, 10, 11]]"] - RM --> RMMatrix - end - - subgraph "Column-Major Layout (Fortran-style)" - CM["Memory: [0,3,6,9,1,4,7,10,2,5,8,11]
Shape: (3,4)
Strides: (1,3)"] - CMMatrix["[[0, 1, 2, 3]
[4, 5, 6, 7]
[8, 9, 10, 11]]"] - CM --> CMMatrix - end - - subgraph "Custom Stride (Transposed View)" - TV["Memory: [0,1,2,3,4,5,6,7,8,9,10,11]
Shape: (4,3)
Strides: (1,4)"] - TVMatrix["[[0, 4, 8]
[1, 5, 9]
[2, 6, 10]
[3, 7, 11]]"] - TV --> TVMatrix - end - - style RM fill:#e0f2fe,stroke:#0284c7,stroke-width:2px - style CM fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - style TV fill:#f3e8ff,stroke:#9333ea,stroke-width:2px - - - - - +.. mermaid:: + + graph TB + subgraph "Row-Major Layout (C-style)" + RM["Memory: [0,1,2,3,4,5,6,7,8,9,10,11]
Shape: (3,4)
Strides: (4,1)"] + RMMatrix["[[0, 1, 2, 3]
[4, 5, 6, 7]
[8, 9, 10, 11]]"] + RM --> RMMatrix + end -.. image:: diagrams/tensor_views_3.svg - :alt: Diagram - :align: center + subgraph "Column-Major Layout (Fortran-style)" + CM["Memory: [0,3,6,9,1,4,7,10,2,5,8,11]
Shape: (3,4)
Strides: (1,3)"] + CMMatrix["[[0, 1, 2, 3]
[4, 5, 6, 7]
[8, 9, 10, 11]]"] + CM --> CMMatrix + end + + subgraph "Custom Stride (Transposed View)" + TV["Memory: [0,1,2,3,4,5,6,7,8,9,10,11]
Shape: (4,3)
Strides: (1,4)"] + TVMatrix["[[0, 4, 8]
[1, 5, 9]
[2, 6, 10]
[3, 7, 11]]"] + TV --> TVMatrix + end Row-Major vs Column-Major Layouts ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -326,45 +265,26 @@ Memory Access Patterns The efficiency of TensorView operations depends on memory access patterns. Understanding these patterns is important for achieving optimal performance. See :ref:`ck_tile_gpu_basics` for hardware considerations. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Memory Access Patterns" - Seq["Sequential Access
(Good cache usage)"] - Stride["Strided Access
(May cause cache misses)"] - Random["Random Access
(Poor cache usage)"] - end - - subgraph "Optimization Strategies" - Opt1["Use row-major for row iteration"] - Opt2["Use col-major for column iteration"] - Opt3["Minimize stride between accesses"] - Opt4["Vectorize when possible"] - end - - Seq --> Opt1 - Stride --> Opt2 - Stride --> Opt3 - Random --> Opt4 - - style Seq fill:#d1fae5,stroke:#10b981,stroke-width:2px - style Stride fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - style Random fill:#fee2e2,stroke:#ef4444,stroke-width:2px - - - - - +.. mermaid:: + + graph LR + subgraph "Memory Access Patterns" + Seq["Sequential Access
(Good cache usage)"] + Stride["Strided Access
(May cause cache misses)"] + Random["Random Access
(Poor cache usage)"] + end + + subgraph "Optimization Strategies" + Opt1["Use row-major for row iteration"] + Opt2["Use col-major for column iteration"] + Opt3["Minimize stride between accesses"] + Opt4["Vectorize when possible"] + end -.. image:: diagrams/tensor_views_4.svg - :alt: Diagram - :align: center + Seq --> Opt1 + Stride --> Opt2 + Stride --> Opt3 + Random --> Opt4 Compile-Time Optimization ~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -391,48 +311,30 @@ TensorView vs BufferView Understanding when to use TensorView versus BufferView is crucial for writing efficient code: -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "BufferView" - BV1["Linear indexing only"] - BV2["buffer[5]"] - BV3["No shape information"] - BV4["Direct memory access"] - end - - subgraph "TensorView" - TV1["Multi-dimensional indexing"] - TV2["tensor(1, 2)"] - TV3["Shape-aware operations"] - TV4["Coordinate transformations"] - end - - subgraph "Use Cases" - UC1["BufferView: Low-level memory ops"] - UC2["TensorView: Matrix/tensor algorithms"] - end - - BV1 --> UC1 - TV1 --> UC2 - - style BV1 fill:#dbeafe,stroke:#3b82f6,stroke-width:2px - style TV1 fill:#fce7f3,stroke:#ec4899,stroke-width:2px - - - - - - -.. image:: diagrams/tensor_views_5.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "BufferView" + BV1["Linear indexing only"] + BV2["buffer[5]"] + BV3["No shape information"] + BV4["Direct memory access"] + end + + subgraph "TensorView" + TV1["Multi-dimensional indexing"] + TV2["tensor(1, 2)"] + TV3["Shape-aware operations"] + TV4["Coordinate transformations"] + end + + subgraph "Use Cases" + UC1["BufferView: Low-level memory ops"] + UC2["TensorView: Matrix/tensor algorithms"] + end + + BV1 --> UC1 + TV1 --> UC2 BufferView excels at raw memory operations where linear access is natural or where the overhead of coordinate calculation would be prohibitive. TensorView is best suited for algorithms that operate in terms of multi-dimensional coordinates, such as matrix operations, image processing, or tensor contractions. diff --git a/docs/conceptual/ck_tile/thread_mapping.rst b/docs/conceptual/ck_tile/thread_mapping.rst index 361912ba9f6..055b94eb404 100644 --- a/docs/conceptual/ck_tile/thread_mapping.rst +++ b/docs/conceptual/ck_tile/thread_mapping.rst @@ -78,61 +78,46 @@ Composable Kernel abstracts thread identification into partition indices, buildi }; -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "GPU Device" - subgraph "Thread Block" - subgraph "Warp 0" - T0["Thread 0
lane_id=0"] - T1["Thread 1
lane_id=1"] - T2["..."] - T31["Thread 31
lane_id=31"] - end - - subgraph "Warp 1" - T32["Thread 32
lane_id=0"] - T33["Thread 33
lane_id=1"] - T34["..."] - T63["Thread 63
lane_id=31"] - end - - W2["Warp 2"] - W3["..."] - W7["Warp 7"] - end - end - - subgraph "Thread Identification" - TID["Thread ID = blockIdx.x * blockDim.x + threadIdx.x"] - WID["Warp ID = threadIdx.x / 32"] - LID["Lane ID = threadIdx.x % 32"] - end - - subgraph "P-space Mapping" - P["P-coordinates
NDimP=1: [thread_id]
NDimP=2: [warp_id, lane_id]"] - end - - T0 --> TID - TID --> WID - TID --> LID - WID --> P - LID --> P - - style T0 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style T32 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style P fill:#fff3e0,stroke:#f57c00,stroke-width:3px - - - - - -.. image:: diagrams/thread_mapping_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "GPU Device" + subgraph "Thread Block" + subgraph "Warp 0" + T0["Thread 0
lane_id=0"] + T1["Thread 1
lane_id=1"] + T2["..."] + T31["Thread 31
lane_id=31"] + end + + subgraph "Warp 1" + T32["Thread 32
lane_id=0"] + T33["Thread 33
lane_id=1"] + T34["..."] + T63["Thread 63
lane_id=31"] + end + + W2["Warp 2"] + W3["..."] + W7["Warp 7"] + end + end + + subgraph "Thread Identification" + TID["Thread ID = blockIdx.x * blockDim.x + threadIdx.x"] + WID["Warp ID = threadIdx.x / 32"] + LID["Lane ID = threadIdx.x % 32"] + end + + subgraph "P-space Mapping" + P["P-coordinates
NDimP=1: [thread_id]
NDimP=2: [warp_id, lane_id]"] + end + + T0 --> TID + TID --> WID + TID --> LID + WID --> P + LID --> P Thread Hierarchy Structure @@ -179,51 +164,36 @@ Thread-to-Data Mapping Once threads know their IDs, they need to map those IDs to specific data elements. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Thread to Data Mapping" - subgraph "Thread Grid" - T00["Thread[0,0]
Warp 0"] - T01["Thread[0,1]
Warp 0"] - T10["Thread[1,0]
Warp 1"] - T11["Thread[1,1]
Warp 1"] - end - - subgraph "Data Tiles" - D00["Data[0:4, 0:4]
16 elements"] - D01["Data[0:4, 4:8]
16 elements"] - D10["Data[4:8, 0:4]
16 elements"] - D11["Data[4:8, 4:8]
16 elements"] - end - - subgraph "Memory Access" - MA["Coalesced Access
Adjacent threads → Adjacent memory"] - end - end - - T00 --> D00 - T01 --> D01 - T10 --> D10 - T11 --> D11 - - D00 --> MA - D01 --> MA - - style T00 fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style D00 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style MA fill:#fff3e0,stroke:#f57c00,stroke-width:2px - - - - - -.. image:: diagrams/thread_mapping_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Thread to Data Mapping" + subgraph "Thread Grid" + T00["Thread[0,0]
Warp 0"] + T01["Thread[0,1]
Warp 0"] + T10["Thread[1,0]
Warp 1"] + T11["Thread[1,1]
Warp 1"] + end + + subgraph "Data Tiles" + D00["Data[0:4, 0:4]
16 elements"] + D01["Data[0:4, 4:8]
16 elements"] + D10["Data[4:8, 0:4]
16 elements"] + D11["Data[4:8, 4:8]
16 elements"] + end + + subgraph "Memory Access" + MA["Coalesced Access
Adjacent threads → Adjacent memory"] + end + end + + T00 --> D00 + T01 --> D01 + T10 --> D10 + T11 --> D11 + + D00 --> MA + D01 --> MA Data Distribution Pattern ------------------------- diff --git a/docs/conceptual/ck_tile/tile_distribution.rst b/docs/conceptual/ck_tile/tile_distribution.rst index 3c016318bfb..24417c70c78 100644 --- a/docs/conceptual/ck_tile/tile_distribution.rst +++ b/docs/conceptual/ck_tile/tile_distribution.rst @@ -19,98 +19,66 @@ This design adapts to diverse computational scenarios without manual interventio Complete Tile Distribution System Overview ------------------------------------------ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Logical View" - T["Tensor
Multi-dimensional data"] - TD["TileDistribution
Work assignment"] - TW["TileWindow
Data view"] - end - - subgraph "Coordinate Spaces" - X["X: Physical tensor coords"] - Y["Y: Tile pattern coords"] - P["P: Processing element coords"] - R["R: Replication coords (optional)"] - end - - subgraph "GPU Execution" - W["Warps
32 threads each"] - L["Lanes
Thread within warp"] - REG["Registers
Thread-local storage"] - end - - T --> TD - TD --> TW - - TD --> X - TD --> Y - TD --> P - TD --> R - - P --> W - P --> L - TW --> REG - - style TD fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style P fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style REG fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - - -.. image:: diagrams/tile_distribution_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Logical View" + T["Tensor
Multi-dimensional data"] + TD["TileDistribution
Work assignment"] + TW["TileWindow
Data view"] + end + + subgraph "Coordinate Spaces" + X["X: Physical tensor coords"] + Y["Y: Tile pattern coords"] + P["P: Processing element coords"] + R["R: Replication coords (optional)"] + end + + subgraph "GPU Execution" + W["Warps
32 threads each"] + L["Lanes
Thread within warp"] + REG["Registers
Thread-local storage"] + end + + T --> TD + TD --> TW + + TD --> X + TD --> Y + TD --> P + TD --> R + + P --> W + P --> L + TW --> REG Coordinate System Architecture ------------------------------ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart LR - subgraph "Input" - TC["Thread Coordinates
(warpId, laneId)"] - end - - subgraph "Transformation Pipeline" - P2Y["P → Y
Thread to pattern"] - Y2X["Y → X
Pattern to physical"] - Y2D["Y → D
Pattern to register"] - end - - subgraph "Output" - MC["Memory Coordinates
Global addresses"] - RI["Register Indices
Local storage"] - end - - TC --> P2Y - P2Y --> Y2X - P2Y --> Y2D - Y2X --> MC - Y2D --> RI - - style TC fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - style MC fill:#d1fae5,stroke:#10b981,stroke-width:2px - style RI fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - - - - - - -.. image:: diagrams/tile_distribution_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + flowchart LR + subgraph "Input" + TC["Thread Coordinates
(warpId, laneId)"] + end + + subgraph "Transformation Pipeline" + P2Y["P → Y
Thread to pattern"] + Y2X["Y → X
Pattern to physical"] + Y2D["Y → D
Pattern to register"] + end + + subgraph "Output" + MC["Memory Coordinates
Global addresses"] + RI["Register Indices
Local storage"] + end + + TC --> P2Y + P2Y --> Y2X + P2Y --> Y2D + Y2X --> MC + Y2D --> RI What is Tile Distribution? -------------------------- @@ -152,50 +120,34 @@ TileDistribution abstracts the mapping between logical problem coordinates and p Problem Space Mapping --------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - - graph TB - subgraph "Problem Space (256×256 Matrix)" - M["Full Matrix
65,536 elements"] - T1["Tile 1
32×32"] - T2["Tile 2
32×32"] - TN["Tile N
32×32"] - end - - subgraph "Thread Assignment" - W0["Warp 0
32 threads"] - W1["Warp 1
32 threads"] - L0["Lane 0-31
Individual threads"] - end - - subgraph "Memory Pattern" - MP["Coalesced Access
Sequential addresses
No bank conflicts"] - end - - M --> T1 - M --> T2 - M --> TN - - T1 --> W0 - T1 --> W1 - W0 --> L0 - L0 --> MP - - style M fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style MP fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - - -.. image:: diagrams/tile_distribution_3.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Problem Space (256×256 Matrix)" + M["Full Matrix
65,536 elements"] + T1["Tile 1
32×32"] + T2["Tile 2
32×32"] + TN["Tile N
32×32"] + end + + subgraph "Thread Assignment" + W0["Warp 0
32 threads"] + W1["Warp 1
32 threads"] + L0["Lane 0-31
Individual threads"] + end + + subgraph "Memory Pattern" + MP["Coalesced Access
Sequential addresses
No bank conflicts"] + end + + M --> T1 + M --> T2 + M --> TN + + T1 --> W0 + T1 --> W1 + W0 --> L0 + L0 --> MP Creating a TileDistribution --------------------------- @@ -369,47 +321,31 @@ Creating and using a TileDistribution: Hierarchical Decomposition -------------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Level 1: Block Distribution" - B["Thread Block
256 threads"] - BT1["Block Tile 1
64×64"] - BT2["Block Tile 2
64×64"] - end - - subgraph "Level 2: Warp Distribution" - W["Warp
32 threads"] - WT1["Warp Tile 1
16×16"] - WT2["Warp Tile 2
16×16"] - end - - subgraph "Level 3: Thread Distribution" - T["Thread"] - TT["Thread Tile
2×2"] - end - - B --> BT1 - BT1 --> W - W --> WT1 - WT1 --> T - T --> TT - - style B fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style W fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style T fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - - -.. image:: diagrams/tile_distribution_4.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Level 1: Block Distribution" + B["Thread Block
256 threads"] + BT1["Block Tile 1
64×64"] + BT2["Block Tile 2
64×64"] + end + + subgraph "Level 2: Warp Distribution" + W["Warp
32 threads"] + WT1["Warp Tile 1
16×16"] + WT2["Warp Tile 2
16×16"] + end + + subgraph "Level 3: Thread Distribution" + T["Thread"] + TT["Thread Tile
2×2"] + end + + B --> BT1 + BT1 --> W + W --> WT1 + WT1 --> T + T --> TT Advanced Example: Matrix Multiplication Distribution ---------------------------------------------------- @@ -462,45 +398,28 @@ Advanced Example: Matrix Multiplication Distribution Work Distribution Pattern ------------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart TB - subgraph "Matrix C (128×128)" - C["16,384 elements"] - end - - subgraph "Thread Grid (32×32)" - TG["1,024 threads"] - end - - subgraph "Per Thread" - PT["4×4 tile
16 elements"] - end - - subgraph "Memory Access" - MA["Coalesced reads
Efficient writes
No conflicts"] - end - - C --> TG - TG --> PT - PT --> MA - - style C fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style TG fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style PT fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style MA fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - - -.. image:: diagrams/tile_distribution_5.svg - :alt: Diagram - :align: center +.. mermaid:: + + flowchart TB + subgraph "Matrix C (128×128)" + C["16,384 elements"] + end + + subgraph "Thread Grid (32×32)" + TG["1,024 threads"] + end + + subgraph "Per Thread" + PT["4×4 tile
16 elements"] + end + + subgraph "Memory Access" + MA["Coalesced reads
Efficient writes
No conflicts"] + end + + C --> TG + TG --> PT + PT --> MA Memory Access Patterns ---------------------- @@ -514,97 +433,67 @@ One of the key benefits of TileDistribution is generating optimal memory access Transformation Pipeline ----------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Input" - TID["Thread ID
(0-1023)"] - end - - subgraph "Stage 1" - P["P-coordinates
(warp, lane)"] - end - - subgraph "Stage 2" - Y["Y-coordinates
(tile position)"] - end - - subgraph "Stage 3" - X["X-coordinates
(tensor indices)"] - end - - subgraph "Output" - ADDR["Memory addresses
Register indices"] - end - - TID --> P - P --> Y - Y --> X - X --> ADDR - - style TID fill:#e0e7ff,stroke:#4338ca,stroke-width:2px - style ADDR fill:#d1fae5,stroke:#10b981,stroke-width:2px - - - - - -.. image:: diagrams/tile_distribution_6.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Input" + TID["Thread ID
(0-1023)"] + end + + subgraph "Stage 1" + P["P-coordinates
(warp, lane)"] + end + + subgraph "Stage 2" + Y["Y-coordinates
(tile position)"] + end + + subgraph "Stage 3" + X["X-coordinates
(tensor indices)"] + end + + subgraph "Output" + ADDR["Memory addresses
Register indices"] + end + + TID --> P + P --> Y + Y --> X + X --> ADDR Performance Comparison ---------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Manual Implementation" - M1["Calculate indices manually"] - M2["Handle boundary conditions"] - M3["Ensure coalescing"] - M4["Manage bank conflicts"] - M5["~200 lines of code"] - end - - subgraph "With TileDistribution" - T1["make_tile_distribution()"] - T2["Automatic optimization"] - T3["~10 lines of code"] - end - - subgraph "Performance" - P1["Same performance"] - P2["Fewer bugs"] - P3["Portable across GPUs"] - end - - M1 --> M5 - T1 --> T3 - - M5 --> P1 - T3 --> P1 - P1 --> P2 - P2 --> P3 - - style M5 fill:#fee2e2,stroke:#ef4444,stroke-width:2px - style T3 fill:#d1fae5,stroke:#10b981,stroke-width:2px - style P3 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px - - - - - - -.. image:: diagrams/tile_distribution_7.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Manual Implementation" + M1["Calculate indices manually"] + M2["Handle boundary conditions"] + M3["Ensure coalescing"] + M4["Manage bank conflicts"] + M5["~200 lines of code"] + end + + subgraph "With TileDistribution" + T1["make_tile_distribution()"] + T2["Automatic optimization"] + T3["~10 lines of code"] + end + + subgraph "Performance" + P1["Same performance"] + P2["Fewer bugs"] + P3["Portable across GPUs"] + end + + M1 --> M5 + T1 --> T3 + + M5 --> P1 + T3 --> P1 + P1 --> P2 + P2 --> P3 Summary ------- diff --git a/docs/conceptual/ck_tile/tile_window.rst b/docs/conceptual/ck_tile/tile_window.rst index 23c006d972b..272526611c2 100644 --- a/docs/conceptual/ck_tile/tile_window.rst +++ b/docs/conceptual/ck_tile/tile_window.rst @@ -13,56 +13,43 @@ TileWindow implements a distribution-aware windowing mechanism that views a subs TileWindow Architecture ----------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Components" - TV["TensorView
Data source"] - TD["TileDistribution
Thread mapping"] - TW["TileWindow
Access gateway"] - LT["LoadStoreTraits
Access optimizer"] - DT["DistributedTensor
Register storage"] - end - - subgraph "Operations" - Load["Load
Global → Registers"] - Compute["Compute
In registers"] - Store["Store
Registers → Global"] - end - - subgraph "Optimizations" - Coal["Coalescing
Adjacent access"] - Vec["Vectorization
Multi-element ops"] - Bank["Bank conflict
avoidance"] - SFC["Space-filling
curve traversal"] - end - - TV --> TW - TD --> TW - TW --> LT - LT --> DT - - TW --> Load - Load --> Compute - Compute --> Store - - Load --> Coal - Load --> Vec - Load --> SFC - Store --> Bank - - style TW fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style LT fill:#fff3e0,stroke:#f57c00,stroke-width:2px - style DT fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - -.. image:: diagrams/tile_window_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Components" + TV["TensorView
Data source"] + TD["TileDistribution
Thread mapping"] + TW["TileWindow
Access gateway"] + LT["LoadStoreTraits
Access optimizer"] + DT["DistributedTensor
Register storage"] + end + + subgraph "Operations" + Load["Load
Global → Registers"] + Compute["Compute
In registers"] + Store["Store
Registers → Global"] + end + + subgraph "Optimizations" + Coal["Coalescing
Adjacent access"] + Vec["Vectorization
Multi-element ops"] + Bank["Bank conflict
avoidance"] + SFC["Space-filling
curve traversal"] + end + + TV --> TW + TD --> TW + TW --> LT + LT --> DT + + TW --> Load + Load --> Compute + Compute --> Store + + Load --> Coal + Load --> Vec + Load --> SFC + Store --> Bank What is a TileWindow? --------------------- @@ -170,42 +157,30 @@ Space-Filling Curves for Memory Access TileWindow uses :ref:`space-filling curves ` to determine the order in which memory is accessed. Space-filling curves provide cache-friendly traversal patterns that help maximize hardware utilization. The "snake" pattern minimizes the distance between consecutive accesses, keeping data in cache longer. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Linear Access Pattern" - L1["0,1,2,3"] - L2["4,5,6,7"] - L3["8,9,10,11"] - L4["12,13,14,15"] - end - - subgraph "Snake Access Pattern" - S1["0,1,2,3"] - S2["7,6,5,4"] - S3["8,9,10,11"] - S4["15,14,13,12"] - end - - L1 --> L2 - L2 --> L3 - L3 --> L4 - - S1 --> S2 - S2 --> S3 - S3 --> S4 - - style S1 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style S2 fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - +.. mermaid:: -.. image:: diagrams/tile_window_2.svg - :alt: Diagram - :align: center + graph LR + subgraph "Linear Access Pattern" + L1["0,1,2,3"] + L2["4,5,6,7"] + L3["8,9,10,11"] + L4["12,13,14,15"] + end + + subgraph "Snake Access Pattern" + S1["0,1,2,3"] + S2["7,6,5,4"] + S3["8,9,10,11"] + S4["15,14,13,12"] + end + + L1 --> L2 + L2 --> L3 + L3 --> L4 + + S1 --> S2 + S2 --> S3 + S3 --> S4 **C++ Space-Filling Curve Implementation:** @@ -237,44 +212,32 @@ TileWindow uses :ref:`space-filling curves ` to det TileWindow Data Flow -------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - flowchart LR - subgraph "Step 1: Create Window" - T["Tensor
[256, 256]"] - O["Origin
(64, 64)"] - W["Window Size
[32, 32]"] - end - - subgraph "Step 2: Apply Distribution" - TD["TileDistribution
Thread mapping"] - TW["TileWindow
Created"] - end - - subgraph "Step 3: Load Data" - GM["Global Memory
Window region"] - REG["Registers
Distributed tensor"] - end - - T --> TW - O --> TW - W --> TW - TD --> TW - - TW --> GM - GM -->|"load()"| REG - - style TW fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style REG fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - +.. mermaid:: + + flowchart LR + subgraph "Step 1: Create Window" + T["Tensor
[256, 256]"] + O["Origin
(64, 64)"] + W["Window Size
[32, 32]"] + end + + subgraph "Step 2: Apply Distribution" + TD["TileDistribution
Thread mapping"] + TW["TileWindow
Created"] + end -.. image:: diagrams/tile_window_3.svg - :alt: Diagram - :align: center + subgraph "Step 3: Load Data" + GM["Global Memory
Window region"] + REG["Registers
Distributed tensor"] + end + + T --> TW + O --> TW + W --> TW + TD --> TW + + TW --> GM + GM -->|"load()"| REG Creating and Using TileWindow ----------------------------- @@ -366,50 +329,37 @@ Calls to ``window.load()`` trigger the following sequence of operations: Load Operation Architecture --------------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Load Analysis" - Analyze["Analyze access pattern
Detect coalescing opportunities"] - end - - subgraph "Vectorization" - V1["Scalar: 4 loads"] - V2["Vector2: 2 loads"] - V4["Vector4: 1 load"] - end - - subgraph "Memory Transaction" - Coal["Coalesced access
32 threads → 1 transaction"] - NonCoal["Non-coalesced
32 threads → 32 transactions"] - end - - subgraph "Result" - Reg["Thread registers
Local data"] - end - - Analyze --> V1 - Analyze --> V2 - Analyze --> V4 - - V4 --> Coal - V1 --> NonCoal - - Coal --> Reg - NonCoal --> Reg - - style V4 fill:#d1fae5,stroke:#10b981,stroke-width:2px - style Coal fill:#d1fae5,stroke:#10b981,stroke-width:2px - style NonCoal fill:#fee2e2,stroke:#ef4444,stroke-width:2px - - +.. mermaid:: + + graph TB + subgraph "Load Analysis" + Analyze["Analyze access pattern
Detect coalescing opportunities"] + end + + subgraph "Vectorization" + V1["Scalar: 4 loads"] + V2["Vector2: 2 loads"] + V4["Vector4: 1 load"] + end + + subgraph "Memory Transaction" + Coal["Coalesced access
32 threads → 1 transaction"] + NonCoal["Non-coalesced
32 threads → 32 transactions"] + end -.. image:: diagrams/tile_window_4.svg - :alt: Diagram - :align: center + subgraph "Result" + Reg["Thread registers
Local data"] + end + + Analyze --> V1 + Analyze --> V2 + Analyze --> V4 + + V4 --> Coal + V1 --> NonCoal + + Coal --> Reg + NonCoal --> Reg Memory Access Patterns ---------------------- @@ -586,39 +536,26 @@ Complete Load-Compute-Store Pipeline Performance Characteristics --------------------------- -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph LR - subgraph "Memory Access Optimization" - V["Vectorization
4x fewer transactions"] - C["Coalescing
32x bandwidth efficiency"] - P["Precomputation
Zero overhead addressing"] - S["Space-filling
Optimal cache usage"] - end - - subgraph "Hardware Utilization" - BW["Memory Bandwidth
Near 100% utilization"] - L["Latency Hiding
Overlapped operations"] - R["Register Reuse
Minimal spills"] - end - - V --> BW - C --> BW - P --> L - S --> R - - style V fill:#e3f2fd,stroke:#1976d2,stroke-width:2px - style C fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - style BW fill:#d1fae5,stroke:#10b981,stroke-width:3px - - - -.. image:: diagrams/tile_window_5.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph LR + subgraph "Memory Access Optimization" + V["Vectorization
4x fewer transactions"] + C["Coalescing
32x bandwidth efficiency"] + P["Precomputation
Zero overhead addressing"] + S["Space-filling
Optimal cache usage"] + end + + subgraph "Hardware Utilization" + BW["Memory Bandwidth
Near 100% utilization"] + L["Latency Hiding
Overlapped operations"] + R["Register Reuse
Minimal spills"] + end + + V --> BW + C --> BW + P --> L + S --> R Best Practices diff --git a/docs/conceptual/ck_tile/transforms.rst b/docs/conceptual/ck_tile/transforms.rst index 3dfea276cbb..64f95cfe132 100644 --- a/docs/conceptual/ck_tile/transforms.rst +++ b/docs/conceptual/ck_tile/transforms.rst @@ -29,36 +29,21 @@ Zero-Copy Logical Operations - **Data Storage**: The actual tensor data remains stored in memory in linear fashion, exactly as specified by the original tensor shape and strides at creation time. See :ref:`ck_tile_buffer_views` for more information about raw memory access. - **Logical Mapping**: Transforms create different logical views of the same underlying data and only change how access coordinates are interpreted. See :ref:`ck_tile_tensor_views` for more information about tensor views. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "Tensor Coordinate Transformation" - US["Lower Dimension Space
Source coordinate system"] - LS["Upper Dimension Space
Target coordinate system"] - - DATA["Linear Data in Memory
Layout determined by tensor
shape & strides"] - end - - US -->|"Forward Transform"| LS - LS -->|"Inverse Transform"| US - - DATA -.->|"Same data,
different views"| US - DATA -.->|"Same data,
different views"| LS - - style US fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style LS fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_1.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "Tensor Coordinate Transformation" + US["Lower Dimension Space
Source coordinate system"] + LS["Upper Dimension Space
Target coordinate system"] + + DATA["Linear Data in Memory
Layout determined by tensor
shape & strides"] + end + + US -->|"Forward Transform"| LS + LS -->|"Inverse Transform"| US + + DATA -.->|"Same data,
different views"| US + DATA -.->|"Same data,
different views"| LS Index Calculation Operations ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -74,82 +59,54 @@ These operations enable bidirectional navigation between different coordinate re Transform System Architecture ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - - subgraph "Transform Types" - EMB["EmbedTransform
Linear → Multi-D Strided"] - UNM["MergeTransform
Multi-D → Linear"] - MRG["UnmergeTransform
Linear → Multi-D"] - REP["ReplicateTransform
0D → Multi-D Broadcast"] - OFF["OffsetTransform
Translation"] - PAS["PassThroughTransform
Identity"] - PAD["PadTransform
Boundaries"] - end - - subgraph "Operations" - FWD["Forward
calculate_lower_index()"] - BWD["Backward
calculate_upper_index()"] - UPD["Update
update_lower_index()"] - end - - EMB --> FWD - UNM --> FWD - MRG --> FWD - REP --> FWD - OFF --> FWD - PAS --> FWD - PAD --> FWD - - style FWD fill:#e8f5e9,stroke:#388e3c,stroke-width:2px - - - - - -.. image:: diagrams/transforms_2.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + + subgraph "Transform Types" + EMB["EmbedTransform
Linear → Multi-D Strided"] + UNM["MergeTransform
Multi-D → Linear"] + MRG["UnmergeTransform
Linear → Multi-D"] + REP["ReplicateTransform
0D → Multi-D Broadcast"] + OFF["OffsetTransform
Translation"] + PAS["PassThroughTransform
Identity"] + PAD["PadTransform
Boundaries"] + end + + subgraph "Operations" + FWD["Forward
calculate_lower_index()"] + BWD["Backward
calculate_upper_index()"] + UPD["Update
update_lower_index()"] + end + + EMB --> FWD + UNM --> FWD + MRG --> FWD + REP --> FWD + OFF --> FWD + PAS --> FWD + PAD --> FWD MergeTransform -------------- MergeTransform collapses multiple dimensions from the lower coordinate space into a single dimension in the upper coordinate space, effectively reducing the dimensionality of the tensor representation while preserving data relationships. This transform is fundamental to the :ref:`tile distribution system `. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "MergeTransform: Multi-D → Linear" - LS["Lower Coordinate Space
2D: [4, 5]
Coord: (2, 3)"] - US["Upper Coordinate Space
1D Linear
Index: 13"] - - DATA["Same Tensor Data
Layout: row-major
Size: 20 elements"] - end - - LS -->|"Forward Transform
2×5 + 3 = 13"| US - US -->|"Inverse Transform
13÷5=2, 13%5=3"| LS - - DATA -.->|"Multi-dimensional
view"| LS - DATA -.->|"Linear
view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_3.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "MergeTransform: Multi-D → Linear" + LS["Lower Coordinate Space
2D: [4, 5]
Coord: (2, 3)"] + US["Upper Coordinate Space
1D Linear
Index: 13"] + + DATA["Same Tensor Data
Layout: row-major
Size: 20 elements"] + end + + LS -->|"Forward Transform
2×5 + 3 = 13"| US + US -->|"Inverse Transform
13÷5=2, 13%5=3"| LS + + DATA -.->|"Multi-dimensional
view"| LS + DATA -.->|"Linear
view"| US **C++ Implementation:** @@ -187,36 +144,21 @@ UnmergeTransform UnmergeTransform expands coordinates from a single dimension in the lower coordinate space into multiple dimensions in the upper coordinate space, effectively increasing the dimensionality of the tensor representation while preserving all data relationships. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "UnmergeTransform: Linear → Multi-D" - LS["Lower Coordinate Space
1D Linear
Index: 14"] - US["Upper Coordinate Space
3D: [3, 4, 2]
Coord: (1, 3, 0)"] - - DATA["Same Tensor Data
Layout: row-major
Size: 24 elements"] - end - - LS -->|"Forward Transform
14 = 1×8 + 3×2 + 0"| US - US -->|"Inverse Transform
linearize back"| LS - - DATA -.->|"Linear
view"| LS - DATA -.->|"Multi-dimensional
view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_4.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "UnmergeTransform: Linear → Multi-D" + LS["Lower Coordinate Space
1D Linear
Index: 14"] + US["Upper Coordinate Space
3D: [3, 4, 2]
Coord: (1, 3, 0)"] + + DATA["Same Tensor Data
Layout: row-major
Size: 24 elements"] + end + + LS -->|"Forward Transform
14 = 1×8 + 3×2 + 0"| US + US -->|"Inverse Transform
linearize back"| LS + + DATA -.->|"Linear
view"| LS + DATA -.->|"Multi-dimensional
view"| US **C++ Implementation:** @@ -264,36 +206,21 @@ EmbedTransform EmbedTransform expands linear indices from the lower coordinate space into multi-dimensional coordinates in the upper coordinate space using configurable strides, enabling flexible strided tensor layouts and sub-tensor views within larger buffers. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "EmbedTransform: Linear → Multi-D Strided" - LS["Lower Coordinate Space
1D Linear
Index: 14"] - US["Upper Coordinate Space
2D: [2, 3]
Coord: (1, 2)"] - - DATA["Linear Buffer in Memory"] - end - - LS -->|"Forward Transform
Strides: [12, 1]
14 ÷ 12 = 1, 14 % 12 = 2"| US - US -->|"Inverse Transform
1×12 + 2×1 = 14"| LS - - DATA -.->|"Linear
index view"| LS - DATA -.->|"Multi-dimensional
strided view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_5.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "EmbedTransform: Linear → Multi-D Strided" + LS["Lower Coordinate Space
1D Linear
Index: 14"] + US["Upper Coordinate Space
2D: [2, 3]
Coord: (1, 2)"] + + DATA["Linear Buffer in Memory"] + end + + LS -->|"Forward Transform
Strides: [12, 1]
14 ÷ 12 = 1, 14 % 12 = 2"| US + US -->|"Inverse Transform
1×12 + 2×1 = 14"| LS + + DATA -.->|"Linear
index view"| LS + DATA -.->|"Multi-dimensional
strided view"| US **C++ Implementation:** @@ -329,36 +256,21 @@ ReplicateTransform ReplicateTransform creates a higher-dimensional tensor by replicating (broadcasting) a lower-dimensional tensor. It's essentially a broadcasting operation that takes a tensor with fewer dimensions and logically replicates it across new dimensions without data duplication. An example is taking a scalar (0-dimensional) input and broadcasting it across multiple dimensions, enabling efficient broadcasting patterns where a single value appears at every position in a multi-dimensional coordinate space. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "ReplicateTransform: 0D → Multi-D Broadcasting" - LS["Lower Coordinate Space
0D: Scalar
Empty coordinate []"] - US["Upper Coordinate Space
2D: [3, 4]
All coords: (i, j)"] - - DATA["Single Scalar Value"] - end - - LS -->|"Forward Transform
[] → (i,j) for any i,j"| US - US -->|"Inverse Transform
(i,j) → [] for any i,j"| LS - - DATA -.->|"One scalar
value"| LS - DATA -.->|"Broadcasted view
at all positions"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_6.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "ReplicateTransform: 0D → Multi-D Broadcasting" + LS["Lower Coordinate Space
0D: Scalar
Empty coordinate []"] + US["Upper Coordinate Space
2D: [3, 4]
All coords: (i, j)"] + + DATA["Single Scalar Value"] + end + + LS -->|"Forward Transform
[] → (i,j) for any i,j"| US + US -->|"Inverse Transform
(i,j) → [] for any i,j"| LS + + DATA -.->|"One scalar
value"| LS + DATA -.->|"Broadcasted view
at all positions"| US **C++ Implementation:** @@ -406,36 +318,21 @@ OffsetTransform OffsetTransform shifts coordinates by a fixed offset, creating a translated view of the coordinate space. It performs translation operations where each coordinate in the upper space is mapped to a coordinate in the lower space by adding a constant offset. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "OffsetTransform: 1D → 1D Translation" - LS["Lower Coordinate Space
1D: [0, 63]
Coord: index + offset"] - US["Upper Coordinate Space
1D: [0, 47]
Coord: index"] - - DATA["Linear Buffer in Memory"] - end - - LS -->|"Forward Transform
idx → idx + 16"| US - US -->|"Inverse Transform
idx + 16 → idx"| LS - - DATA -.->|"Lower
view"| LS - DATA -.->|"Upper
view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_7.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "OffsetTransform: 1D → 1D Translation" + LS["Lower Coordinate Space
1D: [0, 63]
Coord: index + offset"] + US["Upper Coordinate Space
1D: [0, 47]
Coord: index"] + + DATA["Linear Buffer in Memory"] + end + + LS -->|"Forward Transform
idx → idx + 16"| US + US -->|"Inverse Transform
idx + 16 → idx"| LS + + DATA -.->|"Lower
view"| LS + DATA -.->|"Upper
view"| US **C++ Implementation:** @@ -483,36 +380,21 @@ PassThroughTransform - Identity No-op transform that passes coordinates unchanged. The PassThrough transform is the simplest coordinate transformation in CK Tile, implementing a perfect identity mapping where input coordinates are passed through unchanged to the output. This transform is essential as a placeholder in transformation chains and for dimensions that require no modification. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "PassThroughTransform: 1D → 1D Identity" - LS["Lower Coordinate Space
1D: [0, 59]
Coord: index"] - US["Upper Coordinate Space
1D: [0, 59]
Coord: index"] - - DATA["Linear Buffer in Memory"] - end - - LS -.->|"Perfect Identity
idx → idx"| US - US -.->|"Perfect Identity
idx → idx"| LS - - DATA -->|"Same buffer
same view"| LS - DATA -->|"Same buffer
same view"| US - - style LS fill:#e8f5e8,stroke:#2e7d32,stroke-width:3px - style US fill:#e8f5e8,stroke:#2e7d32,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_8.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "PassThroughTransform: 1D → 1D Identity" + LS["Lower Coordinate Space
1D: [0, 59]
Coord: index"] + US["Upper Coordinate Space
1D: [0, 59]
Coord: index"] + + DATA["Linear Buffer in Memory"] + end + + LS -.->|"Perfect Identity
idx → idx"| US + US -.->|"Perfect Identity
idx → idx"| LS + + DATA -->|"Same buffer
same view"| LS + DATA -->|"Same buffer
same view"| US **C++ Implementation:** @@ -555,39 +437,21 @@ PadTransform PadTransform adds padding to tensor dimensions, mapping coordinates from upper dimension space (with padding) to lower dimension space (original data). -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "PadTransform: 1D → 1D with Padding" - LS["Lower Coordinate Space
1D: [0, 2] (original data)"] - US["Upper Coordinate Space
1D: [0, 4] (with padding)"] - - DATA["Tensor Data in Memory"] - end - - LS -->|"Forward Transform
idx + left_pad"| US - US -->|"Inverse Transform
idx - left_pad"| LS - - DATA -.->|"Original view"| LS - DATA -.->|"Padded view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_9.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "PadTransform: 1D → 1D with Padding" + LS["Lower Coordinate Space
1D: [0, 2] (original data)"] + US["Upper Coordinate Space
1D: [0, 4] (with padding)"] + + DATA["Tensor Data in Memory"] + end + + LS -->|"Forward Transform
idx + left_pad"| US + US -->|"Inverse Transform
idx - left_pad"| LS + + DATA -.->|"Original view"| LS + DATA -.->|"Padded view"| US **C++ Implementation:** @@ -633,106 +497,63 @@ XorTransform XorTransform applies a 2D XOR mapping for specialized memory access patterns. It performs XOR operations on coordinates to create transformed memory layouts for specific algorithmic optimizations, particularly useful for avoiding :ref:`LDS bank conflicts `. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "XorTransform: 2D → 2D XOR Mapping" - LS["Lower Coordinate Space
2D: [4, 8]
XOR-transformed coords"] - US["Upper Coordinate Space
2D: [4, 8]
Normal coords"] - - DATA["Same Tensor Data"] - end - - LS -->|"Forward Transform
apply XOR reverse"| US - US -->|"Inverse Transform
apply XOR mapping"| LS - - DATA -.->|"XOR pattern
view"| LS - DATA -.->|"Normal
view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_10.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "XorTransform: 2D → 2D XOR Mapping" + LS["Lower Coordinate Space
2D: [4, 8]
XOR-transformed coords"] + US["Upper Coordinate Space
2D: [4, 8]
Normal coords"] + + DATA["Same Tensor Data"] + end + + LS -->|"Forward Transform
apply XOR reverse"| US + US -->|"Inverse Transform
apply XOR mapping"| LS + + DATA -.->|"XOR pattern
view"| LS + DATA -.->|"Normal
view"| US SliceTransform ~~~~~~~~~~~~~~ SliceTransform extracts a sub-region from a tensor dimension. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "SliceTransform: 1D → 1D Sub-region" - LS["Lower Coordinate Space
1D: [0, 9] (original range)"] - US["Upper Coordinate Space
1D: [0, 4] (slice range)"] - - DATA["Tensor Data in Memory"] - end - - LS -->|"Forward Transform
idx + slice_begin"| US - US -->|"Inverse Transform
idx - slice_begin"| LS - - DATA -.->|"Full tensor
view"| LS - DATA -.->|"Sub-region
view"| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - - - -.. image:: diagrams/transforms_11.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "SliceTransform: 1D → 1D Sub-region" + LS["Lower Coordinate Space
1D: [0, 9] (original range)"] + US["Upper Coordinate Space
1D: [0, 4] (slice range)"] + + DATA["Tensor Data in Memory"] + end + + LS -->|"Forward Transform
idx + slice_begin"| US + US -->|"Inverse Transform
idx - slice_begin"| LS + + DATA -.->|"Full tensor
view"| LS + DATA -.->|"Sub-region
view"| US ModuloTransform ~~~~~~~~~~~~~~~ ModuloTransform applies cyclic wrapping to coordinates using modulo operations. -.. - Original mermaid diagram (edit here, then run update_diagrams.py) - - .. mermaid:: - - graph TB - subgraph "ModuloTransform: 1D → 1D Cyclic" - LS["Lower Coordinate Space
1D: [0, 3] (modulus range)"] - US["Upper Coordinate Space
1D: [0, 15] (full range)"] - - DATA["Tensor Data in Memory"] - end - - LS -->|"Forward Transform
idx * cycle_count"| US - US -->|"Inverse Transform
idx % modulus"| LS - - DATA -.->|" "| LS - DATA -.->|" "| US - - style LS fill:#e3f2fd,stroke:#1976d2,stroke-width:3px - style US fill:#fff3e0,stroke:#f57c00,stroke-width:3px - style DATA fill:#f0f9ff,stroke:#0284c7,stroke-width:2px,stroke-dasharray: 5 5 - - - -.. image:: diagrams/transforms_12.svg - :alt: Diagram - :align: center +.. mermaid:: + + graph TB + subgraph "ModuloTransform: 1D → 1D Cyclic" + LS["Lower Coordinate Space
1D: [0, 3] (modulus range)"] + US["Upper Coordinate Space
1D: [0, 15] (full range)"] + + DATA["Tensor Data in Memory"] + end + + LS -->|"Forward Transform
idx * cycle_count"| US + US -->|"Inverse Transform
idx % modulus"| LS + + DATA -.->|" "| LS + DATA -.->|" "| US Summary ------- diff --git a/docs/conf.py b/docs/conf.py index bb7847e1d6d..e5de44fcf6d 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -45,7 +45,12 @@ for sphinx_var in ROCmDocs.SPHINX_VARS: globals()[sphinx_var] = getattr(docs_core, sphinx_var) -extensions += ['sphinxcontrib.bibtex'] -bibtex_bibfiles = ['refs.bib'] +extensions += [ + "sphinxcontrib.mermaid", + "sphinxcontrib.bibtex", +] + +mermaid_output_format = "raw" +bibtex_bibfiles = ["refs.bib"] cpp_id_attributes = ["__global__", "__device__", "__host__"] \ No newline at end of file diff --git a/docs/sphinx/requirements.in b/docs/sphinx/requirements.in index f2fb27e2b94..015693bb057 100644 --- a/docs/sphinx/requirements.in +++ b/docs/sphinx/requirements.in @@ -1,2 +1,3 @@ rocm-docs-core[api_reference]==1.31.3 sphinxcontrib-bibtex==2.6.5 +sphinxcontrib-mermaid>=1.0 diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index e36975219d8..fb3e6f80234 100644 --- a/docs/sphinx/requirements.txt +++ b/docs/sphinx/requirements.txt @@ -10,12 +10,12 @@ alabaster==1.0.0 # via sphinx asttokens==3.0.1 # via stack-data -attrs==25.4.0 +attrs==26.1.0 # via # jsonschema # jupyter-cache # referencing -babel==2.17.0 +babel==2.18.0 # via # pydata-sphinx-theme # sphinx @@ -23,15 +23,15 @@ beautifulsoup4==4.14.3 # via pydata-sphinx-theme breathe==4.36.0 # via rocm-docs-core -certifi==2026.1.4 +certifi==2026.5.20 # via requests cffi==2.0.0 # via # cryptography # pynacl -charset-normalizer==3.4.4 +charset-normalizer==3.4.7 # via requests -click==8.3.1 +click==8.4.1 # via # click-log # doxysphinx @@ -41,11 +41,11 @@ click-log==0.4.0 # via doxysphinx comm==0.2.3 # via ipykernel -cryptography==46.0.3 +cryptography==48.0.0 # via pyjwt -debugpy==1.8.19 +debugpy==1.8.21 # via ipykernel -decorator==5.2.1 +decorator==5.3.1 # via ipython docutils==0.21.2 # via @@ -66,30 +66,31 @@ fastjsonschema==2.21.2 # rocm-docs-core gitdb==4.0.12 # via gitpython -gitpython==3.1.46 +gitpython==3.1.50 # via rocm-docs-core -greenlet==3.3.0 +greenlet==3.5.1 # via sqlalchemy -idna==3.15 +idna==3.18 # via requests -imagesize==1.4.1 +imagesize==2.0.0 # via sphinx -importlib-metadata==8.7.1 +importlib-metadata==9.0.0 # via # jupyter-cache # myst-nb -ipykernel==7.1.0 +ipykernel==7.2.0 # via myst-nb -ipython==8.38.0 +ipython==8.39.0 # via # ipykernel # myst-nb -jedi==0.19.2 +jedi==0.20.0 # via ipython jinja2==3.1.6 # via # myst-parser # sphinx + # sphinxcontrib-mermaid jsonschema==4.26.0 # via nbformat jsonschema-specifications==2025.9.1 @@ -118,17 +119,17 @@ markdown-it-py==3.0.0 # myst-parser markupsafe==3.0.3 # via jinja2 -matplotlib-inline==0.2.1 +matplotlib-inline==0.2.2 # via # ipykernel # ipython -mdit-py-plugins==0.5.0 +mdit-py-plugins==0.6.1 # via myst-parser mdurl==0.1.2 # via markdown-it-py mpire==2.10.2 # via doxysphinx -myst-nb==1.3.0 +myst-nb==1.4.0 # via rocm-docs-core myst-parser==4.0.1 # via myst-nb @@ -143,40 +144,40 @@ nbformat==5.10.4 # nbclient nest-asyncio==1.6.0 # via ipykernel -packaging==25.0 +packaging==26.2 # via # ipykernel # pydata-sphinx-theme # sphinx -parso==0.8.5 +parso==0.8.7 # via jedi pexpect==4.9.0 # via ipython -platformdirs==4.5.1 +platformdirs==4.10.0 # via jupyter-core prompt-toolkit==3.0.52 # via ipython -psutil==7.2.1 +psutil==7.2.2 # via ipykernel ptyprocess==0.7.0 # via pexpect pure-eval==0.2.3 # via stack-data -pybtex==0.25.1 +pybtex==0.26.1 # via # pybtex-docutils # sphinxcontrib-bibtex pybtex-docutils==1.0.3 # via sphinxcontrib-bibtex -pycparser==2.23 +pycparser==3.0 # via cffi pydata-sphinx-theme==0.15.4 # via # rocm-docs-core # sphinx-book-theme -pygithub==2.8.1 +pygithub==2.9.1 # via rocm-docs-core -pygments==2.19.2 +pygments==2.20.0 # via # accessible-pygments # ipython @@ -185,7 +186,7 @@ pygments==2.19.2 # sphinx pyjson5==1.6.9 # via doxysphinx -pyjwt[crypto]==2.10.1 +pyjwt[crypto]==2.13.0 # via pygithub pynacl==1.6.2 # via pygithub @@ -201,6 +202,7 @@ pyyaml==6.0.3 # pybtex # rocm-docs-core # sphinx-external-toc + # sphinxcontrib-mermaid pyzmq==27.1.0 # via # ipykernel @@ -209,7 +211,7 @@ referencing==0.37.0 # via # jsonschema # jsonschema-specifications -requests==2.33.0 +requests==2.34.2 # via # pygithub # sphinx @@ -221,11 +223,11 @@ rpds-py==0.30.0 # referencing six==1.17.0 # via python-dateutil -smmap==5.0.2 +smmap==5.0.3 # via gitdb -snowballstemmer==3.0.1 +snowballstemmer==3.1.1 # via sphinx -soupsieve==2.8.1 +soupsieve==2.8.4 # via beautifulsoup4 sphinx==8.1.3 # via @@ -238,16 +240,20 @@ sphinx==8.1.3 # sphinx-copybutton # sphinx-design # sphinx-external-toc + # sphinx-multitoc-numbering # sphinx-notfound-page # sphinxcontrib-bibtex + # sphinxcontrib-mermaid sphinx-book-theme==1.1.4 # via rocm-docs-core sphinx-copybutton==0.5.2 # via rocm-docs-core sphinx-design==0.6.1 # via rocm-docs-core -sphinx-external-toc==1.0.1 +sphinx-external-toc==1.1.0 # via rocm-docs-core +sphinx-multitoc-numbering==0.1.3 + # via sphinx-external-toc sphinx-notfound-page==1.1.0 # via rocm-docs-core sphinxcontrib-applehelp==2.0.0 @@ -260,25 +266,27 @@ sphinxcontrib-htmlhelp==2.1.0 # via sphinx sphinxcontrib-jsmath==1.0.1 # via sphinx +sphinxcontrib-mermaid==2.0.2 + # via -r requirements.in sphinxcontrib-qthelp==2.0.0 # via sphinx sphinxcontrib-serializinghtml==2.0.0 # via sphinx -sqlalchemy==2.0.45 +sqlalchemy==2.0.50 # via jupyter-cache stack-data==0.6.3 # via ipython -tabulate==0.9.0 +tabulate==0.10.0 # via jupyter-cache -tomli==2.4.0 +tomli==2.4.1 # via sphinx -tornado==6.5.5 +tornado==6.5.6 # via # ipykernel # jupyter-client tqdm==4.67.3 # via mpire -traitlets==5.14.3 +traitlets==5.15.1 # via # ipykernel # ipython @@ -296,13 +304,14 @@ typing-extensions==4.15.0 # myst-nb # pydata-sphinx-theme # pygithub + # pyjwt # referencing # sqlalchemy urllib3==2.7.0 # via # pygithub # requests -wcwidth==0.2.14 +wcwidth==0.7.0 # via prompt-toolkit -zipp==3.23.0 +zipp==4.1.0 # via importlib-metadata diff --git a/example/01_gemm/CMakeLists.txt b/example/01_gemm/CMakeLists.txt index bc2e6a78e7c..83aea5a33b9 100644 --- a/example/01_gemm/CMakeLists.txt +++ b/example/01_gemm/CMakeLists.txt @@ -25,8 +25,6 @@ add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp16) add_example_executable(example_gemm_xdl_fp16_v2 gemm_xdl_fp16_v2.cpp) add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp16_v2) -add_example_executable(example_gemm_xdl_fp16_streamk_v3 gemm_xdl_fp16_streamk_v3.cpp) -add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp16_streamk_v3) add_example_executable(example_gemm_xdl_fp16_v3 gemm_xdl_fp16_v3.cpp) add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp16_v3) add_example_executable(example_gemm_xdl_fp8_v3 gemm_xdl_fp8_v3.cpp) @@ -35,10 +33,6 @@ add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp8_v3) add_example_executable(example_gemm_xdl_fp16_fp8_v3 gemm_xdl_fp16_fp8_v3.cpp) add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp16_fp8_v3) - -add_example_executable(example_gemm_xdl_fp16_fp8_streamk_v3 gemm_xdl_fp16_fp8_streamk_v3.cpp) -add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp16_fp8_streamk_v3) - add_example_executable(example_gemm_xdl_bf16_v3 gemm_xdl_bf16_v3.cpp) add_example_dependencies(example_gemm_xdl example_gemm_xdl_bf16_v3) @@ -80,8 +74,6 @@ endif(USE_BITINT_EXTENSION_INT4) add_example_executable(example_gemm_xdl_fp64 gemm_xdl_fp64.cpp) add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp64) -add_example_executable(example_gemm_xdl_streamk gemm_xdl_streamk.cpp) - list(APPEND gpu_list gfx90a gfx942 gfx950 gfx1250) set(target 0) foreach(gpu IN LISTS GPU_TARGETS) @@ -95,19 +87,6 @@ foreach(gpu IN LISTS GPU_TARGETS) endif() endforeach() -list(APPEND gpu_list gfx90a gfx942 gfx950 gfx1200 gfx1201 gfx12-generic gfx1250) -set(target 0) -foreach(gpu IN LISTS GPU_TARGETS) - if(gpu IN_LIST gpu_list AND target EQUAL 0) - add_example_executable(example_gemm_xdl_bf16_streamk_v3 gemm_xdl_bf16_streamk_v3.cpp) - add_example_dependencies(example_gemm_xdl example_gemm_xdl_bf16_streamk_v3) - - add_example_executable(example_gemm_xdl_fp8_streamk_v3 gemm_xdl_fp8_streamk_v3.cpp) - add_example_dependencies(example_gemm_xdl example_gemm_xdl_fp8_streamk_v3) - set(target 1) - endif() -endforeach() - list(APPEND gpu_list_tf32 gfx942 gfx950) set(target 0) foreach(gpu IN LISTS GPU_TARGETS) diff --git a/example/01_gemm/README.md b/example/01_gemm/README.md index ae0e918b8d0..7e82067df2c 100644 --- a/example/01_gemm/README.md +++ b/example/01_gemm/README.md @@ -162,8 +162,6 @@ Split-K is supported (requires zeroing output buffer if splitK > 1). - **DeviceGemmMultipleDLayernorm**: GEMM fused with layernorm - **DeviceGemmMultipleDMultipleR**: GEMM fused with reductions and custom global reductions - **DeviceGemmReduce**: GEMM fused with reduction -- **DeviceGemm_Streamk_V2**: Stream K with reduction instead of AtomicAdd -- **DeviceGemmStreamK**: Stream K using AtomicAdd --- diff --git a/example/01_gemm/common.hpp b/example/01_gemm/common.hpp index a7dca891fd5..07eed72aa50 100644 --- a/example/01_gemm/common.hpp +++ b/example/01_gemm/common.hpp @@ -42,33 +42,6 @@ struct ProblemSize final ck::index_t StrideC = -1; }; -struct ProblemSizeStreamK final -{ - ck::index_t M = 3840; - ck::index_t N = 4096; - ck::index_t K = 4096; - - ck::index_t StrideA = -1; - ck::index_t StrideB = -1; - ck::index_t StrideC = -1; - - ck::index_t NumSKBlocks = -1; // number of stream-k blocks -}; -struct ProblemSizeStreamK_universal final -{ - ck::index_t M = 3840; - ck::index_t N = 4096; - ck::index_t K = 4096; - - ck::index_t StrideA = -1; - ck::index_t StrideB = -1; - ck::index_t StrideC = -1; - - ck::index_t Grid_size = -1; // defaults to max occupancy - ck::index_t Streamk_sel = 1; // defaults to 1-tile SK - ck::StreamKReductionStrategy reduction_strategy = ck::StreamKReductionStrategy::Atomic; -}; - struct ProblemSizeSplitK final { ck::index_t M = 3840; @@ -148,123 +121,6 @@ bool parse_cmd_args(int argc, return true; } -template <> -bool parse_cmd_args(int argc, - char* argv[], - ProblemSizeStreamK_universal& problem_size, - ExecutionConfig& config) -{ - if(argc == 1) - { - // use default case - } - else if(argc == 4) - { - config.do_verification = std::stoi(argv[1]); - config.init_method = std::stoi(argv[2]); - config.time_kernel = std::stoi(argv[3]); - } - else if(argc >= 10) - { - config.do_verification = std::stoi(argv[1]); - config.init_method = std::stoi(argv[2]); - config.time_kernel = std::stoi(argv[3]); - - problem_size.M = std::stoi(argv[4]); - problem_size.N = std::stoi(argv[5]); - problem_size.K = std::stoi(argv[6]); - - problem_size.StrideA = std::stoi(argv[7]); - problem_size.StrideB = std::stoi(argv[8]); - problem_size.StrideC = std::stoi(argv[9]); - - if(argc >= 11) - { - problem_size.Streamk_sel = std::stoi(argv[10]); - - if(argc >= 12) - { - problem_size.Grid_size = std::stoi(argv[11]); - - if(argc >= 13) - { - int reduction_strategy = std::stoi(argv[12]); - problem_size.reduction_strategy = reduction_strategy == 0 - ? ck::StreamKReductionStrategy::Atomic - : ck::StreamKReductionStrategy::Reduction; - } - } - } - } - else - { - std::cerr - << "arg1: verification (0=no, 1=CPU, 2=GPU, 3=CPU and GPU)" << std::endl - << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)" << std::endl - << "arg3: time kernel (0=no, 1=yes)" << std::endl - << "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC (default: -1 or 0)" - << std::endl - << "arg10: stream-k select (-1: default config, 0: all DP, 1: 1-tile SK, 2: 2-tile SK)" - << std::endl - << "arg11: Grid_size(-1 for max occupancy)" << std::endl - << "arg12: Reduction strategy (0: Atomic, 1: Reduction)" << std::endl; - return false; - } - - return true; -} - -template <> -bool parse_cmd_args(int argc, - char* argv[], - ProblemSizeStreamK& problem_size, - ExecutionConfig& config) -{ - if(argc == 1) - { - // use default case - } - else if(argc == 4) - { - config.do_verification = std::stoi(argv[1]); - config.init_method = std::stoi(argv[2]); - config.time_kernel = std::stoi(argv[3]); - } - else if(argc >= 10) - { - config.do_verification = std::stoi(argv[1]); - config.init_method = std::stoi(argv[2]); - config.time_kernel = std::stoi(argv[3]); - - problem_size.M = std::stoi(argv[4]); - problem_size.N = std::stoi(argv[5]); - problem_size.K = std::stoi(argv[6]); - - problem_size.StrideA = std::stoi(argv[7]); - problem_size.StrideB = std::stoi(argv[8]); - problem_size.StrideC = std::stoi(argv[9]); - - if(argc >= 11) - { - problem_size.NumSKBlocks = std::stoi(argv[10]); - } - } - else - { - std::cerr - << "arg1: verification (0=no, 1=CPU, 2=GPU, 3=CPU and GPU)" << std::endl - << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)" << std::endl - << "arg3: time kernel (0=no, 1=yes)" << std::endl - << "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC (default: -1 or 0)" - << std::endl - << "arg10: stream-k select (0: all DP, 1: 1-tile SK, 2: 2-tile SK)" - << "\narg11: Grid_size(-1 for max occupancy)" << std::endl; - return false; - } - - return true; -} - template <> bool parse_cmd_args(int argc, char* argv[], diff --git a/example/01_gemm/gemm_xdl_bf16_streamk_v3.cpp b/example/01_gemm/gemm_xdl_bf16_streamk_v3.cpp deleted file mode 100644 index 754cc8f6f53..00000000000 --- a/example/01_gemm/gemm_xdl_bf16_streamk_v3.cpp +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#include "common.hpp" - -#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_streamk_v3.hpp" - -using ADataType = ck::bhalf_t; -using BDataType = ck::bhalf_t; -using CDataType = ck::bhalf_t; -using AccDataType = float; -using CShuffleDataType = ck::bhalf_t; - -using ALayout = Row; -using BLayout = Col; -using CLayout = Row; - -using AElementOp = PassThrough; -using BElementOp = PassThrough; -using CElementOp = PassThrough; - -static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default; - -// clang-format off -using DeviceGemmV2_Streamk_Instance = - ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle_Streamk_V3< - ALayout, BLayout, CLayout, - ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, - PassThrough, PassThrough, PassThrough, GemmDefault, - 256, - 128, 128, - 64, 8, 8, - 16, 16, - 4, 4, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 8, 8, 0, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 8, 8, 0, - 1, 2, S<1, 32, 1, 8>, 8, - ck::BlockGemmPipelineScheduler::Intrawave,ck::BlockGemmPipelineVersion::v3>; -// clang-format on - -using ReferenceGemmInstance = ck::tensor_operation::host:: - ReferenceGemm; - -using ReferenceGemmInstanceGPU = ck::tensor_operation::device::ReferenceGemm; - -#include "run_gemm_example_streamk_v2.inc" - -int main(int argc, char* argv[]) { return !run_gemm_universal_streamk_example(argc, argv); } diff --git a/example/01_gemm/gemm_xdl_fp16_fp8_streamk_v3.cpp b/example/01_gemm/gemm_xdl_fp16_fp8_streamk_v3.cpp deleted file mode 100644 index e7c00610742..00000000000 --- a/example/01_gemm/gemm_xdl_fp16_fp8_streamk_v3.cpp +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#include "common.hpp" - -#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_streamk_v3.hpp" - -using ADataType = ck::half_t; -using BDataType = ck::f8_t; -using AccDataType = float; -using CShuffleDataType = ck::half_t; -using CDataType = ck::half_t; - -using ALayout = Row; -using BLayout = Col; -using CLayout = Row; - -using AElementOp = PassThrough; -using BElementOp = PassThrough; -using CElementOp = PassThrough; - -static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default; - -// clang-format off -using DeviceGemmV2_Streamk_Instance = - ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle_Streamk_V3< - ALayout, BLayout, CLayout, - ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, - AElementOp, BElementOp, CElementOp, GemmDefault, - 64, - 32, 32, - 256, 8, 16, - 16, 16, - 2, 2, - S<32, 2, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 8, 8, 0, - S<16, 4, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 16, 16, 0, - 1, 1, S<1, 16, 1, 4>, 4, - ck::BlockGemmPipelineScheduler::Interwave, ck::BlockGemmPipelineVersion::v1>; -// clang-format on - -using ReferenceGemmInstanceGPU = ck::tensor_operation::device::ReferenceGemm; - -using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm; - -#include "run_gemm_example_streamk_v2.inc" - -int main(int argc, char* argv[]) { return !run_gemm_universal_streamk_example(argc, argv); } diff --git a/example/01_gemm/gemm_xdl_fp16_streamk_v3.cpp b/example/01_gemm/gemm_xdl_fp16_streamk_v3.cpp deleted file mode 100644 index 0997afcdca1..00000000000 --- a/example/01_gemm/gemm_xdl_fp16_streamk_v3.cpp +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#include "common.hpp" - -#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_streamk_v3.hpp" - -using ADataType = ck::half_t; -using BDataType = ck::half_t; -using AccDataType = float; -using CShuffleDataType = float; -using CDataType = ck::half_t; - -using ALayout = Row; -using BLayout = Row; -using CLayout = Row; - -using AElementOp = PassThrough; -using BElementOp = PassThrough; -using CElementOp = PassThrough; - -static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::MNPadding; - -// clang-format off -using DeviceGemmV2_Streamk_Instance = - ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle_Streamk_V3< - ALayout, BLayout, CLayout, - ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, - PassThrough, PassThrough, PassThrough, GemmDefault, - 256, - 224, 256, - 64, 8, 2, - 16, 16, - 7, 8, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 8, 8, 0, - S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, - 1, 8, 2, 0, - 1, 2, S<1, 32, 1, 8>, 8, - ck::BlockGemmPipelineScheduler::Intrawave,ck::BlockGemmPipelineVersion::v3>; -// clang-format on - -using ReferenceGemmInstance = ck::tensor_operation::host:: - ReferenceGemm; - -using ReferenceGemmInstanceGPU = ck::tensor_operation::device::ReferenceGemm; - -#include "run_gemm_example_streamk_v2.inc" - -int main(int argc, char* argv[]) { return !run_gemm_universal_streamk_example(argc, argv); } diff --git a/example/01_gemm/gemm_xdl_fp8_streamk_v3.cpp b/example/01_gemm/gemm_xdl_fp8_streamk_v3.cpp deleted file mode 100644 index e4a01c2c13d..00000000000 --- a/example/01_gemm/gemm_xdl_fp8_streamk_v3.cpp +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#include "common.hpp" - -#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_streamk_v3.hpp" - -using ADataType = ck::f8_t; -using BDataType = ck::f8_t; -using AccDataType = float; -using CShuffleDataType = ck::half_t; -using CDataType = ck::half_t; - -using ALayout = Row; -using BLayout = Col; -using CLayout = Row; - -using AElementOp = PassThrough; -using BElementOp = PassThrough; -using CElementOp = PassThrough; - -static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default; - -// clang-format off -using DeviceGemmV2_Streamk_Instance = - ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle_Streamk_V3< - ALayout, BLayout, CLayout, - ADataType, BDataType, CDataType, AccDataType, CShuffleDataType, - PassThrough, PassThrough, PassThrough, GemmDefault, - 256, - 128, 256, - 128, 16, 16, - 16, 16, - 4, 8, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 16, 16, 1, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, - 2, 16, 16, 1, - 1, 2, S<1, 32, 1, 8>, 8, - ck::BlockGemmPipelineScheduler::Intrawave,ck::BlockGemmPipelineVersion::v3, ck::f8_t>; -// clang-format on - -using ReferenceGemmInstance = ck::tensor_operation::host:: - ReferenceGemm; -using ReferenceGemmInstanceGPU = ck::tensor_operation::device::ReferenceGemm; - -#include "run_gemm_example_streamk_v2.inc" - -int main(int argc, char* argv[]) { return !run_gemm_universal_streamk_example(argc, argv); } diff --git a/example/01_gemm/gemm_xdl_streamk.cpp b/example/01_gemm/gemm_xdl_streamk.cpp deleted file mode 100644 index caf98c1cff3..00000000000 --- a/example/01_gemm/gemm_xdl_streamk.cpp +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#include "common.hpp" - -#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_streamk.hpp" - -using ADataType = ck::half_t; -using BDataType = ck::half_t; -using AccDataType = float; -using CShuffleDataType = float; -using CDataType = ck::half_t; - -using F16 = ck::half_t; - -using ALayout = Row; -using BLayout = Row; -using CLayout = Row; - -using AElementOp = PassThrough; -using BElementOp = PassThrough; -using CElementOp = PassThrough; - -// clang-format off -using DeviceGemmStreamK = ck::tensor_operation::device::DeviceGemmXdlStreamK -// ######| AData| BData| CData| AccData| ALayout| BLayout| CLayout| A| B| C| Block| MPer| NPer| K0Per| K1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| -// ######| Type| Type| Type| Type| | | | Elementwise| Elementwise| Elementwise| Size| Block| Block| Block| | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| -// ######| | | | | | | | Operation| Operation| Operation| | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| -// ######| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, 256, 128, 128, 4, 8, 32, 32, 2, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>; - - // < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, 256, 256, 128, 4, 8, 32, 32, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 1, 1, 1, S<1, 32, 1, 8>, 8>; - // < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, 128, 32, 64, 4, 8, 32, 32, 1, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 16, 1, 8>, 8>; - // < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, 128, 32, 128, 4, 8, 32, 32, 1, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 2, 1, 1, 1, S<1, 32, 1, 4>, 8>; - -// instance for double rate mfma instruction on gfx950 -using DeviceGemmStreamK2 = ck::tensor_operation::device::DeviceGemmXdlStreamK -// ######| AData| BData| CData| AccData| ALayout| BLayout| CLayout| A| B| C| Block| MPer| NPer| K0Per| K1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| -// ######| Type| Type| Type| Type| | | | Elementwise| Elementwise| Elementwise| Size| Block| Block| Block| | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| -// ######| | | | | | | | Operation| Operation| Operation| | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| -// ######| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, 256, 256, 128, 4, 16, 32, 32, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>; - -// clang-format on - -using DeviceGemmInstance = DeviceGemmStreamK; -using DeviceGemmInstance2 = DeviceGemmStreamK2; - -using ReferenceGemmInstance = ck::tensor_operation::host:: - ReferenceGemm; - -using ReferenceGemmInstanceGPU = ck::tensor_operation::device::ReferenceGemm; - -#include "run_gemm_example_streamk.inc" - -int main(int argc, char* argv[]) { return !run_gemm_streamk_example(argc, argv); } diff --git a/example/01_gemm/run_gemm_example_streamk.inc b/example/01_gemm/run_gemm_example_streamk.inc deleted file mode 100644 index 2761ce28e92..00000000000 --- a/example/01_gemm/run_gemm_example_streamk.inc +++ /dev/null @@ -1,270 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#pragma once - -#include "ck/host_utility/device_prop.hpp" -#include "ck/tensor_operation/gpu/device/device_gemm_streamk.hpp" - -template -bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) -{ -#if defined(BUILD_INT4_EXAMPLE) && defined(CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4) - static_assert(sizeof(ck::int4_t) == sizeof(int8_t)); -#endif - - using namespace ck::literals; - - auto M = problem_size.M; - auto N = problem_size.N; - auto K = problem_size.K; - auto StrideA = problem_size.StrideA; - auto StrideB = problem_size.StrideB; - auto StrideC = problem_size.StrideC; - - auto f_host_tensor_descriptor = - [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { - if constexpr(std::is_same_v) - { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); - } - else - { - return HostTensorDescriptor({row, col}, {1_uz, stride}); - } - }; - - auto f_get_default_stride = - [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { - if(stride == -1 || stride == 0) - { - // give a chance if stride is -1, return a default packed stride - if constexpr(std::is_same_v) - { - return static_cast(col); - } - else - { - return static_cast(row); - } - } - else - return static_cast(stride); - }; - - StrideA = f_get_default_stride(M, K, StrideA, ALayout{}); - StrideB = f_get_default_stride(K, N, StrideB, BLayout{}); - StrideC = f_get_default_stride(M, N, StrideC, CLayout{}); - - Tensor a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); - Tensor b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{})); - - switch(config.init_method) - { - case 0: - ck::utils::FillConstant{ck::type_convert(1.f)}(a_m_k); - ck::utils::FillConstant{ck::type_convert(1.f)}(b_k_n); - break; - case 1: - ck::utils::FillUniformDistributionIntegerValue{-5.f, 5.f}(a_m_k); - ck::utils::FillUniformDistributionIntegerValue{-5.f, 5.f}(b_k_n); - break; - case 2: - ck::utils::FillUniformDistribution{-1.f, 1.f}(a_m_k); - ck::utils::FillUniformDistribution{-1.f, 1.f}(b_k_n); - break; - case 3: - ck::utils::FillUniformDistributionIntegerValue{1.f, 1.f}(a_m_k); - ck::utils::FillUniformDistributionIntegerValue{-5.f, 5.f}(b_k_n); - break; - case 4: - ck::utils::FillUniformDistributionIntegerValue{-5.f, 5.f}(a_m_k); - ck::utils::FillUniformDistributionIntegerValue{1.f, 1.f}(b_k_n); - break; - case 5: - ck::utils::FillUniformDistributionIntegerValue{-2.f, 2.f}(a_m_k); - ck::utils::FillUniformDistributionIntegerValue{-2.f, 2.f}(b_k_n); - break; - default: - ck::utils::FillUniformDistribution{-0.1f, 0.1f}(a_m_k); - ck::utils::FillUniformDistribution{-0.1f, 0.1f}(b_k_n); - } - - Tensor c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); - Tensor c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); - Tensor c_m_n_device_ref_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); - - std::cout << "a_m_k: " << a_m_k.mDesc << std::endl; - std::cout << "b_k_n: " << b_k_n.mDesc << std::endl; - std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl; - -#ifdef BUILD_INT4_EXAMPLE - DeviceMem a_m_k_device_buf(sizeof(KernelADataType) * a_m_k.mDesc.GetElementSpaceSize()); - DeviceMem b_k_n_device_buf(sizeof(KernelBDataType) * b_k_n.mDesc.GetElementSpaceSize()); - DeviceMem c_m_n_device_buf(sizeof(KernelCDataType) * - c_m_n_device_result.mDesc.GetElementSpaceSize()); - - const Tensor a_m_k_converted(a_m_k); - const Tensor b_k_n_converted(b_k_n); - - a_m_k_device_buf.ToDevice(a_m_k_converted.mData.data()); - b_k_n_device_buf.ToDevice(b_k_n_converted.mData.data()); -#else - DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize()); - DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize()); - DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize()); - DeviceMem c_m_n_device_ref_buf(sizeof(CDataType) * - c_m_n_device_ref_result.mDesc.GetElementSpaceSize()); - - a_m_k_device_buf.ToDevice(a_m_k.mData.data()); - b_k_n_device_buf.ToDevice(b_k_n.mData.data()); -#endif - DeviceMem workspace; - - auto a_element_op = AElementOp{}; - auto b_element_op = BElementOp{}; - auto c_element_op = CElementOp{}; - - using BaseStreamK = ck::tensor_operation::device::DeviceGemmStreamK; - - // do GEMM - static_assert(std::is_base_of::value && - std::is_base_of::value); - auto gemm = DeviceGemmInstance{}; - auto gemm2 = DeviceGemmInstance2{}; // instance for double rate mfma instruction - BaseStreamK* op_ptr = (ck::get_device_name() == "gfx950") ? static_cast(&gemm2) - : static_cast(&gemm); - - float ave_time = 0; - auto invoker_ptr = op_ptr->MakeInvokerPointer(); - - auto argument_ptr = op_ptr->MakeArgumentPointer( -#ifdef BUILD_INT4_EXAMPLE - static_cast(a_m_k_device_buf.GetDeviceBuffer()), - static_cast(b_k_n_device_buf.GetDeviceBuffer()), - static_cast(c_m_n_device_buf.GetDeviceBuffer()), -#else - static_cast(a_m_k_device_buf.GetDeviceBuffer()), - static_cast(b_k_n_device_buf.GetDeviceBuffer()), - static_cast(c_m_n_device_buf.GetDeviceBuffer()), -#endif - M, - N, - K, - StrideA, - StrideB, - StrideC, - a_element_op, - b_element_op, - c_element_op, - problem_size.NumSKBlocks); - - if(!op_ptr->IsSupportedArgument(argument_ptr.get())) - { - std::cerr << op_ptr->GetTypeString() << " does not support this problem" << std::endl; - - return true; - } - - auto argument = argument_ptr.get(); - std::size_t workspace_size = op_ptr->GetWorkSpaceSize(argument); - if(workspace_size != 0) - { - workspace.Realloc(workspace_size); - op_ptr->SetWorkSpacePointer(argument, workspace.GetDeviceBuffer()); - } - - ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, config.time_kernel}); - - std::size_t flop = 2_uz * M * N * K; - std::size_t num_btype = - sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N; - - float tflops = static_cast(flop) / 1.E9 / ave_time; - - float gb_per_sec = num_btype / 1.E6 / ave_time; - - std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " - << op_ptr->GetTypeString() << std::endl; - - bool pass = true; - - if((config.do_verification == 1) || (config.do_verification == 3)) - { - // CPU verification - auto ref_gemm = ReferenceGemmInstance{}; - auto ref_invoker = ref_gemm.MakeInvoker(); - - auto ref_argument = ref_gemm.MakeArgument( - a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op); - - std::cout << "Running verification on CPU." << std::endl; - ref_invoker.Run(ref_argument); - -#ifdef BUILD_INT4_EXAMPLE - Tensor c_m_n_device_result_converted(c_m_n_host_result.mDesc); - - c_m_n_device_buf.FromDevice(c_m_n_device_result_converted.mData.data()); - - c_m_n_device_result = c_m_n_device_result_converted.CopyAsType(); - - return ck::utils::check_err(c_m_n_device_result_converted, c_m_n_host_result); -#else - c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data()); - - pass &= ck::utils::check_err(c_m_n_device_result, - c_m_n_host_result, - "Error: Incorrect results!", - get_rtol(), - get_atol()); -#endif - } - - if((config.do_verification == 2) || (config.do_verification == 3)) - { - // GPU verification - auto ref_gemm_gpu = ReferenceGemmInstanceGPU{}; - auto ref_invoker_gpu = ref_gemm_gpu.MakeInvoker(); - - auto ref_argument_gpu = ref_gemm_gpu.MakeArgument( - static_cast(a_m_k_device_buf.GetDeviceBuffer()), - static_cast(b_k_n_device_buf.GetDeviceBuffer()), - static_cast(c_m_n_device_ref_buf.GetDeviceBuffer()), - M, - N, - K, - a_element_op, - b_element_op, - c_element_op); - - std::cout << "Running verification on GPU." << std::endl; - ref_invoker_gpu.Run(ref_argument_gpu, StreamConfig{}); - - c_m_n_device_ref_buf.FromDevice(c_m_n_device_ref_result.mData.data()); - c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data()); - - pass &= ck::utils::check_err(c_m_n_device_result, - c_m_n_device_ref_result, - "Error: Incorrect results!", - get_rtol(), - get_atol()); - } - - return pass == true; -} - -bool run_gemm_streamk_example(int argc, char* argv[]) -{ - ProblemSizeStreamK problem_size; - ExecutionConfig config; - - return !parse_cmd_args(argc, argv, problem_size, config) || run_gemm(problem_size, config); -} diff --git a/example/01_gemm/run_gemm_example_streamk_v2.inc b/example/01_gemm/run_gemm_example_streamk_v2.inc deleted file mode 100644 index 4416c601785..00000000000 --- a/example/01_gemm/run_gemm_example_streamk_v2.inc +++ /dev/null @@ -1,270 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#pragma once - -template -bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) -{ -#if defined(BUILD_INT4_EXAMPLE) && defined(CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4) - static_assert(sizeof(ck::int4_t) == sizeof(int8_t)); -#endif - - using namespace ck::literals; - - auto M = problem_size.M; - auto N = problem_size.N; - auto K = problem_size.K; - auto StrideA = problem_size.StrideA; - auto StrideB = problem_size.StrideB; - auto StrideC = problem_size.StrideC; - auto Grid_size = problem_size.Grid_size; - auto Streamk_sel = problem_size.Streamk_sel; - - auto reduction_strategy = problem_size.reduction_strategy; - if(reduction_strategy == ck::StreamKReductionStrategy::Atomic) - { - std::cout << "Using Atomic reduction strategy" << std::endl; - } - else - { - std::cout << "Using Parallel reduction strategy" << std::endl; - } - - auto f_host_tensor_descriptor = - [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { - if constexpr(std::is_same_v) - { - return HostTensorDescriptor({row, col}, {stride, 1_uz}); - } - else - { - return HostTensorDescriptor({row, col}, {1_uz, stride}); - } - }; - - auto f_get_default_stride = - [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { - if(stride == -1 || stride == 0) - { - // give a chance if stride is -1, return a default packed stride - if constexpr(std::is_same_v) - { - return static_cast(col); - } - else - { - return static_cast(row); - } - } - else - return static_cast(stride); - }; - - auto f_get_default_streamk_policy = [](ck::index_t streamk_sel) { - if(streamk_sel == -1) - { - return static_cast(4); - } - else - return static_cast(streamk_sel); - }; - - StrideA = f_get_default_stride(M, K, StrideA, ALayout{}); - StrideB = f_get_default_stride(K, N, StrideB, BLayout{}); - StrideC = f_get_default_stride(M, N, StrideC, CLayout{}); - - Streamk_sel = f_get_default_streamk_policy(Streamk_sel); - - Tensor a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); - Tensor b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{})); - - switch(config.init_method) - { - case 0: - a_m_k.GenerateTensorValue(GeneratorTensor_1{1}); - b_k_n.GenerateTensorValue(GeneratorTensor_1{1}); - break; - case 1: - a_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 2}); - b_k_n.GenerateTensorValue(GeneratorTensor_2{-2, 2}); - break; - case 2: - a_m_k.GenerateTensorValue(GeneratorTensor_1{1}); - b_k_n.GenerateTensorValue(GeneratorTensor_2{-2, 2}); - break; - case 3: - a_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 2}); - b_k_n.GenerateTensorValue(GeneratorTensor_1{1}); - break; - default: - a_m_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); - b_k_n.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); - } - - Tensor c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); - Tensor c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); - Tensor c_m_n_device_ref_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{})); - - std::cout << "a_m_k: " << a_m_k.mDesc << std::endl; - std::cout << "b_k_n: " << b_k_n.mDesc << std::endl; - std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl; - -#ifdef BUILD_INT4_EXAMPLE - DeviceMem a_m_k_device_buf(sizeof(KernelADataType) * a_m_k.mDesc.GetElementSpaceSize()); - DeviceMem b_k_n_device_buf(sizeof(KernelBDataType) * b_k_n.mDesc.GetElementSpaceSize()); - DeviceMem c_m_n_device_buf(sizeof(KernelCDataType) * - c_m_n_device_result.mDesc.GetElementSpaceSize()); - - const Tensor a_m_k_converted(a_m_k); - const Tensor b_k_n_converted(b_k_n); - - a_m_k_device_buf.ToDevice(a_m_k_converted.mData.data()); - b_k_n_device_buf.ToDevice(b_k_n_converted.mData.data()); -#else - DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize()); - DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize()); - DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize()); - DeviceMem c_m_n_device_ref_buf(sizeof(CDataType) * - c_m_n_device_ref_result.mDesc.GetElementSpaceSize()); - - a_m_k_device_buf.ToDevice(a_m_k.mData.data()); - b_k_n_device_buf.ToDevice(b_k_n.mData.data()); -#endif - DeviceMem workspace; - - auto a_element_op = AElementOp{}; - auto b_element_op = BElementOp{}; - auto c_element_op = CElementOp{}; - - // do GEMM - auto gemm = DeviceGemmV2_Streamk_Instance{}; - auto invoker = gemm.MakeInvoker(); - float ave_time = 0; - - auto argument = gemm.MakeArgument( -#ifdef BUILD_INT4_EXAMPLE - static_cast(a_m_k_device_buf.GetDeviceBuffer()), - static_cast(b_k_n_device_buf.GetDeviceBuffer()), - static_cast(c_m_n_device_buf.GetDeviceBuffer()), -#else - static_cast(a_m_k_device_buf.GetDeviceBuffer()), - static_cast(b_k_n_device_buf.GetDeviceBuffer()), - static_cast(c_m_n_device_buf.GetDeviceBuffer()), -#endif - M, - N, - K, - StrideA, - StrideB, - StrideC, - Streamk_sel, - Grid_size, - a_element_op, - b_element_op, - c_element_op, - reduction_strategy); - - if(!gemm.IsSupportedArgument(argument)) - { - std::cerr << gemm.GetTypeString() << " does not support this problem" << std::endl; - - return true; - } - - std::size_t workspace_size = gemm.GetWorkSpaceSize(&argument); - if(workspace_size != 0) - { - workspace.Realloc(workspace_size); - gemm.SetWorkSpacePointer(&argument, workspace.GetDeviceBuffer()); - } - - bool pass = true; - if((config.do_verification == 1) || (config.do_verification == 3)) - { - auto ref_gemm = ReferenceGemmInstance{}; - auto ref_invoker = ref_gemm.MakeInvoker(); - - auto ref_argument = ref_gemm.MakeArgument( - a_m_k, b_k_n, c_m_n_host_result, PassThrough{}, PassThrough{}, PassThrough{}); - - ref_invoker.Run(ref_argument); - - ave_time = invoker.Run(argument, StreamConfig{nullptr, false, 1}); -#ifdef BUILD_INT4_EXAMPLE - Tensor c_m_n_device_result_converted(c_m_n_host_result.mDesc); - - c_m_n_device_buf.FromDevice(c_m_n_device_result_converted.mData.data()); - - c_m_n_device_result = c_m_n_device_result_converted.CopyAsType(); - - return ck::utils::check_err(c_m_n_device_result_converted, c_m_n_host_result); -#else - c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data()); - - pass &= ck::utils::check_err(c_m_n_device_result, - c_m_n_host_result, - "Error: Incorrect results!", - get_rtol(), - get_atol()); -#endif - } - - if((config.do_verification == 2) || (config.do_verification == 3)) - { - // GPU verification - auto ref_gemm_gpu = ReferenceGemmInstanceGPU{}; - auto ref_invoker_gpu = ref_gemm_gpu.MakeInvoker(); - - auto ref_argument_gpu = ref_gemm_gpu.MakeArgument( - static_cast(a_m_k_device_buf.GetDeviceBuffer()), - static_cast(b_k_n_device_buf.GetDeviceBuffer()), - static_cast(c_m_n_device_ref_buf.GetDeviceBuffer()), - M, - N, - K, - a_element_op, - b_element_op, - c_element_op); - - std::cout << "Running verification on GPU." << std::endl; - ref_invoker_gpu.Run(ref_argument_gpu, StreamConfig{}); - - c_m_n_device_ref_buf.FromDevice(c_m_n_device_ref_result.mData.data()); - c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data()); - - pass &= ck::utils::check_err(c_m_n_device_result, - c_m_n_device_ref_result, - "Error: Incorrect results!", - get_rtol(), - get_atol()); - } - - if(config.time_kernel) - { - ave_time = invoker.Run(argument, StreamConfig{nullptr, config.time_kernel}); - - std::size_t flop = 2_uz * M * N * K; - std::size_t num_btype = - sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N; - - float tflops = static_cast(flop) / 1.E9 / ave_time; - - float gb_per_sec = num_btype / 1.E6 / ave_time; - - std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec - << " GB/s, " << gemm.GetTypeString() - << (reduction_strategy == ck::StreamKReductionStrategy::Atomic ? " (Atomic)" - : " (Reduction)") - << std::endl; - } - return pass; -} - -bool run_gemm_universal_streamk_example(int argc, char* argv[]) -{ - ProblemSizeStreamK_universal problem_size; - ExecutionConfig config; - - return !parse_cmd_args(argc, argv, problem_size, config) || run_gemm(problem_size, config); -} diff --git a/example/32_batched_gemm_scale_softmax_gemm/grouped_query_attention_forward_wmma_fp16.cpp b/example/32_batched_gemm_scale_softmax_gemm/grouped_query_attention_forward_wmma_fp16.cpp index 66b2aa8508f..4b714a5f9eb 100644 --- a/example/32_batched_gemm_scale_softmax_gemm/grouped_query_attention_forward_wmma_fp16.cpp +++ b/example/32_batched_gemm_scale_softmax_gemm/grouped_query_attention_forward_wmma_fp16.cpp @@ -3,7 +3,7 @@ /* Grouped Query Attention, -Ainslie, Joshua, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit +Ainslie, Joshua, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints." arXiv, May 22, 2023. https://doi.org/10.48550/arXiv.2305.13245. diff --git a/example/64_fpAintB_gemm/fp16int8_gemm_wmma.cpp b/example/64_fpAintB_gemm/fp16int8_gemm_wmma.cpp index 450d1b643f9..dd7b43fff64 100644 --- a/example/64_fpAintB_gemm/fp16int8_gemm_wmma.cpp +++ b/example/64_fpAintB_gemm/fp16int8_gemm_wmma.cpp @@ -6,7 +6,7 @@ #include "ck/tensor_operation/gpu/device/impl/device_fpAintB_gemm_wmma.hpp" // Implementation follows the paper: -// Kim, Young Jin, Rawn Henry, Raffy Fahim, and Hany Hassan Awadalla. "Who Says Elephants Can’t Run: +// Kim, Young Jin, Rawn Henry, Raffy Fahim, and Hany Hassan Awadalla. "Who Says Elephants Can't Run: // Bringing Large Scale MoE Models into Cloud Scale Production." arXiv, November 17, 2022. // https://doi.org/10.48550/arXiv.2211.10017. Assume weight (Matrix B) is add preprocess to // unsigned. diff --git a/example/ck_tile/01_fmha/CMakeLists.txt b/example/ck_tile/01_fmha/CMakeLists.txt index 0650bd3de01..2d44b449968 100644 --- a/example/ck_tile/01_fmha/CMakeLists.txt +++ b/example/ck_tile/01_fmha/CMakeLists.txt @@ -2,8 +2,8 @@ # SPDX-License-Identifier: MIT set(INST_TARGETS ${SUPPORTED_GPU_TARGETS}) -# Currently only gfx9 and gfx12 archs are supported by FMHA -list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx12") +# Currently only gfx9, gfx11, and gfx12 archs are supported by FMHA +list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx1[12]") if(NOT INST_TARGETS) message(WARNING "Skipping Tile Engine FMHA compilation: No supported GPU targets (gfx9, gfx11, gfx12) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") return() diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index 8079b3d8581..dae78e243ce 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -28,6 +28,13 @@ FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n // auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && \\ + (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || \\ + defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || \\ + defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__)) +#undef CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT +#define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1 +#endif #include "fmha_bwd.hpp" """ diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index 849f463afa0..ed025dcf5fc 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -128,7 +128,7 @@ namespace {{ template void run_instance(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) {{ - if constexpr ({F_hdim} == 128 && {F_bias} == ck_tile::BlockAttentionBiasEnum::NO_BIAS + if constexpr ({F_bias} == ck_tile::BlockAttentionBiasEnum::NO_BIAS && (std::is_same_v<{F_mask}, ck_tile::SimplifiedGenericAttentionMask> || std::is_same_v<{F_mask}, FmhaMasks::NoMask>)) {{ if (a.max_seqlen_q == 1 && a.nhead_k < a.nhead_q) {{ @@ -283,7 +283,7 @@ """ FMHA_FWD_SPLITKV_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) && - ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ + ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_seqtune}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv},{F_sink}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; // get combine kernel tile sizes @@ -364,6 +364,14 @@ def scheck(self) -> str: else: assert False + def seqtune(self, max_bm0: int) -> str: + if self.bm0 == max_bm0: + return "true/*fall back to largest tile*/" + else: + if self.mode == "group": + return f"a.max_seqlen_q <= {self.bm0}" + return f"a.seqlen_q <= {self.bm0}" + @property def skcheck(self) -> str: if self.mode == "group": @@ -561,6 +569,7 @@ def api(self) -> str: for i_dtype, (dtype, pool_by_dtype) in enumerate(pool_by_arch.items()): per_hdim_case = str() for i_hdim, (hdim, pool_by_hdim) in enumerate(pool_by_dtype.items()): + max_bm0 = max((t.bm0 for t in pool_by_hdim), default=0) inners = str() for i_trait, trait in enumerate(pool_by_hdim): inners += FMHA_FWD_SPLITKV_API_INNER_DISPATCH.format( @@ -579,6 +588,7 @@ def api(self) -> str: F_pagedkv=BOOL_MAP[trait.pagedkv], F_sink=BOOL_MAP[trait.sink], F_scheck=trait.scheck, + F_seqtune=trait.seqtune(max_bm0), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, @@ -763,6 +773,7 @@ def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdSplitKVPipeline]: pipelines.append(Pipeline("qr", "row", "t", "f", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip pipelines.append(Pipeline("qr", "row", "t", "t", "t", "t", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip + pipelines.append(Pipeline("qr_nwarp_sshuffle", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip elif dtype in ["fp8", "bf8"]: for logits, mask, bias in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys() @@ -846,11 +857,15 @@ class KernelComponentFactoryGfx11(KernelComponentFactoryBase): def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: if dtype in ["fp16", "bf16"]: return { - # bm0, bn0, bk0, bn1, bk1, - "32" : FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "64" : FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # bm0, bn0, bk0, bn1, bk1, + "32" : [FmhaFwdTileSize( 16, 64, 16, 32, 32, 32, 1, 2, 1, 1, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "64" : [FmhaFwdTileSize( 16, 64, 32, 64, 32, 64, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "128": [FmhaFwdTileSize( 16, 64, 32, 128, 32, 128, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "256": [FmhaFwdTileSize( 16, 64, 32, 256, 32, 256, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip else: return None @@ -863,11 +878,15 @@ class KernelComponentFactoryGfx12(KernelComponentFactoryBase): def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: if dtype in ["fp16", "bf16"]: return { - # bm0, bn0, bk0, bn1, bk1, - "32" : FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "64" : FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # bm0, bn0, bk0, bn1, bk1, + "32" : [FmhaFwdTileSize( 16, 64, 16, 32, 32, 32, 1, 2, 1, 1, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "64" : [FmhaFwdTileSize( 16, 64, 32, 64, 32, 64, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "128": [FmhaFwdTileSize( 16, 128, 32, 128, 32, 128, 1, 8, 1, 1, 8, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "256": [FmhaFwdTileSize( 16, 128, 32, 256, 32, 256, 1, 8, 1, 1, 8, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip elif dtype in ["fp8", "bf8"]: return { @@ -930,11 +949,17 @@ def get_fwd_splitkv_blobs( d = factory.get_hdim_tile_size_dict(dtype) if d is None: continue - # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()): - tile = d[hdim_str] + tiles = d[hdim_str] + if not isinstance(tiles, list): + tiles = [tiles] hdim = int(hdim_str) - for pipeline in factory.get_pipelines(dtype, hdim, mask_impl): + for tile, pipeline in itertools.product( + tiles, factory.get_pipelines(dtype, hdim, mask_impl) + ): + # Use qr_nwarp_sshuffle with multiple N warps and qr otherwise + if (tile.F_rn0 != 1) != (pipeline.tag == "qr_nwarp_sshuffle"): + continue if mode == "group": if pipeline.F_spad != "t" or pipeline.F_skpad != "t": # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 0b51dffa466..243ff87faa9 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -165,8 +165,10 @@ int override_num_splits_if_necessary( if(num_splits < 1 && p_drop == 0.0f) { + // props.multiProcessorCount for >=gfx10 is the number of WGPs (each has 2 CUs) + const int num_blocks_per_SM = props.warpSize == 32 ? 4 : 2; return num_splits_heuristic( - batch * nhead * num_m_blocks, props.multiProcessorCount * 2, 128); + batch * nhead * num_m_blocks, props.multiProcessorCount * num_blocks_per_SM, 128); } return num_splits; @@ -648,8 +650,18 @@ fwd_result fmha_fwd_run(mode_enum mode, // legalize num_splits according to other options if(num_splits < 1) { + int nhead_merged = nhead; + int max_seqlen_q_merged = max_seqlen_q; + // When max_seqlen_q == 1 and multiple head groups are merged (kMergeNumHeadGroupsSeqLenQ) + // then more splits are required + if(bias.type == bias_enum::no_bias && mask.type == mask_enum::no_mask && + max_seqlen_q == 1 && nhead_k < nhead) + { + nhead_merged = nhead_k; + max_seqlen_q_merged = max_seqlen_q * (nhead / nhead_k); + } num_splits = override_num_splits_if_necessary( - batch, nhead, max_seqlen_q, hdim_v, p_drop, num_splits); + batch, nhead_merged, max_seqlen_q_merged, hdim_v, p_drop, num_splits); } if(128 < num_splits) { diff --git a/example/ck_tile/03_gemm/CMakeLists.txt b/example/ck_tile/03_gemm/CMakeLists.txt index 85094df6770..81949dd00a2 100644 --- a/example/ck_tile/03_gemm/CMakeLists.txt +++ b/example/ck_tile/03_gemm/CMakeLists.txt @@ -21,8 +21,8 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx125") list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -mllvm -enable-noalias-to-md-conversion=0) list(APPEND EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS -Wno-unused-local-typedef) list(APPEND EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS -Wno-gnu-line-marker) - #list(APPEND EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS --save-temps) - list(APPEND EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1 -mllvm -enable-noalias-to-md-conversion=0") + # list(APPEND EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS --save-temps) + list(APPEND EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1 -mllvm -enable-noalias-to-md-conversion=1") target_compile_options(tile_example_gemm_basic PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) target_compile_options(tile_example_gemm_universal PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) target_compile_options(tile_example_gemm_weight_preshuffle PRIVATE ${EXAMPLE_WEIGHT_PRESHUFFLE_COMPILE_OPTIONS}) diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index 7d6a2adc384..413b30fec8d 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -44,7 +44,6 @@ int run_gemm_example(ck_tile::ArgParser& arg_parser) #ifdef CK_GFX950_SUPPORT else if(data_type == "tf32") { - // Pass tf32_t as A/B types - epilogue auto-detects and maps to float for data operations return run_gemm_example_prec_type static float gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) { - // ADataTypeCompute: compute type (tf32_t for TF32 mode, used for warp gemm selection) - // ADataTypeBuf: buffer/storage type (fp32 when tf32) - using ADataTypeCompute = ADataType_; - using BDataTypeCompute = BDataType_; - using ADataTypeBuf = ck_tile::if_select_t; - using BDataTypeBuf = ck_tile::if_select_t; - - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { - static_assert(std::is_same_v, - "ADataTypeCompute and BDataTypeCompute must be the same"); + static_assert(std::is_same_v, + "ADataType and BDataType must be the same"); } if constexpr(Persistent) @@ -37,12 +30,13 @@ struct BasicInvoker std::cout << "WARNING: Ignoring persistent kernel option for basic gemm." << std::endl; } - constexpr bool is_fp32_input = std::is_same_v; - constexpr bool is_tf32_compute = std::is_same_v; + constexpr bool is_fp32_or_tf32_input = + std::is_same_v || std::is_same_v; + constexpr bool is_tf32_compute = std::is_same_v; // This part comes from the Codegen - constexpr ck_tile::index_t M_Tile = is_fp32_input ? 128 : 256; - constexpr ck_tile::index_t N_Tile = is_fp32_input ? 128 : 256; + constexpr ck_tile::index_t M_Tile = is_fp32_or_tf32_input ? 128 : 256; + constexpr ck_tile::index_t N_Tile = is_fp32_or_tf32_input ? 128 : 256; constexpr ck_tile::index_t K_Tile = 64; #if CK_TILE_USE_WMMA @@ -53,17 +47,19 @@ struct BasicInvoker constexpr ck_tile::index_t M_Warp_Tile = 16; constexpr ck_tile::index_t N_Warp_Tile = 16; constexpr ck_tile::index_t K_Warp_Tile = - ck_tile::get_k_warp_tile(); + ck_tile::get_k_warp_tile(); ck_tile::ignore = is_tf32_compute; #else // gfx950: fp32 uses 16x16x16 tile (native MFMA) // tf32 uses 32x32x16 tile (3x bf16 32x32x16 MFMA emulation) - constexpr ck_tile::index_t M_Warp = (is_fp32_input && !is_tf32_compute) ? 4 : 2; - constexpr ck_tile::index_t N_Warp = (is_fp32_input && !is_tf32_compute) ? 4 : 2; + constexpr ck_tile::index_t M_Warp = (is_fp32_or_tf32_input && !is_tf32_compute) ? 4 : 2; + constexpr ck_tile::index_t N_Warp = (is_fp32_or_tf32_input && !is_tf32_compute) ? 4 : 2; constexpr ck_tile::index_t K_Warp = 1; - constexpr ck_tile::index_t M_Warp_Tile = (is_fp32_input && !is_tf32_compute) ? 16 : 32; - constexpr ck_tile::index_t N_Warp_Tile = (is_fp32_input && !is_tf32_compute) ? 16 : 32; + constexpr ck_tile::index_t M_Warp_Tile = + (is_fp32_or_tf32_input && !is_tf32_compute) ? 16 : 32; + constexpr ck_tile::index_t N_Warp_Tile = + (is_fp32_or_tf32_input && !is_tf32_compute) ? 16 : 32; constexpr ck_tile::index_t K_Warp_Tile = 16; #endif @@ -81,15 +77,15 @@ struct BasicInvoker BLayout, CLayout>; - using AComputeDataType = std:: - conditional_t, BDataType_, ADataType_>; + using AComputeDataType = + std::conditional_t, BDataType, ADataType>; using BComputeDataType = - std::conditional_t || - std::is_same_v, - ADataType_, - BDataType_>; - using CodegenPipelineProblem = ck_tile::GemmPipelineProblem || + std::is_same_v, + ADataType, + BDataType>; + using CodegenPipelineProblem = ck_tile::GemmPipelineProblem; using GemmEpilogue = ck_tile::CShuffleEpilogue< - ck_tile::CShuffleEpilogueProblem, AccDataType, CDataType, @@ -141,7 +137,7 @@ struct BasicInvoker } // Declare rotating_mem_ptr here so it stays in scope until it is needed - std::unique_ptr> rotating_mem_ptr; + std::unique_ptr> rotating_mem_ptr; std::function preprocess; auto clear_gemm_output = [&]() { @@ -154,21 +150,16 @@ struct BasicInvoker { std::cout << "Flushing cache..." << std::endl; - ck_tile::HostTensor a_m(ck_tile::host_tensor_descriptor( + ck_tile::HostTensor a_m(ck_tile::host_tensor_descriptor( args.M, args.K, args.stride_A, is_row_major(ALayout{}))); - ck_tile::HostTensor b_n(ck_tile::host_tensor_descriptor( + ck_tile::HostTensor b_n(ck_tile::host_tensor_descriptor( args.K, args.N, args.stride_B, is_row_major(BLayout{}))); auto size_a_buffer = a_m.get_element_space_size_in_bytes(); auto size_b_buffer = b_n.get_element_space_size_in_bytes(); - rotating_mem_ptr = - std::make_unique>( - kargs.as_ptr[0], - kargs.bs_ptr[0], - s.rotating_count_, - size_a_buffer, - size_b_buffer); + rotating_mem_ptr = std::make_unique>( + kargs.as_ptr[0], kargs.bs_ptr[0], s.rotating_count_, size_a_buffer, size_b_buffer); rotating_mem_ptr->Print(); preprocess = [&]() { diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index 2574b1dbc42..4c2b09c1ddd 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -3,16 +3,15 @@ #pragma once -#include -#include - #include "ck_tile/core.hpp" -#include "ck_tile/core/numeric/pk_fp4.hpp" #include "ck_tile/host/kernel_launch.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" #include "ck_tile/utility/json_dump.hpp" +#include +#include + struct GemmConfigBase { static constexpr bool kPadM = false; @@ -43,10 +42,6 @@ struct GemmConfigBase ck_tile::DataCachePrefetchKind::None; }; -// Type trait for tf32 storage type (tf32 uses float for memory layout calculations) -template -using prec_storage_type = ck_tile::if_select_t; - template struct GemmConfigMemoryInterwave : public GemmConfigBase { @@ -93,7 +88,7 @@ struct GemmConfigComputeV3 : public GemmConfigBase // Compute V3 only support Intrawave scheduler static constexpr ck_tile::index_t M_Tile = 16; static constexpr ck_tile::index_t N_Tile = 64; - static constexpr ck_tile::index_t K_Tile = 256 / sizeof(prec_storage_type); + static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType); static constexpr ck_tile::index_t M_Warp = 1; static constexpr ck_tile::index_t N_Warp = 4; @@ -133,7 +128,7 @@ struct GemmConfigComputeV3_2 : public GemmConfigBase { static constexpr ck_tile::index_t M_Tile = 128; static constexpr ck_tile::index_t N_Tile = 128; - static constexpr ck_tile::index_t K_Tile = 128 / sizeof(prec_storage_type); + static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); static constexpr ck_tile::index_t M_Warp = 2; static constexpr ck_tile::index_t N_Warp = 2; @@ -313,7 +308,7 @@ struct GemmConfigPreshufflePrefill : public GemmConfigBase { static constexpr ck_tile::index_t M_Tile = 128; static constexpr ck_tile::index_t N_Tile = 128; - static constexpr ck_tile::index_t K_Tile = 128 / sizeof(prec_storage_type); + static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); static constexpr ck_tile::index_t M_Warp = 1; static constexpr ck_tile::index_t N_Warp = 4; @@ -322,7 +317,7 @@ struct GemmConfigPreshufflePrefill : public GemmConfigBase static constexpr ck_tile::index_t M_Warp_Tile = 16; static constexpr ck_tile::index_t N_Warp_Tile = 16; static constexpr ck_tile::index_t K_Warp_Tile = - ck_tile::get_k_warp_tile, M_Warp_Tile, true>(); + ck_tile::get_k_warp_tile(); static constexpr int kBlockPerCu = 2; static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; @@ -331,6 +326,19 @@ struct GemmConfigPreshufflePrefill : public GemmConfigBase static constexpr bool DoubleSmemBuffer = true; static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; + + static constexpr bool Async = false; +}; + +template +struct GemmConfigPreshufflePrefillAsync : public GemmConfigPreshufflePrefill +{ + static constexpr ck_tile::index_t N_Tile = 256; + + // N_Repeat is even in this config + static constexpr bool TiledMMAPermuteN = true; + + static constexpr bool Async = true; }; template @@ -360,8 +368,8 @@ struct GemmTypeConfig; template <> struct GemmTypeConfig { - using ADataType = float; - using BDataType = float; + using ADataType = ck_tile::tf32_t; + using BDataType = ck_tile::tf32_t; using AccDataType = float; using CDataType = float; }; diff --git a/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp b/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp index fcd9243bebf..cb335607bc0 100644 --- a/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp +++ b/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp @@ -33,6 +33,8 @@ struct WeightPreshuffleInvoker GemmConfig::TileParitionerGroupNum, GemmConfig::TileParitionerM01>; + static constexpr ck_tile::index_t VectorSize = 16; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + GemmConfig::DataCachePrefetchB, + GemmConfig::Async>; + constexpr auto scheduler = GemmConfig::Scheduler; using AComputeDataType = diff --git a/example/ck_tile/03_gemm/gemm_weight_preshuffle_tdm_data_cache_prefetch.cpp b/example/ck_tile/03_gemm/gemm_weight_preshuffle_tdm_data_cache_prefetch.cpp index 59b938f92d3..6034b84b576 100644 --- a/example/ck_tile/03_gemm/gemm_weight_preshuffle_tdm_data_cache_prefetch.cpp +++ b/example/ck_tile/03_gemm/gemm_weight_preshuffle_tdm_data_cache_prefetch.cpp @@ -182,6 +182,8 @@ struct GemmConfigWeightPreshuffleTDMPrefetch : public GemmConfigBase static constexpr ck_tile::DataCachePrefetchKind DataCachePrefetchB = DataCachePrefetchB_; static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; + + static constexpr bool Async = false; }; int main(int argc, char* argv[]) diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index 2698d594889..54313336a29 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -209,8 +209,6 @@ std::tuple inline parse_ge return std::make_tuple(M, N, K); } -// ADataType_ and BDataType_ are original types (e.g., tf32_t for TF32 mode) -// They are passed through invoke_gemm to invoker for tf32 auto-detection template float mapping for host tensors and device buffers using TypeConfig = GemmTypeConfig; using ADataTypeBuf = typename TypeConfig::ADataType; using BDataTypeBuf = typename TypeConfig::BDataType; @@ -357,8 +349,8 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, float ave_time = invoke_gemm, AccDataType, CDataType, @@ -411,12 +403,12 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, if(arg_parser.get_int("v") == 1) { - ck_tile::reference_gemm( + ck_tile::reference_gemm( a_m_k, b_k_n, c_m_n_ref); const float max_accumulated_value = *std::max_element(c_m_n_ref.mData.begin(), c_m_n_ref.mData.end()); const auto rtol_atol = - calculate_rtol_atol( + calculate_rtol_atol( K, kbatch, max_accumulated_value); pass = do_verify(c_m_n_dev_result, c_m_n_ref, rtol_atol, "CPU"); } @@ -440,8 +432,8 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, BDataTypeBuf* d_B = static_cast(b_k_n_dev_buf.GetDeviceBuffer()); CDataType* d_C = static_cast(c_m_n_gpu_buf_ref.GetDeviceBuffer()); - ck_tile::reference_gemm_gpu( + calculate_rtol_atol( K, kbatch, max_accumulated_value); pass = do_verify(c_m_n_dev_result, c_m_n_ref, rtol_atol, "GPU"); } diff --git a/example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp b/example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp index 2384dc2aa58..29eedb57dbc 100644 --- a/example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp +++ b/example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp @@ -245,7 +245,7 @@ bool run(const ck_tile::ArgParser& arg_parser) if(pass_op) { - std::cout << "✅ valid results for this operation" << std::endl; + std::cout << "[OK] valid results for this operation" << std::endl; } pass &= pass_op; }); diff --git a/example/ck_tile/20_grouped_convolution/CMakeLists.txt b/example/ck_tile/20_grouped_convolution/CMakeLists.txt index 18e71c255d4..62f36b4fff5 100644 --- a/example/ck_tile/20_grouped_convolution/CMakeLists.txt +++ b/example/ck_tile/20_grouped_convolution/CMakeLists.txt @@ -8,9 +8,6 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") add_executable(tile_example_grouped_conv_fwd grouped_convolution_forward.cpp) target_compile_options(tile_example_grouped_conv_fwd PRIVATE ${EXAMPLE_CONV_COMPILE_OPTIONS}) - add_executable(tile_example_grouped_conv_fwd_large_tensor grouped_convolution_forward_large_tensor.cpp) - target_compile_options(tile_example_grouped_conv_fwd_large_tensor PRIVATE ${EXAMPLE_CONV_COMPILE_OPTIONS}) - add_executable(tile_example_grouped_conv_fwd_bias_clamp grouped_convolution_forward_bias_clamp.cpp) target_compile_options(tile_example_grouped_conv_fwd_bias_clamp PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data_invoker.hpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data_invoker.hpp index 1679fec7dfc..cdecdac9146 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data_invoker.hpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data_invoker.hpp @@ -45,6 +45,8 @@ struct GroupedConvolutionBackwardDataInvoker GroupedConvTraitsType::FixedGemmParams::TilePartitionerGroupNum, GroupedConvTraitsType::FixedGemmParams::TilePartitionerM01>; + constexpr bool LargeTensors = false; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits< GroupedConvTraitsType::FixedGemmParams::kPadM, GroupedConvTraitsType::FixedGemmParams::kPadN, @@ -56,7 +58,13 @@ struct GroupedConvolutionBackwardDataInvoker GroupedConvTraitsType::FixedGemmParams::TransposeC, GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity, GroupedConvTraitsType::FixedGemmParams::Persistent, - ConvConfig::NumWaveGroups>; + ConvConfig::NumWaveGroups, + GroupedConvTraitsType::FixedGemmParams::Preshuffle, + GroupedConvTraitsType::FixedGemmParams::LDSVectorSize, + ck_tile::DataCachePrefetchKind::None, + ck_tile::DataCachePrefetchKind::None, + false, /*Async*/ + LargeTensors>; constexpr auto scheduler = ConvConfig::Scheduler; using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem< diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp index 533abdd3391..2b38e68650a 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp @@ -64,6 +64,8 @@ struct GroupedConvolutionBackwardWeightInvoker using TilePartitioner = typename PartitionerPolicy::template type; + constexpr bool LargeTensors = false; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits< GroupedConvTraitsType::FixedGemmParams::kPadM, GroupedConvTraitsType::FixedGemmParams::kPadN, @@ -75,7 +77,13 @@ struct GroupedConvolutionBackwardWeightInvoker GroupedConvTraitsType::FixedGemmParams::TransposeC, GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity, GroupedConvTraitsType::FixedGemmParams::Persistent, - ConvConfig::NumWaveGroups>; + ConvConfig::NumWaveGroups, + GroupedConvTraitsType::FixedGemmParams::Preshuffle, + GroupedConvTraitsType::FixedGemmParams::LDSVectorSize, + ck_tile::DataCachePrefetchKind::None, + ck_tile::DataCachePrefetchKind::None, + false, /*Async*/ + LargeTensors>; constexpr auto scheduler = ConvConfig::Scheduler; using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem< diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp index 68c85e9495f..e277e1fd31a 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp @@ -50,6 +50,8 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker GroupedConvTraitsType::FixedGemmParams::TilePartitionerGroupNum, GroupedConvTraitsType::FixedGemmParams::TilePartitionerM01>; + constexpr bool LargeTensors = false; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits< GroupedConvTraitsType::FixedGemmParams::kPadM, GroupedConvTraitsType::FixedGemmParams::kPadN, @@ -61,7 +63,13 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker GroupedConvTraitsType::FixedGemmParams::TransposeC, GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity, GroupedConvTraitsType::FixedGemmParams::Persistent, - ConvConfig::NumWaveGroups>; + ConvConfig::NumWaveGroups, + GroupedConvTraitsType::FixedGemmParams::Preshuffle, + GroupedConvTraitsType::FixedGemmParams::LDSVectorSize, + ck_tile::DataCachePrefetchKind::None, + ck_tile::DataCachePrefetchKind::None, + false, /*Async*/ + LargeTensors>; constexpr auto scheduler = ConvConfig::Scheduler; diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp index a396dd82cbf..7fb246f5ab6 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp @@ -1,12 +1,6 @@ // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT -// Regular grouped convolution invoker (no split-image) -// This invoker demonstrates regular convolution without split-image. -// It always uses Kernel (split-image disabled). -// For large images that require split-image, use -// grouped_convolution_forward_split_image_invoker.hpp - #pragma once #include "grouped_convolution_utils.hpp" @@ -53,6 +47,8 @@ struct GroupedConvolutionForwardInvoker GroupedConvTraitsType::FixedGemmParams::TilePartitionerGroupNum, GroupedConvTraitsType::FixedGemmParams::TilePartitionerM01>; + constexpr bool LargeTensors = false; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits< GroupedConvTraitsType::FixedGemmParams::kPadM, GroupedConvTraitsType::FixedGemmParams::kPadN, @@ -64,7 +60,13 @@ struct GroupedConvolutionForwardInvoker GroupedConvTraitsType::FixedGemmParams::TransposeC, GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity, GroupedConvTraitsType::FixedGemmParams::Persistent, - ConvConfig::NumWaveGroups>; + ConvConfig::NumWaveGroups, + GroupedConvTraitsType::FixedGemmParams::Preshuffle, + GroupedConvTraitsType::FixedGemmParams::LDSVectorSize, + ck_tile::DataCachePrefetchKind::None, + ck_tile::DataCachePrefetchKind::None, + false, /*Async*/ + LargeTensors>; constexpr auto scheduler = ConvConfig::Scheduler; // ===================================================================== diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_large_tensor.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_large_tensor.cpp deleted file mode 100644 index 9a7eb7082a9..00000000000 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward_large_tensor.cpp +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -// Large tensor grouped convolution example -// This example demonstrates convolution for large tensors that exceed memory limits. -// It uses automatic tensor splitting when needed to handle large images. -// For regular convolution without tensor splitting, use grouped_convolution_forward.cpp - -#include - -#include -#include -#include -#include -#include - -#include "ck_tile/host.hpp" -#include "grouped_convolution_utils.hpp" -#include "grouped_convolution_forward_large_tensor_invoker.hpp" -#include "run_grouped_convolution_fwd_example.inc" - -template