Separate checkpoint type execution paths

Checkpointing code is pretty unreadable because of a lot of branching in low-level functions for different checkpoint types (local, global, dist ckpt types). We should move the branching to the higher stack level, i.e. have a separate function/path for local checkpoint, global checkpoint, etc. For logic that is definitely common across paths, use helpers.

To be clear, this is not a feature request. More of a readability refactor.