From 83c1654768d41d4642d2828108199511345459d0 Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Tue, 30 May 2023 16:40:58 +0200
Subject: [PATCH] fix(compiler): Move operations in scf.for reinstantiation
 pattern before replacement

The reinstantianting rewrite pattern for `scf.for` operations,
`TypeConvertingReinstantiationPattern<scf::ForOp, false>`, calls
`mlir::ConversionPatternRewriter::replaceOpWithNewOp()` before moving
the operations of the original loop to the newly created loop. Since
`replaceOpWithNewOp()` indirectly marks all operations of the old loop
as ignored for dialect conversion, the dialect converter never
descends recursively into the newly created loop.

This causes operations that are illegal to be preserved, which results
in illegal IR after dialect conversion.

This commit splits the replacement into three steps:

  1. Creation of the new loop via
     mlir::ConversionPatternRewriter::create()`
  2. Moving operations from the old loop to the newly created one
  3. Replacement of the original loop with the results of the new one
     via `mlir::ConversionPatternRewriter::replaceOp()`

This causes the operations of the loops not to be ignored and fixes
dialect conversion.
---
 .../compiler/lib/Conversion/Utils/Dialects/SCF.cpp          | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/compilers/concrete-compiler/compiler/lib/Conversion/Utils/Dialects/SCF.cpp b/compilers/concrete-compiler/compiler/lib/Conversion/Utils/Dialects/SCF.cpp
index 8f077ee8e..a8beae3cc 100644
--- a/compilers/concrete-compiler/compiler/lib/Conversion/Utils/Dialects/SCF.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Conversion/Utils/Dialects/SCF.cpp
@@ -14,8 +14,8 @@ TypeConvertingReinstantiationPattern<scf::ForOp, false>::matchAndRewrite(
     scf::ForOp oldOp, mlir::OpConversionPattern<scf::ForOp>::OpAdaptor adaptor,
     mlir::ConversionPatternRewriter &rewriter) const {
   // Create new for loop with empty body, but converted iter args
-  scf::ForOp newForOp = rewriter.replaceOpWithNewOp<scf::ForOp>(
-      oldOp, adaptor.getLowerBound(), adaptor.getUpperBound(),
+  scf::ForOp newForOp = rewriter.create<scf::ForOp>(
+      oldOp.getLoc(), adaptor.getLowerBound(), adaptor.getUpperBound(),
       adaptor.getStep(), adaptor.getInitArgs(),
       [&](OpBuilder &builder, Location loc, Value iv, ValueRange args) {});
 
@@ -35,6 +35,8 @@ TypeConvertingReinstantiationPattern<scf::ForOp, false>::matchAndRewrite(
                                newForOp.getRegion());
   }
 
+  rewriter.replaceOp(oldOp, newForOp.getResults());
+
   return mlir::success();
 }
 } // namespace concretelang