From 068ad9c6f8bc7d8c7ad3806fd148492e323bc4b1 Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Fri, 13 Dec 2024 13:08:37 -0600
Subject: [PATCH] MLIR: improve num returns error (#2196)

---
 enzyme/Enzyme/MLIR/Passes/EnzymeMLIRPass.cpp | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/enzyme/Enzyme/MLIR/Passes/EnzymeMLIRPass.cpp b/enzyme/Enzyme/MLIR/Passes/EnzymeMLIRPass.cpp
index c3fe53a7c4e..d83532db35a 100644
--- a/enzyme/Enzyme/MLIR/Passes/EnzymeMLIRPass.cpp
+++ b/enzyme/Enzyme/MLIR/Passes/EnzymeMLIRPass.cpp
@@ -157,6 +157,11 @@ struct DifferentiatePass : public DifferentiatePassBase<DifferentiatePass> {
     OpBuilder builder(CI);
     auto dCI = builder.create<func::CallOp>(CI.getLoc(), newFunc.getName(),
                                             newFunc.getResultTypes(), args);
+    if (dCI.getNumResults() != CI.getNumResults()) {
+      CI.emitError() << "Incorrect number of results for enzyme operation: "
+                     << *CI << " expected " << *dCI;
+      return failure();
+    }
     CI.replaceAllUsesWith(dCI);
     CI->erase();
     return success();