Skip to content

[WIP] Enabling Kokkos RDC #33161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
NamjaeChoi wants to merge 1 commit into
base: next
Choose a base branch
from
Draft

[WIP] Enabling Kokkos RDC #33161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 24 additions & 10 deletions framework/app.mk
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,19 +396,26 @@ $(app_KOKKOS_LIB): curr_objs := $(app_KOKKOS_OBJECTS)
$(app_KOKKOS_LIB): $(app_KOKKOS_OBJECTS)
@echo "Linking Kokkos Library "$@"..."
@bash -c '$(libmesh_LIBTOOL) --tag=CXX $(LIBTOOLFLAGS) --mode=link --quiet \
$(KOKKOS_CXX) -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) -rpath $(curr_dir)/lib ${SILENCE_SOME_WARNINGS}'
$(KOKKOS_CXX) -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS) -rpath $(curr_dir)/lib ${SILENCE_SOME_WARNINGS}'
@$(libmesh_LIBTOOL) --mode=install --quiet install -c $@ $(curr_dir)/lib

else
else ifeq ($(KOKKOS_COMPILER),NVCC)

$(app_KOKKOS_LIB): curr_dir := $(APPLICATION_DIR)
$(app_KOKKOS_LIB): curr_objs := $(app_KOKKOS_OBJECTS)
$(app_KOKKOS_LIB): $(app_KOKKOS_OBJECTS)
@mkdir -p $(curr_dir)/lib
@echo "Linking Kokkos Library "$@"..."
@ar rcs $@ $(curr_objs)

# libtool ignores nvcc and just uses mpicxx to link, so cannot be used
else

$(app_KOKKOS_LIB): curr_dir := $(APPLICATION_DIR)
$(app_KOKKOS_LIB): curr_objs := $(app_KOKKOS_OBJECTS)
$(app_KOKKOS_LIB): $(app_KOKKOS_OBJECTS)
@mkdir -p $(curr_dir)/lib
@echo "Linking Kokkos Library "$@"..."
@$(KOKKOS_CXX) --shared -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS)
@$(KOKKOS_CXX) --shared -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS)

endif

Expand Down Expand Up @@ -452,19 +459,26 @@ $(app_KOKKOS_TEST_LIB): curr_objs := $(app_KOKKOS_TEST_OBJECTS)
$(app_KOKKOS_TEST_LIB): $(app_KOKKOS_TEST_OBJECTS)
@echo "Linking Kokkos Test Library "$@"..."
@bash -c '$(libmesh_LIBTOOL) --tag=CXX $(LIBTOOLFLAGS) --mode=link --quiet \
$(KOKKOS_CXX) -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) -rpath $(curr_dir)/lib ${SILENCE_SOME_WARNINGS}'
$(KOKKOS_CXX) -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS) -rpath $(curr_dir)/lib ${SILENCE_SOME_WARNINGS}'
@$(libmesh_LIBTOOL) --mode=install --quiet install -c $@ $(curr_dir)/lib

else
else ifeq ($(KOKKOS_COMPILER),NVCC)

$(app_KOKKOS_TEST_LIB): curr_dir := $(APPLICATION_DIR)/test
$(app_KOKKOS_TEST_LIB): curr_objs := $(app_KOKKOS_TEST_OBJECTS)
$(app_KOKKOS_TEST_LIB): $(app_KOKKOS_TEST_OBJECTS)
@mkdir -p $(curr_dir)/lib
@echo "Linking Kokkos Test Library "$@"..."
@ar rcs $@ $(curr_objs)

# libtool ignores nvcc and just uses mpicxx to link, so cannot be used
else

$(app_KOKKOS_TEST_LIB): curr_dir := $(APPLICATION_DIR)/test
$(app_KOKKOS_TEST_LIB): curr_objs := $(app_KOKKOS_TEST_OBJECTS)
$(app_KOKKOS_TEST_LIB): $(app_KOKKOS_TEST_OBJECTS)
@mkdir -p $(curr_dir)/lib
@echo "Linking Kokkos Test Library "$@"..."
@$(KOKKOS_CXX) --shared -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS)
@$(KOKKOS_CXX) --shared -o $@ $(curr_objs) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS)

endif

