🐛 Describe the bug

pack and unpack in c10/core/Stream.h are concatenating the DeviceType, DeviceIndex and StreamId into an int64_t.
The assumption for this packing mechanism is that the virtual address will use 48 bit and will move the DeviceType and DeviceIndex into the upper bytes while reading the lower 48 bits from id():

uint64_t bits = static_cast<uint64_t>(static_cast<uint8_t>(device_type()))
            << 56 |
        static_cast<uint64_t>(static_cast<uint8_t>(device_index())) << 48 |
        // Remove the sign extension part of the 64-bit address because
        // the id might be used to hold a pointer.
        (static_cast<uint64_t>(id()) & ((1ull << 48) - 1));

The same method then verifies this packing and raises asserts in case something failed.

While this assumption seems to be correct for the majority of CPU architectures, we are hitting an error on an ARM Neoverse N1 CPU.
After debugging this implementation, we noticed that ARMv8.2 implements "LargeVA" support, which uses up to 52 bits.
The recreation of the actual address will thus fail and we will run into the assert (it was great to see asserts there as a silent memory corruption might have been much harder to debug).

Some references:

• https://en.wikichip.org/w/images/b/b3/arm_neoverse_n1_trm.pdf
• https://opensource.com/article/20/12/52-bit-arm64-kernel
• https://developer.arm.com/documentation/101811/0102/Translation-granule
• https://developer.arm.com/documentation/101811/0101/Address-spaces-in-AArch64

Searching a bit more for this VA extension, also yields x86 implementations:
Intel 5-level paging:

[...] It extends the size of virtual addresses from 48 bits to 57 bits, increasing the addressable virtual memory from 256 TB to 128 PB. The extension was first implemented in the Ice Lake processors,[2] and the 4.14 Linux kernel adds support for it. [...]

which seems to extend it even further to 57 bits.

A workaround might be to shift the payload and reduce the range for DeviceType and DeviceIndex.
However, if we want to use the full 57 bits, it would leave only 7 bits for the type and id, which sounds too limited.

Extending the StreamId to a wider type or using another data structure might be alternatives if we want to support these CPU architectures.

CC @ezyang as the original author of this implementation, @mruberry, @mcarilli, @zasdfgbnm to discuss potential approaches.

Versions

Current master build with CUDA11.6 on ARM Neoverse N1.

cc @ngimel