Vectorized histogramming

[jp-ga]

Add vectorization for Screen.reading histogram (see #351 )

Allows fast computation of vectorised histograms when automatic differentiation is not required (model evaluation, for example).

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist

• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code and checked that formatting passes (required).
• I have have fixed all issues found by flake8 (required).
• I have ensured that all pytest tests pass (required).
• I have run pytest on a machine with a CUDA GPU and made sure all tests pass (required).
• I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line.

[jank324]

I ran some benchmarks just to check if this makes sense given the presence of the new Cloud-in-Cell option and compared to the PyTorch-native implementation (in terms of speed).

Compare PyTorch-native histogram to the implementation from this PR

Compare the different methods in Cheetah at the state of this PR

[cr-xu]

I'm confused.. the native pytorch histogram crashed for large batching or the custom vectorized version?

otherwise the speed comparison is pretty good. It's faster than cloud-in-cell for large screen, which justifies its usage?

[jp-ga]

I'm confused.. the native pytorch histogram crashed for large batching or the custom vectorized version?

otherwise the speed comparison is pretty good. It's faster than cloud-in-cell for large screen, which justifies its usage?

It seems that the pytorch histogramdd was the one crashing on GPU because the numbers of the custom vectorized version match with the 'Compare the different methods in Cheetah at the state of this PR' bottom plot. There might be an issue with the histogramdd runs on GPU because the ones that did run were significantly slower than running on CPU. Either way, the custom vectorized histogramming seems to perform better in most usable cases

[jank324]

Yeah, so the PyTorch-native implementation crashed on MPS. This is an issue I've run into before. Very specifically the MPS implementation seems to require unusually large amounts of memory, so it's quite easy to make it crash on large inputs. But this isn't an issue with CUDA, so not a general GPU issue and probably not so relevant for most use cases.

I agree that generally it looks like the new vectorised implementation outperforms the PyTorch-native one. Tow things still confuse me:

• I would have thought the PyTorch-native one is faster because it should be just a single C++ kernel. But I guess it could be a somewhat neglected implementation. It was only added to PyTorch very late (after Cheetah), the MPS implementation has severe issues, it doesn't support vectorisation, and overall it is not exactly something widely used in neural network training implementations.
• More confusing is the "Non-vectorized 1M" case, which is an outlier in that the new implementation is suddenly and only here a lot slower than the original. I don't understand why that is, but it is reproducible.

I did some scans to kind try and understand that outlier. Interestingly, they look more mixed than I would have expected from the bar plots above.

[jp-ga]

The main reason why histogramdd is significantly slower in the vectorized case is because it isn't actually vectorized. It flattens the batch dimensions. That was the motivation to do the custom code.

The results you show are interesting though. I suggest we merge and then report these observations as an open issue for the future.

[jank324]

The main reason why histogramdd is significantly slower in the vectorized case is because it isn't actually vectorized. It flattens the batch dimensions. That was the motivation to do the custom code.

The results you show are interesting though. I suggest we merge and then report these observations as an open issue for the future.

Yes, of course I am aware of that. The case that confuses me is specifically a non-vectorised one.

[jp-ga]

The main reason why histogramdd is significantly slower in the vectorized case is because it isn't actually vectorized. It flattens the batch dimensions. That was the motivation to do the custom code.
The results you show are interesting though. I suggest we merge and then report these observations as an open issue for the future.

Yes, of course I am aware of that. The case that confuses me is specifically a non-vectorised one.

I see. I'm not sure what could be causing that behavior for high number of particles.

The tradeoff of the new implementation would be: 'native' vectorization which is significantly faster for large batch size (>10) but significantly slower for small batches when N>~50k. I still think it is worth it since the evaluation time is small when the batch size is small anyways, and we gain significant speedup for large batches (where the nominal evaluation time is large) + robustness in MPS GPU evaluations. Further optimization of the vectorized histogramming could be left as a separate future issue.