1: /*
  2:     Macros for preloading a block of code so that startup costs (such as first-touch
  3:     allocation) are not attributed to the measured run. See `PetscPreLoadBegin()`.
  4: */
  5: #pragma once

  7: #include <petscsys.h>

  9: /* MANSEC = Sys */
 10: /* SUBMANSEC = Log */

 12: /*MC
 13:   PetscPreLoadBegin - Begin a block of code that is timed twice so that startup costs (such as JIT or first-touch allocation) are not attributed to the measured run

 15:   Synopsis:
 16: #include <petscpreload.h>
 17:   PetscPreLoadBegin(PetscBool flag, const char name[])

 19:   Not Collective; No Fortran Support

 21:   Input Parameters:
 22: + flag - whether preloading is desired (may be overridden by the command-line option `-preload`)
 23: - name - name to use for the `PetscLogStage` created for the measured run

 25:   Level: intermediate

 27:   Note:
 28:   Use this macro in the form
 29: .vb
 30:     PetscPreLoadBegin(PETSC_TRUE, "Compute");
 31:     // ... code to be timed
 32:     PetscPreLoadEnd();
 33: .ve
 34:   It expands to a loop that runs the enclosed code twice when preloading is enabled and once otherwise.
 35:   Pair with `PetscPreLoadEnd()`. Use `PetscPreLoadStage()` to advance to a new named stage between phases.

 37: .seealso: `PetscPreLoadEnd()`, `PetscPreLoadStage()`, `PetscLogStageRegister()`, `PetscLogStagePush()`
 38: M*/
 39: #define PetscPreLoadBegin(flag, name) \
 40:   do { \
 41:     PetscBool     PetscPreLoading = flag; \
 42:     int           PetscPreLoadMax, PetscPreLoadIt; \
 43:     PetscLogStage _stageNum; \
 44:     PetscCall(PetscOptionsGetBool(NULL, NULL, "-preload", &PetscPreLoading, NULL)); \
 45:     PetscPreLoadMax     = (int)(PetscPreLoading); \
 46:     PetscPreLoadingUsed = PetscPreLoading ? PETSC_TRUE : PetscPreLoadingUsed; \
 47:     PetscCall(PetscLogStageGetId(name, &_stageNum)); \
 48:     for (PetscPreLoadIt = (_stageNum == -1) ? 0 : PetscPreLoadMax; PetscPreLoadIt <= PetscPreLoadMax; PetscPreLoadIt++) { \
 49:       PetscPreLoadingOn = (PetscBool)(PetscPreLoadIt < PetscPreLoadMax); \
 50:       PetscCall(PetscBarrier(NULL)); \
 51:       if (_stageNum == -1) PetscCall(PetscLogStageRegister(name, &_stageNum)); \
 52:       PetscCall(PetscLogStageSetActive(_stageNum, (PetscBool)(PetscPreLoadIt == PetscPreLoadMax))); \
 53:       PetscCall(PetscLogStagePush(_stageNum))

 55: /*MC
 56:   PetscPreLoadEnd - Close a preload block started with `PetscPreLoadBegin()`

 58:   Synopsis:
 59: #include <petscpreload.h>
 60:   PetscPreLoadEnd()

 62:   Not Collective; No Fortran Support

 64:   Level: intermediate

 66: .seealso: `PetscPreLoadBegin()`, `PetscPreLoadStage()`
 67: M*/
 68: #define PetscPreLoadEnd() \
 69:   PetscCall(PetscLogStagePop()); \
 70:   } \
 71:   } \
 72:   while (0)

 74: /*MC
 75:   PetscPreLoadStage - Advance to a new named stage inside a `PetscPreLoadBegin()` / `PetscPreLoadEnd()` block

 77:   Synopsis:
 78: #include <petscpreload.h>
 79:   PetscPreLoadStage(const char name[])

 81:   Not Collective; No Fortran Support

 83:   Input Parameter:
 84: . name - name for the new `PetscLogStage`

 86:   Level: intermediate

 88:   Note:
 89:   Pops the previous stage, registers (or reuses) the new stage, and pushes it. Only valid between
 90:   `PetscPreLoadBegin()` and `PetscPreLoadEnd()`.

 92: .seealso: `PetscPreLoadBegin()`, `PetscPreLoadEnd()`, `PetscLogStagePush()`, `PetscLogStagePop()`
 93: M*/
 94: #define PetscPreLoadStage(name) \
 95:   do { \
 96:     PetscCall(PetscLogStagePop()); \
 97:     PetscCall(PetscLogStageGetId(name, &_stageNum)); \
 98:     if (_stageNum == -1) PetscCall(PetscLogStageRegister(name, &_stageNum)); \
 99:     PetscCall(PetscLogStageSetActive(_stageNum, (PetscBool)(!PetscPreLoadMax || PetscPreLoadIt))); \
100:     PetscCall(PetscLogStagePush(_stageNum)); \
101:   } while (0)

103: PETSC_EXTERN PetscBool PetscPreLoadingUsed; /* true if we are or have done preloading */
104: PETSC_EXTERN PetscBool PetscPreLoadingOn;   /* true if we are currently in a preloading calculation */