Skip to content

aes64ks1i and aes64ks2 have wrong target feature #1765

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sayantn opened this issue Apr 6, 2025 · 5 comments
Open

aes64ks1i and aes64ks2 have wrong target feature #1765

sayantn opened this issue Apr 6, 2025 · 5 comments

Comments

@sayantn
Copy link
Contributor

sayantn commented Apr 6, 2025

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
Copy link
Contributor Author

sayantn commented Jun 5, 2025

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?

@Amanieu
Copy link
Member

Amanieu commented Jun 5, 2025

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

@a4lg
Copy link
Contributor

a4lg commented Jun 5, 2025

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).

@sayantn
Copy link
Contributor Author

sayantn commented Jun 6, 2025

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
Copy link
Contributor Author

sayantn commented Jun 7, 2025

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 😿

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants