Fix halo again

[ASKabalan]

This PR comes after this one #19
Instead of allocating using mlir.full_like_aval I allocate using cudecompMalloc but I keep a fallback to XLA if needed using JD_ALLOCATE_WITH_XLA=1
@aboucaud @EiffL Good for review after merging #19

[aboucaud]

Out of my league..

[ASKabalan]

Still much simpler than JAX distributed 🥲

[EiffL]

Hum... but I don't understand what was the issue here, as long as the workspace is created by XLA in each call the primitive and that the primitive uses the provided workspace in the input stream, what goes wrong?

[ASKabalan]

The workspace is not created at each call.
It is only created at the lowering phase as ir.constants

If I do an LPT it is fine.
If I do a nbody sim with a mesh size of 256 it is also fine (on a V100 32G)
512 and above I get a deadlock, with some further inspection it apeared to me that there is an illegal access of memory by some of the processes (which caused the deadlock)

I inspected the memory and found :

Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 0
Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 1
Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 2

Calling cudecompUpdateHalosX : 0x151003004300 On axis 0
Calling cudecompUpdateHalosX : 0x151003004300 On axis 1
Calling cudecompUpdateHalosX : 0x151003004300 On axis 2

XLA changes the memory address which should not happen (again the lowering happends only once)
This means that the memory have been repurposed.

Using cudecompmalloc this does not happend.

You can still reproduce if you get a 512 or 1024 mesh on no more than 4 GPUs with the flag JD_ALLOCATE_WITH_XLA

[ASKabalan]

To get the error to show, you need to remove all block_until_ready() statements

[EiffL]

Hum, but the thing that worries me is that the same could happen for the other ops no? And I don't see why the memory provided in the stream to the kernel would not be ok.

…

On Tue, Jul 9, 2024, 3:04 PM Wassim KABALAN ***@***.***> wrote: The workspace is not created at each call. It is only created at the lowering phase as ir.constants If I do an LPT it is fine. If I do a nbody sim with a mesh size of 256 it is also fine (on a V100 32G) 512 and above I get a deadlock, with some further inspection it apeared to me that there is an illegal access of memory by some of the processes (which caused the deadlock) I inspected the memory and found : Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 0 Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 1 Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 2 Calling cudecompUpdateHalosX : 0x151003004300 On axis 0 Calling cudecompUpdateHalosX : 0x151003004300 On axis 1 Calling cudecompUpdateHalosX : 0x151003004300 On axis 2 XLA changes the memory address which should not happen (again the lowering happends only once) This means that the memory have been repurposed. Using cudecompmalloc this does not happend. You can still reproduce if you get a 512 or 1024 on no more than 4 GPUs with the flag JD_ALLOCATE_WITH_XLA — Reply to this email directly, view it on GitHub <#23 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAGSLFZYSND6GGANMNAAPADZLQX43AVCNFSM6AAAAABKM7FAX6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMJYGQ2DAMZZGY> . You are receiving this because you were mentioned.Message ID: ***@***.*** com>

[EiffL]

Are you sure it's not something different like not having the right size of workspace? Maybe we have a different meaning for halo size

…

On Tue, Jul 9, 2024, 3:25 PM Eiffel ***@***.***> wrote: Hum, but the thing that worries me is that the same could happen for the other ops no? And I don't see why the memory provided in the stream to the kernel would not be ok. On Tue, Jul 9, 2024, 3:04 PM Wassim KABALAN ***@***.***> wrote: > The workspace is not created at each call. > It is only created at the lowering phase as ir.constants > > If I do an LPT it is fine. > If I do a nbody sim with a mesh size of 256 it is also fine (on a V100 > 32G) > 512 and above I get a deadlock, with some further inspection it apeared > to me that there is an illegal access of memory by some of the processes > (which caused the deadlock) > > I inspected the memory and found : > > Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 0 > Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 1 > Calling cudecompUpdateHalosX : 0x14a75f004300 On axis 2 > > Calling cudecompUpdateHalosX : 0x151003004300 On axis 0 > Calling cudecompUpdateHalosX : 0x151003004300 On axis 1 > Calling cudecompUpdateHalosX : 0x151003004300 On axis 2 > > XLA changes the memory address which should not happen (again the > lowering happends only once) > This means that the memory have been repurposed. > > Using cudecompmalloc this does not happend. > > You can still reproduce if you get a 512 or 1024 on no more than 4 GPUs > with the flag JD_ALLOCATE_WITH_XLA > > — > Reply to this email directly, view it on GitHub > <#23 (comment)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/AAGSLFZYSND6GGANMNAAPADZLQX43AVCNFSM6AAAAABKM7FAX6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMJYGQ2DAMZZGY> > . > You are receiving this because you were mentioned.Message ID: > ***@***.*** > com> >

[ASKabalan]

I am going to try to reproduce and tell you.
But the halo_extent here is half of the halo size but what I see is there is enough space
The size of the input is mesh + half of the halo and the other half is the extent

this should be enough

So yeah, for the FFT, the allocation is handled by us and not cufft because we use cufftsetautoallocation and the work is shared for the FFT and Transpose

The only difference I see is that the FFT is inplace (it uses some temporary parameters but nothing dangerous)
The transpositions are inplace aswell

The supplement memory required by a halo exchange is more?
Anyway it cannot be anything else, since when I use cudecomp the illegal access is not reproduced

[EiffL]

I think it's better to let XLA handle memory, because it can decide to free it when not necessary.

[EiffL]

We should remove the cuda memory allocation thing for halos, better to be consistent in memory handling throughout

[ASKabalan]

I just tested.
It is ok for me.
Can I merge?