Add memcpy for 32 (#1133)

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -5887,7 +5887,8 @@ class AdjointGenerator
 
         auto dmemcpy = getOrInsertDifferentialFloatMemcpy(
             *Builder2.GetInsertBlock()->getParent()->getParent(), secretty,
-            /*dstalign*/ 1, /*srcalign*/ 1, dstaddr, srcaddr);
+            /*dstalign*/ 1, /*srcalign*/ 1, dstaddr, srcaddr,
+            cast<IntegerType>(length->getType())->getBitWidth());
 
         Builder2.CreateCall(dmemcpy, args, ReverseDefs);
       }

enzyme/Enzyme/GradientUtils.cpp

@@ -8207,7 +8207,8 @@ void SubTransferHelper(GradientUtils *gutils, DerivativeMode mode,
                               ? getOrInsertDifferentialFloatMemcpy
                               : getOrInsertDifferentialFloatMemmove)(
               *MTI->getParent()->getParent()->getParent(), secretty, dstalign,
-              srcalign, dstaddr, srcaddr);
+              srcalign, dstaddr, srcaddr,
+              cast<IntegerType>(length->getType())->getBitWidth());
           Builder2.CreateCall(dmemcpy, args);
         }
       }

enzyme/Enzyme/Utils.cpp

@@ -539,20 +539,24 @@ void ErrorIfRuntimeInactive(llvm::IRBuilder<> &B, llvm::Value *primal,
 Function *getOrInsertDifferentialFloatMemcpy(Module &M, Type *elementType,
                                              unsigned dstalign,
                                              unsigned srcalign,
-                                             unsigned dstaddr,
-                                             unsigned srcaddr) {
+                                             unsigned dstaddr, unsigned srcaddr,
+                                             unsigned bitwidth) {
   assert(elementType->isFloatingPointTy());
-  std::string name = "__enzyme_memcpyadd_" + tofltstr(elementType) + "da" +
-                     std::to_string(dstalign) + "sa" + std::to_string(srcalign);
+  std::string name = "__enzyme_memcpy";
+  if (bitwidth != 64)
+    name += std::to_string(bitwidth);
+  name += "add_" + tofltstr(elementType) + "da" + std::to_string(dstalign) +
+          "sa" + std::to_string(srcalign);
   if (dstaddr)
     name += "dadd" + std::to_string(dstaddr);
   if (srcaddr)
     name += "sadd" + std::to_string(srcaddr);
-  FunctionType *FT = FunctionType::get(Type::getVoidTy(M.getContext()),
-                                       {PointerType::get(elementType, dstaddr),
-                                        PointerType::get(elementType, srcaddr),
-                                        Type::getInt64Ty(M.getContext())},
-                                       false);
+  FunctionType *FT =
+      FunctionType::get(Type::getVoidTy(M.getContext()),
+                        {PointerType::get(elementType, dstaddr),
+                         PointerType::get(elementType, srcaddr),
+                         IntegerType::get(M.getContext(), bitwidth)},
+                        false);
 
 #if LLVM_VERSION_MAJOR >= 9
   Function *F = cast<Function>(M.getOrInsertFunction(name, FT).getCallee());
@@ -742,15 +746,14 @@ Function *getOrInsertMemcpyStrided(Module &M, PointerType *T, Type *IT,
 }
 
 // TODO implement differential memmove
-Function *getOrInsertDifferentialFloatMemmove(Module &M, Type *T,
-                                              unsigned dstalign,
-                                              unsigned srcalign,
-                                              unsigned dstaddr,
-                                              unsigned srcaddr) {
+Function *
+getOrInsertDifferentialFloatMemmove(Module &M, Type *T, unsigned dstalign,
+                                    unsigned srcalign, unsigned dstaddr,
+                                    unsigned srcaddr, unsigned bitwidth) {
   llvm::errs() << "warning: didn't implement memmove, using memcpy as fallback "
                   "which can result in errors\n";
   return getOrInsertDifferentialFloatMemcpy(M, T, dstalign, srcalign, dstaddr,
-                                            srcaddr);
+                                            srcaddr, bitwidth);
 }
 
 Function *getOrInsertCheckedFree(Module &M, CallInst *call, Type *Ty,

enzyme/Enzyme/Utils.h

@@ -603,10 +603,9 @@ static inline bool isCertainPrint(const llvm::StringRef name) {
 
 /// Create function for type that performs the derivative memcpy on floating
 /// point memory
-llvm::Function *
-getOrInsertDifferentialFloatMemcpy(llvm::Module &M, llvm::Type *T,
-                                   unsigned dstalign, unsigned srcalign,
-                                   unsigned dstaddr, unsigned srcaddr);
+llvm::Function *getOrInsertDifferentialFloatMemcpy(
+    llvm::Module &M, llvm::Type *T, unsigned dstalign, unsigned srcalign,
+    unsigned dstaddr, unsigned srcaddr, unsigned bitwidth);
 
 /// Create function for type that performs memcpy with a stride
 llvm::Function *getOrInsertMemcpyStrided(llvm::Module &M, llvm::PointerType *T,
@@ -615,10 +614,9 @@ llvm::Function *getOrInsertMemcpyStrided(llvm::Module &M, llvm::PointerType *T,
 
 /// Create function for type that performs the derivative memmove on floating
 /// point memory
-llvm::Function *
-getOrInsertDifferentialFloatMemmove(llvm::Module &M, llvm::Type *T,
-                                    unsigned dstalign, unsigned srcalign,
-                                    unsigned dstaddr, unsigned srcaddr);
+llvm::Function *getOrInsertDifferentialFloatMemmove(
+    llvm::Module &M, llvm::Type *T, unsigned dstalign, unsigned srcalign,
+    unsigned dstaddr, unsigned srcaddr, unsigned bitwidth);
 
 llvm::Function *getOrInsertCheckedFree(llvm::Module &M, llvm::CallInst *call,
                                        llvm::Type *Type, unsigned width);