Rollup merge of #146798 - a4lg:riscv-intrinsics-zkne_or_zknd, r=Amanieu

RISC-V: Implement (Zkne or Zknd) intrinsics correctly

On rust-lang/stdarch#1765, it has been pointed out that two RISC-V (64-bit only) intrinsics to perform AES key scheduling have wrong target feature.
`aes64ks1i` and `aes64ks2` instructions require *either* Zkne (scalar cryptography: AES encryption) or Zknd (scalar cryptography: AES decryption) extension (or both) but corresponding Rust intrinsics (in `core::arch::riscv64`) required *both* Zkne and Zknd extensions.

An excerpt from the original intrinsics:

```rust
#[target_feature(enable = "zkne", enable = "zknd")]
```

To fix that, we need to:

1.  Represent a condition where *either* Zkne or Zknd is available and
2.  Workaround an issue: `llvm.riscv.aes64ks1i` / `llvm.riscv.aes64ks2` LLVM intrinsics require either Zkne or Zknd extension.

This PR attempts to resolve them by:

1.  Adding a perma-unstable RISC-V target feature: `zkne_or_zknd` (implied from both `zkne` and `zknd`) and
2.  Using inline assembly to construct machine code directly (because `zkne_or_zknd` alone cannot imply neither Zkne nor Zknd, we cannot use LLVM intrinsics).

The author confirmed that we can construct an AES key scheduling function with decent performance using fixed `aes64ks1i` and `aes64ks2` intrinsics (with optimization enabled).

compiler/rustc_codegen_llvm/src/llvm_util.rs

@@ -262,6 +262,11 @@ pub(crate) fn to_llvm_features<'a>(sess: &Session, s: &'a str) -> Option<LLVMFea
             "power8-crypto" => Some(LLVMFeature::new("crypto")),
             s => Some(LLVMFeature::new(s)),
         },
+        Arch::RiscV32 | Arch::RiscV64 => match s {
+            // Filter out Rust-specific *virtual* target feature
+            "zkne_or_zknd" => None,
+            s => Some(LLVMFeature::new(s)),
+        },
         Arch::Sparc | Arch::Sparc64 => match s {
             "leoncasa" => Some(LLVMFeature::new("hasleoncasa")),
             s => Some(LLVMFeature::new(s)),

compiler/rustc_target/src/target_features.rs

@@ -690,8 +690,9 @@ static RISCV_FEATURES: &[(&str, Stability, ImpliedFeatures)] = &[
     ("zimop", Unstable(sym::riscv_target_feature), &[]),
     ("zk", Stable, &["zkn", "zkr", "zkt"]),
     ("zkn", Stable, &["zbkb", "zbkc", "zbkx", "zkne", "zknd", "zknh"]),
-    ("zknd", Stable, &[]),
-    ("zkne", Stable, &[]),
+    ("zknd", Stable, &["zkne_or_zknd"]),
+    ("zkne", Stable, &["zkne_or_zknd"]),
+    ("zkne_or_zknd", Unstable(sym::riscv_target_feature), &[]), // Not an extension
     ("zknh", Stable, &[]),
     ("zkr", Stable, &[]),
     ("zks", Stable, &["zbkb", "zbkc", "zbkx", "zksed", "zksh"]),

library/stdarch/crates/core_arch/src/riscv64/zk.rs

@@ -1,6 +1,8 @@
 #[cfg(test)]
 use stdarch_test::assert_instr;
 
+use crate::arch::asm;
+
 unsafe extern "unadjusted" {
     #[link_name = "llvm.riscv.aes64es"]
     fn _aes64es(rs1: i64, rs2: i64) -> i64;
@@ -14,12 +16,6 @@ unsafe extern "unadjusted" {
     #[link_name = "llvm.riscv.aes64dsm"]
     fn _aes64dsm(rs1: i64, rs2: i64) -> i64;
 
-    #[link_name = "llvm.riscv.aes64ks1i"]
-    fn _aes64ks1i(rs1: i64, rnum: i32) -> i64;
-
-    #[link_name = "llvm.riscv.aes64ks2"]
-    fn _aes64ks2(rs1: i64, rs2: i64) -> i64;
-
     #[link_name = "llvm.riscv.aes64im"]
     fn _aes64im(rs1: i64) -> i64;
 
@@ -133,15 +129,26 @@ pub fn aes64dsm(rs1: u64, rs2: u64) -> u64 {
 /// # Note
 ///
 /// The `RNUM` parameter is expected to be a constant value inside the range of `0..=10`.
-#[target_feature(enable = "zkne", enable = "zknd")]
+#[target_feature(enable = "zkne_or_zknd")]
 #[rustc_legacy_const_generics(1)]
-#[cfg_attr(test, assert_instr(aes64ks1i, RNUM = 0))]
 #[inline]
 #[unstable(feature = "riscv_ext_intrinsics", issue = "114544")]
 pub fn aes64ks1i<const RNUM: u8>(rs1: u64) -> u64 {
     static_assert!(RNUM <= 10);
-
-    unsafe { _aes64ks1i(rs1 as i64, RNUM as i32) as u64 }
+    unsafe {
+        let rd: u64;
+        asm!(
+            ".option push",
+            ".option arch, +zkne",
+            "aes64ks1i {}, {}, {}",
+            ".option pop",
+            lateout(reg) rd,
+            in(reg) rs1,
+            const RNUM,
+            options(pure, nomem, nostack, preserves_flags)
+        );
+        rd
+    }
 }
 
 /// This instruction implements part of the KeySchedule operation for the AES Block cipher.
@@ -155,12 +162,24 @@ pub fn aes64ks1i<const RNUM: u8>(rs1: u64) -> u64 {
 /// Version: v1.0.1
 ///
 /// Section: 3.11
-#[target_feature(enable = "zkne", enable = "zknd")]
-#[cfg_attr(test, assert_instr(aes64ks2))]
+#[target_feature(enable = "zkne_or_zknd")]
 #[inline]
 #[unstable(feature = "riscv_ext_intrinsics", issue = "114544")]
 pub fn aes64ks2(rs1: u64, rs2: u64) -> u64 {
-    unsafe { _aes64ks2(rs1 as i64, rs2 as i64) as u64 }
+    unsafe {
+        let rd: u64;
+        asm!(
+            ".option push",
+            ".option arch, +zkne",
+            "aes64ks2 {}, {}, {}",
+            ".option pop",
+            lateout(reg) rd,
+            in(reg) rs1,
+            in(reg) rs2,
+            options(pure, nomem, nostack, preserves_flags)
+        );
+        rd
+    }
 }
 
 /// This instruction accelerates the inverse MixColumns step of the AES Block Cipher, and is used to aid creation of

tests/ui/check-cfg/target_feature.stderr

@@ -424,6 +424,7 @@ LL |     cfg!(target_feature = "_UNEXPECTED_VALUE");
 `zkn`
 `zknd`
 `zkne`
+`zkne_or_zknd`
 `zknh`
 `zkr`
 `zks`