s390x: Add try_call / try_call_indirect support (#10557)

Following-up on #10510
this adds the corresponding support to the s390x target.

Due to the separate ABI implementation, this currently needs to
duplicate some of the underlying logic.  I hope to be able to
refactor ABI handling in the future to be able to share more of
that code across all targets.

Author: uweigand

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

@@ -906,8 +906,8 @@ impl ABIMachineSpec for S390xMachineDeps {
         call_conv_of_callee: isa::CallConv,
         is_exception: bool,
     ) -> PRegSet {
-        assert!(!is_exception);
         match call_conv_of_callee {
+            _ if is_exception => ALL_CLOBBERS,
             isa::CallConv::Tail => TAIL_CLOBBERS,
             _ => SYSV_CLOBBERS,
         }
@@ -1009,6 +1009,14 @@ impl ABIMachineSpec for S390xMachineDeps {
     fn retval_temp_reg(_call_conv_of_callee: isa::CallConv) -> Writable<Reg> {
         panic!("Should not be called");
     }
+
+    fn exception_payload_regs(call_conv: isa::CallConv) -> &'static [Reg] {
+        const PAYLOAD_REGS: &'static [Reg] = &[gpr(6), gpr(7)];
+        match call_conv {
+            isa::CallConv::SystemV | isa::CallConv::Tail => PAYLOAD_REGS,
+            _ => &[],
+        }
+    }
 }
 
 impl S390xMachineDeps {
@@ -1378,6 +1386,59 @@ const fn tail_clobbers() -> PRegSet {
 }
 const TAIL_CLOBBERS: PRegSet = tail_clobbers();
 
+const fn all_clobbers() -> PRegSet {
+    PRegSet::empty()
+        .with(gpr_preg(0))
+        .with(gpr_preg(1))
+        .with(gpr_preg(2))
+        .with(gpr_preg(3))
+        .with(gpr_preg(4))
+        .with(gpr_preg(5))
+        .with(gpr_preg(6))
+        .with(gpr_preg(7))
+        .with(gpr_preg(8))
+        .with(gpr_preg(9))
+        .with(gpr_preg(10))
+        .with(gpr_preg(11))
+        .with(gpr_preg(12))
+        .with(gpr_preg(13))
+        .with(gpr_preg(14))
+        .with(gpr_preg(15))
+        .with(vr_preg(0))
+        .with(vr_preg(1))
+        .with(vr_preg(2))
+        .with(vr_preg(3))
+        .with(vr_preg(4))
+        .with(vr_preg(5))
+        .with(vr_preg(6))
+        .with(vr_preg(7))
+        .with(vr_preg(8))
+        .with(vr_preg(9))
+        .with(vr_preg(10))
+        .with(vr_preg(11))
+        .with(vr_preg(12))
+        .with(vr_preg(13))
+        .with(vr_preg(14))
+        .with(vr_preg(15))
+        .with(vr_preg(16))
+        .with(vr_preg(17))
+        .with(vr_preg(18))
+        .with(vr_preg(19))
+        .with(vr_preg(20))
+        .with(vr_preg(21))
+        .with(vr_preg(22))
+        .with(vr_preg(23))
+        .with(vr_preg(24))
+        .with(vr_preg(25))
+        .with(vr_preg(26))
+        .with(vr_preg(27))
+        .with(vr_preg(28))
+        .with(vr_preg(29))
+        .with(vr_preg(30))
+        .with(vr_preg(31))
+}
+const ALL_CLOBBERS: PRegSet = all_clobbers();
+
 fn sysv_create_machine_env() -> MachineEnv {
     MachineEnv {
         preferred_regs_by_class: [

cranelift/codegen/src/isa/s390x/inst.isle

@@ -2669,10 +2669,9 @@
         dst))
 
 ;; Helper for emitting `MInst.Call` instructions.
-(decl call_impl (WritableReg CallSiteInfo) InstOutput)
-(rule (call_impl reg (call_site_info info output))
-      (let ((_ Unit (emit (MInst.Call reg info))))
-        output))
+(decl call_impl (WritableReg BoxCallInfo) SideEffectNoResult)
+(rule (call_impl reg info)
+      (SideEffectNoResult.Inst (MInst.Call reg info)))
 
 ;; Helper for emitting `MInst.ReturnCall` instructions.
 (decl return_call_impl (BoxReturnCallInfo) SideEffectNoResult)
@@ -3478,6 +3477,9 @@
 (decl abi_call_site_info (Sig CallInstDest CallArgList) CallSiteInfo)
 (extern constructor abi_call_site_info abi_call_site_info)
 
+(decl abi_try_call_info (Sig CallInstDest CallArgList ExceptionTable MachLabelSlice) BoxCallInfo)
+(extern constructor abi_try_call_info abi_try_call_info)
+
 (decl abi_return_call_info (Sig CallInstDest CallArgList) BoxReturnCallInfo)
 (extern constructor abi_return_call_info abi_return_call_info)
 
@@ -3495,13 +3497,20 @@
 
 (decl abi_call (Sig CallInstDest CallArgList) InstOutput)
 (rule (abi_call abi dest uses)
-      (let ((info CallSiteInfo (abi_call_site_info abi dest uses)))
-        (call_impl (writable_link_reg) info)))
+      (abi_call_impl (abi_call_site_info abi dest uses)))
+(decl abi_call_impl (CallSiteInfo) InstOutput)
+(rule (abi_call_impl (call_site_info info output))
+      (let ((_ Unit (emit_side_effect (call_impl (writable_link_reg) info))))
+        output))
 
 (decl abi_return_call (Sig CallInstDest CallArgList) SideEffectNoResult)
 (rule (abi_return_call abi dest uses)
       (return_call_impl (abi_return_call_info abi dest uses)))
 
+(decl abi_try_call (Sig CallInstDest CallArgList ExceptionTable MachLabelSlice) SideEffectNoResult)
+(rule (abi_try_call abi dest uses et targets)
+      (call_impl (writable_link_reg) (abi_try_call_info abi dest uses et targets)))
+
 (decl abi_lane_order (Sig) LaneOrder)
 (extern constructor abi_lane_order abi_lane_order)
 

cranelift/codegen/src/isa/s390x/inst/emit.rs

@@ -3218,11 +3218,28 @@ impl Inst {
                 put(sink, enc);
                 sink.add_call_site();
 
+                // Add exception info, if any, at this point (which will
+                // be the return address on stack).
+                if let Some(try_call) = info.try_call_info.as_ref() {
+                    for &(tag, label) in &try_call.exception_dests {
+                        sink.add_exception_handler(tag, label);
+                    }
+                }
+
                 state.nominal_sp_offset -= info.callee_pop_size;
 
                 for inst in S390xMachineDeps::gen_retval_loads(info) {
                     inst.emit(sink, emit_info, state);
                 }
+
+                // If this is a try-call, jump to the continuation
+                // (normal-return) block.
+                if let Some(try_call) = info.try_call_info.as_ref() {
+                    let jmp = Inst::Jump {
+                        dest: try_call.continuation,
+                    };
+                    jmp.emit(sink, emit_info, state);
+                }
             }
             &Inst::ReturnCall { ref info } => {
                 let (epilogue_insts, temp_dest) = S390xMachineDeps::gen_tail_epilogue(

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

@@ -1076,6 +1076,7 @@ impl MachInst for Inst {
             &Inst::CondBr { .. } => MachTerminator::Branch,
             &Inst::IndirectBr { .. } => MachTerminator::Branch,
             &Inst::JTSequence { .. } => MachTerminator::Branch,
+            &Inst::Call { ref info, .. } if info.try_call_info.is_some() => MachTerminator::Branch,
             _ => MachTerminator::None,
         }
     }
@@ -3119,18 +3120,30 @@ impl Inst {
                 if !retval_loads.is_empty() {
                     retval_loads = " ; ".to_string() + &retval_loads;
                 }
+                let try_call = if let Some(try_call_info) = &info.try_call_info {
+                    let dests = try_call_info
+                        .exception_dests
+                        .iter()
+                        .map(|(tag, label)| format!("{tag:?}: {label:?}"))
+                        .collect::<Vec<_>>()
+                        .join(", ");
+                    format!("; jg {:?}; catch [{dests}]", try_call_info.continuation)
+                } else {
+                    "".to_string()
+                };
                 let callee_pop_size = if info.callee_pop_size > 0 {
                     format!(" ; callee_pop_size {}", info.callee_pop_size)
                 } else {
                     "".to_string()
                 };
                 format!(
-                    "{} {}, {}{}{}",
+                    "{} {}, {}{}{}{}",
                     opcode,
                     show_reg(link),
                     dest,
                     callee_pop_size,
-                    retval_loads
+                    retval_loads,
+                    try_call
                 )
             }
             &Inst::ReturnCall { ref info } => {

cranelift/codegen/src/isa/s390x/inst/regs.rs

@@ -10,8 +10,8 @@ use crate::machinst::*;
 // Registers, the Universe thereof, and printing
 
 /// Get a reference to a GPR (integer register).
-pub fn gpr(num: u8) -> Reg {
-    Reg::from(gpr_preg(num))
+pub const fn gpr(num: u8) -> Reg {
+    Reg::from_real_reg(gpr_preg(num))
 }
 
 pub(crate) const fn gpr_preg(num: u8) -> PReg {
@@ -25,8 +25,8 @@ pub fn writable_gpr(num: u8) -> Writable<Reg> {
 }
 
 /// Get a reference to a VR (vector register).
-pub fn vr(num: u8) -> Reg {
-    Reg::from(vr_preg(num))
+pub const fn vr(num: u8) -> Reg {
+    Reg::from_real_reg(vr_preg(num))
 }
 
 pub(crate) const fn vr_preg(num: u8) -> PReg {

cranelift/codegen/src/isa/s390x/lower.isle

@@ -3926,6 +3926,30 @@
         (args_builder_finish uses)))
 
 
+;;;; Rules for `try_call` and `try_call_indirect` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+;; Direct call to an in-range function.
+(rule 1 (lower_branch (try_call (func_ref_data sig_ref name (reloc_distance_near)) args et) targets)
+      (let ((abi Sig (abi_sig sig_ref))
+            (uses CallArgList (lower_call_args abi (range 0 (abi_num_args abi)) args)))
+        (emit_side_effect (abi_try_call abi (CallInstDest.Direct name) uses et targets))))
+
+;; Direct call to an out-of-range function (implicitly via pointer).
+(rule (lower_branch (try_call (func_ref_data sig_ref name _) args et) targets)
+      (let ((abi Sig (abi_sig sig_ref))
+            (uses CallArgList (lower_call_args abi (range 0 (abi_num_args abi)) args))
+            (target Reg (load_symbol_reloc (SymbolReloc.Absolute name 0))))
+        (emit_side_effect (abi_try_call abi (CallInstDest.Indirect target) uses et targets))))
+
+;; Indirect call.
+(rule (lower_branch (try_call_indirect ptr args et) targets)
+      (if-let (exception_sig sig_ref) et)
+      (let ((abi Sig (abi_sig sig_ref))
+            (target Reg (put_in_reg ptr))
+            (uses CallArgList (lower_call_args abi (range 0 (abi_num_args abi)) args)))
+        (emit_side_effect (abi_try_call abi (CallInstDest.Indirect target) uses et targets))))
+
+
 ;;;; Common helpers for argument lowering ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 ;; Lower function arguments (part 1): prepare buffer copies.