Fix argument extension on riscv64 for Wasmtime builtins (bytecodealliance#10069)

* Fix argument extension on riscv64 for Wasmtime builtins

This commit fixes a crash found in the CI of bytecodealliance#10040. That PR itself
isn't the fault per-se but rather uncovered a preexisting issue on
riscv64. According to riscv64's ABI docs it looks like arguments are all
expected to be sign-extended, whereas currently in Wasmtime all host
signatures are zero-extended on all platforms. This commit applies two
changes to fix this:

* A new `TargetIsa::default_argument_extension` method was added which
  is now used instead of unconditionally using `uext` to apply to all
  arguments/results.

* While I was here I went ahead and split apart things where the wasm
  signature of a builtin doesn't use `uext` or `sext`. The host
  signature still does, however, which means that any extension
  necessary happens in the trampoline we generate per-libcall as opposed
  to at all callsites.

I'm not certain that all platforms require `uext` but I've left the
`TargetIsa` implementation as `uext` for now with a comment explaining
why. Currently the only non-`uext` platforms are riscv64, which is
`sext` to fix the issue from bytecodealliance#10040, and Pulley which is "none" as
things work differently there.

* Update test expectations

Author: alexcrichton

cranelift/codegen/src/isa/aarch64/mod.rs

@@ -1,7 +1,7 @@
 //! ARM 64-bit Instruction Set Architecture.
 
 use crate::dominator_tree::DominatorTree;
-use crate::ir::{Function, Type};
+use crate::ir::{self, Function, Type};
 use crate::isa::aarch64::settings as aarch64_settings;
 #[cfg(feature = "unwind")]
 use crate::isa::unwind::systemv;
@@ -238,6 +238,17 @@ impl TargetIsa for AArch64Backend {
     fn has_x86_pmaddubsw_lowering(&self) -> bool {
         false
     }
+
+    fn default_argument_extension(&self) -> ir::ArgumentExtension {
+        // This is copied/carried over from a historical piece of code in
+        // Wasmtime:
+        //
+        // https://github.com/bytecodealliance/wasmtime/blob/a018a5a9addb77d5998021a0150192aa955c71bf/crates/cranelift/src/lib.rs#L366-L374
+        //
+        // Whether or not it is still applicable here is unsure, but it's left
+        // the same as-is for now to reduce the likelihood of problems arising.
+        ir::ArgumentExtension::Uext
+    }
 }
 
 impl fmt::Display for AArch64Backend {

cranelift/codegen/src/isa/mod.rs

@@ -400,6 +400,16 @@ pub trait TargetIsa: fmt::Display + Send + Sync {
     /// Returns whether the CLIF `x86_pmaddubsw` instruction is implemented for
     /// this ISA.
     fn has_x86_pmaddubsw_lowering(&self) -> bool;
+
+    /// Returns the mode of extension used for integer arguments smaller than
+    /// the pointer width in function signatures.
+    ///
+    /// Some platform ABIs require that smaller-than-pointer-width values are
+    /// either zero or sign-extended to the full register width. This value is
+    /// propagated to the `AbiParam` value created for signatures. Note that not
+    /// all ABIs for all platforms require extension of any form, so this is
+    /// generally only necessary for the `default_call_conv`.
+    fn default_argument_extension(&self) -> ir::ArgumentExtension;
 }
 
 /// Function alignment specifications as required by an ISA, returned by

cranelift/codegen/src/isa/pulley_shared/mod.rs

@@ -243,6 +243,10 @@ where
     fn has_x86_pmaddubsw_lowering(&self) -> bool {
         false
     }
+
+    fn default_argument_extension(&self) -> ir::ArgumentExtension {
+        ir::ArgumentExtension::None
+    }
 }
 
 /// Create a new Pulley ISA builder.

cranelift/codegen/src/isa/riscv64/mod.rs

@@ -216,6 +216,19 @@ impl TargetIsa for Riscv64Backend {
     fn has_x86_pmaddubsw_lowering(&self) -> bool {
         false
     }
+
+    fn default_argument_extension(&self) -> ir::ArgumentExtension {
+        // According to https://riscv.org/wp-content/uploads/2024/12/riscv-calling.pdf
+        // it says:
+        //
+        // > In RV64, 32-bit types, such as int, are stored in integer
+        // > registers as proper sign extensions of their 32-bit values; that
+        // > is, bits 63..31 are all equal. This restriction holds even for
+        // > unsigned 32-bit types.
+        //
+        // leading to `sext` here.
+        ir::ArgumentExtension::Sext
+    }
 }
 
 impl fmt::Display for Riscv64Backend {

cranelift/codegen/src/isa/s390x/mod.rs

@@ -1,7 +1,7 @@
 //! IBM Z 64-bit Instruction Set Architecture.
 
 use crate::dominator_tree::DominatorTree;
-use crate::ir::{Function, Type};
+use crate::ir::{self, Function, Type};
 use crate::isa::s390x::settings as s390x_settings;
 #[cfg(feature = "unwind")]
 use crate::isa::unwind::systemv::RegisterMappingError;
@@ -198,6 +198,17 @@ impl TargetIsa for S390xBackend {
     fn has_x86_pmaddubsw_lowering(&self) -> bool {
         false
     }
+
+    fn default_argument_extension(&self) -> ir::ArgumentExtension {
+        // This is copied/carried over from a historical piece of code in
+        // Wasmtime:
+        //
+        // https://github.com/bytecodealliance/wasmtime/blob/a018a5a9addb77d5998021a0150192aa955c71bf/crates/cranelift/src/lib.rs#L366-L374
+        //
+        // Whether or not it is still applicable here is unsure, but it's left
+        // the same as-is for now to reduce the likelihood of problems arising.
+        ir::ArgumentExtension::Uext
+    }
 }
 
 impl fmt::Display for S390xBackend {

cranelift/codegen/src/isa/x64/mod.rs

@@ -4,7 +4,7 @@ pub use self::inst::{args, AtomicRmwSeqOp, EmitInfo, EmitState, Inst};
 
 use super::{OwnedTargetIsa, TargetIsa};
 use crate::dominator_tree::DominatorTree;
-use crate::ir::{types, Function, Type};
+use crate::ir::{self, types, Function, Type};
 #[cfg(feature = "unwind")]
 use crate::isa::unwind::systemv;
 use crate::isa::x64::settings as x64_settings;
@@ -186,6 +186,17 @@ impl TargetIsa for X64Backend {
     fn has_x86_pmaddubsw_lowering(&self) -> bool {
         self.x64_flags.use_ssse3()
     }
+
+    fn default_argument_extension(&self) -> ir::ArgumentExtension {
+        // This is copied/carried over from a historical piece of code in
+        // Wasmtime:
+        //
+        // https://github.com/bytecodealliance/wasmtime/blob/a018a5a9addb77d5998021a0150192aa955c71bf/crates/cranelift/src/lib.rs#L366-L374
+        //
+        // Whether or not it is still applicable here is unsure, but it's left
+        // the same as-is for now to reduce the likelihood of problems arising.
+        ir::ArgumentExtension::Uext
+    }
 }
 
 /// Emit unwind info for an x86 target.

crates/cranelift/src/lib.rs

@@ -343,6 +343,7 @@ struct BuiltinFunctionSignatures {
 
     host_call_conv: CallConv,
     wasm_call_conv: CallConv,
+    argument_extension: ir::ArgumentExtension,
 }
 
 impl BuiltinFunctionSignatures {
@@ -351,6 +352,7 @@ impl BuiltinFunctionSignatures {
             pointer_type: compiler.isa().pointer_type(),
             host_call_conv: CallConv::triple_default(compiler.isa().triple()),
             wasm_call_conv: wasm_call_conv(compiler.isa(), compiler.tunables()),
+            argument_extension: compiler.isa().default_argument_extension(),
         }
     }
 
@@ -363,16 +365,7 @@ impl BuiltinFunctionSignatures {
     }
 
     fn i32(&self) -> AbiParam {
-        // Some platform ABIs require i32 values to be zero- or sign-
-        // extended to the full register width.  We need to indicate
-        // this here by using the appropriate .uext or .sext attribute.
-        // The attribute can be added unconditionally; platforms whose
-        // ABI does not require such extensions will simply ignore it.
-        // Note that currently all i32 arguments or return values used
-        // by builtin functions are unsigned, so we always use .uext.
-        // If that ever changes, we will have to add a second type
-        // marker here.
-        AbiParam::new(ir::types::I32).uext()
+        AbiParam::new(ir::types::I32)
     }
 
     fn i64(&self) -> AbiParam {
@@ -418,6 +411,16 @@ impl BuiltinFunctionSignatures {
     fn host_signature(&self, builtin: BuiltinFunctionIndex) -> Signature {
         let mut sig = self.wasm_signature(builtin);
         sig.call_conv = self.host_call_conv;
+
+        // Once we're declaring the signature of a host function we must
+        // respect the default ABI of the platform which is where argument
+        // extension of params/results may come into play.
+        for arg in sig.params.iter_mut().chain(sig.returns.iter_mut()) {
+            if arg.value_type.is_int() {
+                arg.extension = self.argument_extension;
+            }
+        }
+
         sig
     }
 }

tests/disas/gc/drc/array-new-fixed.wat

@@ -14,7 +14,7 @@
 ;;     gv1 = load.i64 notrap aligned readonly gv0+8
 ;;     gv2 = load.i64 notrap aligned gv1+16
 ;;     gv3 = vmctx
-;;     sig0 = (i64 vmctx, i32 uext, i32 uext, i32 uext, i32 uext) -> i64 tail
+;;     sig0 = (i64 vmctx, i32, i32, i32, i32) -> i64 tail
 ;;     fn0 = colocated u1:27 sig0
 ;;     stack_limit = gv2
 ;;

tests/disas/gc/drc/array-new.wat

@@ -14,7 +14,7 @@
 ;;     gv1 = load.i64 notrap aligned readonly gv0+8
 ;;     gv2 = load.i64 notrap aligned gv1+16
 ;;     gv3 = vmctx
-;;     sig0 = (i64 vmctx, i32 uext, i32 uext, i32 uext, i32 uext) -> i64 tail
+;;     sig0 = (i64 vmctx, i32, i32, i32, i32) -> i64 tail
 ;;     fn0 = colocated u1:27 sig0
 ;;     stack_limit = gv2
 ;;

tests/disas/gc/drc/br-on-cast-fail.wat

@@ -22,7 +22,7 @@
 ;;     gv1 = load.i64 notrap aligned readonly gv0+8
 ;;     gv2 = load.i64 notrap aligned gv1+16
 ;;     gv3 = vmctx
-;;     sig0 = (i64 vmctx, i32 uext, i32 uext) -> i32 uext tail
+;;     sig0 = (i64 vmctx, i32, i32) -> i32 tail
 ;;     sig1 = (i64 vmctx, i64) tail
 ;;     sig2 = (i64 vmctx, i64) tail
 ;;     fn0 = colocated u1:35 sig0