From f7e409deec715aa96c4f1ebb412fc70d37052519 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Fri, 13 Dec 2024 12:32:55 -0500
Subject: [PATCH] MLIR: improve num returns error

---
 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();