MLIR: improve num returns error (#2196)

EnzymeAD · Dec 13, 2024 · 068ad9c · 068ad9c
1 parent 9fe27b8
commit 068ad9c
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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();