stdarch icon indicating copy to clipboard operation
stdarch copied to clipboard

`aes64ks1i` and `aes64ks2` have wrong target feature

Open sayantn opened this issue 8 months ago • 7 comments

The RiscV ZK spec states that aes64ks1i and aes64ks2 is available whenever either one of zkne and zknd is available (see here and here), but the current implementation uses #[target_feature(enable = "zkne", enable = "zknd")], which means the 2 functions require both of zkne and zknd.

I know Rust doesn't have any way of "OR"-ing target features, but can we do in this case? (Also these intrinsics are unstable, so no problems changing their signatures or anything 😄)

This was brought to my attention in rust-lang/rust#137293

sayantn avatar Apr 06 '25 21:04 sayantn

It would be nice if we can introduce a perma-unstable target feature zkne_or_zknd that is implied by both zkne and zknd, then we can make these #[target_feature(enable = "zkne_or_zknd")]. This seems like the only way to emulate ORing target features.

What do you think about this @Amanieu? Should I go ahead and send a rustc PR?

sayantn avatar Jun 05 '25 07:06 sayantn

I'm inclined to agree, but would like a second opinion from @a4lg who has been working on RISC-V target features a lot.

Amanieu avatar Jun 05 '25 10:06 Amanieu

I'm not sure how to support (Zkne or Zknd) but not strongly against it (though I feel an uneasy feelings to zkne_or_zknd).

Using underscores is a great idea since RISC-V extension names does not allow this symbol for extension names.

~~I came up with an alternative name: zkn_aes_ks (AES key schedule components inside the Zkn extension) but I also feel an uneasy feelings to my own idea.~~

Edit: forget it. My idea gives wrong impression that aes64im (performs a part of AES key schedule on decryption) is a part of that (but this is only a part of the Zknd extension).

a4lg avatar Jun 05 '25 14:06 a4lg

InvMixColumns (what aes64im does) is technically not part of the key schedule, although it is typically used to get the decryption round keys from the encryption round keys (which are produced by the key schedule).

Anyway, I don't think the name matters that much, as it would be perma-unstable. I will try it out, and lyk if it works

sayantn avatar Jun 06 '25 03:06 sayantn

I tried it out, sadly it doesn't work in debug mode. The following code (compiled with a patched compiler that implied zkne_or_zknd by zkne and zknd, and with zkne_or_zknd corresponding to no LLVM feature)

#![feature(riscv_target_feature, link_llvm_intrinsics, abi_unadjusted)]
#![allow(internal_features)]
#![crate_type = "lib"]

extern "unadjusted" {
    #[link_name = "llvm.riscv.aes64ks2"]
    fn foo(a: u64, b: u64) -> u64;
}

#[inline]
#[target_feature(enable = "zkne_or_zknd")]
pub fn bar(a: u64, b: u64) -> u64 {
    unsafe { foo(a, b) }
}

#[inline(never)]
#[target_feature(enable = "zkne")]
pub fn baz(a: u64, b: u64) -> u64 {
    bar(a, b)
}

It compiles fine with opt-level 1 and above, but with opt-level 0, it tries to generate asm for bar too, which doesn't compile due to the missing target feature (rustc-LLVM ERROR: Cannot select: intrinsic %llvm.riscv.aes64ks2). So I guess this approach is not viable.

CLang handles this kind of intrinsics by just defining a macro, but that is not a viable option for us too.

Edit: this code would have worked if bar was #[inline(always)], but target features don't allow that 😿

sayantn avatar Jun 07 '25 14:06 sayantn

I know that this is a really, really stupid idea but... we have an option to split e.g. aes64im into aes64im_d and aes64im_e with the same internal contents but different target_feature (should only be considered if there's no alternatives).

a4lg avatar Aug 30 '25 05:08 a4lg

@sayantn Draft #1917 attempts to resolve this issue by substituting LLVM instructions to inline assembly. If it works well with your zkne_or_zknd (perma-unstable) target feature proposal, that would be nice.

a4lg avatar Sep 11 '25 06:09 a4lg