From 851b9732aaf6b5779a8ab900c7e805fdccf46877 Mon Sep 17 00:00:00 2001 From: Jithun Nair Date: Thu, 26 Sep 2024 18:05:51 +0000 Subject: [PATCH] Download pre-compiled AOTriton from GitHub unless AOTRITON_INSTALL_FROM_SOURCE=1 is set (#136603) 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](https://github.com/pytorch/pytorch/blob/a8ed873ba2aa13437a336863ae9d73c235798acc/.ci/docker/ubuntu-rocm/Dockerfile#L110)/[manywheel](https://github.com/pytorch/pytorch/blob/a8ed873ba2aa13437a336863ae9d73c235798acc/.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 --- cmake/External/aotriton.cmake | 46 +++++++++++++++++++++++++++++------ 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake index cb6080f5f7a..d9d99959f3e 100644 --- a/cmake/External/aotriton.cmake +++ b/cmake/External/aotriton.cmake @@ -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)