Reimplement commit ce0c77b472d

The ce0c77b472d was broken for contexts containing mixed CPU/GPU devices.