LOKI: use configurable custom pipeline

ecmwf-ifs · Nov 4, 2024 · f878cf3 · f878cf3
1 parent c52bbef
commit f878cf3
Show file tree

Hide file tree

Showing 3 changed files with 172 additions and 3 deletions.
diff --git a/src/ecwam/CMakeLists.txt b/src/ecwam/CMakeLists.txt
@@ -504,14 +504,12 @@ endif()
 if( HAVE_LOKI )
 
   set( LOKI_FRONTEND "fp" CACHE STRING "Frontend parser for Loki source transformations" )
-  set( SCC_ARGS "")
 
   if( LOKI_MODE MATCHES "idem|idem-stack" )
      set( LOKI_CONFIG_FILE ${CMAKE_CURRENT_SOURCE_DIR}/ecwam_loki.config )
   else()
      set( LOKI_CONFIG_FILE ${CMAKE_CURRENT_SOURCE_DIR}/ecwam_loki_gpu.config )
      target_compile_definitions( ${ecwam} PRIVATE WAM_GPU )
-     set( SCC_ARGS "TRIM_VECTOR_SECTIONS;GLOBAL_VAR_OFFLOAD")
   endif()
 
   # Apply Loki source file transformation to lib target
@@ -521,7 +519,6 @@ if( HAVE_LOKI )
     FRONTEND ${LOKI_FRONTEND}
     CONFIG ${LOKI_CONFIG_FILE}
     PLAN ${CMAKE_CURRENT_BINARY_DIR}/loki_plan_ecwam.cmake
-    ${SCC_ARGS}
     CPP
     DEFINITIONS WAM_GPU
     INCLUDES ${ecwam_intfb_includes}

diff --git a/src/ecwam/ecwam_loki.config b/src/ecwam/ecwam_loki.config
@@ -31,6 +31,58 @@ block = ['ec_parkind', 'parkind_wave']
   intrinsic_names = ['write(iu06']
   kernel_only = true
 
+# Inline transformation
+[transformations.InlineTransformation]
+  module = "loki.transformations"
+[transformations.InlineTransformation.options]
+  allowed_aliases = "IJ"
+  inline_elementals = false
+
+# Split-read-write transformation
+[transformations.SplitReadWriteTransformation]
+  module = "loki.transformations"
+[transformations.SplitReadWriteTransformation.options]
+  dimensions = "%dimensions.horizontal%"
+
+# TemporariesPoolAllocatorTransformation
+[transformations.TemporariesPoolAllocatorTransformation]
+  module = "loki.transformations"
+[transformations.TemporariesPoolAllocatorTransformation.options]
+  horizontal = "%dimensions.horizontal%"
+  directive = "openmp"
+  block_dim = "%dimensions.block_dim%"
+  check_bounds = true
+
+# ModuleWrap transformation
+[transformations.ModuleWrapTransformation]
+  module = "loki.transformations"
+[transformations.ModuleWrapTransformation.options]
+  module_suffix = "_MOD"
+
+# Dependency transformation
+[transformations.DependencyTransformation]
+  module = "loki.transformations"
+[transformations.DependencyTransformation.options]
+  module_suffix = "_MOD"
+  suffix = "_LOKI"
+
+# Idem transformation
+[transformations.IdemTransformation]
+  module = "loki.transformations"
+
+# loki pipelines
+[pipelines.idem]
+  transformations = [
+  'RemoveCodeTransformation', 'SplitReadWriteTransformation', 'InlineTransformation', 
+  'IdemTransformation', 'ModuleWrapTransformation', 'DependencyTransformation'
+]
+
+[pipelines.idem-stack]
+  transformations = [
+  'RemoveCodeTransformation', 'SplitReadWriteTransformation', 'InlineTransformation', 'IdemTransformation', 
+  'TemporariesPoolAllocatorTransformation', 'ModuleWrapTransformation', 'DependencyTransformation'
+]
+
 # Define entry point for call-tree transformation
 [routines.wamintgr]
   role = "driver"
@@ -46,6 +98,40 @@ block = ['ec_parkind', 'parkind_wave']
 [routines.transf]
 [routines.aki_ice]
 
