saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Download pre-compiled AOTriton from GitHub unless AOTRITON_INSTALL_FROM_SOURCE=1 is set (#136603)

Browse source 
PyTorch community members have reported issues with building PyTorch from source for ROCm in an environment that doesn't have aotriton pre-installed, because aotriton is only installed in the [CI](a8ed873ba2/.ci/docker/manywheel/Dockerfile (L197)) docker images. Building aotriton from source can take ~45 minutes.

This PR fixes the issue by downloading the aotriton tarball in such scenarios, *unless the user explicitly wants to build aotriton from source using the AOTRITON_INSTALL_FROM_SOURCE=1 env var*

Pull Request resolved: https://github.com/pytorch/pytorch/pull/136603
Approved by: https://github.com/atalman

Co-authored-by: Xinya Zhang <Xinya.Zhang@amd.com>
This commit is contained in:
Jithun Nair 2024-09-26 18:05:51 +00:00 committed by PyTorch MergeBot
parent f0a92541fe
commit 851b9732aa
1 changed files with 38 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
cmake/External/aotriton.cmake vendored
View file

@ -1,25 +1,27 @@
if(NOT __AOTRITON_INCLUDED)
set(__AOTRITON_INCLUDED TRUE)
set(__AOTRITON_SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/aotriton/src")
set(__AOTRITON_BUILD_DIR "${CMAKE_CURRENT_BINARY_DIR}/aotriton/build")
set(__AOTRITON_EXTERN_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/aotriton")
set(__AOTRITON_INSTALL_DIR "${PROJECT_SOURCE_DIR}/torch")
add_library(__caffe2_aotriton INTERFACE)
# Note it is INSTALL"ED"
if(DEFINED ENV{AOTRITON_INSTALLED_PREFIX})
install(DIRECTORY
$ENV{AOTRITON_INSTALLED_PREFIX}/lib
$ENV{AOTRITON_INSTALLED_PREFIX}/include
DESTINATION ${__AOTRITON_INSTALL_DIR})
set(__AOTRITON_INSTALL_DIR "$ENV{AOTRITON_INSTALLED_PREFIX}")
message(STATUS "Using Preinstalled AOTriton at ${__AOTRITON_INSTALL_DIR}")
else()
elseif(DEFINED ENV{AOTRITON_INSTALL_FROM_SOURCE})
file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/.ci/docker/aotriton_version.txt" __AOTRITON_CI_INFO)
list(GET __AOTRITON_CI_INFO 3 __AOTRITON_CI_COMMIT)
ExternalProject_Add(aotriton_external
GIT_REPOSITORY https://github.com/ROCm/aotriton.git
GIT_TAG ${__AOTRITON_CI_COMMIT}
SOURCE_DIR ${__AOTRITON_SOURCE_DIR}
BINARY_DIR ${__AOTRITON_BUILD_DIR}
PREFIX ${__AOTRITON_INSTALL_DIR}
PREFIX ${__AOTRITON_EXTERN_PREFIX}
INSTALL_DIR ${__AOTRITON_INSTALL_DIR}
CMAKE_ARGS -DCMAKE_INSTALL_PREFIX:PATH=${__AOTRITON_INSTALL_DIR}
-DAOTRITON_COMPRESS_KERNEL=OFF
-DAOTRITON_COMPRESS_KERNEL=ON
-DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
-DAOTRITON_NO_PYTHON=ON
-DAOTRITON_NO_SHARED=OFF
@ -33,7 +35,35 @@ if(NOT __AOTRITON_INCLUDED)
# INSTALL_COMMAND ${MAKE_COMMAND} install
)
add_dependencies(__caffe2_aotriton aotriton_external)
message(STATUS "Using AOTriton compiled from source directory ${__AOTRITON_SOURCE_DIR}")
message(STATUS "Using AOTriton compiled from source directory ${__AOTRITON_EXTERN_PREFIX}")
else()
file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/.ci/docker/aotriton_version.txt" __AOTRITON_CI_INFO)
list(GET __AOTRITON_CI_INFO 0 __AOTRITON_VER)
list(GET __AOTRITON_CI_INFO 1 __AOTRITON_MANYLINUX)
list(GET __AOTRITON_CI_INFO 2 __AOTRITON_ROCM)
list(GET __AOTRITON_CI_INFO 3 __AOTRITON_COMMIT)
list(GET __AOTRITON_CI_INFO 4 __AOTRITON_SHA256)
set(__AOTRITON_ARCH ${CMAKE_HOST_SYSTEM_PROCESSOR})
string(CONCAT __AOTRITON_FILE "aotriton-"
"${__AOTRITON_VER}-${__AOTRITON_MANYLINUX}"
"_${__AOTRITON_ARCH}-${__AOTRITON_ROCM}"
"-shared.tar.gz")
string(CONCAT __AOTRITON_URL "https://github.com/ROCm/aotriton/releases/download/"
"${__AOTRITON_VER}/${__AOTRITON_FILE}")
ExternalProject_Add(aotriton_external
URL "${__AOTRITON_URL}"
URL_HASH SHA256=${__AOTRITON_SHA256}
SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/aotriton_tarball
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
INSTALL_COMMAND ${CMAKE_COMMAND} -E copy_directory
"${CMAKE_CURRENT_BINARY_DIR}/aotriton_tarball"
"${__AOTRITON_INSTALL_DIR}"
BUILD_BYPRODUCTS "${__AOTRITON_INSTALL_DIR}/lib/libaotriton_v2.so"
)
add_dependencies(__caffe2_aotriton aotriton_external)
message(STATUS "Using AOTriton from pre-compiled binary ${__AOTRITON_URL}.\
Set env variables AOTRITON_INSTALL_FROM_SOURCE=1 to build from source.")
endif()
target_link_libraries(__caffe2_aotriton INTERFACE ${__AOTRITON_INSTALL_DIR}/lib/libaotriton_v2.so)
target_include_directories(__caffe2_aotriton INTERFACE ${__AOTRITON_INSTALL_DIR}/include)