recursion test

EnzymeAD · Jan 17, 2025 · 0154874 · 0154874
1 parent d3aab16
commit 0154874
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/enzyme/test/MLIR/Batch/recursion.mlir b/enzyme/test/MLIR/Batch/recursion.mlir
@@ -0,0 +1,21 @@
+// RUN: %eopt -enzyme-batch %s | FileCheck %s
+
+module {
+  func.func private @f(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>) -> tensor<16xf32> {
+    %0 = func.call @f(%arg0, %arg1) : (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>
+    return %0 : tensor<16xf32>
+  }
+  func.func @main(%arg0: tensor<4x16xf32>, %arg1: tensor<4x16xf32>) {
+    %0 = enzyme.batch @f(%arg0, %arg1) {batch_shape = array<i64: 4>} : (tensor<4x16xf32>, tensor<4x16xf32>) -> tensor<4x16xf32>
+    return
+  }
+}
+
+// CHECK:   func.func @main(%[[arg0:.+]]: tensor<4x16xf32>, %[[arg1:.+]]: tensor<4x16xf32>) {
+// CHECK-NEXT:     %[[v0:.+]] = call @batched_f(%[[arg0]], %[[arg1]]) : (tensor<4x16xf32>, tensor<4x16xf32>) -> tensor<4x16xf32>
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+// CHECK:   func.func private @batched_f(%[[arg0:.+]]: tensor<4x16xf32>, %[[arg1:.+]]: tensor<4x16xf32>) -> tensor<4x16xf32> {
+// CHECK-NEXT:     %[[v0:.+]] = call @batched_f(%[[arg0]], %[[arg1]]) : (tensor<4x16xf32>, tensor<4x16xf32>) -> tensor<4x16xf32>
+// CHECK-NEXT:     return %[[v0]] : tensor<4x16xf32>
+// CHECK-NEXT:   }