Add more GPU registrations needed for while_v2.

These registrations are needed because the placer's too stupid to
always place resource handles and accesses on the same device, if
other colocation constraints dictate they be on different
devices. Ideally when soft placement is enabled, the placer would
prioritize colocating resource handles and accesses over other
constraints to avoid runtime errors.

PiperOrigin-RevId: 226359107