From e3f549d7a8118a17203bc8fab3a7a4a36c5dfe67 Mon Sep 17 00:00:00 2001
From: Laurence Tratt <laurie@tratt.net>
Date: Sun, 26 Jan 2025 09:30:38 +0000
Subject: [PATCH] Implement codegen for 64 bit llvm.smax.

This is effectively an IR intrinsic, but is encoded as a call.
---
 ykrt/src/compile/jitc_yk/codegen/x64/mod.rs | 45 +++++++++++++++++++++
 ykrt/src/compile/jitc_yk/jit_ir/jit_ir.l    |  4 +-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs b/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs
index 9d9c16aa5..4d95b1ab3 100644
--- a/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs
+++ b/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs
@@ -1503,6 +1503,11 @@ impl<'a> Assemble<'a> {
             "llvm.assume" => Ok(()),
             "llvm.lifetime.start.p0" => Ok(()),
             "llvm.lifetime.end.p0" => Ok(()),
+            x if x.starts_with("llvm.smax") => {
+                let [lhs_op, rhs_op] = args.try_into().unwrap();
+                self.cg_smax(iidx, lhs_op, rhs_op);
+                Ok(())
+            }
             x => {
                 let va = symbol_to_ptr(x).map_err(|e| CompilationError::General(e.to_string()))?;
                 self.emit_call(iidx, fty, Some(va), None, &args)
@@ -1687,6 +1692,25 @@ impl<'a> Assemble<'a> {
         }
     }
 
+    fn cg_smax(&mut self, iidx: InstIdx, lhs: Operand, rhs: Operand) {
+        assert_eq!(lhs.bitw(self.m), rhs.bitw(self.m));
+        let bitw = lhs.bitw(self.m);
+        let [lhs_reg, rhs_reg] = self.ra.assign_gp_regs(
+            &mut self.asm,
+            iidx,
+            [RegConstraint::InputOutput(lhs), RegConstraint::Input(rhs)],
+        );
+        match bitw {
+            64 => {
+                dynasm!(self.asm
+                    ; cmp Rq(lhs_reg.code()), Rq(rhs_reg.code())
+                    ; cmovl Rq(lhs_reg.code()), Rq(rhs_reg.code())
+                );
+            }
+            x => todo!("{x}"),
+        }
+    }
+
     /// Return the [VarLocation] an [Operand] relates to.
     fn op_to_var_location(&self, op: Operand) -> VarLocation {
         match op {
@@ -3899,6 +3923,27 @@ mod tests {
         );
     }
 
+    #[test]
+    fn cg_call_smax() {
+        codegen_and_test(
+            "
+             func_decl llvm.smax.i64 (i64, i64) -> i64
+             entry:
+               %0: i64 = param 0
+               %1: i64 = param 1
+               %2: i64 = call @llvm.smax.i64(%0, %1)
+               black_box %2
+            ",
+            "
+               ...
+               ; %2: i64 = call @llvm.smax.i64(%0, %1)
+               cmp r.64.a, r.64.b
+               cmovl r.64.a, r.64.b
+            ",
+            false,
+        );
+    }
+
     #[test]
     fn cg_eq() {
         codegen_and_test(
diff --git a/ykrt/src/compile/jitc_yk/jit_ir/jit_ir.l b/ykrt/src/compile/jitc_yk/jit_ir/jit_ir.l
index 4bce969fa..c1a2f736f 100644
--- a/ykrt/src/compile/jitc_yk/jit_ir/jit_ir.l
+++ b/ykrt/src/compile/jitc_yk/jit_ir/jit_ir.l
@@ -1,5 +1,5 @@
 %%
-@[a-zA-Z_.][a-zA-Z_0-9]* "GLOBAL"
+@[a-zA-Z_.][a-zA-Z_0-9.]* "GLOBAL"
 %[0-9]+ "LOCAL_OPERAND"
 i[0-9]+ "INT_TYPE"
 float "FLOAT_TYPE"
@@ -80,7 +80,7 @@ f_true "F_TRUE"
 urem "UREM"
 xor "XOR"
 [a-zA_Z_]+: "LABEL"
-[a-zA_Z_][a-zA_Z_0-9]* "ID"
+[a-zA_Z_][a-zA_Z_0-9.]* "ID"
 volatile "VOLATILE"
 \< "<"
 \> ">"