What is "the correct" way to use TBI/UAI/LAM in Rust

community
1

Given that TBI(armv8.5a+)/UAI(amd zen4)/LAM(intel) are coming to real life, I wonder how should a user actually use these features in Rust. (Let's consider use level usage rather than compiler optimizers).

For instance, if you have Apple Silicon devices (or any other device fulfills armv8.5a profile with TBI enabled kernel), the following code should work:

use std::mem::{size_of, ManuallyDrop};
use std::ops::{Deref, DerefMut};
use std::rc::Rc;


pub trait PointerLike: Sized {}

impl<T> PointerLike for &'_ T {}

impl<T> PointerLike for &'_ mut T {}

impl<T> PointerLike for Box<T> {}

impl<T> PointerLike for Rc<T> {}

union TBITag<T: PointerLike> {
    tag: [u8; size_of::<*const u8>()],
    pointer: std::mem::ManuallyDrop<T>,
}

impl<T: PointerLike> TBITag<T> {
    fn mask(&self) -> u8 {
        unsafe {
            if cfg!(target_endian = "big") {
                self.tag[0]
            } else {
                self.tag[core::mem::size_of::<*const u8>() - 1]
            }
        }
    }
    fn set_mask(&mut self, x: u8) {
        unsafe {
            if cfg!(target_endian = "big") {
                self.tag[0] = x;
            } else {
                self.tag[core::mem::size_of::<*const u8>() - 1] = x;
            }
        }
    }
    fn new(data: T) -> Self {
        Self {
            pointer: ManuallyDrop::new(data),
        }
    }
}

impl<T: PointerLike> Drop for TBITag<T> {
    fn drop(&mut self) {
        self.set_mask(0);
        unsafe {
            ManuallyDrop::drop(&mut self.pointer);
        }
    }
}

impl<T: PointerLike + Deref> Deref for TBITag<T> {
    type Target = T::Target;
    fn deref(&self) -> &Self::Target {
        unsafe { &self.pointer }
    }
}

impl<T: PointerLike + DerefMut> DerefMut for TBITag<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut self.pointer }
    }
}

fn main() {
    let mut x = TBITag::new(Box::new(5));
    x.set_mask(170);
    println!("x = {}", *x);
    println!("tag = {}", x.mask());
}

However, if you run this with miri, it will complain about UB (of cuz!).

Leaving alone the miri problem, I wonder if TBI is actually "safe" with current implementation. To me, wrapping the pointer deeply into Union and ManuallyDrop should have already forbidden the compiler from optimizing the code in unexpected ways, but, emmm, it is "UB" after all!

2

I can't answer your actual question, sorry. This comment is more of a meta-comment.

Regarding

For instance, if you have Apple Silicon devices (or any other device fulfills armv8.5a profile with TBI enabled kernel), the following code should work:
[...]
To me, wrapping the pointer deeply into Union and ManuallyDrop should have already forbidden the compiler from optimizing the code in unexpected ways

The code is in Rust and UB is defined at the language level, not the target hardware level. And the compiler and optimizers operate accordingly. You can't rely on assuming any particular relationship between the Rust virtual machine and what hardware you're running on, unless such a relationship is made part of the Rust language.

Perhaps you knew all that --

emmm, it is "UB" after all!

-- but I got conflicting messages from your phrasing.

So maybe just want to know what you can "get away with" -- what you can write that will compile down how you want it to, even though you recognize it's UB:

Leaving alone the miri problem, I wonder if TBI is actually "safe" with current implementation.

That can be a perfectly valid technical pursuit for personal projects and exploration IMO.[1] However, you will likely still hit a cultural barrier when asking for advice about such a pursuit. Not only is there a cultural stance in Rust to recognize unsound (much less UB) code as invalid generally, but there is much less interest in "getting away with it", so not many people will know or even care when it's possible.

Also there's a new compiler every 6 weeks, so it's pretty fleeting knowledge.

(Outside of personal projects and exploration, for a publicly used crate, the ecosystem tends to strongly reject projects that are cavalier about unsoundness and UB.)

  1. Others disagree. ↩ī¸Ž

5 Likes
3

If miri does not know that the tagged pointer is equally valid as the non-tagged pointer (does the arch give the tag any semantics? or is that up to the user?), then it reports UB, no surprise there. If the compiler does not know that these pointers are valid, then some UB-exploiting optimisations can kick in. I am less worried about rustc itself, but LLVM does heavy optimisations on pointers, sometimes too much (different issue, but still).

I have doubts that "UB-hiding by obscurity" is a solid appraoch. In particular, hiding things from LLVM is more difficult than using union. I guess the following approach can feasible:

  1. crate a struct #[repr(transparent)] struct TaggedPointer([u8; sizeof::<*const u8>()]), or something similar, not sure.
  2. convert any pointers you get to that type.
  3. perform tagging on that type only
  4. when accessing the pointer, dont use vanilla ptr-read or *. Use some intrinsic that performs ptr::read that LLVM has for tagged pointer. Or: create a temporary "un-tagged" pointer and read that normally.

If miri complains about provenance, it might be necessary to call ptr::expose_addr.