amrex::GpuComplex not ideal for GPU

[AlexanderSinn]

Specifically when accessing an array of amrex::GpuComplex on the GPU, it would be better if the alignment of amrex::GpuComplex was 2*alignof(T) instead of alignof(T). This is because one GPU thread can write both real and imaginary part in a single memory transaction if the memory region is aligned to its size. If this is not the case the memory access from the GPU warp is non-coalesced (https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-memory-accesses).

struct alignas(2*sizeof(amrex::Real)) AlignedComplex {
    amrex::Real real;
    amrex::Real imag;
};

Using AlignedComplex was more than 10 times faster than amrex::GpuComplex<amrex::Real> (Real = double) when writing to a large array in pinned memory form an A100 (pinned memory is especially sensitive to non-coalesced memory access).

Note: increasing the alignment would break reinterpret_cast from std::complex to amrex::GpuComplex.

[BenWibking]

reinterpret_cast between different types (except when casting to char, or unsigned vs. signed variants of the same type) is always undefined behavior. That should never be done, and abuse of reintepret_cast has caused incorrect results with our AMReX code on AMD GPUs.

[WeiqunZhang]

https://en.cppreference.com/w/cpp/numeric/complex Casting from std::complex to T[2] or T is allowed.

[WeiqunZhang]

NVIDIA/libcudacxx#151 Looks like cuda::std::complex aligns to 2*sizeof(T) now.

[WeiqunZhang]

@AlexanderSinn How did you create the array of GpuComplex? I am asking because the alignmen of amrex::Arena allocation is 16.

[AlexanderSinn]

amrex::BaseFab<AlignedComplex> or amrex::BaseFab<amrex::GpuComplex<amrex::Real>>
fab.resize(domain, 1, amrex::The_Pinned_Arena());
fab.setVal<amrex::RunOn::Host>({0,0});

[WeiqunZhang]

That probably means cudaMalloc and cudaHostMalloc allocate memory with smaller alignment size, I guess, because in the Arean we make sure the blocks we give out are multiples of 16 bytes.

[WeiqunZhang]

So we might be able to fix the alignment in Arena without changing GpuComplex.

[AlexanderSinn]

The alignment needs to be known at compile time so nvcc can do the 16-byte read/write. cudaMalloc is actually 256 bytes aligned.

[AlexanderSinn]

It generates different types of load and store instructions for aligned and non-aligned types.
https://godbolt.org/z/hMv9TWxzz

[WeiqunZhang]

Yes, you are right. Since the basic types like unsigned long long need 16 bytes alignment, malloc is at least 16 bytes aligned.

[WeiqunZhang]

I think we should add the required alignment to amrex::GpuComplex. @AlexanderSinn Could you submit a PR?

[AlexanderSinn]

Yes

[AlexanderSinn]

Fixed in #3691