From 0d6b5ddf71903949db276601979f612bf985521f Mon Sep 17 00:00:00 2001
From: beetrees <b@beetr.ee>
Date: Sat, 15 Jun 2024 22:30:25 +0100
Subject: [PATCH] Add `f16` inline ASM support for RISC-V

---
 compiler/rustc_codegen_llvm/src/asm.rs | 15 ++++++++++++++
 compiler/rustc_target/src/asm/riscv.rs |  7 ++++---
 tests/assembly/asm/riscv-types.rs      | 27 +++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_codegen_llvm/src/asm.rs b/compiler/rustc_codegen_llvm/src/asm.rs
index 60e63b956db6e..12759c96f766a 100644
--- a/compiler/rustc_codegen_llvm/src/asm.rs
+++ b/compiler/rustc_codegen_llvm/src/asm.rs
@@ -1029,6 +1029,13 @@ fn llvm_fixup_input<'ll, 'tcx>(
                 _ => value,
             }
         }
+        (InlineAsmRegClass::RiscV(RiscVInlineAsmRegClass::freg), Abi::Scalar(s)) if s.primitive() == Primitive::Float(Float::F16) => {
+            // Smaller floats are always "NaN-boxed" inside larger floats on RISC-V.
+            let value = bx.bitcast(value, bx.type_i16());
+            let value = bx.zext(value, bx.type_i32());
+            let value = bx.or(value, bx.const_u32(0xFFFF_0000));
+            bx.bitcast(value, bx.type_f32())
+        }
         _ => value,
     }
 }
@@ -1140,6 +1147,11 @@ fn llvm_fixup_output<'ll, 'tcx>(
                 _ => value,
             }
         }
+        (InlineAsmRegClass::RiscV(RiscVInlineAsmRegClass::freg), Abi::Scalar(s)) if s.primitive() == Primitive::Float(Float::F16) => {
+            let value = bx.bitcast(value, bx.type_i32());
+            let value = bx.trunc(value, bx.type_i16());
+            bx.bitcast(value, bx.type_f16())
+        }
         _ => value,
     }
 }
@@ -1242,6 +1254,9 @@ fn llvm_fixup_output_type<'ll, 'tcx>(
                 _ => layout.llvm_type(cx),
             }
         }
+        (InlineAsmRegClass::RiscV(RiscVInlineAsmRegClass::freg), Abi::Scalar(s)) if s.primitive() == Primitive::Float(Float::F16) => {
+            cx.type_f32()
+        }
         _ => layout.llvm_type(cx),
     }
 }
diff --git a/compiler/rustc_target/src/asm/riscv.rs b/compiler/rustc_target/src/asm/riscv.rs
index 3845a0e14af09..02a4a5e2ece59 100644
--- a/compiler/rustc_target/src/asm/riscv.rs
+++ b/compiler/rustc_target/src/asm/riscv.rs
@@ -40,12 +40,13 @@ impl RiscVInlineAsmRegClass {
         match self {
             Self::reg => {
                 if arch == InlineAsmArch::RiscV64 {
-                    types! { _: I8, I16, I32, I64, F32, F64; }
+                    types! { _: I8, I16, I32, I64, F16, F32, F64; }
                 } else {
-                    types! { _: I8, I16, I32, F32; }
+                    types! { _: I8, I16, I32, F16, F32; }
                 }
             }
-            Self::freg => types! { f: F32; d: F64; },
+            // FIXME(f16_f128): Add `q: F128;` once LLVM support the `Q` extension.
+            Self::freg => types! { f: F16, F32; d: F64; },
             Self::vreg => &[],
         }
     }
diff --git a/tests/assembly/asm/riscv-types.rs b/tests/assembly/asm/riscv-types.rs
index 0d1f8305d375e..3cfd0c7dec20a 100644
--- a/tests/assembly/asm/riscv-types.rs
+++ b/tests/assembly/asm/riscv-types.rs
@@ -6,7 +6,7 @@
 //@[riscv32] needs-llvm-components: riscv
 //@ compile-flags: -C target-feature=+d
 
-#![feature(no_core, lang_items, rustc_attrs)]
+#![feature(no_core, lang_items, rustc_attrs, f16)]
 #![crate_type = "rlib"]
 #![no_core]
 #![allow(asm_sub_register)]
@@ -33,6 +33,7 @@ type ptr = *mut u8;
 
 impl Copy for i8 {}
 impl Copy for i16 {}
+impl Copy for f16 {}
 impl Copy for i32 {}
 impl Copy for f32 {}
 impl Copy for i64 {}
@@ -103,6 +104,12 @@ macro_rules! check_reg {
 // CHECK: #NO_APP
 check!(reg_i8 i8 reg "mv");
 
+// CHECK-LABEL: reg_f16:
+// CHECK: #APP
+// CHECK: mv {{[a-z0-9]+}}, {{[a-z0-9]+}}
+// CHECK: #NO_APP
+check!(reg_f16 f16 reg "mv");
+
 // CHECK-LABEL: reg_i16:
 // CHECK: #APP
 // CHECK: mv {{[a-z0-9]+}}, {{[a-z0-9]+}}
@@ -141,6 +148,12 @@ check!(reg_f64 f64 reg "mv");
 // CHECK: #NO_APP
 check!(reg_ptr ptr reg "mv");
 
+// CHECK-LABEL: freg_f16:
+// CHECK: #APP
+// CHECK: fmv.s f{{[a-z0-9]+}}, f{{[a-z0-9]+}}
+// CHECK: #NO_APP
+check!(freg_f16 f16 freg "fmv.s");
+
 // CHECK-LABEL: freg_f32:
 // CHECK: #APP
 // CHECK: fmv.s f{{[a-z0-9]+}}, f{{[a-z0-9]+}}
@@ -165,6 +178,12 @@ check_reg!(a0_i8 i8 "a0" "mv");
 // CHECK: #NO_APP
 check_reg!(a0_i16 i16 "a0" "mv");
 
+// CHECK-LABEL: a0_f16:
+// CHECK: #APP
+// CHECK: mv a0, a0
+// CHECK: #NO_APP
+check_reg!(a0_f16 f16 "a0" "mv");
+
 // CHECK-LABEL: a0_i32:
 // CHECK: #APP
 // CHECK: mv a0, a0
@@ -197,6 +216,12 @@ check_reg!(a0_f64 f64 "a0" "mv");
 // CHECK: #NO_APP
 check_reg!(a0_ptr ptr "a0" "mv");
 
+// CHECK-LABEL: fa0_f16:
+// CHECK: #APP
+// CHECK: fmv.s fa0, fa0
+// CHECK: #NO_APP
+check_reg!(fa0_f16 f16 "fa0" "fmv.s");
+
 // CHECK-LABEL: fa0_f32:
 // CHECK: #APP
 // CHECK: fmv.s fa0, fa0