Using LoopVectorization to compute the commutator of two Hermitian matrices

[christian-hh]

I'm trying to leverage LoopVectorization to implement a fast way to calculate the commutator on two Hermitian matrices H and rho (times a factor im):

im * [H, rho] = im * (H * rho - rho * H)

This would then be used to model the evolution of the so-called density matrix rho, a common construct for modelling quantum systems, see e.g., von Neumann equation.

Since the result im * [H, rho] itself is guaranteed to be Hermitian, I was wondering if there might be a way to utilize @avx that only calculates the upper triangular part of the result, and/or potentially calculates both of the products H * rho and rho * H in a single pass. To illustrate my thinking, looking at the following function from #https://discourse.julialang.org/t/ann-loopvectorization/32843 to perform matrix multiplication,

@inline function gemmavx!(C, A, B)
    @avx for i ∈ 1:size(A,1), j ∈ 1:size(B,2)
        Cᵢⱼ = 0.0
        for k ∈ 1:size(A,2)
            Cᵢⱼ += A[i,k] * B[k,j]
        end
        C[i,j] = Cᵢⱼ
    end
end

one could imagine simply changing the indexing on j to start at i rather than 1, i.e., j ∈ i:size(B,2) in the second line. This clearly breaks the assumptions of @avx in that the order of indexing matters, of course.

For context, my current approach to computing this commutator is via defining StructArrays à la #https://discourse.julialang.org/t/ann-loopvectorization/32843/34 to perform the two complex matrix multiplications H * rho and rho * H separately. While this is certainly fast, I was simply wondering whether there is any obvious additional speed-up that could be achieved given the strict structure of this problem.

Thanks for any help!

[chriselrod]

one could imagine simply changing the indexing on j to start at i rather than 1, i.e., j ∈ i:size(B,2) in the second line. This clearly breaks the assumptions of @avx in that the order of indexing matters, of course.

I have been planning on adding support for this for a while. I think I've figured out a lot of the details. Hopefully fairly soon I'll start at least dedicating a few hours every night to making progress here.

and/or potentially calculates both of the products H * rho and rho * H in a single pass.

This would probably be difficult to make optimal, because in A * B, it is normally optimal to load vectors from A and scalars from B.

[christian-hh]

I will look forward to it! I'd love to be able to contribute to this development in the future, albeit I don't have much familiarity with vectorization (or, more generally, how compilers work). Are there any references you might suggest for an introduction to the sort of knowledge required for building a package such as LoopVectorization.jl?

[chriselrod]

Hmm, it's not really about vectorization specifically, but I found Agner Fog's C++ optimization guide very helpful. I ought to make the time to reread it again.

Something else that'd probably help is Anatomy of High-Performance Matrix Multiplication. It focuses in depth on a specific example.
Register tiling is one of the most important optimizations LoopVectorization does, and an important part of creating a matmul microkernel.

That's probably enough to get you started. I'd also suggest experimenting a lot and looking at asm.
I'm also always happy to answer questions.