From 5845e17341cd5633af8308c114ac426d92d436b2 Mon Sep 17 00:00:00 2001 From: Samuel Omlin Date: Tue, 18 Jul 2023 20:23:28 +0200 Subject: [PATCH] set Enzyme runtime activity to true for threads by default --- src/AD.jl | 2 +- src/ParallelKernel/AD.jl | 2 +- src/ParallelKernel/init_parallel_kernel.jl | 1 + src/ParallelKernel/shared.jl | 7 ------- 4 files changed, 3 insertions(+), 9 deletions(-) diff --git a/src/AD.jl b/src/AD.jl index 6c7d39f9..024d302b 100644 --- a/src/AD.jl +++ b/src/AD.jl @@ -44,7 +44,7 @@ Provides GPU-compatible wrappers for automatic differentiation functions of the main() !!! note "Enzyme runtime activity default" - If ParallelStencil is initialized with Threads, then `Enzyme.API.runtimeActivity!(true)` is called at module load time to ensure correct behavior of Enzyme. If you want to disable this behavior, then call `Enzyme.API.runtimeActivity!(false)` after loading ParallelStencil. + If ParallelStencil is initialized with Threads, then `Enzyme.API.runtimeActivity!(true)` is called to ensure correct behavior of Enzyme. If you want to disable this behavior, then call `Enzyme.API.runtimeActivity!(false)` after loading ParallelStencil. To see a description of a function type `?`. """ diff --git a/src/ParallelKernel/AD.jl b/src/ParallelKernel/AD.jl index e40a950e..f03978d2 100644 --- a/src/ParallelKernel/AD.jl +++ b/src/ParallelKernel/AD.jl @@ -11,7 +11,7 @@ Provides GPU-compatible wrappers for automatic differentiation functions of the - `autodiff_deferred_thunk!`: wraps function `autodiff_deferred_thunk`. !!! note "Enzyme runtime activity default" - If ParallelKernel is initialized with Threads, then `Enzyme.API.runtimeActivity!(true)` is called at module load time to ensure correct behavior of Enzyme. If you want to disable this behavior, then call `Enzyme.API.runtimeActivity!(false)` after loading ParallelStencil. + If ParallelKernel is initialized with Threads, then `Enzyme.API.runtimeActivity!(true)` is called to ensure correct behavior of Enzyme. If you want to disable this behavior, then call `Enzyme.API.runtimeActivity!(false)` after loading ParallelStencil. To see a description of a function type `?`. """ diff --git a/src/ParallelKernel/init_parallel_kernel.jl b/src/ParallelKernel/init_parallel_kernel.jl index cacb3c10..7479ccf9 100644 --- a/src/ParallelKernel/init_parallel_kernel.jl +++ b/src/ParallelKernel/init_parallel_kernel.jl @@ -35,6 +35,7 @@ function init_parallel_kernel(caller::Module, package::Symbol, numbertype::DataT pkg_import_cmd = :() end ad_import_cmd = :(import ParallelStencil.ParallelKernel.Enzyme) + if (package == PKG_THREADS) ad_import_cmd = :(import ParallelStencil.ParallelKernel.Enzyme; Enzyme.API.runtimeActivity!(true)) end # NOTE: Enzyme requires this currently to work correctly with threads. if !isdefined(caller, :Data) || (@eval(caller, isa(Data, Module)) && length(symbols(caller, :Data)) == 1) # Only if the module Data does not exist in the caller or is empty, create it. if (datadoc_call==:()) if (numbertype == NUMBERTYPE_NONE) datadoc_call = :(@doc ParallelStencil.ParallelKernel.DATA_DOC_NUMBERTYPE_NONE Data) diff --git a/src/ParallelKernel/shared.jl b/src/ParallelKernel/shared.jl index 1544db4c..41818f3d 100644 --- a/src/ParallelKernel/shared.jl +++ b/src/ParallelKernel/shared.jl @@ -82,13 +82,6 @@ macro ranges() esc(RANGES_VARNAME) end macro rangelengths() esc(:(($(RANGELENGTHS_VARNAMES...),))) end -## MODULE INITIALIZATION FUNCTION - -function __init__() - if (get_package() == PKG_THREADS) Enzyme.API.runtimeActivity!(true) end # NOTE: Enzyme requires this currently to work correctly with threads. -end - - ## FUNCTIONS TO GET CREATE AND MANAGE CUDA STREAMS, AMDGPU QUEUES AND "ROCSTREAMS" @static if ENABLE_CUDA