Support int calling for cublas (#2685)

* Support int calling for cublas

* fix cublasv2

Author: wsmoses

enzyme/Enzyme/ActivityAnalysis.cpp

@@ -1362,16 +1362,24 @@ bool ActivityAnalyzer::isConstantValue(TypeResults const &TR, Value *Val) {
                                << " dependant on " << *LoadReval << "\n";
                 ReEvaluateValueIfInactiveInst[LoadReval].insert(Val);
               }
-              if (StoreReval && EnzymeEnableRecursiveHypotheses)
+              if (StoreReval && EnzymeEnableRecursiveHypotheses) {
+                if (EnzymePrintActivity)
+                  llvm::errs() << " global activity of " << *Val
+                               << " dependant on " << *StoreReval << "\n";
                 ReEvaluateValueIfInactiveInst[StoreReval].insert(Val);
+              }
               if (ValLoadReval && EnzymeEnableRecursiveHypotheses) {
                 if (EnzymePrintActivity)
                   llvm::errs() << " global activity of " << *Val
                                << " dependant on " << *ValLoadReval << "\n";
                 ReEvaluateValueIfInactiveValue[ValLoadReval].insert(Val);
               }
-              if (ValStoreReval && EnzymeEnableRecursiveHypotheses)
+              if (ValStoreReval && EnzymeEnableRecursiveHypotheses) {
+                if (EnzymePrintActivity)
+                  llvm::errs() << " global activity of " << *Val
+                               << " dependant on " << *ValStoreReval << "\n";
                 ReEvaluateValueIfInactiveValue[ValStoreReval].insert(Val);
+              }
             }
           }
         }

enzyme/Enzyme/DiffeGradientUtils.cpp

@@ -818,14 +818,18 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig,
 
   bool needsCast = false;
 #if LLVM_VERSION_MAJOR < 17
-  if (origptr->getContext().supportsTypedPointers()) {
+  if (isa<PointerType>(origptr->getType()) &&
+      origptr->getContext().supportsTypedPointers()) {
     needsCast = origptr->getType()->getPointerElementType() != addingType;
   }
 #endif
 
   assert(ptr);
-  if (start != 0 || needsCast) {
+  if (start != 0 || needsCast || !isa<PointerType>(origptr->getType())) {
     auto rule = [&](Value *ptr) {
+      if (!isa<PointerType>(origptr->getType())) {
+        ptr = BuilderM.CreateIntToPtr(ptr, getUnqual(addingType));
+      }
       if (start != 0) {
         auto i8 = Type::getInt8Ty(ptr->getContext());
         ptr = BuilderM.CreatePointerCast(
@@ -845,7 +849,9 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig,
     ptr = applyChainRule(
         PointerType::get(
             addingType,
-            cast<PointerType>(origptr->getType())->getAddressSpace()),
+            isa<PointerType>(origptr->getType())
+                ? cast<PointerType>(origptr->getType())->getAddressSpace()
+                : 0),
         BuilderM, rule, ptr);
   }
 

enzyme/tools/enzyme-tblgen/blas-tblgen.cpp

@@ -1473,7 +1473,13 @@ void rev_call_args(bool forward, Twine argName, const TGPattern &pattern,
     // Distinguish later trough byRef if it is cblas (thus has layout)
     os << "        if (cblas) " << argName << ".push_back(arg_layout);\n";
   }
-  os << "        if (cublas) " << argName << ".push_back(arg_handle);\n";
+  os << "        if (cublas) " << argName << ".push_back(";
+  if (!forward)
+    os << "lookup(";
+  os << "arg_handle";
+  if (!forward)
+    os << ", Builder2)";
+  os << ");\n";
 
   for (size_t pos = fncHasLayout ? 1 : 0; pos < numArgs; pos++) {
     os << "        for (auto item : ";

enzyme/tools/enzyme-tblgen/blasDeclUpdater.h

@@ -117,6 +117,13 @@ inline void emit_attributeBLAS(const TGPattern &pattern, raw_ostream &os) {
           "llvm::Attribute::get(F->getContext(), llvm::Attribute::ZExt));\n";
   }
   os << "  }\n";
+
+  if (has_active_return(name)) {
+    os << "const bool cublasv2 = blas.prefix == "
+          "\"cublas\" && llvm::StringRef(blas.suffix).contains(\"v2\");\n";
+    os << "  if (cublasv2) argTys.push_back(getUnqual(fpType));\n";
+  }
+
   os << "  auto nextFT = llvm::FunctionType::get(prevFT->getReturnType(), "
         "argTys, false);\n";
   os << "  if (nextFT != prevFT && F->empty()) {\n";