WIP Binomial checkpoint

[sriharikrishna]

I am creating this WIP PR to get figure out how to to explicit cache values during the reverse pass.
This is related to #2039 where it is pointed out that we should/can perform an x = cache[i / M]

// reverse pass


for i = reverse(1:N) {
	x = cache[i / M]
	// mini fwd pass
	for i2 = i / M .. i {
		x = foo(x)
	}

	grad_foo(x)
}

I am not sure how to go about it in src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp reverseWithCheckpointingRevolve in the portion related to rvStore.

[sriharikrishna]

Bump @Pangoraw

[Pangoraw]

I am not sure how to go about it in src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp reverseWithCheckpointingRevolve in the portion related to rvStore.

Could you give more details about your blocker here?

[sriharikrishna]

Apologies for the delayed response.
I follow how, for periodic checkpointing
cache[i / M] = x
happens in the single forward sweep followed by
x = cache[i / M]
in the single reverse sweep.

In binomial checkpointing, the reverse sweep is not a monolith. Instead it consists of
new mini-forward sweeps and mini-reverse sweeps. Additionally, more values are cached and read.

So, the question is whether the infrastructure to implement the caching and reading exists once the reverse-sweep has begun.

[Pangoraw]

i had a go at another implementation in #2485 if you want to take a look.

[Pangoraw]

Closing as #2485 as been merged, feel free to open new issues / pr on top of the merged implementation.