WIP: Experimental/corbett/jit #1333

corbett5 · 2021-02-23T01:29:01Z

This is a proof of concept for a dynamic library based JIT. It takes the constitutive and element dispatch out of regionBasedKernelApplication and instead compiles and links in the kernel at runtime. I confirmed that it is passing the SSLE_sedov_01 problem on Quartz with Clang debug and GCC release and on Lassen with Clang debug.

It should pass all the one rank problems if I were to #include the appropriate headers for all the various kernels at the bottom of KernelBase.hpp like I did for SolidMechanicsSmallStrainExplicitNewmarkKernel.hpp. As is there are some compilation errors (with the main code not the JIT) if I do this, but structuring things like this is a terrible idea so I just proceeded with my hack.

Each run re-compiles the kernels it needs, this is wasteful but would be easily fixed.

Multiple rank runs fail because each rank tries to compile the kernel, easily fixed with some synchronization.

Related to
GEOS-DEV/LvArray#225

corbett5 · 2021-02-23T01:32:27Z

@rrsettgast @klevzoff @francoishamon take a look.

TotoGaz · 2021-03-03T08:16:01Z

Hi @corbett5 ! Would this jit approach means that instead of deploying GEOSX we need to provide the end-user with a complete GEOSX build environment? Will I still be able to work with models that I do not want to share (i.e. without accessing the sources)?

corbett5 · 2021-03-03T18:51:08Z

@TotoGaz yes you will be able to pre-compile things.

…d jit compilation, testing and configuration still needed

…ss to avoid cumbersome preprocessor usage, only managed to remove the kernel name string from the arglist, which at least enforces the name being correct

…sedov 1-rank passing

…uartz

…due to file conflicts if multiple runs try to compile a kernel simultaneously

wrtobin · 2021-09-01T23:11:15Z

src/coreComponents/finiteElement/kernelInterface/KernelBase.hpp

+  #define KernelDispatch KernelDispatchJIT
+#else
+  #define KernelDispatch KernelDispatchTemplate
+#endif


Per the comments, this is the part of the current implementation I want to change the most. It should be possible, but wound up eating hours while I was trying to get it implemented the way I would prefer.

It should be possible to make the two dispatcher classes similar by replacing NAME and HEADER parameters in KernelDispatchJIT with with just KERNEL_TYPE parameter (i.e. template template parameter, like in the non-JIT dispatcher).

NAME could then be extracted as (I'm trying to decide if this is 100% robust):

string const fullName = LvArray::system::demangleType< KERNEL_TYPE< SUBREGION_TYPE, FE_TYPE, CONSTITUTIVE_TYPE > >(); string const name = fullName.substr( 0, fullName.find( '<' ) );

HEADER could be replaced by adding inside each "leaf" kernel class something like

template< typename SUBREGION_TYPE, typename FE_TYPE, typename CONSTITUTIVE_TYPE > class MyKernel { static constexpr char const source_location[] = __FILE__; ... }; // Sad part: this is required in C++14 to avoid linker errors template< typename SUBREGION_TYPE, typename FE_TYPE, typename CONSTITUTIVE_TYPE > constexpr char const MyKernel< SUBREGION_TYPE, FE_TYPE, CONSTITUTIVE_TYPE >::source_location[];

and accessing as

string const header = KERNEL_TYPE< SUBREGION_TYPE, FE_TYPE, CONSTITUTIVE_TYPE >::source_location;

The downside is having this extra stuff to remember to put in kernels. Maybe there's a better way to associate kernel class with source file name?

These ideas aren't too pretty, just thought I'd mention them. They allow us to gets rid of JITTI_DECL/JITTI_TPARAM macros entirely and unify the two dispatchers, possibly even making them specializations as the comment above suggests, and allowing for a per-kernel JITting decision.

I think the end goal is to always use JITTI, although that doesn't mean that we need to always jit things at run-time. We need to have the capability to pre-jit everything at build-time, and I think that would suffice.

wrtobin · 2021-09-01T23:14:07Z

I'll need to merge from dev and revert the submodule change prior to merge, but this is functional and in a good place to review. I still need to check some things on lassen, but was able to get 100% of the regression tests passing on quartz. It did require multiple passes as it is possible for multiple tests to try to compile the same kernel simultaneously resulting in file system errors that kill the runs.

This had to be avoiding internally via some mpi checks, but when multiple executables are running simultaneously it might require adding sentinel files or something similar, I'm open to ideas to mitigate that issue.

klevzoff

Looks great! Just have one suggestion (not necessarily a good one) re: unifying the two dispatchers.

klevzoff · 2021-09-09T23:05:22Z

src/coreComponents/finiteElement/CMakeLists.txt

   )
 #
 # Specify all sources
 #
 set( finiteElement_sources
     FiniteElementDiscretization.cpp
     FiniteElementDiscretizationManager.cpp
+     kernelInterface/kernelJIT.cpp


This file is not included in the PR, and I can't see it being generated anywhere either

klevzoff · 2021-09-09T23:12:40Z

src/coreComponents/finiteElement/kernelInterface/KernelBase.hpp

+    camp::tuple< ARGS ... > m_args;
+  };
+
+  jitti::CompilationInfo getKernelCompilationInfo( const string & header );


