Support batching scalar types (#2175)

* support batching scalar types * formatting
EnzymeAD · Nov 28, 2024 · e8ba058 · e8ba058
1 parent 0a8fdc7
commit e8ba058
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 2 deletions.
diff --git a/enzyme/Enzyme/CallDerivatives.cpp b/enzyme/Enzyme/CallDerivatives.cpp
@@ -3312,6 +3312,9 @@ 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,7 @@ 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,9 @@ 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 +3527,9 @@ 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))

diff --git a/enzyme/Enzyme/MLIR/Passes/EnzymeBatchPass.cpp b/enzyme/Enzyme/MLIR/Passes/EnzymeBatchPass.cpp
@@ -30,7 +30,11 @@ namespace {
 
 static mlir::TensorType applyBatchSizes(mlir::Type Ty,
                                         llvm::ArrayRef<int64_t> batchSizes) {
-  auto T = cast<TensorType>(Ty);
+  auto T = dyn_cast<TensorType>(Ty);
+  if (!T) {
+    return RankedTensorType::get(batchSizes, Ty);
+  }
+
   SmallVector<int64_t> shape(batchSizes.begin(), batchSizes.end());
   shape.append(T.getShape().begin(), T.getShape().end());
   auto T2 = T.clone(shape);

diff --git a/enzyme/test/MLIR/Batch/batched_scalar.mlir b/enzyme/test/MLIR/Batch/batched_scalar.mlir
@@ -0,0 +1,21 @@
+// RUN: %eopt --enzyme-batch %s | FileCheck %s
+
+module {
+  func.func @square(%x : f64) -> f64 {
+    %y = math.sin %x : f64
+    return %y : f64
+  }
+  func.func @dsq(%x : tensor<10x2xf64>) -> tensor<10x2xf64> {
+    %r = enzyme.batch @square(%x) { batch_shape=array<i64: 10, 2> } : (tensor<10x2xf64>) -> (tensor<10x2xf64>)
+    return %r : tensor<10x2xf64>
+  }
+}
+
+// CHECK:  func.func @dsq(%arg0: tensor<10x2xf64>) -> tensor<10x2xf64> {
+// CHECK-NEXT:    %0 = call @batched_square(%arg0) : (tensor<10x2xf64>) -> tensor<10x2xf64>
+// CHECK-NEXT:    return %0 : tensor<10x2xf64>
+// CHECK-NEXT:  }
+// CHECK:  func.func private @batched_square(%arg0: tensor<10x2xf64>) -> tensor<10x2xf64> {
+// CHECK-NEXT:    %0 = math.sin %arg0 : tensor<10x2xf64>
+// CHECK-NEXT:    return %0 : tensor<10x2xf64>
+// CHECK-NEXT:  }