Support batching scalar types (#2175) (#2180)

* support batching scalar types * formatting Co-authored-by: jumerckx <[email protected]>
EnzymeAD · Nov 28, 2024 · 095ee7e · 095ee7e
1 parent b59ab66
commit 095ee7e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/enzyme/Enzyme/CallDerivatives.cpp b/enzyme/Enzyme/CallDerivatives.cpp
@@ -3312,6 +3312,10 @@ bool AdjointGenerator::handleKnownCallDerivatives(
       }
 #endif
       Value *replacement = B.CreateAlloca(elTy, Size);
+      for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type",
+                      "enzymejl_allocart"})
+        if (auto M = call.getMetadata(MD))
+          cast<AllocaInst>(replacement)->setMetadata(MD, M);
       if (I)
         replacement->takeName(I);
       else

diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp
@@ -508,7 +508,8 @@ UpgradeAllocasToMallocs(Function *NewF, DerivativeMode mode,
                     {ConstantAsMetadata::get(ConstantInt::get(
                         IntegerType::get(AI->getContext(), 64), align))}));
 
-    for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type"})
+    for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type",
+                    "enzymejl_allocart"})
       if (auto M = AI->getMetadata(MD))
         CI->setMetadata(MD, M);
 

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -3280,6 +3280,10 @@ BasicBlock *GradientUtils::prepRematerializedLoopEntry(LoopContext &lc) {
             auto replacement = NB.CreateAlloca(
                 Type::getInt8Ty(I.getContext()),
                 lookupM(getNewFromOriginal(I.getOperand(0)), NB, available));
+            for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type",
+                            "enzymejl_allocart"})
+              if (auto M = I.getMetadata(MD))
+                replacement->setMetadata(MD, M);
             auto Alignment =
                 cast<ConstantInt>(
                     cast<ConstantAsMetadata>(MD->getOperand(0))->getValue())
@@ -3524,6 +3528,10 @@ BasicBlock *GradientUtils::prepRematerializedLoopEntry(LoopContext &lc) {
                 auto rule = [&](Value *anti) {
                   AllocaInst *replacement = NB.CreateAlloca(
                       Type::getInt8Ty(orig->getContext()), args[0]);
+                  for (auto MD : {"enzyme_active", "enzyme_inactive",
+                                  "enzyme_type", "enzymejl_allocart"})
+                    if (auto M = I.getMetadata(MD))
+                      replacement->setMetadata(MD, M);
                   replacement->takeName(anti);
                   auto Alignment = cast<ConstantInt>(cast<ConstantAsMetadata>(
                                                          MD->getOperand(0))