fix(compiler): Preserve original values of extra conversions in TFHE …

…circuit parametrization

Since conversions at the edge of partitions determined by the
optimizer are specified at the producer of a value,
`MaterializePartitionBoundaryPattern` may eagerly replace references
to the original value before conversion, resulting in conflicts.

This commit adds a check to the conflict handling routine that checks
if the producer of an operation with a conflicting operand is a
keyswitch operation that converts the original value of an extra
conversion with the required, non-conflicting type. If this is the
case, the conflict is handled by simply forwarding the original value
from the operand of the producing keyswitch operation.

Resolves Issue #538
(zama-ai/concrete-internal#538).

compilers/concrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp

@@ -593,9 +593,26 @@ class TFHECircuitSolutionRewriter : public TypeInferenceRewriter {
                                                    resolvedType, producerValue);
     }
 
+    mlir::Operation *producer = producerValue.getDefiningOp();
+
+    // Since conversions at the edge of partitions are specified for
+    // the producer of the value,
+    // `MaterializePartitionBoundaryPattern` may have replaced
+    // references to the original value before conversion too eagerly,
+    // resulting in a conflict. Check if the current producer is a
+    // keyswitch operation that converts the original value needed
+    // here. If so, this is a spurious conflict that can be resolved
+    // by simply reusing the original value before conversion.
+    if (TFHE::KeySwitchGLWEOp ancestor =
+            llvm::dyn_cast_or_null<TFHE::KeySwitchGLWEOp>(producer)) {
+      if (ancestor.getCiphertext().getType() == resolvedType) {
+        return ancestor.getCiphertext();
+      }
+    }
+
     // Place keyswitch operation near the producer of the value to
     // avoid nesting it too depply into loops
-    if (mlir::Operation *producer = producerValue.getDefiningOp())
+    if (producer)
       rewriter.setInsertionPointAfter(producer);
 
     assert(cttFrom->getKey().getParameterized().has_value());