Implement checkpointing for random GPU operators

[szkarpinski]

Category:

New feature (non-breaking change which adds functionality)

Description:

In this PR I add checkpointing support to fn.random operators on the GPU.

Additional information:

Affected modules and functionalities:

curand_states (from randomizer.cuh) now has get_states and set_states methods, which allow to export internal state or restore it.

RNGBase (from rng_base.h) can now save/restore states also on the GPU. The checkpoint of GPU random operator is a pointer to states (as returned by get_states). An ability to serialize/deserialize such checkpoints was added.

Key points relevant for the review:

Checkpoints are generally kept on the device

Note that the states returned by get_states are still on the GPU and set_states also expects its input to be on the GPU. RNGBase's Save and Restore also keep the checkpoint as a pointer to the GPU memory. The checkpoint is copied to the host at the very last moment during Serialize. This limits the device->host transfers significantly as we only copy the checkpoints that the user intends to save, while the checkpoints which we just trace and then delete never leave the device.

Serialization of curandState

During serialization, we treat curandState as a sequence of bytes, not caring about its structure. I consider it an implementation detail of curand.

Tests:

• Existing tests apply
• New tests added
  • Python tests
  • GTests
  • Benchmark
  • Other
• N/A

Checklist

Documentation

• Existing documentation applies
• Documentation updated
  • Docstring
  • Doxygen
  • RST
  • Jupyter
  • Other
• N/A

DALI team only

Requirements

• Implements new requirements
• Affects existing requirements
• N/A

REQ IDs: N/A

JIRA TASK: DALI-3532

[szkarpinski]

Found some performance problems, converting back to draft

[szkarpinski]

There are 1024 curand states kept per sample (48B each), serializing them one by one was not the best idea. Now a4c521e I serialize array of states to one big bytestring and it reduces the overhead by ~4x. Now the overhead is comparable to that of CPU random operators, so I'd leave it as it is and try to optimize them both in a follow-up.

[klecki]

Looks ok, one small question.

I also have more general question regarding synchronous operations. Are there any plans to make the checkpointing to run on some stream? Now we do the synchronous copies in few places. (I don't see this as a bad thing, the change is simpler thanks to that, but it may have some impact in cases where the checkpoints are created more often).

[stiepan]

Looks nice! Please remeber to add the more fine-grained synchronization to the roadmap.

[klecki]

One more place with state vs curand_state, otherwise ok.

[szkarpinski]

!build

[dali-automaton]

CI MESSAGE: [10787351]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [10787351]: BUILD FAILED

[szkarpinski]

!build

[dali-automaton]

CI MESSAGE: [10839737]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [10839737]: BUILD PASSED