mirror of
https://github.com/pseXperiments/icicle.git
synced 2026-01-08 20:48:06 -05:00
14997566ff
## Describe the changes This PR fixes an issue when `RunOnDevice` is called for multi-gpu while other goroutines calling device operations are run outside of `RunOnDevice`. The issue comes from setting a device other than the default device (device 0) on a host thread within `RunOnDevice` and not unsetting that host threads device when `RunOnDevice` finishes. When `RunOnDevice` locks a host thread to ensure that all other calls in the go routine are on the same device, it never unsets that thread’s device. Once the thread is unlocked, other go routines can get scheduled to it but it still has the device set to whatever it was before while it was locked so its possible that the following sequence happens: 1. NTT domain is initialized on thread 2 via a goroutine on device 0 2. MSM multiGPU test runs and is locked on thread 3 setting its device to 1 3. Other tests run concurrently on threads other than 3 (since it is locked) 4. MSM multiGPU test finishes and release thread 3 back to the pool but its device is still 1 5. NTT test runs and is assigned to thread 3 --> this will fail because the thread’s device wasn’t released back We really only want to set a thread's device while the thread is locked. But once we unlock a thread, it’s device should return to whatever it was set at originally. In theory, it should always be 0 if `SetDevice` is never used outside of `RunOnDevice` - which it shouldn’t be in most situations