+# we add loki inlined routines here to force them to be created
+[routines.femeanws]
+[routines.frcutindex]
+[routines.omegagc]
+[routines.tau_phi_hf]
+[routines.stresso]
+[routines.wsigstar]
+[routines.sinput]
+[routines.sinput_ard]
+[routines.sinput_jan]
+[routines.taut_z0]
+[routines.z0wave]
+[routines.airsea]
+[routines.femean]
+[routines.meansqs_lf]
+[routines.halphap]
+[routines.wnfluxes]
+[routines.sdiwbk]
+[routines.sbottom]
+[routines.fkmean]
+[routines.imphftail]
+[routines.setice]
+[routines.stokestrn]
+[routines.stokesdrift]
+[routines.cimsstrn]
+[routines.semean]
+[routines.sdepthlim]
+[routines.ciwabr]
+[routines.sinflx]
+[routines.sdissip_ard]
+[routines.sdissip_jan]
+[routines.sdissip]
+[routines.peak_ang]
+
 # Define indices and bounds for array dimensions
 [dimensions.horizontal]
   size = "KIJL"

diff --git a/src/ecwam/ecwam_loki_gpu.config b/src/ecwam/ecwam_loki_gpu.config
@@ -35,6 +35,92 @@ block = ['ec_parkind', 'parkind_wave', 'yowdrvtype']
   intrinsic_names = ['write(iu06']
   kernel_only = true
 
+# Inline transformation
+[transformations.InlineTransformation]
+  module = "loki.transformations"
+[transformations.InlineTransformation.options]
+  allowed_aliases = "IJ"
+  inline_elementals = false
+
+# Split-read-write transformation
+[transformations.SplitReadWriteTransformation]
+  module = "loki.transformations"
+[transformations.SplitReadWriteTransformation.options]
+  dimensions = "%dimensions.horizontal%"
+
+# GlobalVariableAnalysis
+[transformations.GlobalVariableAnalysis]
+  module = "loki.transformations"
+[transformations.GlobalVariableAnalysis.options]
+
+# GlobalVarOffloadTransformation
+[transformations.GlobalVarOffloadTransformation]
+  module = "loki.transformations"
+[transformations.GlobalVarOffloadTransformation.options]
+
+# SCC transformations
+[transformations.SCCVectorPipeline]
+  module = "loki.transformations.single_column"
+[transformations.SCCVectorPipeline.options]
+  horizontal = "%dimensions.horizontal%"
+  directive = "openacc"
+  trim_vector_sections = true
+
+[transformations.SCCStackPipeline]
+  module = "loki.transformations.single_column"
+[transformations.SCCStackPipeline.options]
+  horizontal = "%dimensions.horizontal%"
+  directive = "openacc"
+  trim_vector_sections = true
+  block_dim = "%dimensions.block_dim%"
+  check_bounds = false
+
+[transformations.SCCHoistPipeline]
+  module = "loki.transformations.single_column"
+[transformations.SCCHoistPipeline.options]
+  horizontal = "%dimensions.horizontal%"
+  directive = "openacc"
+  trim_vector_sections = true
+  block_dim = "%dimensions.block_dim%"
+
+# ModuleWrap transformation
+[transformations.ModuleWrapTransformation]
+  module = "loki.transformations"
+[transformations.ModuleWrapTransformation.options]
+  module_suffix = "_MOD"
+
+# Dependency transformation
+[transformations.DependencyTransformation]
+  module = "loki.transformations"
+[transformations.DependencyTransformation.options]
+  module_suffix = "_MOD"
+  suffix = "_LOKI"
+
+# FileWrite transformation
+[transformations.FileWriteTransformation]
+  module = "loki.transformations"
+[transformations.FileWriteTransformation.options]
+  include_module_var_imports = true
+
+# loki pipelines
+[pipelines.scc]
+  transformations = [
+  'RemoveCodeTransformation', 'SplitReadWriteTransformation', 'InlineTransformation', 'GlobalVariableAnalysis',
+  'GlobalVarOffloadTransformation', 'SCCVectorPipeline', 'ModuleWrapTransformation', 'DependencyTransformation'
+]
+
+[pipelines.scc-stack]
+  transformations = [
+  'RemoveCodeTransformation', 'SplitReadWriteTransformation', 'InlineTransformation', 'GlobalVariableAnalysis',
+  'GlobalVarOffloadTransformation', 'SCCStackPipeline', 'ModuleWrapTransformation', 'DependencyTransformation'
+]
+
+[pipelines.scc-hoist]
+  transformations = [
+  'RemoveCodeTransformation', 'SplitReadWriteTransformation', 'InlineTransformation', 'GlobalVariableAnalysis',
+  'GlobalVarOffloadTransformation', 'SCCHoistPipeline', 'ModuleWrapTransformation', 'DependencyTransformation'
+]
+
 # Define entry point for call-tree transformation
 [routines.wamintgr_loki_gpu]
   role = "driver"