Expand All @@ -480,7 +494,7 @@ $(KOKKOS_DEVICE_LINK_OBJECT): curr_dir := $(APPLICATION_DIR)
$(KOKKOS_DEVICE_LINK_OBJECT): $(KOKKOS_OBJECTS)
@mkdir -p $(curr_dir)/lib
@echo "Device Linking Kokkos Objects..."
@$(KOKKOS_CXX) -dlink -o $@ $(KOKKOS_LDFLAGS) $(KOKKOS_OBJECTS) $(KOKKOS_LIBS)
@$(KOKKOS_CXX) -dlink -o $@ $(KOKKOS_LDFLAGS) $(KOKKOS_OBJECTS) $(KOKKOS_LIBS) $(libmesh_LIBS)

endif

Expand Down Expand Up @@ -583,7 +597,7 @@ endif
$(app_EXEC): $(app_LIBS) $(MOOSE_KOKKOS_LIB) $(app_KOKKOS_LIBS) $(app_KOKKOS_TEST_LIB) $(KOKKOS_DEVICE_LINK_OBJECT) $(mesh_library) $(main_object) $(app_test_LIB) $(depend_test_libs) $(ADDITIONAL_EXEC_OBJECTS)
@echo "Linking Executable "$@"..."
@bash -c '$(libmesh_LIBTOOL) --tag=CXX $(LIBTOOLFLAGS) --mode=link --quiet \
$(libmesh_CXX) $(libmesh_CXXFLAGS) -o $@ $(main_object) $(if $(filter yes,$(HAVE_NO_AS_NEEDED)),$(NO_AS_NEEDED_FLAG)) $(depend_test_libs_flags) $(applibs) $(KOKKOS_DEVICE_LINK_OBJECT) $(ADDITIONAL_LIBS) $(ADDITIONAL_EXEC_OBJECTS) $(LDFLAGS) $(libmesh_LDFLAGS) $(libmesh_LIBS) $(EXTERNAL_FLAGS) ${SILENCE_SOME_WARNINGS}'
$(libmesh_CXX) $(libmesh_CXXFLAGS) -o $@ $(main_object) $(if $(filter yes,$(HAVE_NO_AS_NEEDED)),$(NO_AS_NEEDED_FLAG)) $(depend_test_libs_flags) $(applibs) $(KOKKOS_DEVICE_LINK_OBJECT) $(ADDITIONAL_LIBS) $(ADDITIONAL_EXEC_OBJECTS) $(LDFLAGS) $(libmesh_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS) $(EXTERNAL_FLAGS) ${SILENCE_SOME_WARNINGS}'
@$(codesign)

###### install stuff #############
Expand Down
73 changes: 41 additions & 32 deletions framework/kokkos.mk
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,38 +69,42 @@ ifeq ($(PETSC_HAVE_SYCL),1)
endif

