1 use super::regs; 2 use crate::RegIndexEnv; 3 use crate::abi::{ABI, ABIOperand, ABIParams, ABIResults, ABISig, ParamsOrReturns, align_to}; 4 use crate::codegen::CodeGenError; 5 use crate::isa::{CallingConvention, reg::Reg}; 6 use anyhow::{Result, bail}; 7 use wasmtime_environ::{WasmHeapType, WasmValType}; 8 9 #[derive(Default)] 10 pub(crate) struct Aarch64ABI; 11 12 /// The x28 register serves as the shadow stack pointer. For further details, 13 /// please refer to [`crate::isa::aarch64::regs::shadow_sp`]. 14 /// 15 /// This register is designated as callee-saved to prevent corruption during 16 /// function calls. This is especially important for Wasm-to-Wasm calls in 17 /// Winch-generated code, as Winch's default calling convention does not define 18 /// any callee-saved registers. 19 /// 20 /// Note that 16 bytes are used to save the shadow stack pointer register even 21 /// though only 8 are needed. 16 is used for simplicity to ensure that the 22 /// 16-byte alignment requirement for memory addressing is met at the function's 23 /// prologue and epilogue. 24 pub const SHADOW_STACK_POINTER_SLOT_SIZE: u8 = 16; 25 26 impl ABI for Aarch64ABI { 27 // TODO change to 16 once SIMD is supported 28 fn stack_align() -> u8 { 29 8 30 } 31 32 fn call_stack_align() -> u8 { 33 16 34 } 35 36 fn arg_base_offset() -> u8 { 37 // Two 8-byte slots: 38 // * One for link register 39 // * One for the frame pointer 40 // 41 // ┌──────────┬───────── Argument base 42 // │ LR │ 43 // │ │ 44 // ├──────────┼ 45 // │ │ 46 // │ FP │ 47 // └──────────┴ -> 16 48 16 49 } 50 51 fn initial_frame_size() -> u8 { 52 // The initial frame size is composed of 4 8-byte slots: 53 // * Two slots for the link register and the frame pointer. See 54 // [`Self::arg_base_offset`] 55 // * Two for the shadow stack pointer register. See 56 // [`SHADOW_STACK_POINTER_SIZE`] 57 // 58 // ┌──────────┬───────── Argument base 59 // │ LR │ 60 // │ │ 61 // ├──────────┼ 62 // │ │ 63 // │ FP │ 64 // ┌──────────┬ 65 // │ │ 66 // │ │ 67 // │ │ 68 // │ x28 │ 69 // └──────────┴ -> 32 70 Self::arg_base_offset() + SHADOW_STACK_POINTER_SLOT_SIZE 71 } 72 73 fn word_bits() -> u8 { 74 64 75 } 76 77 fn sig_from( 78 params: &[WasmValType], 79 returns: &[WasmValType], 80 call_conv: &CallingConvention, 81 ) -> Result<ABISig> { 82 assert!(call_conv.is_apple_aarch64() || call_conv.is_default()); 83 // The first element tracks the general purpose register index, capped at 7 (x0-x7). 84 // The second element tracks the floating point register index, capped at 7 (v0-v7). 85 // Follows 86 // https://github.com/ARM-software/abi-aa/blob/2021Q1/aapcs64/aapcs64.rst#64parameter-passing 87 let mut params_index_env = RegIndexEnv::with_limits_per_class(8, 8); 88 let results = Self::abi_results(returns, call_conv)?; 89 let params = 90 ABIParams::from::<_, Self>(params, 0, results.on_stack(), |ty, stack_offset| { 91 Self::to_abi_operand( 92 ty, 93 stack_offset, 94 &mut params_index_env, 95 call_conv, 96 ParamsOrReturns::Params, 97 ) 98 })?; 99 100 Ok(ABISig::new(*call_conv, params, results)) 101 } 102 103 fn abi_results(returns: &[WasmValType], call_conv: &CallingConvention) -> Result<ABIResults> { 104 assert!(call_conv.is_apple_aarch64() || call_conv.is_default()); 105 // Use absolute count for results given that for Winch's 106 // default CallingConvention only one register is used for results 107 // independent of the register class. 108 // In the case of 2+ results, the rest are passed in the stack, 109 // similar to how Wasmtime handles multi-value returns. 110 let mut returns_index_env = RegIndexEnv::with_absolute_limit(1); 111 112 ABIResults::from(returns, call_conv, |ty, stack_offset| { 113 Self::to_abi_operand( 114 ty, 115 stack_offset, 116 &mut returns_index_env, 117 call_conv, 118 ParamsOrReturns::Returns, 119 ) 120 }) 121 } 122 123 fn vmctx_reg() -> Reg { 124 regs::xreg(9) 125 } 126 127 fn stack_slot_size() -> u8 { 128 Self::word_bytes() 129 } 130 131 fn sizeof(ty: &WasmValType) -> u8 { 132 match ty { 133 WasmValType::Ref(rt) => match rt.heap_type { 134 WasmHeapType::Func => Self::word_bytes(), 135 ht => unimplemented!("Support for WasmHeapType: {ht}"), 136 }, 137 WasmValType::F64 | WasmValType::I64 => Self::word_bytes(), 138 WasmValType::F32 | WasmValType::I32 => Self::word_bytes() / 2, 139 WasmValType::V128 => Self::word_bytes() * 2, 140 } 141 } 142 } 143 144 impl Aarch64ABI { 145 fn to_abi_operand( 146 wasm_arg: &WasmValType, 147 stack_offset: u32, 148 index_env: &mut RegIndexEnv, 149 call_conv: &CallingConvention, 150 params_or_returns: ParamsOrReturns, 151 ) -> Result<(ABIOperand, u32)> { 152 let (reg, ty) = match wasm_arg { 153 ty @ (WasmValType::I32 | WasmValType::I64) => { 154 (index_env.next_gpr().map(regs::xreg), ty) 155 } 156 157 ty @ (WasmValType::F32 | WasmValType::F64) => { 158 (index_env.next_fpr().map(regs::vreg), ty) 159 } 160 161 ty @ WasmValType::Ref(rt) => match rt.heap_type { 162 WasmHeapType::Func | WasmHeapType::Extern => { 163 (index_env.next_gpr().map(regs::xreg), ty) 164 } 165 _ => bail!(CodeGenError::unsupported_wasm_type()), 166 }, 167 168 _ => bail!(CodeGenError::unsupported_wasm_type()), 169 }; 170 171 let ty_size = <Self as ABI>::sizeof(wasm_arg); 172 let default = || { 173 let arg = ABIOperand::stack_offset(stack_offset, *ty, ty_size as u32); 174 let slot_size = Self::stack_slot_size(); 175 // Stack slots for parameters are aligned to a fixed slot size, 176 // in the case of Aarch64, 8 bytes. 177 // For the non-default calling convention, stack slots for 178 // return values are type-sized aligned. 179 // For the default calling convention, we don't type-size align, 180 // given that results on the stack must match spills generated 181 // from within the compiler, which are not type-size aligned. 182 let next_stack = if params_or_returns == ParamsOrReturns::Params { 183 align_to(stack_offset, slot_size as u32) + (slot_size as u32) 184 } else if call_conv.is_default() { 185 stack_offset + (ty_size as u32) 186 } else { 187 align_to(stack_offset, ty_size as u32) + (ty_size as u32) 188 }; 189 (arg, next_stack) 190 }; 191 Ok(reg.map_or_else(default, |reg| { 192 (ABIOperand::reg(reg, *ty, ty_size as u32), stack_offset) 193 })) 194 } 195 } 196 197 #[cfg(test)] 198 mod tests { 199 use super::Aarch64ABI; 200 use crate::{ 201 abi::{ABI, ABIOperand}, 202 isa::CallingConvention, 203 isa::aarch64::regs, 204 isa::reg::Reg, 205 }; 206 use wasmtime_environ::{ 207 WasmFuncType, 208 WasmValType::{self, *}, 209 }; 210 211 use anyhow::Result; 212 213 #[test] 214 fn xreg_abi_sig() -> Result<()> { 215 let wasm_sig = WasmFuncType::new( 216 [I32, I64, I32, I64, I32, I32, I64, I32, I64].into(), 217 [].into(), 218 ); 219 220 let sig = Aarch64ABI::sig(&wasm_sig, &CallingConvention::Default)?; 221 let params = sig.params; 222 223 match_reg_arg(params.get(0).unwrap(), I32, regs::xreg(0)); 224 match_reg_arg(params.get(1).unwrap(), I64, regs::xreg(1)); 225 match_reg_arg(params.get(2).unwrap(), I32, regs::xreg(2)); 226 match_reg_arg(params.get(3).unwrap(), I64, regs::xreg(3)); 227 match_reg_arg(params.get(4).unwrap(), I32, regs::xreg(4)); 228 match_reg_arg(params.get(5).unwrap(), I32, regs::xreg(5)); 229 match_reg_arg(params.get(6).unwrap(), I64, regs::xreg(6)); 230 match_reg_arg(params.get(7).unwrap(), I32, regs::xreg(7)); 231 match_stack_arg(params.get(8).unwrap(), I64, 0); 232 Ok(()) 233 } 234 235 #[test] 236 fn vreg_abi_sig() -> Result<()> { 237 let wasm_sig = WasmFuncType::new( 238 [F32, F64, F32, F64, F32, F32, F64, F32, F64].into(), 239 [].into(), 240 ); 241 242 let sig = Aarch64ABI::sig(&wasm_sig, &CallingConvention::Default)?; 243 let params = sig.params; 244 245 match_reg_arg(params.get(0).unwrap(), F32, regs::vreg(0)); 246 match_reg_arg(params.get(1).unwrap(), F64, regs::vreg(1)); 247 match_reg_arg(params.get(2).unwrap(), F32, regs::vreg(2)); 248 match_reg_arg(params.get(3).unwrap(), F64, regs::vreg(3)); 249 match_reg_arg(params.get(4).unwrap(), F32, regs::vreg(4)); 250 match_reg_arg(params.get(5).unwrap(), F32, regs::vreg(5)); 251 match_reg_arg(params.get(6).unwrap(), F64, regs::vreg(6)); 252 match_reg_arg(params.get(7).unwrap(), F32, regs::vreg(7)); 253 match_stack_arg(params.get(8).unwrap(), F64, 0); 254 Ok(()) 255 } 256 257 #[test] 258 fn mixed_abi_sig() -> Result<()> { 259 let wasm_sig = WasmFuncType::new( 260 [F32, I32, I64, F64, I32, F32, F64, F32, F64].into(), 261 [].into(), 262 ); 263 264 let sig = Aarch64ABI::sig(&wasm_sig, &CallingConvention::Default)?; 265 let params = sig.params; 266 267 match_reg_arg(params.get(0).unwrap(), F32, regs::vreg(0)); 268 match_reg_arg(params.get(1).unwrap(), I32, regs::xreg(0)); 269 match_reg_arg(params.get(2).unwrap(), I64, regs::xreg(1)); 270 match_reg_arg(params.get(3).unwrap(), F64, regs::vreg(1)); 271 match_reg_arg(params.get(4).unwrap(), I32, regs::xreg(2)); 272 match_reg_arg(params.get(5).unwrap(), F32, regs::vreg(2)); 273 match_reg_arg(params.get(6).unwrap(), F64, regs::vreg(3)); 274 match_reg_arg(params.get(7).unwrap(), F32, regs::vreg(4)); 275 match_reg_arg(params.get(8).unwrap(), F64, regs::vreg(5)); 276 Ok(()) 277 } 278 279 #[test] 280 fn int_abi_sig_multi_returns() -> Result<()> { 281 let wasm_sig = WasmFuncType::new( 282 [I32, I64, I32, I64, I32, I32].into(), 283 [I32, I32, I32].into(), 284 ); 285 286 let sig = Aarch64ABI::sig(&wasm_sig, &CallingConvention::Default)?; 287 let params = sig.params; 288 let results = sig.results; 289 290 match_reg_arg(params.get(0).unwrap(), I32, regs::xreg(1)); 291 match_reg_arg(params.get(1).unwrap(), I64, regs::xreg(2)); 292 match_reg_arg(params.get(2).unwrap(), I32, regs::xreg(3)); 293 match_reg_arg(params.get(3).unwrap(), I64, regs::xreg(4)); 294 match_reg_arg(params.get(4).unwrap(), I32, regs::xreg(5)); 295 match_reg_arg(params.get(5).unwrap(), I32, regs::xreg(6)); 296 297 match_stack_arg(results.get(0).unwrap(), I32, 4); 298 match_stack_arg(results.get(1).unwrap(), I32, 0); 299 match_reg_arg(results.get(2).unwrap(), I32, regs::xreg(0)); 300 Ok(()) 301 } 302 303 #[test] 304 fn mixed_abi_sig_multi_returns() -> Result<()> { 305 let wasm_sig = WasmFuncType::new( 306 [F32, I32, I64, F64, I32].into(), 307 [I32, F32, I32, F32, I64].into(), 308 ); 309 310 let sig = Aarch64ABI::sig(&wasm_sig, &CallingConvention::Default)?; 311 let params = sig.params; 312 let results = sig.results; 313 314 match_reg_arg(params.get(0).unwrap(), F32, regs::vreg(0)); 315 match_reg_arg(params.get(1).unwrap(), I32, regs::xreg(1)); 316 match_reg_arg(params.get(2).unwrap(), I64, regs::xreg(2)); 317 match_reg_arg(params.get(3).unwrap(), F64, regs::vreg(1)); 318 match_reg_arg(params.get(4).unwrap(), I32, regs::xreg(3)); 319 320 match_stack_arg(results.get(0).unwrap(), I32, 12); 321 match_stack_arg(results.get(1).unwrap(), F32, 8); 322 match_stack_arg(results.get(2).unwrap(), I32, 4); 323 match_stack_arg(results.get(3).unwrap(), F32, 0); 324 match_reg_arg(results.get(4).unwrap(), I64, regs::xreg(0)); 325 Ok(()) 326 } 327 328 #[track_caller] 329 fn match_reg_arg(abi_arg: &ABIOperand, expected_ty: WasmValType, expected_reg: Reg) { 330 match abi_arg { 331 &ABIOperand::Reg { reg, ty, .. } => { 332 assert_eq!(reg, expected_reg); 333 assert_eq!(ty, expected_ty); 334 } 335 stack => panic!("Expected reg argument, got {stack:?}"), 336 } 337 } 338 339 #[track_caller] 340 fn match_stack_arg(abi_arg: &ABIOperand, expected_ty: WasmValType, expected_offset: u32) { 341 match abi_arg { 342 &ABIOperand::Stack { offset, ty, .. } => { 343 assert_eq!(offset, expected_offset); 344 assert_eq!(ty, expected_ty); 345 } 346 reg => panic!("Expected stack argument, got {reg:?}"), 347 } 348 } 349 } 350