Update calling conventions for OCaml

llvm/include/llvm/AsmParser/LLToken.h

@@ -173,6 +173,11 @@ enum Kind {
   kw_amdgpu_kernel,
   kw_amdgpu_gfx,
   kw_tailcc,
+  kw_oxcaml_fpcc,
+  kw_oxcaml_nofpcc,
+  kw_oxcaml_ccc,
+  kw_oxcaml_c_stackcc,
+  kw_oxcaml_alloccc,
 
   // Attributes:
   kw_attributes,

llvm/include/llvm/IR/CallingConv.h

@@ -241,11 +241,12 @@ namespace CallingConv {
     /// Preserve X2-X15, X19-X29, SP, Z0-Z31, P0-P15.
     AArch64_SME_ABI_Support_Routines_PreserveMost_From_X2 = 103,
 
-    /// Calling convention for OCaml
-    OCaml = 104,
-
-    /// Special case of the C calling convention for calling from OCaml
-    OCaml_C_Call = 105,
+    /// Calling conventions used by OxCaml
+    OxCaml_WithFP = 104, // normal OxCaml functions (with frame pointers)
+    OxCaml_WithoutFP = 105, // normal OxCaml functions (w/o frame pointers)
+    OxCaml_C_Call = 106, // [caml_c_call]
+    OxCaml_C_Call_StackArgs = 107, // [caml_c_call_stack_args_llvm_backend]
+    OxCaml_Alloc = 108, // [caml_call_gc] and [caml_local_realloc]
 
     /// The highest possible ID. Must be some 2^k - 1.
     MaxID = 1023

llvm/lib/AsmParser/LLLexer.cpp

@@ -631,6 +631,11 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(amdgpu_kernel);
   KEYWORD(amdgpu_gfx);
   KEYWORD(tailcc);
+  KEYWORD(oxcaml_fpcc);
+  KEYWORD(oxcaml_nofpcc);
+  KEYWORD(oxcaml_ccc);
+  KEYWORD(oxcaml_c_stackcc);
+  KEYWORD(oxcaml_alloccc);
 
   KEYWORD(cc);
   KEYWORD(c);

llvm/lib/AsmParser/LLParser.cpp

@@ -2019,6 +2019,11 @@ void LLParser::parseOptionalDLLStorageClass(unsigned &Res) {
 ///   ::= 'amdgpu_cs'
 ///   ::= 'amdgpu_kernel'
 ///   ::= 'tailcc'
+///   ::= 'oxcaml_fpcc'
+///   ::= 'oxcaml_nofpcc'
+///   ::= 'oxcaml_ccc'
+///   ::= 'oxcaml_c_stackcc'
+///   ::= 'oxcaml_alloccc'
 ///   ::= 'cc' UINT
 ///
 bool LLParser::parseOptionalCallingConv(unsigned &CC) {
@@ -2077,6 +2082,11 @@ bool LLParser::parseOptionalCallingConv(unsigned &CC) {
   case lltok::kw_amdgpu_cs:      CC = CallingConv::AMDGPU_CS; break;
   case lltok::kw_amdgpu_kernel:  CC = CallingConv::AMDGPU_KERNEL; break;
   case lltok::kw_tailcc:         CC = CallingConv::Tail; break;
+  case lltok::kw_oxcaml_fpcc:    CC = CallingConv::OxCaml_WithFP; break;
+  case lltok::kw_oxcaml_nofpcc:  CC = CallingConv::OxCaml_WithoutFP; break;
+  case lltok::kw_oxcaml_ccc:     CC = CallingConv::OxCaml_C_Call; break;
+  case lltok::kw_oxcaml_c_stackcc:CC = CallingConv::OxCaml_C_Call_StackArgs; break;
+  case lltok::kw_oxcaml_alloccc: CC = CallingConv::OxCaml_Alloc; break;
   case lltok::kw_cc: {
       Lex.Lex();
       return parseUInt32(CC);

llvm/lib/IR/AsmWriter.cpp

@@ -340,7 +340,15 @@ static void PrintCallingConv(unsigned cc, raw_ostream &Out) {
   case CallingConv::AMDGPU_CS:     Out << "amdgpu_cs"; break;
   case CallingConv::AMDGPU_KERNEL: Out << "amdgpu_kernel"; break;
   case CallingConv::AMDGPU_Gfx:    Out << "amdgpu_gfx"; break;
-  case CallingConv::OCaml:         Out << "ocamlcc"; break;
+  case CallingConv::OxCaml_WithFP: Out << "oxcaml_fpcc"; break;
+  case CallingConv::OxCaml_WithoutFP:
+    Out << "oxcaml_nofpcc";
+    break;
+  case CallingConv::OxCaml_C_Call: Out << "oxcaml_ccc"; break;
+  case CallingConv::OxCaml_C_Call_StackArgs:
+    Out << "oxcaml_c_stackcc";
+    break;
+  case CallingConv::OxCaml_Alloc:  Out << "oxcaml_alloccc"; break;
   }
 }
 

llvm/lib/IR/Verifier.cpp

@@ -3581,10 +3581,11 @@ void Verifier::verifyMustTailCall(CallInst &CI) {
   AttributeList CalleeAttrs = CI.getAttributes();
   if (CI.getCallingConv() == CallingConv::SwiftTail ||
       CI.getCallingConv() == CallingConv::Tail ||
-      CI.getCallingConv() == CallingConv::OCaml) {
+      CI.getCallingConv() == CallingConv::OxCaml_WithFP ||
+      CI.getCallingConv() == CallingConv::OxCaml_WithoutFP) {
     StringRef CCName =
         CI.getCallingConv() == CallingConv::Tail ? "tailcc" :
-        CI.getCallingConv() == CallingConv::SwiftTail ? "swifttailcc" : "ocamlcc";
+        CI.getCallingConv() == CallingConv::SwiftTail ? "swifttailcc" : "oxcamlcc";
 
     // - Only sret, byval, swiftself, and swiftasync ABI-impacting attributes
     //   are allowed in swifttailcc call

llvm/lib/Target/X86/X86CallingConv.td

@@ -440,12 +440,12 @@ def RetCC_X86_64_HHVM: CallingConv<[
                                  RAX, R10, R11, R13, R14, R15]>>
 ]>;
 
-def RetCC_X86_64_OCaml : CallingConv<[
+def RetCC_X86_64_OxCaml : CallingConv<[
   // Promote all types to i64
   CCIfType<[i8, i16, i32], CCPromoteToType<i64>>,
 
   // Runtime registers (R14, R15) are threaded through function calls
-  // See the argument calling conventions for OCaml for more details
+  // See the argument calling conventions for OxCaml for more details
   CCIfType<[i64], CCAssignToReg<[R14, R15, RAX, RBX, RDI, RSI, RDX, RCX, R8, R9, R12, R13]>>,
 
   CCIfType<[f32, f64], CCAssignToReg<[
@@ -454,7 +454,13 @@ def RetCC_X86_64_OCaml : CallingConv<[
   ]>>
 ]>;
 
-def RetCC_X86_64_OCaml_C_Call : CallingConv<[
+def RetCC_X86_64_OxCaml_C_Call : CallingConv<[
+  CCIfType<[i64], CCAssignToReg<[R14, R15]>>,
+
+  CCDelegateTo<RetCC_X86_64_C>
+]>;
+
+def RetCC_X86_64_OxCaml_C_Call_StackArgs : CallingConv<[
   CCIfType<[i64], CCAssignToReg<[R14, R15]>>,
 
   CCDelegateTo<RetCC_X86_64_C>
@@ -513,9 +519,15 @@ def RetCC_X86_64 : CallingConv<[
   // Mingw64 and native Win64 use Win64 CC
   CCIfSubtarget<"isTargetWin64()", CCDelegateTo<RetCC_X86_Win64_C>>,
 
-  // Handle OCaml calls
-  CCIfCC<"CallingConv::OCaml", CCDelegateTo<RetCC_X86_64_OCaml>>,
-  CCIfCC<"CallingConv::OCaml_C_Call", CCDelegateTo<RetCC_X86_64_OCaml_C_Call>>,
+  // OxCaml CCs
+  CCIfCC<"CallingConv::OxCaml_WithFP", CCDelegateTo<RetCC_X86_64_OxCaml>>,
+  CCIfCC<"CallingConv::OxCaml_WithoutFP", CCDelegateTo<RetCC_X86_64_OxCaml>>,
+  CCIfCC<"CallingConv::OxCaml_C_Call", CCDelegateTo<RetCC_X86_64_OxCaml_C_Call>>,
+  CCIfCC<"CallingConv::OxCaml_C_Call_StackArgs",
+    CCDelegateTo<RetCC_X86_64_OxCaml_C_Call_StackArgs>>,
+
+  // See arg CC for OxCaml_Alloc for why we use OxCaml here.
+  CCIfCC<"CallingConv::OxCaml_Alloc", CCDelegateTo<RetCC_X86_64_OxCaml>>,
 
   // Otherwise, drop to normal X86-64 CC
   CCDelegateTo<RetCC_X86_64_C>
@@ -723,7 +735,7 @@ def CC_X86_Win64_VectorCall : CallingConv<[
   CCDelegateTo<CC_X86_Win64_C>
 ]>;
 
-def CC_X86_64_OCaml : CallingConv<[
+def CC_X86_64_OxCaml : CallingConv<[
   // Promote i8/i16/i32 arguments to i64.
   CCIfType<[i8, i16, i32], CCPromoteToType<i64>>,
 
@@ -733,6 +745,8 @@ def CC_X86_64_OCaml : CallingConv<[
   // runtime) must be preserved through the run of the program. To ensure LLVM
   // doesn't mess with them, they are explicitly threaded through function calls
   // and returns. These are R14 and R15.
+  // Note that whether frame pointers are enabled or not doesn't impact the
+  // argument or return conventions. It only changes the callee-saved registers.
   CCIfType<[i64], CCAssignToReg<[R14, R15, RAX, RBX, RDI, RSI, RDX, RCX, R8, R9, R12, R13]>>,
 
   CCIfType<[f32, f64], CCAssignToReg<[
@@ -741,9 +755,21 @@ def CC_X86_64_OCaml : CallingConv<[
   ]>>
 ]>;
 
-def CC_X86_64_OCaml_C_Call : CallingConv<[
-  // OCaml wraps C calls with another function that transfers stack arguments.
-  // RAX contains the function address, and R12 contains the size of the stack arguments.
+def CC_X86_64_OxCaml_C_Call : CallingConv<[
+  // Calling conventions followed by [caml_c_call] which wraps non-noalloc C calls.
+  // RAX contains the function address to call.
+  CCIfType<[i64], CCAssignToReg<[R14, R15, RAX]>>,
+
+  // Follow C convention normally otherwise
+  CCDelegateTo<CC_X86_64_C>
+]>;
+
+def CC_X86_64_OxCaml_C_Call_StackArgs : CallingConv<[
+  // Calling conventions followed by [caml_c_call_stack_args] to additionally handle
+  // transfer of stack arguments. Note that this function normally takes a pair of
+  // pointers on the stack, but since LLVM makes it hard to directly meddle with the
+  // stack, this in reality calls yet anothr wrapper which calculates this range given
+  // the number of stack arguments in bytes in R12.
   CCIfType<[i64], CCAssignToReg<[R14, R15, RAX, R12]>>,
 
   // Follow C convention normally otherwise
@@ -1152,8 +1178,19 @@ def CC_X86_64 : CallingConv<[
     CCIfSubtarget<"isTargetWin64()", CCDelegateTo<CC_X86_Win64_RegCall>>>,
   CCIfCC<"CallingConv::X86_RegCall", CCDelegateTo<CC_X86_SysV64_RegCall>>,
   CCIfCC<"CallingConv::X86_INTR", CCCustom<"CC_X86_Intr">>,
-  CCIfCC<"CallingConv::OCaml", CCDelegateTo<CC_X86_64_OCaml>>,
-  CCIfCC<"CallingConv::OCaml_C_Call", CCDelegateTo<CC_X86_64_OCaml_C_Call>>,
+
+  // OxCaml calling conventions
+  CCIfCC<"CallingConv::OxCaml_WithFP", CCDelegateTo<CC_X86_64_OxCaml>>,
+  CCIfCC<"CallingConv::OxCaml_WithoutFP", CCDelegateTo<CC_X86_64_OxCaml>>,
+  CCIfCC<"CallingConv::OxCaml_C_Call", CCDelegateTo<CC_X86_64_OxCaml_C_Call>>,
+  CCIfCC<"CallingConv::OxCaml_C_Call_StackArgs",
+    CCDelegateTo<CC_X86_64_OxCaml_C_Call_StackArgs>>,
+
+  // OxCaml_Alloc follows the same CC as OxCaml, but doesn't pass arguments
+  // (or return) in any non-runtime register. So, as long as it is used that way,
+  // it is safe for it to follow OxCaml's arg and ret CCs and have a different
+  // CSR mask.
+  CCIfCC<"CallingConv::OxCaml_Alloc", CCDelegateTo<CC_X86_64_OxCaml>>,
 
   // Mingw64 and native Win64 use Win64 CC
   CCIfSubtarget<"isTargetWin64()", CCDelegateTo<CC_X86_Win64_C>>,
@@ -1179,9 +1216,6 @@ def CSR_NoRegs : CalleeSavedRegs<(add)>;
 def CSR_32 : CalleeSavedRegs<(add ESI, EDI, EBX, EBP)>;
 def CSR_64 : CalleeSavedRegs<(add RBX, R12, R13, R14, R15, RBP)>;
 
-// R14 and R15 are used as return registers, so they aren't callee saved.
-def CSR_64_OCaml_C_Call : CalleeSavedRegs<(sub CSR_64, R14, R15)>;
-
 def CSR_64_SwiftError : CalleeSavedRegs<(sub CSR_64, R12)>;
 def CSR_64_SwiftTail : CalleeSavedRegs<(sub CSR_64, R13, R14)>;
 
@@ -1281,3 +1315,16 @@ def CSR_SysV64_RegCall_NoSSE : CalleeSavedRegs<(add RBX, RBP,
                                                (sequence "R%u", 12, 15))>;
 def CSR_SysV64_RegCall       : CalleeSavedRegs<(add CSR_SysV64_RegCall_NoSSE,               
                                                (sequence "XMM%u", 8, 15))>;
+
+def CSR_64_OxCaml_WithFP : CalleeSavedRegs<(add RBP)>;
+def CSR_64_OxCaml_WithoutFP : CalleeSavedRegs<(add)>;
+
+// R14 and R15 (and also R12 in the latter) are used as return registers,
+// so they aren't callee saved.
+def CSR_64_OxCaml_C_Call : CalleeSavedRegs<(sub CSR_64, R14, R15)>;
+def CSR_64_OxCaml_C_Call_StackArgs : CalleeSavedRegs<(sub CSR_64, R14, R15, R12)>;
+
+// See [Proc.destroyed_at_alloc_or_poll] for more details:
+// https://github.com/oxcaml/oxcaml/blob/main/backend/amd64/proc.ml#L457
+// Note that R14 and R15 are also not saved since they are passed as arguments.
+def CSR_64_OxCaml_Alloc : CalleeSavedRegs<(sub CSR_64_AllRegs, R10, R11, R14, R15)>;

llvm/lib/Target/X86/X86FastISel.cpp

@@ -3140,7 +3140,8 @@ static unsigned computeBytesPoppedByCalleeForSRet(const X86Subtarget *Subtarget,
     return 0;
   if (CC == CallingConv::Fast || CC == CallingConv::GHC ||
       CC == CallingConv::HiPE || CC == CallingConv::Tail ||
-      CC == CallingConv::SwiftTail || CC == CallingConv::OCaml)
+      CC == CallingConv::SwiftTail || CC == CallingConv::OxCaml_WithFP ||
+      CC == CallingConv::OxCaml_WithoutFP)
     return 0;
 
   if (CB)

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -3670,7 +3670,8 @@ static bool canGuaranteeTCO(CallingConv::ID CC) {
   return (CC == CallingConv::Fast || CC == CallingConv::GHC ||
           CC == CallingConv::X86_RegCall || CC == CallingConv::HiPE ||
           CC == CallingConv::HHVM || CC == CallingConv::Tail ||
-          CC == CallingConv::SwiftTail || CC == CallingConv::OCaml);
+          CC == CallingConv::SwiftTail || CC == CallingConv::OxCaml_WithFP ||
+          CC == CallingConv::OxCaml_WithoutFP);
 }
 
 /// Return true if we might ever do TCO for calls with this calling convention.

llvm/lib/Target/X86/X86RegisterInfo.cpp

@@ -300,10 +300,17 @@ X86RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
   switch (CC) {
   case CallingConv::GHC:
   case CallingConv::HiPE:
-  case CallingConv::OCaml:
     return CSR_NoRegs_SaveList;
-  case CallingConv::OCaml_C_Call:
-    return CSR_64_OCaml_C_Call_SaveList;
+  case CallingConv::OxCaml_WithFP:
+    return CSR_64_OxCaml_WithFP_SaveList;
+  case CallingConv::OxCaml_WithoutFP:
+    return CSR_64_OxCaml_WithoutFP_SaveList;
+  case CallingConv::OxCaml_C_Call:
+    return CSR_64_OxCaml_C_Call_SaveList;
+  case CallingConv::OxCaml_C_Call_StackArgs:
+    return CSR_64_OxCaml_C_Call_StackArgs_SaveList;
+  case CallingConv::OxCaml_Alloc:
+    return CSR_64_OxCaml_Alloc_SaveList;
   case CallingConv::AnyReg:
     if (HasAVX)
       return CSR_64_AllRegs_AVX_SaveList;
@@ -426,10 +433,17 @@ X86RegisterInfo::getCallPreservedMask(const MachineFunction &MF,
   switch (CC) {
   case CallingConv::GHC:
   case CallingConv::HiPE:
-  case CallingConv::OCaml:
     return CSR_NoRegs_RegMask;
-  case CallingConv::OCaml_C_Call:
-    return CSR_64_OCaml_C_Call_RegMask;
+  case CallingConv::OxCaml_WithFP:
+    return CSR_64_OxCaml_WithFP_RegMask;
+  case CallingConv::OxCaml_WithoutFP:
+    return CSR_64_OxCaml_WithoutFP_RegMask;
+  case CallingConv::OxCaml_C_Call:
+    return CSR_64_OxCaml_C_Call_RegMask;
+  case CallingConv::OxCaml_C_Call_StackArgs:
+    return CSR_64_OxCaml_C_Call_StackArgs_RegMask;
+  case CallingConv::OxCaml_Alloc:
+    return CSR_64_OxCaml_Alloc_RegMask;
   case CallingConv::AnyReg:
     if (HasAVX)
       return CSR_64_AllRegs_AVX_RegMask;