Rename and move code

Author: jithunnair-amd

caffe2/CMakeLists.txt

@@ -930,8 +930,8 @@ if(USE_ROCM)
     target_link_libraries(torch_hip PRIVATE __caffe2_aotriton)
   endif()
 # link CK library
-  if(USE_CK_FLASH_ATTENTION)
-    target_link_libraries(torch_hip PRIVATE __ck_lib)
+  if(USE_CK_FLASH_ATTENTION) # TODO: Add condition of not building CK_kernels from source
+    target_link_libraries(torch_hip PRIVATE __ck_kernels_lib)
   endif()
   set(CUDA_LINK_LIBRARIES_KEYWORD)
   torch_compile_options(torch_hip)  # see cmake/public/utils.cmake

cmake/External/ck_kernels.cmake

@@ -9,14 +9,17 @@ if(NOT __ck_kernels_included)
   set(__ck_kernels_version 0.1)
 
   # create INTERFACE target
-  add_library(__ck_lib INTERFACE)
+  add_library(__ck_kernels_lib INTERFACE)
 
   if(DEFINED ENV{CK_KERNELS_INSTALLED_PREFIX})
     # Copy .so from $ENV{CK_KERNELS_INSTALLED_PREFIX} into ${ck_kernels_install_dir}
     install(DIRECTORY
             $ENV{CK_KERNELS_INSTALLED_PREFIX}/
             DESTINATION ${ck_kernels_install_dir}
-            FILES_MATCHING_PATTERN "*.so")
+            FILES_MATCHING_PATTERN "libck_kernels.so")
+    set(ck_kernels_install_path "${ck_kernels_install_dir}/libck_kernels.so")
+    # specify path to CK library
+    target_link_libraries(__ck_kernels_lib INTERFACE ${ck_kernels_install_path})
     message(STATUS "Using Preinstalled CK_kernels from $ENV{CK_KERNELS_INSTALLED_PREFIX}; installed at ${ck_kernels_install_dir}")
   elseif(DEFINED ENV{CK_KERNELS_PACKAGE_BASE_URL})
     # get CK commit hash
@@ -49,12 +52,12 @@ if(NOT __ck_kernels_included)
       "${ck_kernels_install_dir}"
       INSTALL_BYPRODUCTS "${ck_kernels_install_path}"
     )
-    add_dependencies(__ck_lib ck_kernels_external)
+    add_dependencies(__ck_kernels_lib ck_kernels_external)
     message(STATUS "Using CK_kernels from pre-compiled binary ${ck_kernels_package_full_url}; installed at ${ck_kernels_install_dir}")
+    # specify path to CK library
+    target_link_libraries(__ck_kernels_lib INTERFACE ${ck_kernels_install_path})
   else()
-    message(FATAL_ERROR "Unable to find an existing CK_kernels installation or to install CK_kernels library")
-  endif() # DEFINED ENV{AOTRITON_INSTALLED_PREFIX}
+    # TODO: Enable build from source
+  endif() # DEFINED ENV{CK_KERNELS_INSTALLED_PREFIX}
 
-  # specify path to CK library
-  target_link_libraries(__ck_lib INTERFACE ${ck_kernels_install_path})
 endif() # __ck_kernels_included