Implementing OpenMP parallelization in UFS-Fire simulations

[masih-e]

We @pedro-jm @mkavulich are working on integrating the OpenMP-parallelized fire model (https://github.com/NCAR/fire_behavior/tree/develop) into the UFS Weather Model using the UFS SRW App (https://github.com/mkavulich/ufs-srweather-app/tree/feature/ufs_fire_updates), and we are running into issues with how OpenMP threading is handled on the FIRE cap within the ESMF coupling framework.

The fire model itself is already parallelized with OpenMP. In the SRW App, OpenMP is enabled using ufs.configure file where the number of OpenMP threads can be set for the ATM and FIRE:

Here is the ufs.configure for a 2-thread example:

#############################################
####  UFS Run-Time Configuration File  #####
#############################################
# ESMF #
logKindFlag:            ESMF_LOGKIND_MULTI_ON_ERROR
globalResourceControl:  true
# EARTH #
EARTH_component_list: ATM FIRE
EARTH_attributes::
  Verbosity = 0
::
# ATM #
ATM_model:                      fv3
ATM_petlist_bounds:             0 59
ATM_omp_num_threads:            2
ATM_attributes::
  Verbosity = 0
::
# FIRE #
FIRE_model:                     fire_behavior
FIRE_petlist_bounds:            60 61
FIRE_omp_num_threads:            2
FIRE_attributes::
  Verbosity = high
  Diagnostic = 0
::
# Run Sequence #
runSeq::
  @36
    ATM -> FIRE
    FIRE -> ATM :remapmethod=conserve
    ATM
    FIRE
  @
::

If we put 2 threads in the ATM and 1 in the FIRE the UFS-Fire sim works fine. The problem is when FIRE_omp_num_threads > 1 in the fire because the fire cap runs as a parallel region and thus calling initialization and other subs multiple times. Also each thread has is portion of the fire grid (ESMF is doing the domain decomposition for OpenMP).

This behavior conflicts with how the fire model is designed to run. We would like:

1. The FIRE cap to run serially by default.
2. OpenMP parallelism to be controlled explicitly inside the fire model using our own !$omp parallel do regions.

Is it expected that setting *_omp_num_threads > 1 causes the component cap itself to execute as a parallel OpenMP region? Is there a supported way to force the FIRE cap to run serially, while still allowing OpenMP parallelism inside the fire model? Any guidance, documentation pointers, or example configurations would be greatly appreciated.

[jdkublnick]

Hello @masih-e,

We have received your question, and I will be forwarding it to our internal SME's to gather a response.

Thank you,

Josh Kublnick
Epic User Support

[FernandoAndrade-NOAA]

@danrosen25 could you provide any guidance on this?

[danrosen25]

@theurich I don't see the issue with the configuration. Do you have any ideas on what might be going on inside the fire_behavior model?

[theurich]

Without looking into the fire_behavior cap and model itself I can only give limited insight. However, from NUOPC/ESMF level, this should work the way you were hoping. I.e. ESMF does NOT multi-thread your cap, etc. Instead, for the configuration shown

# FIRE #
FIRE_model:                     fire_behavior
FIRE_petlist_bounds:            60 61
FIRE_omp_num_threads:            2
FIRE_attributes::
  Verbosity = high
  Diagnostic = 0
::

i.e. with 2 PETs, and then omp_num_threads set to 2, the expected behavior is that your cap will run on 1 PET (which has access to 2 PEs, i.e. cores) and when user level OpenMP regions are encountered, that is where 2-way threading happens.

Can you confirm that behavior on the top level cap? I.e. 1 PET entering and executing the CAP. You might want to use ESMF_VMLog() to write out the VM info on the fire_behavior cap level (inside any of the specializations), to validate the expected VM setup.

[masih-e]

@theurich Thanks for the suggestion. I am going to rerun the code with additional print statements and share what I find once that run is complete.

One question: the fire cap source doesn’t currently store an ESMF_VM handle, so I’m not sure what object you expect to pass into ESMF_VMLog(), or where to put the call.

https://github.com/NCAR/fire_behavior/blob/98e93a953f30caf4789a52c42796b21eac18906d/nuopc/fire_behavior_nuopc.F90

[danrosen25]

@masih-e The VM is stored within the GridComp (the object extended to create a NUOPC model). The following call will retrieve the ESMF_VM object.
17.4.7 ESMF_GridCompGet - Get GridComp information

[masih-e]

@theurich @danrosen25 Thanks, I added VM logging at the start of Advertise as suggested.

call ESMF_GridCompGet(model, vm=vm, rc=rc) if (ESMF_LogFoundError(rcToCheck=rc, msg=ESMF_LOGERR_PASSTHRU, & line=__LINE__, & file=__FILE__)) & return ! bail out
call ESMF_VMLog(vm, rc=rc) if (ESMF_LogFoundError(rcToCheck=rc, msg=ESMF_LOGERR_PASSTHRU, & line=__LINE__, & file=__FILE__)) & return ! bail out

I believe here is the output from the PET60 and 61:

20260210 155307.813 INFO PET60 --- VMK::log() start -------------------------------------
20260210 155307.813 INFO PET60 vm located at: 0x86057c30
20260210 155307.813 INFO PET60 mpionly=1 threadsflag=0
20260210 155307.813 INFO PET60 ssiCount=1 localSsi=0
20260210 155307.813 INFO PET60 devCount=0 ssiLocalDevCount=0
20260210 155307.813 INFO PET60 petCount=2 ssiLocalPetCount=2
20260210 155307.813 INFO PET60 localPet=0 mypthid=23266495069056 ssiLocalPet=0 currentSsiPe=60
20260210 155307.813 INFO PET60 Current system level affinity pinning for local PET:
20260210 155307.813 INFO PET60 SSIPE=60
20260210 155307.813 INFO PET60 Current system level OMP_NUM_THREADS setting for local PET: 1
20260210 155307.813 INFO PET60 Current system level OMP_NUM_PROCS setting for local PET: 1
20260210 155307.813 INFO PET60 ESMF level logical PET->PE association (but consider system level affinity pinning!):
20260210 155307.813 INFO PET60 PET=0 lpid=0 tid=0 pid=60 peCount=1 accCount=0
20260210 155307.813 INFO PET60 PE=60 SSI=0 SSIPE=60
20260210 155307.813 INFO PET60 PET=1 lpid=1 tid=0 pid=61 peCount=1 accCount=0
20260210 155307.813 INFO PET60 PE=61 SSI=0 SSIPE=61
20260210 155307.813 INFO PET60 --- VMK::log() end ---------------------------------------

20260210 155307.813 INFO PET61 --- VMK::log() start -------------------------------------
20260210 155307.813 INFO PET61 vm located at: 0x94f36aa0
20260210 155307.813 INFO PET61 mpionly=1 threadsflag=0
20260210 155307.813 INFO PET61 ssiCount=1 localSsi=0
20260210 155307.813 INFO PET61 devCount=0 ssiLocalDevCount=0
20260210 155307.813 INFO PET61 petCount=2 ssiLocalPetCount=2
20260210 155307.813 INFO PET61 localPet=1 mypthid=23265372232576 ssiLocalPet=1 currentSsiPe=61
20260210 155307.813 INFO PET61 Current system level affinity pinning for local PET:
20260210 155307.813 INFO PET61 SSIPE=61
20260210 155307.813 INFO PET61 Current system level OMP_NUM_THREADS setting for local PET: 1
20260210 155307.813 INFO PET61 Current system level OMP_NUM_PROCS setting for local PET: 1
20260210 155307.813 INFO PET61 ESMF level logical PET->PE association (but consider system level affinity pinning!):
20260210 155307.813 INFO PET61 PET=0 lpid=0 tid=0 pid=60 peCount=1 accCount=0
20260210 155307.813 INFO PET61 PE=60 SSI=0 SSIPE=60
20260210 155307.813 INFO PET61 PET=1 lpid=1 tid=0 pid=61 peCount=1 accCount=0
20260210 155307.813 INFO PET61 PE=61 SSI=0 SSIPE=61
20260210 155307.813 INFO PET61 --- VMK::log() end ---------------------------------------

I also added prints at the beginning/end of each subroutine (Advertise/Realize/SetClock/Advance/...) and all prints appear twice in the log. The code crashes in advance and does not leave the subroutine. I think it's crashing in this line

https://github.com/NCAR/fire_behavior/blob/98e93a953f30caf4789a52c42796b21eac18906d/nuopc/fire_behavior_nuopc.F90#L810C1-L811C1

grid%fz0(1:grid%nx,1:grid%ny) = ptr_z0(clb(1):cub(1),clb(2):cub(2)) * 0.01

I added the following prints in the advance

`
print *, 'Masih Enter Advance'

print *, 'grid%nx, grid%ny', grid%nx, grid%ny

print *, 'clb(1):cub(1),clb(2):cub(2)', clb(1),cub(1),clb(2),cub(2)
`

Output:
Masih Enter Advance
grid%nx, grid%ny 200 200
clb(1):cub(1),clb(2):cub(2) 1 100 1 200
Masih Enter Advance
grid%nx, grid%ny 200 200
clb(1):cub(1),clb(2):cub(2) 1 100 1 200

Note the size mismatch (200×200 vs 100×200) where we expected to be the same.

[theurich]

@masih-e Thank you for the output. Seeing it I wonder if the fire_behavior component provides a SetVM routine.... okay, I do see it here https://github.com/NCAR/fire_behavior/blob/6311b0104238175f8b8c351380c9ca030c38370c/nuopc/fire_behavior_nuopc.F90#L28, so that is okay.

However, next I wonder if the component is integrated into UFS correctly calling the SetVM routine. Can someone please point me to the ufs driver Fortran code that is used to integrate the component into UFS SRWeather?

[theurich]

Okay I think I see it here: https://github.com/ufs-community/ufs-weather-model/blob/a796a4f616286f1f6a765b53d3c0302986f9a971/driver/UFSDriver.F90#L89C1-L91C7

So that's the problem: there needs to be also FIRE_BEHAVIOR_SV => SetVM, which needs to be used down in line

ufs-weather-model/driver/UFSDriver.F90

Line 491 in a796a4f

call NUOPC_DriverAddComp(driver, trim(prefix), FIRE_BEHAVIOR_SS, &

.

E.g. look at how this is done for WW3 with WW3_SS and WW3_SV, or any other component that uses both the SS and the SV.

[masih-e]

Thanks again for your help. @theurich Following up, I wired SetVM into the UFS driver as you suggested in driver/UFSDriver.F90

#ifdef FRONT_FIRE_BEHAVIOR
      use FRONT_FIRE_BEHAVIOR, only: FIRE_BEHAVIOR_SS => SetServices, &
                                     FIRE_BEHAVIOR_SV => SetVM
#endif

#ifdef FRONT_FIRE_BEHAVIOR
          if (trim(model) == "fire_behavior") then
            call NUOPC_DriverAddComp(driver, trim(prefix), FIRE_BEHAVIOR_SS, &
              FIRE_BEHAVIOR_SV, &
              petList=petList, comp=comp, rc=rc)
            if (ChkErr(rc,__LINE__,u_FILE_u)) return
            found_comp = .true.
          end if
#endif

My question now is what the expected SetVM logic should be inside the fire_behavior NUOPC cap. I looked at WW3 for the pattern; I see it has public SetServices, SetVM and it also use NUOPC_Model, only: NUOPC_ModelGet, SetVM in the cap. I don't know how that SetVM from NUOPC_Model is intended to be used in the fire cap. If you can point me to the exact call you expect inside the component SetVM (if there are WW3-equivalent lines that actually apply the VM), I can implementation in fire_behavior and rerun to see if this resolves the duplicated runs and partition mismatch in the fire advance.

[theurich]

@masih-e nothing further to do, the SetVM comes from NUOPC_Model already. For clarity, you might want to add it explicitly to the use line, as you see in WW3, but since you don't have the only statement in your use NUOPC_Model, you get the use association by default.

What is happening is that NUOPC (though NUOPC_Model) provides you with a default SetVM implementation. And for most typical cases that is a good implementation you can directly use. You could specialize SetVM locally in your model if you needed additional control of how the incoming parent component resources (PETs and PEs) are managed, but that is not typically needed. The default SetVM provided by NUOPC_Model does the expected management wrt the FIRE_omp_num_threads attribute.

So now that you wired this correctly into driver/UFSDriver.F90, you are ready to run the test, and you should see that you arrive with a single PET in the fire_behavior component, but with 2 OpenMP threads available whenever user level OpenMP regions in the component are encountered.

[masih-e]

Thanks for the help @theurich I reran the model after wiring SetVM through UFSDriver.F90, but the logs look unchanged compared to the previous run:

20260212 132745.877 INFO             PET60   --- VMK::log() start -------------------------------------
20260212 132745.877 INFO             PET60   vm located at: 0x907e0c30
20260212 132745.877 INFO             PET60   mpionly=1 threadsflag=0
20260212 132745.877 INFO             PET60   ssiCount=1 localSsi=0
20260212 132745.877 INFO             PET60   devCount=0 ssiLocalDevCount=0
20260212 132745.877 INFO             PET60   petCount=2 ssiLocalPetCount=2
20260212 132745.877 INFO             PET60   localPet=0 mypthid=22760573660032 ssiLocalPet=0 currentSsiPe=60
20260212 132745.877 INFO             PET60   Current system level affinity pinning for local PET:
20260212 132745.877 INFO             PET60    SSIPE=60
20260212 132745.877 INFO             PET60   Current system level OMP_NUM_THREADS setting for local PET: 1
20260212 132745.877 INFO             PET60   Current system level OMP_NUM_PROCS setting for local PET: 1
20260212 132745.878 INFO             PET60   ESMF level logical PET->PE association (but consider system level affinity pinning!):
20260212 132745.878 INFO             PET60   PET=0 lpid=0 tid=0 pid=60 peCount=1 accCount=0
20260212 132745.878 INFO             PET60    PE=60 SSI=0 SSIPE=60
20260212 132745.878 INFO             PET60   PET=1 lpid=1 tid=0 pid=61 peCount=1 accCount=0
20260212 132745.878 INFO             PET60    PE=61 SSI=0 SSIPE=61
20260212 132745.878 INFO             PET60   --- VMK::log() end ---------------------------------------

20260212 132745.877 INFO             PET61   --- VMK::log() start -------------------------------------
20260212 132745.877 INFO             PET61   vm located at: 0x820faaa0
20260212 132745.877 INFO             PET61   mpionly=1 threadsflag=0
20260212 132745.877 INFO             PET61   ssiCount=1 localSsi=0
20260212 132745.877 INFO             PET61   devCount=0 ssiLocalDevCount=0
20260212 132745.877 INFO             PET61   petCount=2 ssiLocalPetCount=2
20260212 132745.877 INFO             PET61   localPet=1 mypthid=22690109410176 ssiLocalPet=1 currentSsiPe=61
20260212 132745.877 INFO             PET61   Current system level affinity pinning for local PET:
20260212 132745.877 INFO             PET61    SSIPE=61
20260212 132745.877 INFO             PET61   Current system level OMP_NUM_THREADS setting for local PET: 1
20260212 132745.877 INFO             PET61   Current system level OMP_NUM_PROCS setting for local PET: 1
20260212 132745.878 INFO             PET61   ESMF level logical PET->PE association (but consider system level affinity pinning!):
20260212 132745.878 INFO             PET61   PET=0 lpid=0 tid=0 pid=60 peCount=1 accCount=0
20260212 132745.878 INFO             PET61    PE=60 SSI=0 SSIPE=60
20260212 132745.878 INFO             PET61   PET=1 lpid=1 tid=0 pid=61 peCount=1 accCount=0
20260212 132745.878 INFO             PET61    PE=61 SSI=0 SSIPE=61
20260212 132745.878 INFO             PET61   --- VMK::log() end ---------------------------------------

[theurich]

Sorry, I should have noticed this earlier when you quoted the code that you added. You also need to pass the info object. This is where the OpenMP relevant information is passed into NUOPC_DriverAddComp(). So the full calling block looks like this:

#ifdef FRONT_FIRE_BEHAVIOR
          if (trim(model) == "fire_behavior") then
            call NUOPC_DriverAddComp(driver, trim(prefix), FIRE_BEHAVIOR_SS, &
              FIRE_BEHAVIOR_SV, info=info, &
              petList=petList, comp=comp, rc=rc)
            if (ChkErr(rc,__LINE__,u_FILE_u)) return
            found_comp = .true.
          end if
#endif

[masih-e]

Just a quick update on this issue. The model is now running as expected after the suggested changes and the original problem has been resolved. Thanks for the suggestions and support!

[theurich]

Perfect! Thank you for the update.