From 234bc58d2d0c342363e923e8477a59fc658f93b7 Mon Sep 17 00:00:00 2001
From: Illia Silin <98187287+illsilin@users.noreply.github.com>
Date: Wed, 4 Sep 2024 07:36:27 -0700
Subject: [PATCH] Add an option to select an alternative python version during
 build. (#1496)

* locate a newwer version of python when -DRHEL=ON flag is set

* allow setting python version on cmake command line

[ROCm/composable_kernel commit: 841009c5ee5a7913c0b19eb541ac39fea9ef6a4c]
---
 CMakeLists.txt                                     | 14 +++++++++++++-
 .../gpu/mha/CMakeLists.txt                         | 14 +++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 697eb3b745..bf42b6aa15 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -26,7 +26,19 @@ set(version 1.1.0)
 project(composable_kernel VERSION ${version} LANGUAGES CXX HIP)
 include(CTest)
 
-find_package(Python3 3.6 COMPONENTS Interpreter REQUIRED)
+if(NOT CK_USE_ALTERNATIVE_PYTHON)
+   find_package(Python3 3.6 COMPONENTS Interpreter REQUIRED)
+else()
+   message("Using alternative python version")
+   set(EXTRA_PYTHON_PATH)
+   string(REPLACE "/bin/python3.8" "" EXTRA_PYTHON_PATH "${CK_USE_ALTERNATIVE_PYTHON}")
+   message("alternative python path is: ${EXTRA_PYTHON_PATH}")
+   find_package(Python3 3.6 COMPONENTS Interpreter REQUIRED)
+   add_definitions(-DPython3_EXECUTABLE="${CK_USE_ALTERNATIVE_PYTHON}")
+   set(Python3_EXECUTABLE "${CK_USE_ALTERNATIVE_PYTHON}")
+   set(PYTHON_EXECUTABLE "${CK_USE_ALTERNATIVE_PYTHON}")
+   set(ENV{LD_LIBRARY_PATH} "${EXTRA_PYTHON_PATH}/lib:$ENV{LD_LIBRARY_PATH}")
+endif()
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
 
diff --git a/library/src/tensor_operation_instance/gpu/mha/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/mha/CMakeLists.txt
index 59ae09b739..8d159662a2 100644
--- a/library/src/tensor_operation_instance/gpu/mha/CMakeLists.txt
+++ b/library/src/tensor_operation_instance/gpu/mha/CMakeLists.txt
@@ -2,7 +2,19 @@ set(FMHA_CPP_FOLDER ${CMAKE_CURRENT_BINARY_DIR})
 set(FMHA_SRC_FOLDER ${CMAKE_SOURCE_DIR}/example/ck_tile/01_fmha/)
 set(CK_TILE_SRC_FOLDER ${CMAKE_SOURCE_DIR}/include/ck_tile/)
 # python stuff
-find_package(PythonInterp 3 REQUIRED)
+if(NOT CK_USE_ALTERNATIVE_PYTHON)
+  find_package(PythonInterp 3 REQUIRED)
+else()
+  message("Using alternative python version")
+  set(EXTRA_PYTHON_PATH)
+  string(REPLACE "/bin/python3.8" "" EXTRA_PYTHON_PATH "${CK_USE_ALTERNATIVE_PYTHON}")
+  message("alternative python path is: ${EXTRA_PYTHON_PATH}")
+  find_package(Python3 3.6 COMPONENTS Interpreter REQUIRED)
+  add_definitions(-DPython3_EXECUTABLE="${CK_USE_ALTERNATIVE_PYTHON}")
+  set(Python3_EXECUTABLE "${CK_USE_ALTERNATIVE_PYTHON}")
+  set(PYTHON_EXECUTABLE "${CK_USE_ALTERNATIVE_PYTHON}")
+  set(ENV{LD_LIBRARY_PATH} "${EXTRA_PYTHON_PATH}/lib:$ENV{LD_LIBRARY_PATH}")
+endif()
 
 rocm_install(DIRECTORY ${CK_TILE_SRC_FOLDER} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/ck_tile)