Implement MA for Base.Rational

[blegat]

@ericphanson started it in https://discourse.julialang.org/t/memory-consumtion-rational-bigint-on-bernoulli-number-example/33812/7

[ericphanson]

I think it would be great to have efficient implementations for rationals in MA!

There were three things that came to mind when I was working on the discourse post that I thought I’d mention here.

First, since rationals aren’t actually immutable (but the bigints they hold are), this is just a special case of the general problem of wanting to use MA within library code that wasn’t written for it. I think it’s a worthwhile special case though, since rationals of bigint are the only way I know of getting exact results from many computations and one can easily imagine users using MA with them. But maybe the general problem can be addressed using Cassette to add @rewrite statements into library code? (In this case, I think one would still need to add a few more primitives for BigInts, like gcd as well as checked_* operations, unless those are already no-ops on BigInts since they don’t overflow). Then I could imagine a rewrite_recursive macro exported from MA that sets up the Cassette context and does the overdubbing to just add @rewrite to all arithmetic it finds.

Secondly, it seems like you can do a fair bit better than just copying Base and adding @rewrite everywhere. In the code in the post, I try to reuse buffers and the output values where possible (though there may be even better tricks I’ve missed). It was fun to do, like a puzzle, but it also seems like something a compiler could be doing better— I think one essentially wants to figure out the dependencies between variables and then use that info to choose the order to do things in and what memory to use for temporary variables. That makes me wonder if there’s anyway to get Julia’s optimizer to do some of these transformations, or maybe somehow set it up so LLVM knows it’s free to do these things.

Or maybe one can even write a fancier version of @rewrite which is given some buffers to use and tries to decide the best way to use them. You’ll notice almost all the code I wrote for rationals just uses the MA API for BigInts, which makes me think the idea makes some sense. In any case, such automatic optimizations might be able to be used in the hypothetical Cassette recursive rewriting.

The third point is related to the word “almost” in the previous paragraph: I found Base.GMP.MPZ.set! very useful, as a kind of copyto! for BigInts. Would it make sense to add such a function to the MA API?

[ericphanson]

One other idea: MA could globally hold some buffers to be used in @rewrite; these could be separate for each thread for thread-safety. Even three BigInts would be enough to speed up the code in the discourse post, and would mean theoretically a simple @rewrite could achieve the same performance as the buffer reuse I do in that post. You only need as many buffers as would be used simultaneously within one MA function call, which I don’t think is ever very much.

I think one would need to be a bit careful how they are documented and used to prevent corruption, and would also want to make sure there’s a hard cap on how many buffers are held globally to prevent memory leaks, but that it could make a really simple user API be very performant.

edit: as a design idea, this could look like

with_buffer(BigInt) do buffer
    ... code using `buffer`
end

where with_buffer(::Type{T}) gets pop!s an element of type T off some global buffer stack (global as in module-level but with separate stacks for each thread), passes it to the function, then push!s it back on the stack when done. And if the stack is empty (all buffers in use) it just creates a local buffer and passes it to the function without putting it back on the stack at the end (to enforce a maximum number of non-gc'd buffers). It could also lazily populate the buffer stack, i.e. to have a buffer stack of size 10, the first 10 times it's called it creates the buffer itself.

Then the mutable_operate_to! methods could simply call mutable_buffered_operate_to! inside a with_buffer do-block, to reuse a global buffer when possible, and otherwise just generate a local one.

I know globals are frowned on as a general rule, but I think it could be done properly with a safe API such as this with_buffer that handles all the global interaction in a sensible way.

[blegat]

First, since rationals aren’t actually immutable (but the bigints they hold are)

I would do MA.mutability(::Type{Rational{T}}) where {T} = MA.mutability(T).

The third point is related to the word “almost” in the previous paragraph: I found Base.GMP.MPZ.set! very useful, as a kind of copyto! for BigInts. Would it make sense to add such a function to the MA API?

Yes, that would make sense to add it.

One other idea: MA could globally hold some buffers to be used in @rewrite

What's tricky is that some types need buffer while other types do not depending on the operation.
For instance, MA.add_mul!(::JuMP.AffExpr, ::Float64, ::JuMP.VariableRef) does not need any buffer while MA.add_mul!(::BigInt, ::BigInt, ::BigInt) does need one.

[ericphanson]

One other idea: MA could globally hold some buffers to be used in @rewrite

What's tricky is that some types need buffer while other types do not depending on the operation.
For instance, MA.add_mul!(::JuMP.AffExpr, ::Float64, ::JuMP.VariableRef) does not need any buffer while MA.add_mul!(::BigInt, ::BigInt, ::BigInt) does need one.

One way to resolve this is to just do things like

function MA.add_mul!(x::BigInt, y::BigInt, z::BigInt)
    with_buffer(BigInt) do buffer
        MA.add_buffered_mul!(buffer, x, y, z)
    end
end

and never call the buffered functions manually. As long as we can grab buffers from and put them back on the global stack efficiently enough, I think we never really need to manually call the buffered methods. Then the generic things like rewrite can just never call the buffered ones, but that's OK because they dispatch to use the supply of global buffers anyway.

[blegat]

Closed by #104