ifneq ($(PETSC_HAVE_CUDA),)
KOKKOS_DEVICE := CUDA
KOKKOS_ARCH := $(KOKKOS_CUDA_ARCH_$(CUDA_ARCH))
KOKKOS_COMPILER := NVCC
KOKKOS_CXX = $(CUDA_COMPILER)
KOKKOS_CXXFLAGS = -arch=sm_$(CUDA_ARCH) --extended-lambda
KOKKOS_CXXFLAGS += --forward-unknown-to-host-compiler --disable-warnings -x cu -ccbin $(word 1, $(libmesh_CXX))
KOKKOS_CXXFLAGS += $(filter-out -Werror=return-type,$(CXXFLAGS) $(libmesh_CXXFLAGS)) # Incompatible with NVCC
KOKKOS_CPPFLAGS = $(subst -Werror,-Werror=all-warnings,$(libmesh_CPPFLAGS) $(ADDITIONAL_CPPFLAGS) ${ADDITIONAL_KOKKOS_CPPFLAGS})
KOKKOS_LDFLAGS = --forward-unknown-to-host-compiler -arch=sm_$(CUDA_ARCH)
KOKKOS_DEVICE := CUDA
KOKKOS_ARCH := $(KOKKOS_CUDA_ARCH_$(CUDA_ARCH))
KOKKOS_COMPILER := NVCC
KOKKOS_CXX = $(CUDA_COMPILER)
KOKKOS_CXXFLAGS = -arch=sm_$(CUDA_ARCH) --extended-lambda --relocatable-device-code=true --dlink-time-opt
KOKKOS_CXXFLAGS += --forward-unknown-to-host-compiler --disable-warnings -x cu -ccbin $(word 1, $(libmesh_CXX))
KOKKOS_CXXFLAGS += $(filter-out -Werror=return-type,$(CXXFLAGS) $(libmesh_CXXFLAGS)) # Incompatible with NVCC
KOKKOS_CPPFLAGS = $(subst -Werror,-Werror=all-warnings,$(libmesh_CPPFLAGS) $(ADDITIONAL_CPPFLAGS) ${ADDITIONAL_KOKKOS_CPPFLAGS})
KOKKOS_LDFLAGS = --forward-unknown-to-host-compiler --relocatable-device-code=true --dlink-time-opt -arch=sm_$(CUDA_ARCH)
KOKKOS_LIBS = -lpetsckokkos
libmesh_PETSC_KOKKOS_LIBS := $(filter %petsckokkos %petsckokkosdlink,$(libmesh_LIBS))
libmesh_LIBS := $(filter-out %petsckokkos %petsckokkosdlink,$(libmesh_LIBS))
else ifneq ($(PETSC_HAVE_HIP),) # To be determined for HIP
KOKKOS_DEVICE := HIP
KOKKOS_ARCH :=
KOKKOS_COMPILER := HIPCC
KOKKOS_CXX = $(HIP_COMPILER)
KOKKOS_CXXFLAGS =
KOKKOS_CPPFLAGS =
KOKKOS_LDFLAGS =
KOKKOS_DEVICE := HIP
KOKKOS_ARCH :=
KOKKOS_COMPILER := HIPCC
KOKKOS_CXX = $(HIP_COMPILER)
KOKKOS_CXXFLAGS =
KOKKOS_CPPFLAGS =
KOKKOS_LDFLAGS =
else ifneq ($(PETSC_HAVE_SYCL),) # To be determined for SYCL
KOKKOS_DEVICE := SYCL
KOKKOS_ARCH :=
KOKKOS_COMPILER := DPCPP
KOKKOS_CXX = $(SYCL_COMPILER)
KOKKOS_CXXFLAGS = -fsycl
KOKKOS_CPPFLAGS =
KOKKOS_LDFLAGS =
KOKKOS_DEVICE := SYCL
KOKKOS_ARCH :=
KOKKOS_COMPILER := DPCPP
KOKKOS_CXX = $(SYCL_COMPILER)
KOKKOS_CXXFLAGS = -fsycl
KOKKOS_CPPFLAGS =
KOKKOS_LDFLAGS =
else
KOKKOS_COMPILER := CPU
KOKKOS_CXX = $(libmesh_CXX)
KOKKOS_CXXFLAGS = $(CXXFLAGS) $(libmesh_CXXFLAGS) -x c++
KOKKOS_CPPFLAGS = $(libmesh_CPPFLAGS) $(ADDITIONAL_CPPFLAGS) ${ADDITIONAL_KOKKOS_CPPFLAGS}
KOKKOS_LDFLAGS =
ENABLE_KOKKOS_GPU := false
KOKKOS_COMPILER := CPU
KOKKOS_CXX = $(libmesh_CXX)
KOKKOS_CXXFLAGS = $(CXXFLAGS) $(libmesh_CXXFLAGS) -x c++
KOKKOS_CPPFLAGS = $(libmesh_CPPFLAGS) $(ADDITIONAL_CPPFLAGS) ${ADDITIONAL_KOKKOS_CPPFLAGS}
KOKKOS_LDFLAGS =
KOKKOS_LIBS =
ENABLE_KOKKOS_GPU := false
endif

ifeq ($(ENABLE_KOKKOS_GPU),false)
Expand All @@ -113,15 +117,21 @@ endif
KOKKOS_CXXFLAGS += -DMOOSE_KOKKOS_SCOPE=1 -DMETAPHYSICL_KOKKOS_COMPILATION=1 -fPIC
KOKKOS_LDFLAGS += $(libmesh_LDFLAGS)
KOKKOS_INCLUDE = $(libmesh_INCLUDE)
KOKKOS_LIBS = $(libmesh_LIBS)

ifeq ($(METHOD),opt)
KOKKOS_CXXFLAGS += -DNDEBUG
endif

ifeq ($(KOKKOS_COMPILER),CPU)
KOKKOS_LIB_SUFFIX := _kokkos-$(METHOD).la
else ifeq ($(KOKKOS_COMPILER),NVCC)
KOKKOS_LIB_SUFFIX := _kokkos-$(METHOD).a
else
KOKKOS_LIB_SUFFIX := _kokkos-$(METHOD).so
endif

