[release/2.13] Consolidate CK GEMM allow-list

CMakeLists.txt

@@ -269,6 +269,13 @@ cmake_dependent_option(
   "USE_CUDA AND LINUX AND BUILD_PYTHON" OFF)
 cmake_dependent_option(USE_ROCM "Use ROCm" ON "LINUX OR WIN32" OFF)
 cmake_dependent_option(USE_ROCM_CK_GEMM "Use ROCm Composable Kernel for GEMMs" ON "USE_ROCM;NOT WIN32" OFF)
+# Single source of truth for the GPU archs Composable Kernel GEMM supports. It
+# drives both the ck_gemm build arch filter (torch_filter_ck_gemm_arches in
+# cmake/public/utils.cmake) and the runtime support check / dispatch, which read
+# it through the generated ATen/ROCmCKGEMMConfig.h. Add a WMMA arch (e.g.
+# gfx1200) here to opt it back in without patching source.
+set(TORCH_CK_GEMM_ARCH_LIST "gfx90a;gfx942;gfx950" CACHE STRING
+    "Semicolon-separated GPU archs Composable Kernel GEMM is built and enabled for")
 cmake_dependent_option(USE_ROCM_CK_SDPA "Use ROCm Composable Kernel for SDPA" ON "USE_ROCM;NOT WIN32" OFF)
 option(CAFFE2_STATIC_LINK_CUDA "Statically link CUDA libraries" OFF)
 cmake_dependent_option(USE_CUDNN "Use cuDNN" ON "USE_CUDA" OFF)

aten/src/ATen/CMakeLists.txt

@@ -55,6 +55,12 @@ set_bool(AT_CUDNN_ENABLED CAFFE2_USE_CUDNN)
 set_bool(AT_CUSPARSELT_ENABLED CAFFE2_USE_CUSPARSELT)
 set_bool(AT_HIPSPARSELT_ENABLED CAFFE2_USE_HIPSPARSELT)
 
+# CK GEMM allowlist (TORCH_CK_GEMM_ARCH_LIST, top-level CMakeLists.txt) surfaced
+# to Context.cpp and the ck_gemm .hip dispatch, so the runtime check stays in
+# step with the build arch filter (torch_filter_ck_gemm_arches).
+string(REPLACE ";" "," TORCH_CK_GEMM_ARCH_LIST_STR "${TORCH_CK_GEMM_ARCH_LIST}")
+configure_file(ROCmCKGEMMConfig.h.in "${CMAKE_BINARY_DIR}/aten/src/ATen/ROCmCKGEMMConfig.h" @ONLY)
+
 configure_file(Config.h.in "${CMAKE_CURRENT_SOURCE_DIR}/Config.h")
 # TODO: Do not generate CUDAConfig.h for ROCm BUILDS
 # At the moment, `jit_macros.h` include CUDAConfig.h for both CUDA and HIP builds
@@ -667,15 +673,10 @@ if(USE_ROCM)
     exclude(ATen_HIP_SRCS "${ATen_HIP_SRCS}"
       ${native_hip_bgemm} ${native_hip_ck})
   else()
-    # composable_kernel has no gfx1250 or gfx1250-strict support yet, so its CK
-    # GEMM kernels fail to compile for those archs. Build CK GEMM as a separate
-    # library with them filtered out of HIP_ARCHITECTURES so every other
-    # requested arch still gets CK GEMM. Remove this filtering once the CK
-    # submodule supports gfx1250 and gfx1250-strict.
+    # Build CK GEMM only for architectures in TORCH_CK_GEMM_ARCH_LIST.
     exclude(ATen_HIP_SRCS "${ATen_HIP_SRCS}"
       ${native_hip_bgemm} ${native_hip_ck})
