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

compiler/rustc_codegen_llvm/src/llvm_util.rs

@@ -251,6 +251,8 @@ pub(crate) fn to_llvm_features<'a>(sess: &Session, s: &'a str) -> Option<LLVMFea
         // Filter out features that are not supported by the current LLVM version
         ("loongarch32" | "loongarch64", "32s") if major < 21 => None,
         ("powerpc", "power8-crypto") => Some(LLVMFeature::new("crypto")),
+        // Filter out Rust-specific *virtual* target features
+        ("riscv32" | "riscv64", "zkne_or_zknd") => None,
         ("sparc", "leoncasa") => Some(LLVMFeature::new("hasleoncasa")),
         ("x86", "sse4.2") => Some(LLVMFeature::with_dependencies(
             "sse4.2",

compiler/rustc_target/src/target_features.rs

@@ -691,8 +691,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`