ifeq ($(KOKKOS_COMPILER),CPU)

KOKKOS_LIB_SUFFIX := _kokkos-$(METHOD).la
KOKKOS_OBJ_SUFFIX := $(libmesh_HOST).$(METHOD).lo

%.$(KOKKOS_OBJ_SUFFIX) : %.K
Expand All @@ -131,7 +141,6 @@ KOKKOS_OBJ_SUFFIX := $(libmesh_HOST).$(METHOD).lo

else

KOKKOS_LIB_SUFFIX := _kokkos-$(METHOD).so
KOKKOS_OBJ_SUFFIX := $(libmesh_HOST).$(METHOD).o

%.$(KOKKOS_OBJ_SUFFIX) : %.K
Expand Down
24 changes: 20 additions & 4 deletions framework/moose.mk
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,15 +520,20 @@ ifeq ($(KOKKOS_COMPILER),CPU)
$(MOOSE_KOKKOS_LIB): $(MOOSE_KOKKOS_OBJECTS)
@echo "Linking Kokkos Library "$@"..."
@bash -c '$(libmesh_LIBTOOL) --tag=CXX $(LIBTOOLFLAGS) --mode=link --quiet \
$(KOKKOS_CXX) -o $@ $(MOOSE_KOKKOS_OBJECTS) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) -rpath $(FRAMEWORK_DIR) ${SILENCE_SOME_WARNINGS}'
$(KOKKOS_CXX) -o $@ $(MOOSE_KOKKOS_OBJECTS) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS) -rpath $(FRAMEWORK_DIR) ${SILENCE_SOME_WARNINGS}'
@$(libmesh_LIBTOOL) --mode=install --quiet install -c $(MOOSE_KOKKOS_LIB) $(FRAMEWORK_DIR)

else ifeq ($(KOKKOS_COMPILER),NVCC)

$(MOOSE_KOKKOS_LIB): $(MOOSE_KOKKOS_OBJECTS)
@echo "Linking Kokkos Library "$@"..."
@ar rcs $@ $(MOOSE_KOKKOS_OBJECTS)

else

# libtool ignores nvcc and just uses mpicxx to link, so cannot be used
$(MOOSE_KOKKOS_LIB): $(MOOSE_KOKKOS_OBJECTS)
@echo "Linking Kokkos Library "$@"..."
@$(KOKKOS_CXX) --shared -o $@ $(MOOSE_KOKKOS_OBJECTS) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS)
@$(KOKKOS_CXX) --shared -o $@ $(MOOSE_KOKKOS_OBJECTS) $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(libmesh_LIBS)

endif

Expand Down Expand Up @@ -626,6 +631,17 @@ exodiff_objects := $(patsubst %.C, %.$(obj-suffix), $(exodiff_srcfiles))
exodiff_includes := -I$(exodiff_DIR) $(libmesh_INCLUDE)
# dependency files
exodiff_deps := $(patsubst %.C, %.$(obj-suffix).d, $(exodiff_srcfiles))
# Kokkos CUDA strips out -lpetsckokkos and -lpetsckokkosdlink from libmesh_LIBS, because
# those static archives contain CUDA registration/device images and need to be linked
# exactly once per program (namely, intermediate shared libraries should not link them,
# and only the final executable should link them). Since exodiff is an independent program,
# those flags should be added back for proper linking.
exodiff_link_group_end := -Wl,--end-group
exodiff_LIBS = $(if $(libmesh_PETSC_KOKKOS_LIBS), \
$(if $(findstring $(exodiff_link_group_end),$(libmesh_LIBS)), \
$(subst $(exodiff_link_group_end),$(libmesh_PETSC_KOKKOS_LIBS) $(exodiff_link_group_end),$(libmesh_LIBS)), \
$(libmesh_LIBS) $(libmesh_PETSC_KOKKOS_LIBS)), \
$(libmesh_LIBS))

all: exodiff

Expand All @@ -637,7 +653,7 @@ $(exodiff_objects): | prebuild
$(exodiff_APP): $(exodiff_objects)
@echo "Linking Executable "$@"..."
@bash -c '$(libmesh_LIBTOOL) --tag=CXX $(LIBTOOLFLAGS) --mode=link --quiet \
$(libmesh_CXX) $(libmesh_CPPFLAGS) $(libmesh_CXXFLAGS) $(libmesh_INCLUDE) $(exodiff_objects) -o $@ $(LDFLAGS) $(libmesh_LDFLAGS) $(libmesh_LIBS) $(EXTERNAL_FLAGS) ${SILENCE_SOME_WARNINGS}'
$(libmesh_CXX) $(libmesh_CPPFLAGS) $(libmesh_CXXFLAGS) $(libmesh_INCLUDE) $(exodiff_objects) -o $@ $(LDFLAGS) $(libmesh_LDFLAGS) $(exodiff_LIBS) $(EXTERNAL_FLAGS) ${SILENCE_SOME_WARNINGS}'

