[DRAFT] ITT380 add option to cache graph structures

[stadlmax]

Description

Repeated operation on edge_index structures (sorting, chunking) etc., technically don't contribute in a major fashion to the overall runtime, however, can add overheads by causing host-device syncs or by mostly being a series of small kernels exposing latencies. Overall, this also can lead to further latencies being exposed down the road due to

What problem does this change solve?

• add env-var which allows to trade-off memory and potentially reduced latencies by caching graph operations
• use index_select instead indexing kernels

What issue or task does this change relate to?

NA

Additional notes

Looking forward, moving operations like those into the start-up / initialization phase likely is more elegant. The index_select also currently stills adds overheads by not exploiting the indices of most masks being unique, i.e. there usually are compaction and other kernels scheduled beforehand to extract information out of index tensors.

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.