polylithic(?) initialization of Kokkos

[streichler]

As described in #2651, Legion wants to initialize different Kokkos execution spaces from different threads. This clearly isn't possible with a single call to Kokkos::initialize, but Legion is almost able to make it work by directly calling the impl_initialize methods in each execution space (observing a few ordering requirements as it does so). The missing piece is that Kokkos::is_initialized() continues to return false, which causes some things to get unhappy (e.g.:

kokkos/core/src/impl/Kokkos_SharedAlloc.cpp

Lines 235 to 248 in e237df7

	if (!Kokkos::is_initialized()) {
	std::stringstream ss;
	ss << "Kokkos allocation \"";
	ss << arg_record->get_label();
	ss << "\" is being deallocated after Kokkos::finalize was called\n";
	auto s = ss.str();
	#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
	std::cerr << s;
	std::cerr << "This behavior is incorrect Kokkos usage, and will crash in "
	"future releases\n";
	#else
	Kokkos::Impl::throw_runtime_exception(s);
	#endif
	}

). Unfortunately, g_is_initialized is hidden inside an anonymous namespace in Kokkos_Core.cpp and I can't get to it from Legion.

Plausible remediations (one of which I strongly prefer) are:

1. Give that namespace a name so Legion code can go in and manually set g_is_initialized = true. This is no more or less kosher than the calls to impl_initialize methods now, but presumes that there's no Kokkos initialization that isn't done by at least one of the execution space initializations.
2. Change code that depends on Kokkos::is_initialized to be more precise about what it needs (e.g. the SharedAllocationRecord being valid in the case above, if I understand the code correctly). At this point, you'd wonder why you had Kokkos::is_initialized at all though.
3. Make it so that all execution spaces report is_initialized correctly in any thread (which would fix Unnecessary(?) check for host execution space initialization from Cuda initialization #2652 as well), and then call Kokkos::initialize after all the impl_initialize calls have been made.
4. Add an boolean flag called skip_execution_space_init to Kokkos::InitArguments that does all init except for execution spaces. A caller that sets this would be expected to "know what they're doing" and then call the execution space impl_initialize methods after the call to Kokkos::initialize.

Although any of these would technically work (and perhaps others may suggest additional approaches), my preference would be for (4). The first is a horrible hack (not that I'm above that!), and the second seems really fragile, and my concern with the third is that it will further strengthen the assumption in Kokkos right now that each execution space itself is "monolithic" (as opposed to, e.g. the Cuda execution space knowing there's two different GPUs, each of which might or might not be initialized yet).

[crtrott]

There is an option to make this a compile time behavior. I.e. that Kokkos::initialize doesn't call the impl_initialize of each execution space. For this use case here as some other use cases where interoperability with larger runtimes is important (i.e. maybe HPX) this could be a valid option. On the other hand we did have a request for the option of not initializing CUDA in a build which has CUDA enabled too.

[crtrott]

Also I would potentially start with adding an "impl" option first. Lets see whether this does all the things we need it to do for Legion support and then go form there.

[crtrott]

Ok we discussed and do it via a different a initialize call in the Impl namespace.

[dalg24]

Do we need anything else than what is already provided by Kokkos::Profiling::initialize();?

[gshipman]

Thanks for moving this along, LANL is very interested in this.

[janciesko]

Would the code below have the desired effect? I have restructured initialize_internal (./impl/Kokkos_Core.cpp) into initialize_backends() and initialize_profiling(). Compiling and running the example shown further down seems to work for (cond = false;).

void initialize_internal(const InitArguments& args) {
  if (args.disable_warnings) {
    g_show_warnings = false;
  }
  if(cond)
     initialize_backends(args); 
  initialize_profiling(args);
  g_is_initialized = true;
}

#include <Kokkos_Core.hpp>

int main(int argc, char* argv[]) {
Kokkos::initialize(argc, argv);
Kokkos::finalize();
return 1;
}

[streichler]

Would the code below have the desired effect?

Yes, that looks reasonable to me. It wasn't clear to me whether it'd be more thematically consistent for cond to be a field of InitArguments or another parameter, but I'm fine with either.

Here is the change that I've been using successfully in a fork: streichler@8b87035

Do we need anything else than what is already provided by Kokkos::Profiling::initialize();?

That may be equivalent at the moment, but my concern with this approach is that if anything else were ever to be added to Kokkos::initialize, this path would require additional changes (and awareness of the need for those changes) to avoid being broken.

[dalg24]

Do we need anything else than what is already provided by Kokkos::Profiling::initialize();?

That may be equivalent at the moment, but my concern with this approach is that if anything else were ever to be added to Kokkos::initialize, this path would require additional changes (and awareness of the need for those changes) to avoid being broken.

Fair enough.

[dalg24]

It wasn't clear to me whether it'd be more thematically consistent for cond to be a field of InitArguments or another parameter, but I'm fine with either.

We see this as an experimental feature at the moment and we don't want to expose it in the Kokkos:: namespace. Therefore it cannot be a new data member of InitArguments yet. This can be re-evaluated in the future.

[crtrott]

OK we got this merged in now. We don't have tests in place yet, will wait for you to check this out and tell us your usage pattern.

The relevant logic looks like this now:

void initialize(int& narg, char* arg[]) {
  InitArguments arguments;
  Impl::parse_command_line_arguments(narg, arg, arguments);
  Impl::parse_environment_variables(arguments);
  Impl::initialize_internal(arguments);
}

void initialize(InitArguments arguments) {
  Impl::parse_environment_variables(arguments);
  Impl::initialize_internal(arguments);
}

void initialize_profiling(const InitArguments& args) {
#if defined(KOKKOS_ENABLE_PROFILING)
  Kokkos::Profiling::initialize();
#else
  if (getenv("KOKKOS_PROFILE_LIBRARY") != nullptr) {
    std::cerr << "Kokkos::initialize() warning: Requested Kokkos Profiling, "
                 "but Kokkos was built without Profiling support"
              << std::endl;
  }
#endif
}

void pre_initialize_internal(const InitArguments& args) {
  if (args.disable_warnings) g_show_warnings = false;
}

void post_initialize_internal(const InitArguments& args) {
  initialize_profiling(args);
  g_is_initialized = true;
}

void initialize_internal(const InitArguments& args) {
  pre_initialize_internal(args);
  initialize_backends(args);
  post_initialize_internal(args);
}

So all in all someone needs to call:

void pre_initialize_internal(const InitArguments& args) ;
// Init all the backends you want
void post_initialize_internal(const InitArguments& args);

And you can do the check of environment variables independently too.

[crtrott]

@streichler Can you folks check whether this works for you?

[streichler]

Yes, I'll try to look at it early this week.

[streichler]

I've modified Legion to use Kokkos::Impl::{pre,post}_initialize and things looks good. Here's the initialization code we're using:
https://gitlab.com/StanfordLegion/legion/-/blob/c7c07f2d4d1df4531a3e03ca7e7ce467ab0cf182/runtime/realm/kokkos_interop.cc#L144-233
With these changes, we're able to pass our (limited, but better than nothing) CI tests for OpenMP and for Serial+CUDA Kokkos configurations.

In order to generate a unit test for Kokkos, I think it'd be sufficient to do something like:

#ifdef KOKKOS_ENABLE_xyz
std::thread([]{ Kokkos::xyz::impl_initialize(...); }).join();
#endif

for each execution space between the pre/post_initialize calls.

As a reminder, there's still one outstanding issue that must be addressed for CUDA+OpenMP execution in Kokkos-enabled Legion (or in the threaded initialization test proposed above): #2652