-include $(wildcard $(exodiff_DIR)/*.d)

Expand Down
4 changes: 0 additions & 4 deletions framework/src/kokkos/problems/KokkosFEProblemBase.K
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,10 +136,6 @@ FEProblemBase::hasKokkosFunction(const std::string & name) const
Moose::Kokkos::Function
FEProblemBase::getKokkosFunction(const std::string & name)
{
#ifdef MOOSE_ENABLE_KOKKOS_GPU
mooseError("Retrieving a Kokkos function as abstract type is currently not supported for GPU.");
#endif

const auto & function = getKokkosFunction<Moose::FunctionBase>(name);

std::shared_ptr<Moose::Kokkos::FunctionWrapperHostBase> wrapper =
Expand Down
4 changes: 2 additions & 2 deletions test/tests/kokkos/functions/constant_function/tests
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
input = 'kokkos_constant_function.i'
exodiff = 'kokkos_constant_function_out.e'
requirement = 'The Kokkos function system shall include a constant function and the ability to retrieve it as an abstract type.'
capabilities = 'kokkos & !cuda'
compute_devices = 'cpu'
capabilities = 'kokkos'
compute_devices = 'cpu cuda'
[]
[]
4 changes: 2 additions & 2 deletions test/tests/kokkos/functions/default_function/tests
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
input = 'kokkos_default_function.i'
exodiff = 'kokkos_default_function_out.e'
requirement = 'The Kokkos function system shall include the ability to set default values for input parameters expecting a function name.'
capabilities = 'kokkos & !cuda'
compute_devices = 'cpu'
capabilities = 'kokkos'
compute_devices = 'cpu cuda'
[]
[]
16 changes: 8 additions & 8 deletions test/tests/kokkos/functions/piecewise_constant/tests
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,35 +7,35 @@
type = 'Exodiff'
input = 'kokkos_piecewise_constant.i'
exodiff = 'kokkos_piecewise_constant_out.e'
capabilities = 'kokkos & !cuda'
compute_devices = 'cpu'
capabilities = 'kokkos'
compute_devices = 'cpu cuda'
detail = 'provided separately as two vectors in the input,'
[]
[xy_data]
type = 'Exodiff'
input = 'kokkos_piecewise_constant.i'
cli_args = 'Kernels/diff/coef=func_xy_data AuxKernels/coef_aux/function=func_xy_data'
exodiff = 'kokkos_piecewise_constant_out.e'
capabilities = 'kokkos & !cuda'
compute_devices = 'cpu'
capabilities = 'kokkos'
compute_devices = 'cpu cuda'
detail = 'provided as a table in the input,'
[]
[csv]
type = 'Exodiff'
input = 'kokkos_piecewise_constant.i'
cli_args = 'Kernels/diff/coef=func_csv AuxKernels/coef_aux/function=func_csv'
exodiff = 'kokkos_piecewise_constant_out.e'
capabilities = 'kokkos & !cuda'
compute_devices = 'cpu'
capabilities = 'kokkos'
compute_devices = 'cpu cuda'
detail = 'provided as a CSV file,'
[]
[json]
type = 'Exodiff'
input = 'kokkos_piecewise_constant.i'
cli_args = 'Kernels/diff/coef=func_json AuxKernels/coef_aux/function=func_json'
exodiff = 'kokkos_piecewise_constant_out.e'
capabilities = 'kokkos & !cuda'
compute_devices = 'cpu'
capabilities = 'kokkos'
compute_devices = 'cpu cuda'
detail = 'provided as a JSON file.'
[]
[]
Expand Down
2 changes: 1 addition & 1 deletion test/tests/kokkos/materials/parsed/tests
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,6 @@
exodiff = 'kokkos_parsed_material_out.e'
requirement = 'The Kokkos system shall support parsed material property calculations.'
capabilities = 'kokkos'
compute_devices = 'cpu'
compute_devices = 'cpu cuda'
[]
[]