diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ed8ebf5 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 7db8711..a973508 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,5 @@ cmake_minimum_required(VERSION 3.12) +set(CMAKE_PROJECT_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/cmake/enable_shared_libs.cmake") project(FTorch_CESM_WRAPPER LANGUAGES Fortran CXX) enable_language(Fortran) # Force use of Fortran compiler for linking @@ -11,8 +12,9 @@ if("${USE_FTORCH}" STREQUAL "ON") add_subdirectory(src) add_library(ftorch_wrapper SHARED FTorch_cesm_interface.F90) target_link_libraries(ftorch_wrapper PRIVATE ftorch) + # Place ftorch.mod in modules/ to match INCLDIR += -I$(FTORCH_DIR)/modules in cime/CIME/Tools/Makefile set_target_properties(ftorch PROPERTIES - Fortran_MODULE_DIRECTORY ${CMAKE_BINARY_DIR}/external/ftorch + Fortran_MODULE_DIRECTORY ${CMAKE_BINARY_DIR}/modules ) # The cpp file in this library will cause cmake to prefer the cxx compiler for the link step, we override and use fc set_target_properties(ftorch PROPERTIES @@ -21,7 +23,7 @@ if("${USE_FTORCH}" STREQUAL "ON") # Expose that module directory to anything that links against ftorch target_include_directories(ftorch PUBLIC - $ + $ ) else() message("Building FTorch wrapper only ${USE_FTORCH}") diff --git a/FTorch_cesm_interface.F90 b/FTorch_cesm_interface.F90 index 7c946d8..76c4665 100644 --- a/FTorch_cesm_interface.F90 +++ b/FTorch_cesm_interface.F90 @@ -1,35 +1,81 @@ module FTorch_cesm_interface #ifdef USE_FTORCH use ftorch, only: torch_kCPU, torch_tensor, torch_model, torch_tensor_from_array - use ftorch, only: torch_model_load, torch_model_forward + use ftorch, only: torch_model_load, torch_model_forward, torch_delete #else use shr_abort_mod, only : shr_abort_abort - use shr_kind_mod, only : r8 => shr_kind_r8 + implicit none type :: torch_model end type torch_model type :: torch_tensor end type torch_tensor - integer :: torch_kcpu + integer, parameter :: torch_kCPU = 0 ! CPU device enum placeholder when FTorch is disabled character(len=*), parameter :: message="ERROR: Using FTorch Interface without USE_FTORCH=TRUE" + + interface torch_tensor_from_array + module procedure torch_tensor_from_array_layout + module procedure torch_tensor_from_array_nolayout + end interface torch_tensor_from_array + + interface torch_delete + module procedure torch_model_delete + module procedure torch_tensor_delete + module procedure torch_tensor_array_delete + end interface torch_delete + contains - subroutine torch_model_load(model, fname, i) - type(torch_model) :: model - character(len=*) :: fname - integer :: i + subroutine torch_model_load(model, filename, device_type, device_index, & + requires_grad, is_training) + type(torch_model), intent(out) :: model + character(len=*), intent(in) :: filename + integer, intent(in) :: device_type + integer, optional, intent(in) :: device_index + logical, optional, intent(in) :: requires_grad + logical, optional, intent(in) :: is_training call shr_abort_abort(message) end subroutine torch_model_load - subroutine torch_tensor_from_array(tensor, data, layout, i) - type(torch_tensor) :: tensor - real(r8) :: data(:,:) - integer :: layout(:) - integer :: i + + subroutine torch_tensor_from_array_layout(tensor, data_in, layout, device_type, & + device_index, requires_grad) + type(torch_tensor), intent(out) :: tensor + type(*), dimension(..), intent(in) :: data_in + integer, intent(in) :: layout(:) + integer, intent(in) :: device_type + integer, optional, intent(in) :: device_index + logical, optional, intent(in) :: requires_grad + call shr_abort_abort(message) + end subroutine torch_tensor_from_array_layout + + subroutine torch_tensor_from_array_nolayout(tensor, data_in, device_type, & + device_index, requires_grad) + type(torch_tensor), intent(out) :: tensor + type(*), dimension(..), intent(in) :: data_in + integer, intent(in) :: device_type + integer, optional, intent(in) :: device_index + logical, optional, intent(in) :: requires_grad call shr_abort_abort(message) - end subroutine torch_tensor_from_array - subroutine torch_model_forward(model, tensor_in, tensor_out) - type(torch_model) :: model - type(torch_tensor) :: tensor_in(:), tensor_out(:) + end subroutine torch_tensor_from_array_nolayout + + subroutine torch_model_forward(model, input_tensors, output_tensors, requires_grad) + type(torch_model), intent(in) :: model + type(torch_tensor), intent(in) :: input_tensors(:), output_tensors(:) + logical, optional, intent(in) :: requires_grad call shr_abort_abort(message) end subroutine torch_model_forward + subroutine torch_model_delete(model) + type(torch_model), intent(in) :: model + call shr_abort_abort(message) + end subroutine torch_model_delete + + subroutine torch_tensor_delete(tensor) + type(torch_tensor), intent(inout) :: tensor + call shr_abort_abort(message) + end subroutine torch_tensor_delete + + subroutine torch_tensor_array_delete(tensor_array) + type(torch_tensor), dimension(:), intent(inout) :: tensor_array + call shr_abort_abort(message) + end subroutine torch_tensor_array_delete #endif end module FTorch_cesm_interface diff --git a/buildlib b/buildlib index ad4c603..bb37e5a 100755 --- a/buildlib +++ b/buildlib @@ -12,12 +12,29 @@ from standard_script_setup import * from CIME.utils import copyifnewer, run_bld_cmd_ensure_logging, expect from CIME.case import Case from CIME.build import get_standard_makefile_args -from CIME.Servers.wget import WGET -import glob -import zipfile logger = logging.getLogger(__name__) +# Actionable error message displayed when USE_FTORCH=TRUE but LibTorch cannot be found +_MISSING_TORCH_MSG = """ +PyTorch / LibTorch library not found (USE_FTORCH=TRUE). + Configured TORCH_DIR (required): {torch_dir} + + On NCAR Derecho, TORCH_DIR is configured automatically. + On user workstations or custom machines, point TORCH_DIR to your PyTorch + or LibTorch installation using one of the following options: + + 1. Use your active Python/Conda environment's PyTorch (e.g., ctsm_pylib): + ./xmlchange TORCH_DIR=$(python3 -c "import torch; print(torch.__path__[0])") + ---- OR ---- + 2. Download and extract standalone LibTorch (CPU) manually: + wget https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-latest.zip -O {installpath}/libtorch.zip + unzip {installpath}/libtorch.zip -d {installpath} + ./xmlchange TORCH_DIR={installpath}/libtorch + + Then rerun ./case.build. +""" + def parse_command_line(args, description): ############################################################################### @@ -60,29 +77,23 @@ def buildlib(bldroot, installpath, case): ftorch_dir = os.path.join(srcroot, "libraries", "FTorch") casetools = case.get_value("CASETOOLS") - # Path to the torch library - usually installed with pyTorch - torch_dir = case.get_value("TORCH_DIR") - torch_library = os.path.join(torch_dir, "lib", "libtorch_cpu.so") - if not os.path.isfile(torch_library): - torch_dir = os.path.join(installpath,"libtorch") - torch_library = os.path.join(torch_dir, "lib", "libtorch_cpu.so") - if not os.path.isfile(torch_library): - logger.info("libtorch not found, will attempt to use the nightly build from pytorch") - wget_server = WGET("https://download.pytorch.org") - zfile = os.path.join(installpath,"libtorch-shared-with-deps-latest.zip") - if not os.path.isfile(zfile): - wget_server.getfile("libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip", zfile) - with zipfile.ZipFile(zfile, 'r') as zip_ref: - zip_ref.extractall(installpath) - case.set_value("TORCH_DIR", torch_dir) - - expect(torch_dir,"TORCH_DIR not set in XML.\n") if case.get_value("USE_FTORCH"): use_ftorch = "ON" + # Path to the torch library - usually installed with PyTorch + torch_dir = case.get_value("TORCH_DIR") + expect( + torch_dir and os.path.isdir(torch_dir), + _MISSING_TORCH_MSG.format( + torch_dir=torch_dir if torch_dir else "(not set)", + installpath=installpath, + ), + ) + cmake_opts = '"-DCMAKE_PREFIX_PATH={} -DUSE_FTORCH={} "'.format( + torch_dir, use_ftorch + ) else: use_ftorch = "OFF" - - cmake_opts = '"-DCMAKE_PREFIX_PATH={} -DTORCH_LIBRARY={} -DUSE_FTORCH={} "'.format(torch_dir, torch_library, use_ftorch) + cmake_opts = '"-DUSE_FTORCH={} "'.format(use_ftorch) gmake_vars = ( "CASEROOT={caseroot} COMP_NAME=FTorch " @@ -108,8 +119,7 @@ def buildlib(bldroot, installpath, case): # This runs the FTorch cmake command from the cime case Makefile cmd = "{} {}".format(gmake_cmd, gmake_opts) - if not os.path.isfile(os.path.join(bldroot,"Makefile")): - run_bld_cmd_ensure_logging(cmd, logger, from_dir=bldroot) + run_bld_cmd_ensure_logging(cmd, logger, from_dir=bldroot) # This runs the pio make command from the cmake generated Makefile run_bld_cmd_ensure_logging( @@ -119,8 +129,6 @@ def buildlib(bldroot, installpath, case): ) - - def _main(argv, documentation): bldroot, installpath, caseroot = parse_command_line(argv, documentation) with Case(caseroot) as case: diff --git a/cmake/enable_shared_libs.cmake b/cmake/enable_shared_libs.cmake new file mode 100644 index 0000000..28f61e2 --- /dev/null +++ b/cmake/enable_shared_libs.cmake @@ -0,0 +1,8 @@ +# Re-enable shared library support (intended for Cray systems where CIME sets CMAKE_SYSTEM_NAME=Catamount, triggering make files that disable shared libs) +# Enable shared libs so we can link in C++ libtorch.so. +set_property(GLOBAL PROPERTY TARGET_SUPPORTS_SHARED_LIBS TRUE) +set(CMAKE_SHARED_LIBRARY_SUFFIX ".so") +set(CMAKE_FIND_LIBRARY_SUFFIXES ".so" ".a") +set(CMAKE_SHARED_LIBRARY_SONAME_C_FLAG "-Wl,-soname,") +set(CMAKE_SHARED_LIBRARY_SONAME_CXX_FLAG "-Wl,-soname,") +set(CMAKE_SHARED_LIBRARY_SONAME_Fortran_FLAG "-Wl,-soname,")