[kshyatt]

Needs the TensorOperations branch merged and tagged first, also happy if people want this allocator to be a kwarg. I think the scheduler should also possibly be a kwarg at least at top level? Anyway starting CI just to see how much works.

[lkdvos]

To give some information about the decision process of what I did in MPSKit:

The idea was that since allocators have state and it is important to not make them long-lived (don't accidentally hold on to the large buffers), nor incorrectly multithread with them (these are absolutely not thread-safe), I basically decided to not really expose that as part of the algorithms, to avoid mis-use. The claim being that there really isn't a good reason to not use the BufferAllocator, and you should be protected from memory corruption etc. In other words, this is not something you can configure in MPSKit.

As a result, within the body of the toplevel/outward-facing functions I simply create the bufferallocators, and then pass that down to all internal functions. In that sense it doesn't really matter all that much if it is keyword or regular argument, and decided this based on what "looked most natural" given the other arguments/keywords.

I agree that in principle you could reuse buffers between different runs of an algorithm, but this is ultimately just the cost of allocating/deallocating a small amount of large vectors, and I don't think we are at the point where that overhead really is dominant.

I do actually still think this might be the correct initial approach, it is not too hard to change afterwards and already gets the benefits, so I'd recommend/suggest following a similar approach here. As far as I can tell, this is exactly what you did, so honestly looks like a great PR to me.

Any chance we can get some benchmarks on this? I'd love to see also how much the CPU version improves from this, and how much of the allocations are still left. Otherwise for me this is good to go, of course with the caveat that we release the TensorOperations version

[kshyatt]

Let me first make sure CI passes 😅 but then yes benchmarks for sure! I think we should allow the scheduler to be settable though, rather than hardcoded.

[kshyatt]

this is ultimately just the cost of allocating/deallocating a small amount of large vectors, and I don't think we are at the point where that overhead really is dominant.

For AMD it might still be because they have no warm pool there, but we can benchmark that

[kshyatt]

Doing some benchmarking now. This needs QuantumKitHub/TensorKit.jl#608 (or something like it) and a TK release, as well as the TO release (which is not yet tagged)

[kshyatt]

One annoying complication I hit was that there's a huge slowdown on consumer cards like the 5090 I'm testing on for double precision. This won't be a problem on real HPC hardware but it is annoying. I'll try to get some fp32 benchmarks so we can get a rough idea.

[lkdvos]

should this scheduler then be passed down to the algorithm? Also, just want to double-check, but the SimultaneousCTMRG is presumably doing some multithreaded stuff so I guess this then cannot use a buffer? In the long term, it would probably be good to figure out a way to use something like a buffer pool to also enable this kind of workflow, but 🤷

[lkdvos]

looking at this, it might make sense to create some new primitive for threaded maps with allocators, since this is quite easy to accidentally pass a wrong allocator that would break. It's also not that simple to assert the allocator is correct so I'd be slightly more comfortable if we could end up fixing that in one spot. I'm not claiming this has to be in this PR though, opening an issue might actually already suffice

[lkdvos]

Suggested change

	MPSKit.expectation_value(
	peps::InfinitePEPS, O::LocalOperator, env;
	scheduler = Defaults.scheduler[],
	allocator = default_allocator(peps, scheduler)
	) = expectation_value(peps, O, peps, env; scheduler, allocator)
	MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, env; kwargs...) =
	expectation_value(peps, O, peps, env; kwargs...)

just to only have one place where these methods are defined.

[kshyatt]

Some benchmarking info: this + using CUDABufferAllocator cuts C4v CTMRG runtimes by up to 22% per iteration (D = 4, chi = 128), even in the worst case (D = 2, chi = 1024) there's a 5% gain.