Fix bipermutations in contract

[ogauthe]

This PR fixes bipermutation inversion in contract. A given bipermutation alone does not carry enough information to be inverted: the bipartition of the output must be specified. It turns out both permutation are needed in contract stack at different times. I managed to pass only one inside each function and to invert it using information from other arguments.

The logic of the code is now correct and able to handle arrays with bipartitions. We need bipermutations for both directions and be explicit which is which. I would be happy to improve the names I used.

Note that BlockArrays tests fail due to JuliaArrays/BlockArrays.jl#295.

[codecov]

Codecov Report

❌ Patch coverage is 97.36842% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 94.16%. Comparing base (3aaed6c) to head (26c9765).

Files with missing lines	Patch %	Lines
src/contract/blockedperms.jl	90.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #75      +/-   ##
==========================================
- Coverage   94.90%   94.16%   -0.74%     
==========================================
  Files          14       14              
  Lines         451      463      +12     
==========================================
+ Hits          428      436       +8     
- Misses         23       27       +4     

Flag	Coverage Δ
docs	0.00% <0.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mtfishman]

Suggested change

	invbiperm(perm, t::AbstractBlockTuple{2}) = invbiperm(perm, Val(first(blocklength(t))))
	invbiperm(perm, t::AbstractBlockTuple{2}) = invbiperm(perm, Val(first(blocklengths(t))))

[mtfishman]

This is a bit strange to me that it just allows anything as the second argument. Maybe this should be invbiperm(perm)? Is this used anywhere right now?

[mtfishman]

Looking at the rest of the code, I see that it is being used in calls like biperm_a12_to_dest = invbiperm(biperm_dest_to_a12, axes(a_dest)), where axes(a_dest) might output a blocked tuple or a flat tuple.

I think the invbiperm function is trying to do too much and therefore makes the code harder to understand. Instead, maybe we could introduce new functions biperm and length_codomain:

function biperm(perm, blocklength1::Int)
  return biperm(perm, Val(blocklength1))
end
function biperm(perm, ::Val{BlockLength1}) where {BlockLength1}
  # Check: BlockLength1 <= length(perm)
  return blockedperm(Tuple(perm),(BlockLength1, length(perm) - BlockLength1))
end

length_codomain(t::AbstractBlockTuple{2}) = first(blocklength(t))
# Assume all dimensions are in the domain by default
length_codomain(t) = 0

and the use a combination of invperm, biperm, and length_codomain in the contract code, for example:

biperm_a12_to_dest = biperm(invperm(biperm_dest_to_a12), length_codomain(axes(a_dest)))

[mtfishman]

I think in the context of this function, the name biperm_dest_to_a12 is more confusing than helpful (that name only makes sense in the context of contract but this function could be called for other purposes). I think it should be clear that biperm means the permutation that should be performed on m after it is reinterpreted as a length(axes_dest)-dimensional array (unless I'm misunderstanding the conventions of this function, in which case we should change it to that convention).

[mtfishman]

Same thing with the name axes_dest, I think axes is clear enough.

[mtfishman]

Same comments as above about the naming.

[mtfishman]

I think I lost track of the discussions we had about the conventions we want to use in this PR, I thought we had discussed that we would change the convention of unmatricize so that the biperm that gets input would be taken "literally", i.e. it wouldn't need to be inverted.

[mtfishman]

Maybe as an alternative, we could change the convention of unmatricize so that the axes that get input are the unpermuted axes, i.e. the axes corresponding directly to the memory ordering of the input matrix m. Then, the bipermutation that gets input is just the bipermutation that needs to be done to get the desired output. I.e. it would be equivalent to:

function unmatricize(
  style::FusionStyle,
  m::AbstractMatrix,
  ax,
  biperm::AbstractBlockPermutation{2},
)
  a = unmatricize(style, m, ax)
  return permutedims(a, biperm)
end

[lkdvos]

I missed a part of the discussion but this also came up when writing the extension to TensorOperations, so let me perhaps elaborate slightly on what TensorOperations does, which could be used as inspiration of what (not) to do.

From the point of TensorOperations, we chose our permutations/inverses such that effectively a contraction for a dense array would be, if simplified:

function tensorcontract!(C, A, (p1A, p2A), B, (p1B, p2B), (p1AB, p2AB), ...)
    A_mat = reshape(permutedims(A, (p1A..., p2A...)), prod(i->size(A, i), p1A), prod(i -> size(A, i), p2A))
    B_mat = reshape(permutedims(B, (p1B..., p2B...)), prod(i -> size(B, i), p1B), prod(i -> size(B, i), p2B))
    AB = reshape(A_mat * B_mat, size.(Ref(A), p1A), size.(Ref(B), p2B))
    C += permutedims(AB, (p1AB..., p2AB...))
    return C
end

Importantly, the partition for pX = (p1X, p2X) is always chosen such that you could write permutedims(X, pX) and have enough information to know what the output of that would be.

Now, in order to avoid forming AB whenever this is not required, we basically have an if statement that detects when permutedims(C, inv_AB(pAB)) (with inv_AB a permutation inverse with partition of AB) results in a stridedview that is still compatible with BLAS, in which case we indeed do mul!(permutedims(C, inv_AB(pAB)), A_mat, B_mat, alfa, beta), otherwise we simply allocate a temporary AB = A_mat * B_mat and then do an in-place addition C .= beta .* C .+ alfa .* permutedims(AB, pAB...)

The main idea being that where we use alfa and beta never really affects performance, and indeed we require both the inverse and regular blocked permutation to have both fast paths.

I think I quite like the naming scheme (might just be habit) of having permX being the thing that corresponds to permutedims(X, permX), which would entail permutedims(AB, permAB) and permC = inv_AB(permAB) for C_ = permutedims(C, permC).
Additionally, using permAB as the argument instead of permC is a bit more convenient since the partition of AB can be inferred from permA and permB alone, while the partition of permAB can only be inferred from the output array C, which for regular arrays just does not store that information.