fix unroll with block

EnzymeAD · Mar 10, 2024 · 514c9e9 · 514c9e9
1 parent e98d0a1
commit 514c9e9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOUnroll.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOUnroll.cpp
@@ -97,12 +97,11 @@ struct WhileUnroll : public OpRewritePattern<mlir::stablehlo::WhileOp> {
 
     auto iters = (*limit.begin()) - (*start.begin());
 
-    IRMapping operandMap;
-
     SmallVector<Value> results(op.getOperands().begin(),
                                op.getOperands().end());
 
     for (size_t iter = 0; iter < iters.getSExtValue(); iter++) {
+      IRMapping operandMap;
       operandMap.map(loopBodyBlock->getArguments(), results);
 
       Block::iterator srcBlockEnd = std::prev(loopBodyBlock->end(), 2);