PMPI recognition (rust-lang#548)

Co-authored-by: Ubuntu <[email protected]>

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -5981,7 +5981,8 @@ class AdjointGenerator
       return;
     }
 
-    if (funcName == "MPI_Send" || funcName == "MPI_Ssend") {
+    if (funcName == "MPI_Send" || funcName == "MPI_Ssend" ||
+        funcName == "PMPI_Send" || funcName == "PMPI_Ssend") {
       if (Mode == DerivativeMode::ReverseModeGradient ||
           Mode == DerivativeMode::ReverseModeCombined ||
           Mode == DerivativeMode::ForwardMode) {
@@ -6008,6 +6009,12 @@ class AdjointGenerator
           statusArg--;
           if (auto PT = dyn_cast<PointerType>(statusArg->getType()))
             statusType = PT->getPointerElementType();
+        } else if (Function *recvfn =
+                       called->getParent()->getFunction("PMPI_Recv")) {
+          auto statusArg = recvfn->arg_end();
+          statusArg--;
+          if (auto PT = dyn_cast<PointerType>(statusArg->getType()))
+            statusType = PT->getPointerElementType();
         }
         if (statusType == nullptr) {
           statusType = ArrayType::get(Type::getInt8Ty(call.getContext()), 24);

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -3627,7 +3627,9 @@ void TypeAnalyzer::visitCallInst(CallInst &call) {
     }
     if (funcName == "MPI_Send" || funcName == "MPI_Ssend" ||
         funcName == "MPI_Bsend" || funcName == "MPI_Recv" ||
-        funcName == "MPI_Brecv") {
+        funcName == "MPI_Brecv" || funcName == "PMPI_Send" ||
+        funcName == "PMPI_Ssend" || funcName == "PMPI_Bsend" ||
+        funcName == "PMPI_Recv" || funcName == "PMPI_Brecv") {
       TypeTree buf = TypeTree(BaseType::Pointer);
 
       if (Constant *C = dyn_cast<Constant>(call.getOperand(2))) {