I can't find the definition of this function anywhere. Is it in kernelJIT.cpp?

klevzoff · 2021-09-10T07:57:47Z

src/coreComponents/finiteElement/kernelInterface/KernelBase.hpp

+  #define KernelDispatch KernelDispatchJIT
+#else
+  #define KernelDispatch KernelDispatchTemplate
+#endif


It should be possible to make the two dispatcher classes similar by replacing NAME and HEADER parameters in KernelDispatchJIT with with just KERNEL_TYPE parameter (i.e. template template parameter, like in the non-JIT dispatcher).

NAME could then be extracted as (I'm trying to decide if this is 100% robust):

string const fullName = LvArray::system::demangleType< KERNEL_TYPE< SUBREGION_TYPE, FE_TYPE, CONSTITUTIVE_TYPE > >(); string const name = fullName.substr( 0, fullName.find( '<' ) );

HEADER could be replaced by adding inside each "leaf" kernel class something like

template< typename SUBREGION_TYPE, typename FE_TYPE, typename CONSTITUTIVE_TYPE > class MyKernel { static constexpr char const source_location[] = __FILE__; ... }; // Sad part: this is required in C++14 to avoid linker errors template< typename SUBREGION_TYPE, typename FE_TYPE, typename CONSTITUTIVE_TYPE > constexpr char const MyKernel< SUBREGION_TYPE, FE_TYPE, CONSTITUTIVE_TYPE >::source_location[];

and accessing as

string const header = KERNEL_TYPE< SUBREGION_TYPE, FE_TYPE, CONSTITUTIVE_TYPE >::source_location;

The downside is having this extra stuff to remember to put in kernels. Maybe there's a better way to associate kernel class with source file name?

These ideas aren't too pretty, just thought I'd mention them. They allow us to gets rid of JITTI_DECL/JITTI_TPARAM macros entirely and unify the two dispatchers, possibly even making them specializations as the comment above suggests, and allowing for a per-kernel JITting decision.

TotoGaz · 2021-09-10T16:49:54Z

src/coreComponents/finiteElement/CMakeLists.txt

+add_custom_command( OUTPUT ${CMAKE_BINARY_DIR}/include/kernelJITCompileCommands.hpp
+                    WORKING_DIRECTORY ${CMAKE_BINARY_DIR} 
+                    COMMAND python ${CMAKE_CURRENT_LIST_DIR}/../LvArray/src/jitti/generateCompileCommandsHeader.py 
+                                   ${CMAKE_BINARY_DIR}/compile_commands.json
+                                   --cpp ${CMAKE_CURRENT_LIST_DIR}/kernelInterface/kernelJIT.cpp 
+                                   --hpp ${CMAKE_BINARY_DIR}/include/kernelJITCompileCommands.hpp
+                                   --include ${CMAKE_BINARY_DIR}/include 
+                                   --linker ${CMAKE_CXX_COMPILER} )


Instead of relying on headers that you generate on the file system, could it have been more robust to store the C++ pre-processing output as a string alongside the compile/link commands directly into the lib?
And that would be more "modern" JIT style imho.

Yeah I think you could do that, although it would require the user to list out the functions that they want to JIT at configuration time in CMake. Then for each of those functions we could pre-process jitti/templateSource.cpp. But only up to a point, since we don't know what template params the user will want. Then at run time when we know the params we could compile the pre-processed source with the additional command line definitions of JITTI_TEMPLATE_PARAMS and JITTI_TEMPLATE_PARAMS_STRING.

However, how we pass a string embedded in our library to the compiler without using the file system is beyond me. Not to mention that at the end of the day the compiler is going to spit out a library on the file system that we then have to open. So this frees us from having to have access to the same headers used to build, but unless we do something really fancy I think file system access is a requirement.

Also at least with CUDA the cost of the compilation time itself will greatly outweigh the cost of opening and pre-processing the source.

Also at least with CUDA the cost of the compilation time itself will greatly outweigh the cost of opening and pre-processing the source.

It's more a problem of having a consistent self contained GEOSX installation than a performance problem imho.
Embedding the pre-processed source files directly into GEOSX would help at non relying on sources that may be modified behind the scenes.
I do not know all the JIT details so maybe it's not a problem anyway: I do not want to solve problems that do not exist.

corbett5 · 2021-09-13T03:24:56Z

src/cmake/GeosxOptions.cmake

+option( ENABLE_JITTI "Build all compute kernels just-in-time at runtime." OFF )
+
+if ( ENABLE_JITTI )
+  message( "JITTI is ENABLED")


I think the way we've been doing this is GEOSX_ENABLE_JITTI, or maybe it's GEOSX_USE_JITTI, I forget. But either way I like defined/not defined instead of 1/0.

corbett5 · 2021-09-13T03:33:14Z

src/coreComponents/finiteElement/kernelInterface/KernelBase.hpp

+  #define KernelDispatch KernelDispatchJIT
+#else
+  #define KernelDispatch KernelDispatchTemplate
+#endif


I think the end goal is to always use JITTI, although that doesn't mean that we need to always jit things at run-time. We need to have the capability to pre-jit everything at build-time, and I think that would suffice.

corbett5 · 2021-10-07T16:40:50Z

src/coreComponents/finiteElement/kernelInterface/kernelJIT.cpp

+  constexpr bool compilerIsNVCC = false;
+#endif
+
+    jitti::CompilationInfo getKernelCompilationInfo( const string & header )


So my idea is that for each kernel header we'd have a source that includes the header and has the method to get thejitti::CompilationInfo(s) for the kernel(s) in the header. That way if the kernel (header) changes and the user rebuilds GEOSX then the compile time in the compilation info gets updated and so things will be re-jitted.

I think with this implementation it won't rejit if say the solid mechanics kernels are modified and GEOSX is rebuilt, only if some dependency of KernelBase.hpp is modified.

Yeah that is probably true in the current implementation.

This is the first thing I'll take a look at changing from the current implementation, and then I'll look at adding in the compile-time pre-jit.

Sounds good.

corbett5 changed the title ~~Experimental/corbett/jit~~ WIP: Experimental/corbett/jit Feb 23, 2021

corbett5 marked this pull request as draft February 23, 2021 01:30

rrsettgast requested review from wrtobin, rrsettgast and klevzoff July 22, 2021 17:39

wrtobin added 9 commits August 26, 2021 08:53

should be full impl for kernels/sparsity kernels for both standard an…

ece63c9

…d jit compilation, testing and configuration still needed

switching to full JIT and doing some small cleanup

599f4ca

attempted unifying jit/template classes into a single specialized cla…

c340fa4

…ss to avoid cumbersome preprocessor usage, only managed to remove the kernel name string from the arglist, which at least enforces the name being correct

adding header info to template params, needed for compilation, basic …

301734c

…sedov 1-rank passing

resolve mpi issues, jitti-active build passing integration tests on q…

8016eeb

…uartz

missed from last commit

5860e90

add full-jit cmake option

ea0ab5c

all regression tests passing, though multiple passes can be required …

16068c3

…due to file conflicts if multiple runs try to compile a kernel simultaneously

submodule update I'll have to revert before the merge can go in

9bc353d

wrtobin force-pushed the experimental/corbett/jit branch from 7916cda to 9bc353d Compare September 1, 2021 23:07

wrtobin self-assigned this Sep 1, 2021

wrtobin added the flag: ready for review label Sep 1, 2021

wrtobin reviewed Sep 1, 2021

View reviewed changes

wrtobin added flag: requires updated submodule(s) type: feature New feature or request labels Sep 1, 2021

rrsettgast requested a review from wrtobin September 9, 2021 17:32

klevzoff approved these changes Sep 10, 2021

View reviewed changes

TotoGaz reviewed Sep 10, 2021

View reviewed changes

corbett5 commented Sep 13, 2021

View reviewed changes

missing file

d610d60

corbett5 commented Oct 7, 2021

View reviewed changes

francoishamon mentioned this pull request Dec 21, 2021

Implemented new kernel interface for compositional properties #1698

Merged

rrsettgast closed this Jul 3, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

WIP: Experimental/corbett/jit #1333

WIP: Experimental/corbett/jit #1333

corbett5 commented Feb 23, 2021

corbett5 commented Feb 23, 2021

TotoGaz commented Mar 3, 2021

corbett5 commented Mar 3, 2021

wrtobin Sep 1, 2021

klevzoff Sep 10, 2021 •

edited

Loading

corbett5 Sep 13, 2021

wrtobin commented Sep 1, 2021 •

edited

Loading

klevzoff left a comment

klevzoff Sep 9, 2021

klevzoff Sep 9, 2021

klevzoff Sep 10, 2021 •

edited

Loading

TotoGaz Sep 10, 2021

corbett5 Sep 13, 2021

TotoGaz Sep 13, 2021

corbett5 Sep 13, 2021

corbett5 Sep 13, 2021

corbett5 Oct 7, 2021

wrtobin Oct 7, 2021

corbett5 Oct 7, 2021

WIP: Experimental/corbett/jit #1333

WIP: Experimental/corbett/jit #1333

Conversation

corbett5 commented Feb 23, 2021

corbett5 commented Feb 23, 2021

TotoGaz commented Mar 3, 2021

corbett5 commented Mar 3, 2021

Choose a reason for hiding this comment

klevzoff Sep 10, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wrtobin commented Sep 1, 2021 • edited Loading

klevzoff left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

klevzoff Sep 10, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

klevzoff Sep 10, 2021 •

edited

Loading

wrtobin commented Sep 1, 2021 •

edited

Loading

klevzoff Sep 10, 2021 •

edited

Loading