Indexing -> torch.gather for better performance?

[squidrice21]

Hi,

Thanks for the great research and sharing the code! I've been using large-steps preconditioning for my project and it's working well. Recently, I worked on improving the performance of my pytorch code and found that changing from pytorch indexing to torch.gather gives me about 30x speed up. I also visualized the computational graph of the part using large-steps preconditioning via torchviz and saw DifferentiableSolveBackward -> IndexBackward0. This indexing seems to come from this repo. I'm pretty satisfied with the speed of the current version but just wondering if it can be further improved by switching to torch.gather.

[bathal1]

Hi,

Thanks for bringing this up ! Do you have more information regarding where indexing happens in the computation graph?

[squidrice21]

Thanks for replying. Here's my computation graph (generated using torchviz):

Note the IndexBackwards following the DifferentiableSolveBackward (I'm not sure why there are two streams of operations...). I think they are from the large-steps code. There are also some copyslices, which may introduce some performance overheads.

[bathal1]

Are you using the renderer implemented in this repo, or only the parameterization provided in the pip package ? If you are only using the parameterization, could you please provide a minimal reproducer for this behavior ?