Add c api inverted bundle (rust-lang#806)

Author: wsmoses

enzyme/Enzyme/CApi.cpp

@@ -632,6 +632,35 @@ const char *EnzymeGradientUtilsInvertedPointersToString(GradientUtils *gutils,
   return cstr;
 }
 
+LLVMValueRef EnzymeGradientUtilsCallWithInvertedBundles(
+    GradientUtils *gutils, LLVMValueRef func, LLVMValueRef *args_vr,
+    uint64_t args_size, LLVMValueRef orig_vr, CValueType *valTys,
+    uint64_t valTys_size, LLVMBuilderRef B, uint8_t lookup) {
+  auto orig = cast<CallInst>(unwrap(orig_vr));
+
+  ArrayRef<ValueType> ar((ValueType *)valTys, valTys_size);
+
+  IRBuilder<> &BR = *unwrap(B);
+
+  auto Defs = gutils->getInvertedBundles(orig, ar, BR, lookup != 0);
+
+  SmallVector<Value *, 1> args;
+  for (size_t i = 0; i < args_size; i++) {
+    args.push_back(unwrap(args_vr[i]));
+  }
+
+  auto callval = unwrap(func);
+
+#if LLVM_VERSION_MAJOR > 7
+  auto res = BR.CreateCall(
+      cast<FunctionType>(callval->getType()->getPointerElementType()), callval,
+      args, Defs);
+#else
+  auto res = BR.CreateCall(callval, args, Defs);
+#endif
+  return wrap(res);
+}
+
 void EnzymeStringFree(const char *cstr) { delete[] cstr; }
 
 void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2,

enzyme/Enzyme/CApi.h

@@ -70,6 +70,13 @@ struct CTypeTree {
 };
 */
 
+typedef enum {
+  VT_None = 0,
+  VT_Primal = 1,
+  VT_Shadow = 2,
+  VT_Both = VT_Primal | VT_Shadow,
+} CValueType;
+
 struct EnzymeTypeTree;
 typedef struct EnzymeTypeTree *CTypeTreeRef;
 CTypeTreeRef EnzymeNewTypeTree();

enzyme/Enzyme/Enzyme.cpp

@@ -206,11 +206,10 @@ handleCustomDerivative(llvm::Module &M, llvm::GlobalVariable &g,
                   if (!byref.count(realidx))
                     args.push_back(arg.getType());
                   else
-                    args.push_back(
-                        cast<PointerType>(arg.getType())->getElementType());
+                    args.push_back(arg.getType()->getPointerElementType());
                   realidx++;
                 } else {
-                  sretTy = cast<PointerType>(arg.getType())->getElementType();
+                  sretTy = arg.getType()->getPointerElementType();
                 }
                 i++;
               }
@@ -680,7 +679,7 @@ static bool replaceOriginalCall(CallInst *CI, Function *fn, Value *diffret,
     } else if (CI->hasStructRetAttr()) {
       Value *sret = CI->getArgOperand(0);
       PointerType *stype = cast<PointerType>(sret->getType());
-      StructType *st = dyn_cast<StructType>(stype->getElementType());
+      StructType *st = dyn_cast<StructType>(stype->getPointerElementType());
 
       // Assign results to struct allocated at the call site.
       if (st && st->isLayoutIdentical(diffretsty)) {
@@ -1285,7 +1284,7 @@ class Enzyme : public ModulePass {
       }
 
       if (byRefSize) {
-        Type *subTy = cast<PointerType>(res->getType())->getElementType();
+        Type *subTy = res->getType()->getPointerElementType();
         auto &DL = fn->getParent()->getDataLayout();
         auto BitSize = DL.getTypeSizeInBits(subTy);
         if (BitSize / 8 != byRefSize) {

enzyme/Enzyme/EnzymeLogic.cpp

@@ -1832,7 +1832,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
           if (!foundcalled->hasParamAttribute(i, Attribute::StructRet))
             args.push_back(arg.getType());
           else {
-            sretTy = cast<PointerType>(arg.getType())->getElementType();
+            sretTy = arg.getType()->getPointerElementType();
             //  sretTy = foundcalled->getParamStructRetType(i);
           }
           i++;

enzyme/Enzyme/Utils.cpp

@@ -130,7 +130,7 @@ Value *CreateAllocation(IRBuilder<> &Builder, llvm::Type *T, Value *Count,
   if (ZeroMem) {
     auto PT = cast<PointerType>(malloccall->getType());
     Value *tozero = malloccall;
-    if (!PT->getElementType()->isIntegerTy(8))
+    if (!PT->getPointerElementType()->isIntegerTy(8))
       tozero = Builder.CreatePointerCast(
           tozero, PointerType::get(Type::getInt8Ty(PT->getContext()),
                                    PT->getAddressSpace()));