broadcast_inputs triggers tensor storage copy, peaks CUDA memory consumption

[hxu296]

Summary

It seems that the following line in def broadcast_inputs(x, y) triggerred a tensor storage copy that caused a CUDA memory overflow when I tried to run a small bundle adjustment dataset with 31843 pixel observations. Both reshape and contiguous could trigger a memory copy. If we can avoid memory copy in broadcast_inputs, we can avoid overflowing CUDA memory at this step.

pypose/pypose/lietensor/operation.py

Line 914 in 6598a84

x = x.expand(shape+(x.shape[-1],)).reshape(-1,x.shape[-1]).contiguous()

Improvements

refactor broadcast_inputs to not use reshape and contiguous.

Risks

TBD

Involved components

• pypose/pypose/lietensor/operation.py

Optional: Intended side effects

TBD

Optional: Missing test coverage

TBD

[wang-chen]

@hxu296 This is a historical issue, the functions were implemented by Cuda and the tensor copy behavior was expected. But now it seems not needed anymore. A similar function could be torch.broadcast_tensors and torch.broadcast_shapes, you may re-check those functions called after broadcast_inputs, if they are implemented by batch, then broadcast_inputs can be safely removed.