-    set(_ck_gemm_hip_arches ${PYTORCH_ROCM_ARCH})
-    list(REMOVE_ITEM _ck_gemm_hip_arches "gfx1250" "gfx1250-strict")
+    torch_filter_ck_gemm_arches(_ck_gemm_hip_arches)
     set_source_files_properties(${native_hip_bgemm} ${native_hip_ck} PROPERTIES LANGUAGE HIP)
     add_library(ck_gemm STATIC ${native_hip_bgemm} ${native_hip_ck})
     set_target_properties(ck_gemm PROPERTIES

aten/src/ATen/Context.cpp

@@ -8,6 +8,17 @@
 #include <array>
 #include <cctype>
 #include <string>
+#include <string_view>
+
+#if defined(USE_ROCM) && defined(__has_include) && \
+    __has_include(<ATen/ROCmCKGEMMConfig.h>)
+#include <ATen/ROCmCKGEMMConfig.h>
+#endif
+#ifndef TORCH_CK_GEMM_ARCH_LIST
+// Non-CMake ROCm builds (e.g. Buck) do not generate the header; fall back to
+// the default CK GEMM allowlist.
+#define TORCH_CK_GEMM_ARCH_LIST "gfx90a,gfx942,gfx950"
+#endif
 
 #include <ATen/cpu/FlushDenormal.h>
 
@@ -561,12 +572,13 @@ at::BlasBackend Context::blasPreferredBackend() {
 bool Context::ckSupported() {
 #ifdef USE_ROCM
   static const std::vector<std::string> supported_archs = {
-    "gfx942", "gfx950"
+      "gfx942", "gfx950",
   };
   for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) {
-    if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) {
+    if (!detail::getCUDAHooks().isGPUArch(supported_archs, index)) {
       TORCH_WARN_ONCE(
-        "Attempting to use CK on an unsupported architecture! Cannot set backend to CK");
+          "Attempting to use CK on an unsupported architecture! Cannot set "
+          "backend to CK");
       return false;
     }
   }
@@ -576,6 +588,51 @@ bool Context::ckSupported() {
 #endif
 }
 
+#ifdef USE_ROCM
+namespace {
+// Split the generated comma-separated CK GEMM architecture list.
+std::vector<std::string> parseArchList(std::string_view archs) {
+  std::vector<std::string> result;
+  while (!archs.empty()) {
+    const auto comma = archs.find(',');
+    result.emplace_back(archs.substr(0, comma));
+    if (comma == std::string_view::npos) {
+      break;
+    }
+    archs.remove_prefix(comma + 1);
+  }
+  return result;
+}
+} // namespace
+#endif
+
+bool Context::ckGemmSupported() {
+#ifdef USE_ROCM
+  // TORCH_CK_GEMM_ARCH_LIST is the allowlist this build was compiled for (the
+  // same list that drives the ck_gemm build arch filter), so the check stays in
+  // step with the build. It is empty if overridden empty. Cached because the
+  // CK GEMM dispatch calls this on every GEMM.
+  static const bool supported = [] {
+    const auto supported_archs = parseArchList(TORCH_CK_GEMM_ARCH_LIST);
+    if (supported_archs.empty()) {
+      return false;
+    }
+    for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) {
+      if (!detail::getCUDAHooks().isGPUArch(supported_archs, index)) {
+        TORCH_WARN_ONCE(
+            "Attempting to use CK GEMM on an unsupported architecture! Cannot "
+            "set blas backend to CK");
+        return false;
+      }
+    }
+    return true;
+  }();
+  return supported;
+#else
+  return false;
+#endif
+}
+
 void Context::setBlasPreferredBackend(at::BlasBackend b) {
 #ifdef _MSC_VER
   TORCH_WARN_ONCE(
@@ -586,11 +643,11 @@ void Context::setBlasPreferredBackend(at::BlasBackend b) {
   TORCH_CHECK((b != at::BlasBackend::Cublaslt) || hasCuBLASLt(),
       "Cannot set preferred backend to cuBLASLt if PyTorch has not been compiled with cuBLASLt.");
 #ifdef USE_ROCM
-  static const bool ckSupportedFlag = ckSupported();
+  static const bool ckGemmSupportedFlag = ckGemmSupported();
   static const bool hasCKGEMMFlag = hasCKGEMM();
-  TORCH_CHECK((b != at::BlasBackend::Ck) || (ckSupportedFlag && hasCKGEMMFlag),
+  TORCH_CHECK((b != at::BlasBackend::Ck) || (ckGemmSupportedFlag && hasCKGEMMFlag),
       "Cannot set preferred blas backend to CK since following conditions are not true: ",
-      "architecture supported for CK: ", ckSupportedFlag,
+      "architecture supported for CK GEMM: ", ckGemmSupportedFlag,
       ", PyTorch built with CK GEMM support: ", hasCKGEMMFlag);
 #endif
   if (b != at::BlasBackend::Default && b != at::BlasBackend::Cublas) {

aten/src/ATen/Context.h

@@ -159,6 +159,7 @@ class TORCH_API Context {
   static bool hasLAPACK();
   static bool hasMKLDNN();
   static bool ckSupported();
+  static bool ckGemmSupported();
   static bool hasEigenSparse();
   static bool hasMAGMA() {
     return detail::getCUDAHooks().hasMAGMA();

aten/src/ATen/ROCmCKGEMMConfig.h.in

@@ -0,0 +1,14 @@
+#pragma once
+
+// Comma-separated allowlist of the architectures Composable Kernel GEMM is built
+// and enabled for. Generated from the TORCH_CK_GEMM_ARCH_LIST CMake cache
+// variable (the single source of truth, declared in the top-level
+// CMakeLists.txt), which also drives the ck_gemm build arch filter
+// (torch_filter_ck_gemm_arches). The runtime support check
+// (at::Context::ckGemmSupported) and the .hip dispatch read this macro so they
+// stay in step with what was built; override the whole set with
+// -DTORCH_CK_GEMM_ARCH_LIST=... at configure time.
+//
+// Keep this generated configuration header out of public headers; include it
+// only from implementation files.
+#define TORCH_CK_GEMM_ARCH_LIST "@TORCH_CK_GEMM_ARCH_LIST_STR@"

aten/src/ATen/native/hip/ck_gemm_bfloat16.hip

@@ -781,7 +781,7 @@ void gemm_internal_ck<at::BFloat16>(CUDABLAS_GEMM_ARGTYPES(at::BFloat16)) {
   if (at::detail::getCUDAHooks().isGPUArch(wmma_archs)) {
     dispatch_bfloat16_gemm_wmma(CUDABLAS_GEMM_ARGS(at::BFloat16));
   }
-  else if (at::detail::getCUDAHooks().isGPUArch({"gfx9"})) {
+  else if (at::globalContext().ckGemmSupported()) {
     dispatch_bfloat16_gemm(CUDABLAS_GEMM_ARGS(at::BFloat16));
   }
   else {

aten/src/ATen/native/hip/ck_gemm_half.hip

@@ -608,7 +608,7 @@ void gemm_internal_ck<at::Half>(CUDABLAS_GEMM_ARGTYPES(at::Half)) {
   if (at::detail::getCUDAHooks().isGPUArch(wmma_archs)) {
     dispatch_half_gemm_wmma(CUDABLAS_GEMM_ARGS(at::Half));
   }
-  else if (at::detail::getCUDAHooks().isGPUArch({"gfx9"})) {
+  else if (at::globalContext().ckGemmSupported()) {
     dispatch_half_gemm(CUDABLAS_GEMM_ARGS(at::Half));
   }
   else {

cmake/Dependencies.cmake

@@ -1051,20 +1051,10 @@ if(USE_ROCM)
     if(HIPBLASLT_VEC_EXT)
       list(APPEND HIP_CXX_FLAGS -DHIPBLASLT_VEC_EXT)
     endif()
-    # composable_kernel has no gfx1250 or gfx1250-strict support, so
-    # aten/src/ATen/CMakeLists.txt filters those archs out of the ck_gemm
-    # target's HIP_ARCHITECTURES. When they are the only requested archs that
-    # list is empty and CMake fails at generate time with 'HIP_ARCHITECTURES
-    # is empty for target "ck_gemm"', so turn CK GEMM off entirely here --
-    # before -DUSE_ROCM_CK_GEMM is added below, since the non-CK TUs guarded by
-    # that define (HIPBlas.cpp, HIPHooks.cpp, GroupedBlas.cpp) would otherwise
-    # reference symbols the unbuilt ck_gemm target never provides. Mirrors the
-    # USE_ROCM_CK_SDPA / USE_MSLK auto-disable pattern in
-    # aten/src/ATen/CMakeLists.txt. Remove once the CK submodule supports
-    # gfx1250 and gfx1250-strict.
+    # Disable CK GEMM when none of the requested architectures are in its
+    # configured allow-list; an empty HIP_ARCHITECTURES list is invalid.
     if(USE_ROCM_CK_GEMM)
-      set(_ck_gemm_supported_arches ${PYTORCH_ROCM_ARCH})
-      list(REMOVE_ITEM _ck_gemm_supported_arches "gfx1250" "gfx1250-strict")
+      torch_filter_ck_gemm_arches(_ck_gemm_supported_arches)
       if(_ck_gemm_supported_arches STREQUAL "")
         message(STATUS "USE_ROCM_CK_GEMM disabled: PYTORCH_ROCM_ARCH (${PYTORCH_ROCM_ARCH}) has no arch supported by composable_kernel")
         caffe2_update_option(USE_ROCM_CK_GEMM OFF)

cmake/EnvVarForwarding.cmake

@@ -39,6 +39,7 @@ set(_ENV_PASSTHROUGH
   STATIC_DISPATCH_BACKEND
   SELECTED_OP_LIST
   TORCH_CUDA_ARCH_LIST
+  TORCH_CK_GEMM_ARCH_LIST
   TORCH_XPU_ARCH_LIST
   TRACING_BASED
   PYTHON_LIB_REL_PATH

cmake/public/utils.cmake

@@ -293,6 +293,28 @@ macro(torch_hip_get_arch_list store_var)
   string(REPLACE " " ";" ${store_var} "${_TMP}")
 endmacro()
 
+##############################################################################
+# Returns the subset of PYTORCH_ROCM_ARCH that CK GEMM supports, i.e. the archs
+# in TORCH_CK_GEMM_ARCH_LIST (the single source of truth, declared in the
+# top-level CMakeLists.txt and also surfaced to C++ via ATen/ROCmCKGEMMConfig.h).
+# Matches on the base arch so feature-suffixed targets (e.g. gfx942:xnack+) are
+# kept, mirroring the substring match in at::cuda::detail::CUDAHooks::isGPUArch.
+# Usage:
+#   torch_filter_ck_gemm_arches(variable_to_store_arches)
+#
+function(torch_filter_ck_gemm_arches out_var)
+  set(_supported "${TORCH_CK_GEMM_ARCH_LIST}")
+  set(_filtered "")
+  foreach(_arch IN LISTS PYTORCH_ROCM_ARCH)
+    string(REGEX REPLACE ":.*$" "" _base "${_arch}")
+    list(FIND _supported "${_base}" _idx)
+    if(NOT _idx EQUAL -1)
+      list(APPEND _filtered "${_arch}")
+    endif()
+  endforeach()
+  set(${out_var} "${_filtered}" PARENT_SCOPE)
+endfunction()
+
 ##############################################################################
 # Get the XPU arch flags specified by TORCH_XPU_ARCH_LIST.
 # Usage: