Support batching scalar types (#2175)

* support batching scalar types * formatting
EnzymeAD · Nov 27, 2024 · 32a2118 · 32a2118
1 parent 0a8fdc7
commit 32a2118
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
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:  }