Hi, I got something very primitive here: https://github.com/trivialfis/raft/tree/fea-restricted-accessor for annotating the access function with __device__ and __host__ tag. It doesn't compile as we access the device mdarray (not span) on host for tests. (see the device_reference class). There are a few options we can proceed with this:

• Don't annotate the access function. The type should provide sufficient information/protection for most use cases. It's c++, we are not going to have 100% safety.
• Remove the device_reference class and disable the operator() for device_mdarray completely. Since we can't access an owning object in the kernel so the operator will be inaccessible anywhere. But less convenient for testing and debugging.
• Create different accessors for mdarray and mdspan and annotate only the latter. The code is more complicated.

My personal preference is for the first option, I think annotating it with __device__ is a great idea for safety, however, I also think the current type system is safe "enough" for most cases.

@gpucibot merge