Repository: zama-ai/concrete (language: python)

Issue:
Type Inference: Add support for multiple functions



Discussion / hints:

Fix patch:
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h b/compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h
index 5768c863f5..13051db973 100644
--- a/compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h
+++ b/compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h
@@ -60,6 +60,7 @@
 #include <llvm/ADT/TypeSwitch.h>
 #include <llvm/Support/Debug.h>
 #include <mlir/Analysis/DataFlow/SparseAnalysis.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
 
 namespace mlir {
 namespace concretelang {
@@ -192,63 +193,31 @@ class LocalInferenceState {
   bool changed;
 };
 
-// Lattice value for type inference analysis, encapsulating a single
-// inferred type
-class TypeInferenceLatticeValue {
+// A type resolver is responsible for the actual inference of the
+// types related to an operation based on the previous inference state
+class TypeResolver {
 public:
-  TypeInferenceLatticeValue(const InferredType &t = {}) : inferredType(t) {}
-
-  static TypeInferenceLatticeValue getPessimisticValueState(mlir::Value value) {
-    return TypeInferenceLatticeValue();
-  }
-
-  bool operator==(const TypeInferenceLatticeValue &rhs) {
-    return rhs.getInferredType() == inferredType;
-  }
-
-  static TypeInferenceLatticeValue join(const TypeInferenceLatticeValue &lhs,
-                                        const TypeInferenceLatticeValue &rhs) {
-    return TypeInferenceLatticeValue(lhs.getInferredType().hasType()
-                                         ? lhs.getInferredType()
-                                         : rhs.getInferredType());
-  }
-
-  static TypeInferenceLatticeValue meet(const TypeInferenceLatticeValue &lhs,
-                                        const TypeInferenceLatticeValue &rhs) {
-    return join(lhs, rhs);
-  }
-
-  const InferredType &getInferredType() const { return inferredType; }
-  void setType(const InferredType &t) { this->inferredType = t; }
-
-  void print(raw_ostream &os) const { os << inferredType; }
-
-protected:
-  InferredType inferredType;
-};
+  virtual ~TypeResolver() {}
 
-// Lattice for forward and backward type inference analysis
-class TypeInferenceLattice
-    : public mlir::dataflow::Lattice<TypeInferenceLatticeValue> {
-public:
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TypeInferenceLattice)
+  // Resolve the types for all values related to `op` based on the
+  // previous inference state `prevState` and the current state
+  // `currState`.
+  virtual LocalInferenceState resolve(mlir::Operation *op,
+                                      const LocalInferenceState &prevState) = 0;
 
-  using Lattice::Lattice;
-};
+  // Returns true for any type that is considered unresolved by the
+  // resolver
+  virtual bool isUnresolvedType(mlir::Type t) const = 0;
 
-// Collection of common operations for forward and backward type
-// inference and subsequent transformations (e.g., rewriting the IR
-// with the inferred types).
-class TypeInferenceUtils {
-public:
   // Invokes the callback function `fn` to all values related to the
   // operation `op`. Iteration over the related value stops if `fn`
   // returns `false`.
   //
   // Returns `true` if `fn` returned `true` for all related values,
   // otherwise `false`.
-  static bool iterateRelatedValues(mlir::Operation *op,
-                                   llvm::function_ref<bool(mlir::Value)> fn) {
+  virtual bool
+  iterateRelatedValues(mlir::Operation *op,
+                       llvm::function_ref<bool(mlir::Value)> fn) const {
     auto applyToOperandsAndResults = [&](mlir::Operation *op) {
       for (mlir::Value v : op->getOperands()) {
         if (!fn(v))
@@ -287,17 +256,67 @@ class TypeInferenceUtils {
 
     return true;
   }
+};
+
+// Lattice value for type inference analysis, encapsulating a single
+// inferred type
+class TypeInferenceLatticeValue {
+public:
+  TypeInferenceLatticeValue(const InferredType &t = {}) : inferredType(t) {}
+
+  static TypeInferenceLatticeValue getPessimisticValueState(mlir::Value value) {
+    return TypeInferenceLatticeValue();
+  }
+
+  bool operator==(const TypeInferenceLatticeValue &rhs) {
+    return rhs.getInferredType() == inferredType;
+  }
+
+  static TypeInferenceLatticeValue join(const TypeInferenceLatticeValue &lhs,
+                                        const TypeInferenceLatticeValue &rhs) {
+    return TypeInferenceLatticeValue(lhs.getInferredType().hasType()
+                                         ? lhs.getInferredType()
+                                         : rhs.getInferredType());
+  }
+
+  static TypeInferenceLatticeValue meet(const TypeInferenceLatticeValue &lhs,
+                                        const TypeInferenceLatticeValue &rhs) {
+    return join(lhs, rhs);
+  }
+
+  const InferredType &getInferredType() const { return inferredType; }
+  void setType(const InferredType &t) { this->inferredType = t; }
 
+  void print(raw_ostream &os) const { os << inferredType; }
+
+protected:
+  InferredType inferredType;
+};
+
+// Lattice for forward and backward type inference analysis
+class TypeInferenceLattice
+    : public mlir::dataflow::Lattice<TypeInferenceLatticeValue> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TypeInferenceLattice)
+
+  using Lattice::Lattice;
+};
+
+// Collection of common operations for forward and backward type
+// inference and subsequent transformations (e.g., rewriting the IR
+// with the inferred types).
+class TypeInferenceUtils {
+public:
   // Looks up the inferred types for all values related to the
   // operation `op`, i.e., the inferred types for operands, results,
   // region arguments and operands of terminators used in directly
   // nested regions.
   static LocalInferenceState getLocalInferenceState(
-      mlir::Operation *op,
+      mlir::Operation *op, const TypeResolver &typeResolver,
       llvm::function_ref<const TypeInferenceLattice *(mlir::Value)> lookup) {
     LocalInferenceState map;
 
-    iterateRelatedValues(op, [&](mlir::Value v) {
+    typeResolver.iterateRelatedValues(op, [&](mlir::Value v) {
       map.set(v, getInferredType(v, lookup));
       return true;
     });
@@ -308,9 +327,11 @@ class TypeInferenceUtils {
   // Retrieves the inferred types for all values related to `op` using
   // `DataFlowSolver::lookupState` invoked on `solver`.
   static LocalInferenceState
-  getLocalInferenceState(const DataFlowSolver &solver, mlir::Operation *op) {
-    LocalInferenceState map =
-        TypeInferenceUtils::getLocalInferenceState(op, [&](mlir::Value v) {
+  getLocalInferenceState(const DataFlowSolver &solver,
+                         const TypeResolver &typeResolver,
+                         mlir::Operation *op) {
+    LocalInferenceState map = TypeInferenceUtils::getLocalInferenceState(
+        op, typeResolver, [&](mlir::Value v) {
           return solver.lookupState<TypeInferenceLattice>(v);
         });
 
@@ -341,27 +362,14 @@ class TypeInferenceUtils {
       llvm::function_ref<const TypeInferenceLattice *(mlir::Value)> lookup) {
     const TypeInferenceLattice *latticeOperand = lookup(v);
 
-    return InferredType(latticeOperand->getValue().getInferredType());
+    if (latticeOperand) {
+      return InferredType(latticeOperand->getValue().getInferredType());
+    } else {
+      return InferredType{};
+    }
   }
 };
 
-// A type resolver is responsible for the actual inference of the
-// types related to an operation based on the previous inference state
-class TypeResolver {
-public:
-  virtual ~TypeResolver() {}
-
-  // Resolve the types for all values related to `op` based on the
-  // previous inference state `prevState` and the current state
-  // `currState`.
-  virtual LocalInferenceState resolve(mlir::Operation *op,
-                                      const LocalInferenceState &prevState) = 0;
-
-  // Returns true for any type that is considered unresolved by the
-  // resolver
-  virtual bool isUnresolvedType(mlir::Type t) const = 0;
-};
-
 // Base class for all type constraints that can be added to a
 // `TypeConstraintSet`.
 class TypeConstraint {
@@ -590,16 +598,25 @@ template <typename YieldT>
 using DynamicSameTypeConstraint =
     SameTypeConstraintBase<DynamicTypeConstraint<YieldT>>;
 
-// Constraint that ensures that the types inferred for two values use
-// the same element type. If different types have been inferred
-// previously for the two values, precedence is given to the element
-// type inferred for the first value yielded by
-// `BaseConstraintT::yield`.
+// Base for constraints that ensure that the types inferred for two
+// values have the same nested type. What is considered a nested type
+// (e.g., an arbitrarily deeply nested scalar type or a shallow type
+// of a pointer) can be customized through `getNestedType()` and
+// `applyNestedType()`.
+//
+// The depth of the nesting for the left and right values can be set
+// independently through the template parameters `leftDepth` and
+// `rightDepth`.
+//
+// If different types have been inferred previously for the two
+// values, precedence is given to the element type inferred for the
+// first ("left") value yielded by `BaseConstraintT::yield`.
 //
 // If the element types for the two values are fixed, but
 // contradictory, an assertion is triggered.
-template <typename BaseConstraintT>
-class SameElementTypeConstraintBase : public BaseConstraintT {
+template <typename BaseConstraintT, unsigned int leftDepth = 1,
+          unsigned int rightDepth = 1>
+class SameNestedTypeConstraintBase : public BaseConstraintT {
 public:
   using BaseConstraintT::BaseConstraintT;
 
@@ -619,14 +636,17 @@ class SameElementTypeConstraintBase : public BaseConstraintT {
     // type for the second value
     if (!leftUnresolved && rightUnresolved) {
       currState.set(leftValue, leftType);
-      currState.set(rightValue, TypeInferenceUtils::applyElementType(
-                                    getElementType(leftType), rightType));
+      currState.set(rightValue,
+                    applyNestedType(getNestedType(leftType, leftDepth),
+                                    rightType, rightDepth));
     } else if (leftUnresolved && !rightUnresolved) {
-      currState.set(leftValue, TypeInferenceUtils::applyElementType(
-                                   getElementType(rightType), leftType));
+      currState.set(leftValue,
+                    applyNestedType(getNestedType(rightType, rightDepth),
+                                    leftType, leftDepth));
       currState.set(rightValue, rightType);
     } else if (!leftUnresolved && !rightUnresolved) {
-      assert(getElementType(leftType) == getElementType(rightType) &&
+      assert(getNestedType(leftType, leftDepth) ==
+                 getNestedType(rightType, rightDepth) &&
              "Constraint cannot be matched, as fixed types of the values"
              "are incompatible types");
       currState.set(leftValue, leftType);
@@ -640,29 +660,72 @@ class SameElementTypeConstraintBase : public BaseConstraintT {
 
       if (leftCurrType.hasType()) {
         currState.set(leftValue, leftCurrType);
-        currState.set(rightValue,
-                      TypeInferenceUtils::applyElementType(
-                          getElementType(leftCurrType.getType()), rightType));
+        currState.set(
+            rightValue,
+            applyNestedType(getNestedType(leftCurrType.getType(), leftDepth),
+                            rightType, rightDepth));
       } else if (rightCurrType.hasType()) {
-        currState.set(leftValue,
-                      TypeInferenceUtils::applyElementType(
-                          getElementType(rightCurrType.getType()), leftType));
+        currState.set(
+            leftValue,
+            applyNestedType(getNestedType(rightCurrType.getType(), rightDepth),
+                            leftType, leftDepth));
         currState.set(rightValue, rightCurrType);
       }
     }
   }
 
 protected:
-  static mlir::Type getElementType(mlir::Type t) {
-    if (mlir::RankedTensorType rtt =
-            llvm::dyn_cast<mlir::RankedTensorType>(t)) {
-      return rtt.getElementType();
-    } else if (mlir::MemRefType mrt = llvm::dyn_cast<mlir::MemRefType>(t)) {
-      return mrt.getElementType();
+  virtual mlir::Type getNestedType(mlir::Type t, unsigned int depth) {
+    mlir::Type res = t;
+
+    for (unsigned int i = 0; i < depth; i++)
+      res = getNestedType(res);
+
+    return res;
+  }
+
+  virtual mlir::Type getNestedType(mlir::Type t) = 0;
+
+  virtual mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t,
+                                     unsigned int depth) {
+    if (depth == 0) {
+      return nestedType;
+    } else {
+      return applyNestedType(
+          applyNestedType(nestedType, getNestedType(t), depth - 1), t);
+    }
+  }
+
+  virtual mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) = 0;
+};
+
+// Constraint that ensures that the types inferred for two values use
+// the same element type. If different types have been inferred
+// previously for the two values, precedence is given to the element
+// type inferred for the first value yielded by
+// `BaseConstraintT::yield`.
+//
+// If the element types for the two values are fixed, but
+// contradictory, an assertion is triggered.
+template <typename BaseConstraintT>
+class SameElementTypeConstraintBase
+    : public SameNestedTypeConstraintBase<BaseConstraintT> {
+public:
+  using SameNestedTypeConstraintBase<
+      BaseConstraintT>::SameNestedTypeConstraintBase;
+
+protected:
+  mlir::Type getNestedType(mlir::Type t) override {
+    if (mlir::ShapedType sht = llvm::dyn_cast<mlir::ShapedType>(t)) {
+      return sht.getElementType();
     } else {
       return t;
     }
   }
+
+  mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) override {
+    return TypeInferenceUtils::applyElementType(nestedType, t);
+  }
 };
 
 // Same element type constraint for static yield operators
@@ -826,8 +889,8 @@ class TypeInferenceAnalysisBase : public AnalysisT {
   // for which no type has been inferred so far, a new lattice with an
   // empty inferred types is created.
   LocalInferenceState getCurrentInferredTypes(mlir::Operation *op) {
-    LocalInferenceState map =
-        TypeInferenceUtils::getLocalInferenceState(op, [&](mlir::Value v) {
+    LocalInferenceState map = TypeInferenceUtils::getLocalInferenceState(
+        op, resolver, [&](mlir::Value v) {
           return this->template getOrCreate<TypeInferenceLattice>(v);
         });
 
@@ -860,22 +923,48 @@ class TypeInferenceAnalysisBase : public AnalysisT {
 
   // Updates the lattice for the inferred type for the value `v` using
   // its inferred type specified by `state`.
-  void updateLatticeValuesFromState(LocalInferenceState &state, mlir::Value v) {
+  void updateLatticeValuesAndDependenciesFromState(LocalInferenceState &state,
+                                                   mlir::Value v,
+                                                   mlir::Operation *op) {
     TypeInferenceLattice *lattice =
         this->template getOrCreate<TypeInferenceLattice>(v);
+
+    this->addDependency(lattice, op);
+
     auto latticeValue = lattice->getValue();
     latticeValue.setType(state.find(v));
     mlir::ChangeResult res = lattice->join(latticeValue);
     this->propagateIfChanged(lattice, res);
   }
 
+  // Visits all operations recursively by invoking `doVisitOperation`
+  // on them. Since `doVisitOperation` registers dependencies between
+  // related values an operations explicitly, this ensures that all
+  // operations are visited upon an update of any of their
+  // dependencies, even those for which the dataflow analysis
+  // framework does invoke `visitOperation`.
+  void initializeRecursively(mlir::Operation *op) {
+    for (Region &region : op->getRegions()) {
+      for (Operation &op : region.getOps()) {
+        doVisitOperation(&op);
+        initializeRecursively(&op);
+      }
+    }
+  }
+
+  LogicalResult initialize(mlir::Operation *op) override {
+    initializeRecursively(op);
+
+    return AnalysisT::initialize(op);
+  }
+
   // Visits a single operation: looks up the lattice values for the
   // values related to the operation, extracts the inferred types,
   // invokes the type resolver and updates the lattice values with the
   // newly inferred types.
   void doVisitOperation(Operation *op) {
     // Skip operations that do not use unresolved types
-    if (TypeInferenceUtils::iterateRelatedValues(op, [&](mlir::Value v) {
+    if (resolver.iterateRelatedValues(op, [&](mlir::Value v) {
           return !resolver.isUnresolvedType(v.getType());
         })) {
       return;
@@ -901,8 +990,8 @@ class TypeInferenceAnalysisBase : public AnalysisT {
         .Default([&](auto op) { state = resolver.resolve(op, inferredTypes); });
 
     // Store the resulting state in the lattice values
-    TypeInferenceUtils::iterateRelatedValues(op, [&](mlir::Value v) {
-      updateLatticeValuesFromState(state, v);
+    resolver.iterateRelatedValues(op, [&](mlir::Value v) {
+      updateLatticeValuesAndDependenciesFromState(state, v, op);
       return true;
     });
   }
@@ -963,15 +1052,7 @@ class TypeInferenceAnalysisBase : public AnalysisT {
   // the name of the operation, followed by the types inferred for
   // each operand, followed by `->`, followed by the types inferred
   // for the results.
-  void dumpAllState(mlir::Operation *op) {
-    mlir::Operation *funcOp = op;
-    while (funcOp && !llvm::isa<mlir::func::FuncOp>(funcOp))
-      funcOp = funcOp->getParentOp();
-
-    assert(funcOp);
-
-    dumpStateForOp(funcOp, 0);
-  }
+  void dumpAllState(mlir::ModuleOp module) { dumpStateForOp(module, 0); }
 
   TypeResolver &resolver;
 };
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Dialect/RT/TypeInference.h b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/RT/TypeInference.h
new file mode 100644
index 0000000000..e2d9da7599
--- /dev/null
+++ b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/RT/TypeInference.h
@@ -0,0 +1,111 @@
+// Part of the Concrete Compiler Project, under the BSD3 License with Zama
+// Exceptions. See
+// https://github.com/zama-ai/concrete/blob/main/LICENSE.txt
+// for license information.
+
+#ifndef CONCRETELANG_DIALECT_RT_TYPEINFERENCE_H
+#define CONCRETELANG_DIALECT_RT_TYPEINFERENCE_H
+
+#include <concretelang/Analysis/TypeInferenceAnalysis.h>
+#include <concretelang/Dialect/RT/IR/RTTypes.h>
+
+namespace mlir {
+namespace concretelang {
+namespace RT {
+
+template <typename BaseConstraintT, unsigned int leftDepth,
+          unsigned int rightDepth>
+class SameNestedTypeContraint
+    : public SameNestedTypeConstraintBase<BaseConstraintT, leftDepth,
+                                          rightDepth> {
+public:
+  using SameNestedTypeConstraintBase<BaseConstraintT, leftDepth,
+                                     rightDepth>::SameNestedTypeConstraintBase;
+
+protected:
+  mlir::Type getNestedType(mlir::Type t) override {
+    if (concretelang::RT::PointerType ptrt =
+            llvm::dyn_cast<concretelang::RT::PointerType>(t)) {
+      return ptrt.getElementType();
+    } else if (concretelang::RT::FutureType futt =
+                   llvm::dyn_cast<concretelang::RT::FutureType>(t)) {
+      return futt.getElementType();
+    } else {
+      return t;
+    }
+  }
+
+  mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) override {
+    if (llvm::isa<concretelang::RT::PointerType>(t)) {
+      return concretelang::RT::PointerType::get(nestedType);
+    } else if (llvm::isa<concretelang::RT::FutureType>(t)) {
+      return concretelang::RT::FutureType::get(nestedType);
+    } else {
+      return t;
+    }
+  }
+};
+
+// Constraint ensuring that two types are identical when the
+// indirection through a pointer is stripped. The parameters
+// `leftIsPointer` and `rightIsPointer` indicate which types should be
+// stripped from the pointer indirection before comparison.
+template <typename BaseConstraintT, bool leftIsPointer, bool rightIsPointer>
+class SamePointerTypeContraint
+    : public SameNestedTypeConstraintBase<
+          BaseConstraintT, leftIsPointer ? 1 : 0, rightIsPointer ? 1 : 0> {
+public:
+  using SameNestedTypeConstraintBase<
+      BaseConstraintT, leftIsPointer ? 1 : 0,
+      rightIsPointer ? 1 : 0>::SameNestedTypeConstraintBase;
+
+protected:
+  mlir::Type getNestedType(mlir::Type t) override {
+    concretelang::RT::PointerType ptrt =
+        llvm::cast<concretelang::RT::PointerType>(t);
+    return ptrt.getElementType();
+  }
+
+  mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) override {
+    if (llvm::isa<concretelang::RT::PointerType>(t)) {
+      return concretelang::RT::PointerType::get(nestedType);
+    } else {
+      return t;
+    }
+  }
+};
+
+// Constraint ensuring that two types are identical when the
+// indirection through a future is stripped. The parameters
+// `leftIsFuture` and `rightIsFuture` indicate which types should be
+// stripped from the future indirection before comparison.
+template <typename BaseConstraintT, bool leftIsFuture, bool rightIsFuture>
+class SameFutureTypeContraint
+    : public SameNestedTypeConstraintBase<BaseConstraintT, leftIsFuture ? 1 : 0,
+                                          rightIsFuture ? 1 : 0> {
+public:
+  using SameNestedTypeConstraintBase<
+      BaseConstraintT, leftIsFuture ? 1 : 0,
+      rightIsFuture ? 1 : 0>::SameNestedTypeConstraintBase;
+
+protected:
+  mlir::Type getNestedType(mlir::Type t) override {
+    concretelang::RT::FutureType ptrt =
+        llvm::cast<concretelang::RT::FutureType>(t);
+    return ptrt.getElementType();
+  }
+
+  mlir::Type applyNestedType(mlir::Type nestedType, mlir::Type t) override {
+    if (llvm::isa<concretelang::RT::FutureType>(t)) {
+      return concretelang::RT::FutureType::get(nestedType);
+    } else {
+      return t;
+    }
+  }
+};
+
+} // namespace RT
+} // namespace concretelang
+} // namespace mlir
+
+#endif // CONCRETELANG_DIALECT_RT_TYPEINFERENCE_H
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Transforms/TypeInferenceRewriter.h b/compilers/concrete-compiler/compiler/include/concretelang/Transforms/TypeInferenceRewriter.h
index bef222cd66..29e9573789 100644
--- a/compilers/concrete-compiler/compiler/include/concretelang/Transforms/TypeInferenceRewriter.h
+++ b/compilers/concrete-compiler/compiler/include/concretelang/Transforms/TypeInferenceRewriter.h
@@ -59,11 +59,81 @@ class TypeInferenceRewriter {
       return mlir::failure();
     }
 
+    // Re-write temporary function names and erase old functions that
+    // have been rewritten
+    llvm::SmallVector<mlir::func::FuncOp> funcsToErase;
+
+    for (auto [fromOp, toOp] : mapping.getOperationMap()) {
+      mlir::func::FuncOp fromFunc = llvm::dyn_cast<mlir::func::FuncOp>(fromOp);
+      mlir::func::FuncOp toFunc = llvm::dyn_cast<mlir::func::FuncOp>(toOp);
+
+      if (fromFunc && toFunc) {
+        std::string fromFuncName = fromFunc.getName().str();
+        toFunc.setName(fromFuncName);
+        funcsToErase.push_back(fromFunc);
+      }
+    }
+
+    for (mlir::func::FuncOp funcToErase : funcsToErase) {
+      funcToErase.erase();
+    }
+
     mlir::OpBuilder::InsertionGuard guard(rewriter);
     return this->postRewriteHook(rewriter, module, module);
   }
 
 protected:
+  // Rewrites all the blocks of all regions of `srcOp` into `tgtOp`
+  mlir::LogicalResult rewriteRegions(mlir::Operation *srcOp,
+                                     mlir::Operation *tgtOp,
+                                     const LocalInferenceState &resolvedTypes,
+                                     mlir::IRRewriter &rewriter) {
+    for (size_t regionIdx = 0; regionIdx < srcOp->getNumRegions();
+         regionIdx++) {
+      mlir::Region &oldRegion = srcOp->getRegion(regionIdx);
+      mlir::Region &newRegion = tgtOp->getRegion(regionIdx);
+
+      // Create empty blocks with the right argument types
+      for (auto [blockIdx, oldBlock] : llvm::enumerate(oldRegion.getBlocks())) {
+        mlir::Block *newBlock = new Block();
+        mapping.map(&oldBlock, newBlock);
+
+        for (auto [argIdx, oldArg] : llvm::enumerate(oldBlock.getArguments())) {
+          mlir::Type newArgType =
+              typeResolver.isUnresolvedType(oldArg.getType())
+                  ? resolvedTypes.lookup(oldArg)
+                  : oldArg.getType();
+
+          if (!newArgType || typeResolver.isUnresolvedType(newArgType)) {
+            srcOp->emitError()
+                << "Type of block argument #" << argIdx << " of block #"
+                << blockIdx << " of region #" << regionIdx << " is unresolved";
+            return mlir::failure();
+          }
+
+          mlir::BlockArgument newArg =
+              newBlock->addArgument(newArgType, oldArg.getLoc());
+
+          mapping.map(oldArg, newArg);
+        }
+
+        newRegion.getBlocks().insert(newRegion.end(), newBlock);
+      }
+
+      // Rewrite the contents of the blocks
+      for (auto [oldBlock, newBlock] :
+           llvm::zip_equal(oldRegion.getBlocks(), newRegion.getBlocks())) {
+        rewriter.setInsertionPointToEnd(&newBlock);
+
+        for (mlir::Operation &oldOp : oldBlock)
+          if (rewrite(&oldOp, rewriter).failed())
+            return mlir::failure();
+      }
+    }
+
+    return mlir::success();
+  }
+
   // Rewrite a single function
   mlir::LogicalResult rewrite(mlir::func::FuncOp func,
                               mlir::IRRewriter &rewriter) {
@@ -73,25 +143,55 @@ class TypeInferenceRewriter {
     // function to get the actual types for the arguments and return
     // values
     LocalInferenceState inferredTypes =
-        TypeInferenceUtils::getLocalInferenceState(solver, func);
+        TypeInferenceUtils::getLocalInferenceState(solver, typeResolver, func);
     LocalInferenceState resolvedTypes =
         typeResolver.resolve(func, inferredTypes);
 
     llvm::SmallVector<mlir::Type> argTypes =
         resolvedTypes.lookup(func.getBody().getArguments());
 
-    llvm::SmallVector<mlir::Type> returnTypes = resolvedTypes.lookup(
-        func.getRegion().getBlocks().front().getTerminator()->getOperands());
-
-    if (checkAllResolved(func, argTypes, "argument").failed() ||
-        checkAllResolved(func, returnTypes, "return value").failed()) {
+    // Check that all argument types have been resolved
+    if (checkAllResolved(func, argTypes, "argument").failed()) {
       return mlir::failure();
     }
 
+    std::optional<llvm::SmallVector<mlir::Type>> returnTypes;
+
+    // Check that the types of the values returned by all func.return
+    // terminators have the same types
+    for (auto &[i, block] : llvm::enumerate(func.getRegion().getBlocks())) {
+      if (mlir::func::ReturnOp returnOp =
+              llvm::dyn_cast<mlir::func::ReturnOp>(block.getTerminator())) {
+        if (!returnTypes.has_value()) {
+          returnTypes = resolvedTypes.lookup(returnOp->getOperands());
+
+          // Check that the types of the values returned through first
+          // terminator have all been resolved
+          if (checkAllResolved(returnOp, *returnTypes, "argument").failed()) {
+            return mlir::failure();
+          }
+        } else {
+          for (auto [j, operand] : llvm::enumerate(returnOp->getOperands())) {
+            mlir::Type thisOperandType = resolvedTypes.lookup(operand);
+            mlir::Type firstReturnOperand = (*returnTypes)[j];
+
+            if (thisOperandType != firstReturnOperand) {
+              func->emitError("Resolved type ")
+                  << thisOperandType << " for operand #" << i
+                  << " of terminator #" << j << " differs from the type "
+                  << firstReturnOperand
+                  << " of the same operand of the first "
+                     "return operation";
+            }
+          }
+        }
+      }
+    }
+
     // Temporarily create a new function with an suffix for the name
     // that later replaces the original function
     mlir::FunctionType ft =
-        FunctionType::get(func.getContext(), argTypes, returnTypes);
+        FunctionType::get(func.getContext(), argTypes, *returnTypes);
 
     std::string funcName =
         (func.getName() + "__rewriting_type_inference_rewriter").str();
@@ -100,22 +200,9 @@ class TypeInferenceRewriter {
 
     mapping.map(func.getOperation(), newFunc.getOperation());
 
-    newFunc.addEntryBlock();
-
-    for (auto [oldOperand, newOperand] : llvm::zip_equal(
-             func.getBody().getArguments(), newFunc.getBody().getArguments())) {
-      mapping.map(oldOperand, newOperand);
-    }
-
-    mlir::Block &entryBlock = func.getBody().getBlocks().front();
-    mlir::Block &newEntryBlock = newFunc.getBody().getBlocks().front();
-    rewriter.setInsertionPointToStart(&newEntryBlock);
-
-    // Recurse into the body of the function
-    for (mlir::Operation &op : entryBlock.getOperations()) {
-      if (rewrite(&op, rewriter).failed())
-        return mlir::failure();
-    }
+    // Recurse into nested regions and blocks
+    if (rewriteRegions(func, newFunc, resolvedTypes, rewriter).failed())
+      return mlir::failure();
 
     {
       mlir::OpBuilder::InsertionGuard guard(rewriter);
@@ -123,11 +210,10 @@ class TypeInferenceRewriter {
         return mlir::failure();
     }
 
-    // Replace original function and remove suffix from the name of the new
-    // function
-    std::string oldFuncName = func.getName().str();
-    rewriter.replaceOp(func, newFunc->getResults());
-    newFunc.setName(oldFuncName);
+    // Replace results of the original function, but do not yet erase
+    // it, since it may still be referenced in functions that have not
+    // been rewritten, yet
+    func->replaceAllUsesWith(newFunc->getResults());
 
     return mlir::success();
   }
@@ -163,7 +249,7 @@ class TypeInferenceRewriter {
       resolvedResultTypes = llvm::to_vector(op->getResultTypes());
     } else {
       LocalInferenceState inferredTypes =
-          TypeInferenceUtils::getLocalInferenceState(solver, op);
+          TypeInferenceUtils::getLocalInferenceState(solver, typeResolver, op);
 
       // Return-like ops are a bit special, since their operand types
       // are tied to the types of their parent op. The types inferred
@@ -217,55 +303,21 @@ class TypeInferenceRewriter {
       }
     }
 
+    SmallVector<Block *> newSuccessors;
+    for (Block *successor : op->getSuccessors()) {
+      newSuccessors.push_back(mapping.lookupOrDefault(successor));
+    }
+
     mlir::Operation *newOp = Operation::create(
         op->getLoc(), op->getName(), resolvedResultTypes, newOperands,
-        op->getAttrs(), op->getSuccessors(), op->getNumRegions());
+        op->getAttrs(), newSuccessors, op->getNumRegions());
 
     rewriter.insert(newOp);
     mapping.map(op, newOp);
 
     // Recurse into nested regions and blocks
-    for (size_t regionIdx = 0; regionIdx < op->getNumRegions(); regionIdx++) {
-      mlir::Region &oldRegion = op->getRegion(regionIdx);
-      mlir::Region &newRegion = newOp->getRegion(regionIdx);
-
-      // Create empty blocks with the right argument types
-      for (auto [blockIdx, oldBlock] : llvm::enumerate(oldRegion.getBlocks())) {
-        mlir::Block *newBlock = new Block();
-        mapping.map(&oldBlock, newBlock);
-
-        for (auto [argIdx, oldArg] : llvm::enumerate(oldBlock.getArguments())) {
-          mlir::Type newArgType =
-              typeResolver.isUnresolvedType(oldArg.getType())
-                  ? resolvedTypes.lookup(oldArg)
-                  : oldArg.getType();
-
-          if (!newArgType || typeResolver.isUnresolvedType(newArgType)) {
-            op->emitError()
-                << "Type of block argument #" << argIdx << " of block #"
-                << blockIdx << " of region #" << regionIdx << " is unresolved";
-            return mlir::failure();
-          }
-
-          mlir::BlockArgument newArg =
-              newBlock->addArgument(newArgType, oldArg.getLoc());
-
-          mapping.map(oldArg, newArg);
-        }
-
-        newRegion.getBlocks().insert(newRegion.end(), newBlock);
-      }
-
-      // Rewrite the contents of the blocks
-      for (auto [oldBlock, newBlock] :
-           llvm::zip_equal(oldRegion.getBlocks(), newRegion.getBlocks())) {
-        rewriter.setInsertionPointToEnd(&newBlock);
-
-        for (mlir::Operation &oldOp : oldBlock)
-          if (rewrite(&oldOp, rewriter).failed())
-            return mlir::failure();
-      }
-    }
+    if (rewriteRegions(op, newOp, resolvedTypes, rewriter).failed())
+      return mlir::failure();
 
     rewriter.setInsertionPointAfter(newOp);
 
diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp
index 9aba320326..0ebc6aefb9 100644
--- a/compilers/concrete-compiler/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp
@@ -64,6 +64,10 @@ static func::FuncOp outlineWorkFunction(RT::DataflowTaskOp DFTOp,
   Region &DFTOpBody = DFTOp.getBody();
   OpBuilder::InsertionGuard guard(builder);
 
+  SetVector<Value> operands;
+  getUsedValuesDefinedAbove(DFTOpBody, operands);
+  DFTOp->setOperands(operands.takeVector());
+
   // Instead of outlining with the same operands/results, we pass all
   // results as operands as well.  For now we preserve the results'
   // types, which will be changed to use an indirection when lowering.
@@ -591,6 +595,23 @@ struct FinalizeTaskCreationPass
         op->setOperand(1, clone);
       }
     });
+
+    module.walk([&](RT::WorkFunctionReturnOp op) {
+      OpBuilder builder(op);
+
+      Value val = op.getOperand(0);
+      if (val.getType().isa<mlir::MemRefType>() &&
+          isa<RT::DerefWorkFunctionArgumentPtrPlaceholderOp>(
+              val.getDefiningOp())) {
+        Value newval =
+            builder
+                .create<mlir::memref::AllocOp>(
+                    val.getLoc(), val.getType().dyn_cast<mlir::MemRefType>())
+                .getResult();
+        builder.create<mlir::memref::CopyOp>(val.getLoc(), val, newval);
+        op->setOperand(0, newval);
+      }
+    });
   }
   FinalizeTaskCreationPass(bool debug) : debug(debug){};
 
diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp
index 3b74d97cb5..e97140b9d2 100644
--- a/compilers/concrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp
@@ -4,6 +4,9 @@
 // for license information.
 
 #include "concrete-optimizer.hpp"
+#include "concretelang/Dialect/RT/IR/RTOps.h"
+#include "concretelang/Dialect/RT/IR/RTTypes.h"
+#include "concretelang/Dialect/RT/TypeInference.h"
 #include "concretelang/Dialect/TFHE/IR/TFHEParameters.h"
 #include "concretelang/Dialect/TypeInference/IR/TypeInferenceOps.h"
 
@@ -20,6 +23,7 @@
 #include <mlir/Analysis/DataFlow/SparseAnalysis.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/Dialect/Bufferization/IR/Bufferization.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Dialect/SCF/IR/SCF.h>
 #include <mlir/Dialect/Tensor/IR/Tensor.h>
 #include <mlir/Transforms/GreedyPatternRewriteDriver.h>
@@ -33,10 +37,13 @@ template <class T> static std::optional<T> tryGetScalarType(mlir::Type t) {
   if (T ctt = t.dyn_cast<T>())
     return ctt;
 
-  if (mlir::RankedTensorType rtt = t.dyn_cast<mlir::RankedTensorType>())
-    return tryGetScalarType<T>(rtt.getElementType());
-  else if (mlir::MemRefType mrt = t.dyn_cast<mlir::MemRefType>())
-    return tryGetScalarType<T>(mrt.getElementType());
+  if (mlir::ShapedType sht = t.dyn_cast<mlir::ShapedType>()) {
+    return tryGetScalarType<T>(sht.getElementType());
+  } else if (RT::PointerType rtptr = llvm::dyn_cast<RT::PointerType>(t)) {
+    return tryGetScalarType<T>(rtptr.getElementType());
+  } else if (RT::FutureType futt = llvm::dyn_cast<RT::FutureType>(t)) {
+    return tryGetScalarType<T>(futt.getElementType());
+  }
 
   return std::nullopt;
 }
@@ -203,15 +210,224 @@ class TFHEParametrizationTypeResolver : public TypeResolver {
       std::optional<CircuitSolutionWrapper> solution)
       : solution(solution) {}
 
+  // Adds constraints that ensure that the argument types of `funcOp`
+  // and the operand types of `createAsyncTaskOp` are compatible and
+  // that the result types of all return operations of `funcOp` are
+  // compatible with the output arguments of `callOp`
+  void
+  addFuncAndCreateTaskConstraints(TypeConstraintSet<> &cs,
+                                  mlir::func::FuncOp funcOp,
+                                  RT::CreateAsyncTaskOp createAsyncTaskOp) {
+    for (auto [createArg, funcArg] :
+         llvm::zip_equal(createAsyncTaskOp.getOperands().drop_front(3),
+                         funcOp.getArguments())) {
+      if (llvm::isa<RT::PointerType>(createArg.getType()) &&
+          llvm::isa<RT::FutureType>(
+              llvm::cast<RT::PointerType>(createArg.getType())
+                  .getElementType())) {
+        // This is an input argument of type
+        // RT.rtptr<!RT.future<...>>, the corresponding function
+        // argument type is !RT.rtptr<...>
+        cs.addConstraint<RT::SameNestedTypeContraint<
+            DynamicTypeConstraint<DynamicFunctorYield>, 2, 1>>(
+            [createArg = createArg]() { return createArg; },
+            [funcArg = funcArg]() { return funcArg; });
+      } else if (llvm::isa<RT::FutureType>(createArg.getType()) &&
+                 llvm::isa<RT::PointerType>(funcArg.getType())) {
+        // This is an output argument of type
+        // !RT.future<...>, the corresponding function
+        // argument type is !RT.rtptr<...>
+        cs.addConstraint<RT::SameNestedTypeContraint<
+            DynamicTypeConstraint<DynamicFunctorYield>, 1, 1>>(
+            [createArg = createArg]() { return createArg; },
+            [funcArg = funcArg]() { return funcArg; });
+      } else {
+        llvm_unreachable("Asynchronous task creation with unknown scheme "
+                         "for argument types");
+      }
+    }
+  }
+
+  // Adds constraints that ensure that the argument types of `funcOp`
+  // and the operand types of `callOp` are equal and that the result
+  // types of all return operations of `funcOp` are equal to the
+  // result types of `callOp`
+  void addFuncAndCallConstraints(TypeConstraintSet<> &cs,
+                                 mlir::func::FuncOp funcOp,
+                                 mlir::func::CallOp callOp) {
+    // Ensure that the operands of the call op have the same
+    // type as the function arguments
+    for (auto [callOperand, funcArgument] :
+         llvm::zip_equal(callOp->getOperands(), funcOp.getArguments())) {
+      cs.addConstraint<DynamicSameTypeConstraint<DynamicFunctorYield>>(
+          [callOperand = callOperand]() { return callOperand; },
+          [funcArgument = funcArgument]() { return funcArgument; });
+    }
+
+    // Ensure that the return types of the function defined via
+    // its return ops are equal to the result types of the call
+    for (mlir::Block &block : funcOp.getRegion().getBlocks()) {
+      if (mlir::func::ReturnOp returnOp =
+              llvm::dyn_cast<mlir::func::ReturnOp>(block.getTerminator())) {
+        for (auto [returnOperand, callResult] :
+             llvm::zip_equal(returnOp->getOperands(), callOp->getResults())) {
+          cs.addConstraint<DynamicSameTypeConstraint<DynamicFunctorYield>>(
+              [returnOperand = returnOperand]() { return returnOperand; },
+              [callResult = callResult]() { return callResult; });
+        }
+      }
+    }
+  }
+
+  // Attempts to assemble the types for arguments and return values of
+  // the function with the name `funcName` from `state`. The (partial)
+  // result if passed in `retArgTypes` and `retResTypes`.
+  void
+  assembleInferredFunctionType(mlir::ModuleOp module, StringRef funcName,
+                               const LocalInferenceState &state,
+                               llvm::SmallVector<mlir::Type> &retArgTypes,
+                               llvm::SmallVector<mlir::Type> &retResTypes) {
+    mlir::func::FuncOp func = lookupFunction(module, funcName);
+    mlir::FunctionType funcType = func.getFunctionType();
+    llvm::SmallVector<mlir::Type> argTypes;
+    llvm::SmallVector<mlir::Type> resTypes(funcType.getNumResults());
+
+    // Assemble argument types
+    for (mlir::Value arg : func.getArguments()) {
+      argTypes.push_back(isUnresolvedType(arg.getType()) ? state.lookup(arg)
+                                                         : arg.getType());
+    }
+
+    // Assemble result types from types inferred for the operands
+    // of return ops
+    for (mlir::Block &block : func.getRegion().getBlocks()) {
+      if (mlir::func::ReturnOp returnOp =
+              llvm::dyn_cast<mlir::func::ReturnOp>(block.getTerminator()))
+        for (auto [i, operand] : llvm::enumerate(returnOp.getOperands())) {
+          if (!resTypes[i]) {
+            resTypes[i] = isUnresolvedType(operand.getType())
+                              ? state.lookup(operand)
+                              : operand.getType();
+          }
+        }
+    }
+
+    retArgTypes = std::move(argTypes);
+    retResTypes = std::move(resTypes);
+  }
+
+  // Attemps to infer the type of a function from `state`. If the
+  // function type cannot be inferred entirely (e.g., the type for at
+  // least one argument or one result cannot be inferred), the
+  // function returns an empty function type.
+  mlir::FunctionType
+  tryGetFullyInferredFunctionType(mlir::ModuleOp module, StringRef funcName,
+                                  const LocalInferenceState &state) {
+    llvm::SmallVector<mlir::Type> argTypes;
+    llvm::SmallVector<mlir::Type> resTypes;
+
+    assembleInferredFunctionType(module, funcName, state, argTypes, resTypes);
+
+    auto typeDefined = [=](mlir::Type t) { return !!t; };
+
+    if (llvm::all_of(argTypes, typeDefined) &&
+        llvm::all_of(resTypes, typeDefined)) {
+      return mlir::FunctionType::get(module.getContext(), argTypes, resTypes);
+    }
+
+    return nullptr;
+  }
+
+  // Returns the function operation for the function with the name
+  // `funcName`.
+  mlir::func::FuncOp lookupFunction(mlir::ModuleOp module,
+                                    llvm::StringRef funcName) {
+    mlir::func::FuncOp ret;
+
+    module.walk([&](mlir::func::FuncOp funcOp) {
+      if (funcOp.getName() == funcName) {
+        ret = funcOp;
+        return mlir::WalkResult::interrupt();
+      }
+
+      return mlir::WalkResult::advance();
+    });
+
+    return ret;
+  }
+
   LocalInferenceState
   resolve(mlir::Operation *op,
           const LocalInferenceState &inferredTypes) override {
     LocalInferenceState state = inferredTypes;
 
     mlir::TypeSwitch<mlir::Operation *>(op)
+        .Case<mlir::func::CallOp>([&](auto op) {
+          TypeConstraintSet<> cs;
+
+          mlir::ModuleOp module =
+              op->template getParentOfType<mlir::ModuleOp>();
+
+          mlir::func::FuncOp funcOp = lookupFunction(module, op.getCallee());
+          addFuncAndCallConstraints(cs, funcOp, op);
+          cs.converge(op, *this, state, inferredTypes);
+        })
+        .Case<mlir::func::ConstantOp>([&](auto op) {
+          // One-way dependency from the argument and return types of
+          // the corresponding function to the constant op
+          mlir::ModuleOp module =
+              op->template getParentOfType<mlir::ModuleOp>();
+          if (mlir::FunctionType funcType = tryGetFullyInferredFunctionType(
+                  module, op.getValue(), state)) {
+            state.set(op.getResult(), funcType);
+          }
+        })
         .Case<mlir::func::FuncOp>([&](auto op) {
           TypeConstraintSet<> cs;
 
+          mlir::ModuleOp module =
+              op->template getParentOfType<mlir::ModuleOp>();
+
+          module.walk([&](RT::CreateAsyncTaskOp createAsyncTaskOp) {
+            mlir::SymbolRefAttr workfnSymbol = createAsyncTaskOp.getWorkfn();
+
+            if (workfnSymbol.getRootReference() == op.getName()) {
+              addFuncAndCreateTaskConstraints(cs, op, createAsyncTaskOp);
+            }
+          });
+
+          module.walk([&](mlir::func::CallOp callOp) {
+            if (callOp.getCallee() == op.getName()) {
+              addFuncAndCallConstraints(cs, op, callOp);
+            }
+          });
+
+          // Make sure that all return operations have the same operand types
+          mlir::func::ReturnOp firstReturnOp;
+
+          for (mlir::Block &block : op.getRegion().getBlocks()) {
+            if (mlir::func::ReturnOp returnOp =
+                    llvm::dyn_cast<mlir::func::ReturnOp>(
+                        block.getTerminator())) {
+              if (!firstReturnOp) {
+                firstReturnOp = returnOp;
+              } else {
+                for (auto [returnOperand, firstReturnOperand] :
+                     llvm::zip_equal(returnOp->getOperands(),
+                                     firstReturnOp->getResults())) {
+                  cs.addConstraint<
+                      DynamicSameTypeConstraint<DynamicFunctorYield>>(
+                      [returnOperand = returnOperand]() {
+                        return returnOperand;
+                      },
+                      [firstReturnOperand = firstReturnOperand]() {
+                        return firstReturnOperand;
+                      });
+                }
+              }
+            }
+          }
+
           if (solution.has_value()) {
             cs.addConstraint<ApplySolverSolutionToFunctionArgsConstraint>(
                 *this, solution.value());
@@ -228,10 +444,16 @@ class TFHEParametrizationTypeResolver : public TypeResolver {
               TFHE::BatchedBootstrapGLWEOp, TFHE::EncodeExpandLutForBootstrapOp,
               TFHE::EncodeLutForCrtWopPBSOp, TFHE::EncodePlaintextWithCrtOp,
               TFHE::WopPBSGLWEOp, mlir::func::ReturnOp,
-              Tracing::TraceCiphertextOp, mlir::tensor::EmptyOp>([&](auto op) {
+              Tracing::TraceCiphertextOp, mlir::tensor::EmptyOp,
+              mlir::tensor::DimOp, RT::BuildReturnPtrPlaceholderOp,
+              RT::RegisterTaskWorkFunctionOp>([&](auto op) {
           converge<NoTypeConstraint>(op, state, inferredTypes);
         })
-
+        .Case<mlir::arith::SelectOp>([&](auto op) {
+          converge<SameOperandTypeConstraint<1, 2>,
+                   SameOperandAndResultTypeConstraint<1, 0>>(op, state,
+                                                             inferredTypes);
+        })
         .Case<TFHE::AddGLWEOp, TFHE::ABatchedAddGLWEOp>([&](auto op) {
           converge<SameOperandTypeConstraint<0, 1>,
                    SameOperandAndResultTypeConstraint<0, 0>>(op, state,
@@ -393,6 +615,40 @@ class TFHEParametrizationTypeResolver : public TypeResolver {
                 DynamicSameElementTypeConstraint<DynamicFunctorYield>>(
                 [=]() { return result; }, [=]() { return terminatorOperand; });
           }
+          cs.converge(op, *this, state, inferredTypes);
+        })
+        .Case<RT::DerefWorkFunctionArgumentPtrPlaceholderOp,
+              RT::DerefReturnPtrPlaceholderOp>([&](auto _op) {
+          converge<RT::SamePointerTypeContraint<
+              StaticTypeConstraint<OperandAndResultValueYield<0, 0>>, true,
+              false>>(op, state, inferredTypes);
+        })
+        .Case<RT::WorkFunctionReturnOp>([&](auto _op) {
+          converge<RT::SamePointerTypeContraint<
+              StaticTypeConstraint<OperandValueYield<0, 1>>, false, true>>(
+              op, state, inferredTypes);
+        })
+        .Case<RT::MakeReadyFutureOp>([&](auto _op) {
+          converge<RT::SameFutureTypeContraint<
+              StaticTypeConstraint<OperandAndResultValueYield<0, 0>>, false,
+              true>>(op, state, inferredTypes);
+        })
+        .Case<RT::AwaitFutureOp>([&](auto _op) {
+          converge<RT::SameFutureTypeContraint<
+              StaticTypeConstraint<OperandAndResultValueYield<0, 0>>, true,
+              false>>(op, state, inferredTypes);
+        })
+        .Case<RT::CreateAsyncTaskOp>([&](auto op) {
+          TypeConstraintSet<> cs;
+
+          mlir::SymbolRefAttr workfnSymbol = op.getWorkfn();
+          mlir::ModuleOp module =
+              op->template getParentOfType<mlir::ModuleOp>();
+
+          mlir::func::FuncOp workfn =
+              module.lookupSymbol<mlir::func::FuncOp>(workfnSymbol);
+
+          addFuncAndCreateTaskConstraints(cs, workfn, op);
 
           cs.converge(op, *this, state, inferredTypes);
         })
@@ -401,8 +657,116 @@ class TFHEParametrizationTypeResolver : public TypeResolver {
     return state;
   }
 
+  static bool staticIsUnresolvedType(mlir::Type t) {
+    if (RT::PointerType rtptr = llvm::dyn_cast<RT::PointerType>(t)) {
+      return staticIsUnresolvedType(rtptr.getElementType());
+    } else if (RT::FutureType futt = llvm::dyn_cast<RT::FutureType>(t)) {
+      return staticIsUnresolvedType(futt.getElementType());
+    } else if (mlir::ShapedType sht = llvm::dyn_cast<mlir::ShapedType>(t)) {
+      return staticIsUnresolvedType(sht.getElementType());
+    } else if (mlir::FunctionType funct =
+                   llvm::dyn_cast<mlir::FunctionType>(t)) {
+      return llvm::any_of(funct.getInputs(), staticIsUnresolvedType) ||
+             llvm::any_of(funct.getResults(), staticIsUnresolvedType);
+    } else {
+      return isUnparametrizedGLWEType(t);
+    }
+  }
+
   bool isUnresolvedType(mlir::Type t) const override {
-    return isUnparametrizedGLWEType(t);
+    return staticIsUnresolvedType(t);
+  }
+
+  bool iterateRelatedValues(
+      mlir::Operation *op,
+      llvm::function_ref<bool(mlir::Value)> fn) const override {
+    if (!TypeResolver::iterateRelatedValues(op, fn))
+      return false;
+
+    return mlir::TypeSwitch<mlir::Operation *, bool>(op)
+        .Case<RT::CreateAsyncTaskOp>([&](auto op) {
+          mlir::SymbolRefAttr workfnSymbol = op.getWorkfn();
+          mlir::ModuleOp module =
+              op->template getParentOfType<mlir::ModuleOp>();
+
+          auto workfn = module.lookupSymbol<mlir::func::FuncOp>(workfnSymbol);
+
+          for (mlir::Value v : workfn.getArguments())
+            if (!fn(v))
+              return false;
+
+          return true;
+        })
+        .Case<mlir::func::CallOp, mlir::func::ConstantOp>(
+            [&](mlir::Operation *op) {
+              mlir::StringRef funcName = "";
+
+              if (mlir::func::CallOp callOp =
+                      llvm::dyn_cast<mlir::func::CallOp>(op)) {
+                funcName = callOp.getCallee();
+              } else if (mlir::func::ConstantOp cstOp =
+                             llvm::dyn_cast<mlir::func::ConstantOp>(op)) {
+                funcName = cstOp.getValue();
+              }
+
+              mlir::ModuleOp module =
+                  op->template getParentOfType<mlir::ModuleOp>();
+
+              if (module
+                      .walk([&](mlir::func::FuncOp func) {
+                        if (func.getName() == funcName) {
+                          for (mlir::Value arg : func.getArguments()) {
+                            if (!fn(arg))
+                              return mlir::WalkResult::interrupt();
+                          }
+
+                          for (mlir::Block &block :
+                               func.getRegion().getBlocks()) {
+                            if (mlir::func::ReturnOp returnOp =
+                                    llvm::dyn_cast<mlir::func::ReturnOp>(
+                                        block.getTerminator())) {
+                              for (mlir::Value res : returnOp->getOperands()) {
+                                if (!fn(res))
+                                  return mlir::WalkResult::interrupt();
+                              }
+                            }
+                          }
+                        }
+
+                        return mlir::WalkResult::advance();
+                      })
+                      .wasInterrupted()) {
+                return false;
+              } else {
+                return true;
+              }
+            })
+        .Case<mlir::func::FuncOp>([&](auto op) {
+          mlir::ModuleOp module =
+              op->template getParentOfType<mlir::ModuleOp>();
+
+          if (module
+                  .walk([&](RT::CreateAsyncTaskOp createAsyncTaskOp) {
+                    mlir::SymbolRefAttr workfnSymbol =
+                        createAsyncTaskOp.getWorkfn();
+
+                    if (workfnSymbol.getRootReference() == op.getName()) {
+                      for (mlir::Value operand :
+                           createAsyncTaskOp.getOperands().drop_front(3)) {
+                        if (!fn(operand))
+                          return mlir::WalkResult::interrupt();
+                      }
+                    }
+
+                    return mlir::WalkResult::advance();
+                  })
+                  .wasInterrupted()) {
+            return false;
+          } else {
+            return true;
+          }
+        })
+        .Default([&](auto _op) { return true; });
   }
 
 protected:
diff --git a/compilers/concrete-compiler/compiler/lib/Support/CompilerEngine.cpp b/compilers/concrete-compiler/compiler/lib/Support/CompilerEngine.cpp
index adb8e6e97c..0eab6a7206 100644
--- a/compilers/concrete-compiler/compiler/lib/Support/CompilerEngine.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Support/CompilerEngine.cpp
@@ -308,16 +308,6 @@ CompilerEngine::compile(mlir::ModuleOp moduleOp, Target target,
 
   auto dataflowParallelize =
       options.autoParallelize || options.dataflowParallelize;
-  if (options.optimizerConfig.strategy == optimizer::Strategy::DAG_MULTI &&
-      dataflowParallelize == true) {
-    // FIXME: DF is not currently compatible with multi-parameters as
-    // the generation of dataflow tasks obfuscates the code before the
-    // analysis can be done. Until this is fixed we cannot allow both.
-    dataflowParallelize = false;
-    warnx("WARNING: dataflow parallelization is not compatible with the "
-          "optimizer strategy [dag-multi]. Continuing with dataflow "
-          "parallelization disabled.");
-  }
   auto loopParallelize = options.autoParallelize || options.loopParallelize;
 
   if (loopParallelize)
diff --git a/compilers/concrete-compiler/compiler/lib/Support/Pipeline.cpp b/compilers/concrete-compiler/compiler/lib/Support/Pipeline.cpp
index a755c1b9f7..67093ee8f6 100644
--- a/compilers/concrete-compiler/compiler/lib/Support/Pipeline.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Support/Pipeline.cpp
@@ -561,6 +561,9 @@ mlir::LogicalResult lowerToStd(mlir::MLIRContext &context,
       pm, mlir::bufferization::createBufferDeallocationPass(), enablePass);
   addPotentiallyNestedPass(pm, mlir::concretelang::createStartStopPass(),
                            enablePass);
+  addPotentiallyNestedPass(pm, mlir::createCanonicalizerPass(), enablePass);
+  addPotentiallyNestedPass(pm, mlir::createBufferizationToMemRefPass(),
+                           enablePass);
   addPotentiallyNestedPass(
       pm, mlir::concretelang::createFixupBufferDeallocationPass(), enablePass);
 
@@ -589,9 +592,13 @@ lowerStdToLLVMDialect(mlir::MLIRContext &context, mlir::ModuleOp &module,
   pipelinePrinting("StdToLLVM", pm, context);
 
   // Convert to MLIR LLVM Dialect
+  addPotentiallyNestedPass(pm, mlir::arith::createArithExpandOpsPass(),
+                           enablePass);
   addPotentiallyNestedPass(
       pm, mlir::concretelang::createConvertMLIRLowerableDialectsToLLVMPass(),
       enablePass);
+  addPotentiallyNestedPass(pm, mlir::createReconcileUnrealizedCastsPass(),
+                           enablePass);
 
   return pm.run(module);
 }


Test patch:
diff --git a/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference-multifunc.mlir b/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference-multifunc.mlir
new file mode 100644
index 0000000000..4c4e716a87
--- /dev/null
+++ b/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference-multifunc.mlir
@@ -0,0 +1,60 @@
+// RUN: concretecompiler --split-input-file --action=dump-parametrized-tfhe --skip-program-info --passes=tfhe-circuit-solution-parametrization %s 2>&1| FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   func.func @main(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>> {
+// CHECK-NEXT:     %0 = call @aux(%arg0, %arg1) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:     return %0 : !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   func.func @aux(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>> {
+// CHECK-NEXT:     %0 = "TFHE.add_glwe"(%arg0, %arg1) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:     return %0 : !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+func.func @main(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>) {
+  %a0 = "TypeInference.propagate_downward"(%arg0) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  %a1 = "TypeInference.propagate_downward"(%arg1) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  %0 = call @aux(%a0, %a1) : (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  return %0: !TFHE.glwe<sk?>
+}
+
+func.func @aux(%arg0: !TFHE.glwe<sk?>, %arg1: !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>) {
+  %0 = "TFHE.add_glwe"(%arg0, %arg1): (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  return %0: !TFHE.glwe<sk?>
+}
+
+// -----
+
+// CHECK:      module {
+// CHECK-NEXT:   func.func @main(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>> {
+// CHECK-NEXT:     %0 = call @intermediate(%arg0, %arg1, %arg0) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:     return %0 : !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   func.func @intermediate(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>, %arg2: !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>> {
+// CHECK-NEXT:     %0 = call @aux(%arg0, %arg1) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:     %1 = "TFHE.add_glwe"(%arg2, %0) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:     return %1 : !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   func.func @aux(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>> {
+// CHECK-NEXT:     %0 = "TFHE.add_glwe"(%arg0, %arg1) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:     return %0 : !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+func.func @main(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>) {
+  %a0 = "TypeInference.propagate_downward"(%arg0) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  %0 = call @intermediate(%a0, %arg1, %arg0) : (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>, !TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  return %0: !TFHE.glwe<sk?>
+}
+
+func.func @intermediate(%arg0: !TFHE.glwe<sk?>, %arg1: !TFHE.glwe<sk?>, %arg2: !TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>) {
+  %a2 = "TypeInference.propagate_downward"(%arg2) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  %0 = call @aux(%arg0, %arg1) : (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  %1 = "TFHE.add_glwe"(%a2, %0): (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  return %1: !TFHE.glwe<sk?>
+}
+
+func.func @aux(%arg0: !TFHE.glwe<sk?>, %arg1: !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>) {
+  %0 = "TFHE.add_glwe"(%arg0, %arg1): (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  return %0: !TFHE.glwe<sk?>
+}
diff --git a/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference.mlir b/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference.mlir
index bd8f22d82a..87b9aa6568 100644
--- a/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference.mlir
+++ b/compilers/concrete-compiler/compiler/tests/check_tests/TypeInference/inference.mlir
@@ -332,3 +332,26 @@ func.func @tiled_2(%arg0: tensor<8x4x!TFHE.glwe<sk?>>, %arg1: tensor<4x2xi7>) ->
   }
   return %5 : tensor<8x2x!TFHE.glwe<sk?>>
 }
+
+// -----
+
+// CHECK:      func.func @multi_block(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>, %arg2: !TFHE.glwe<sk[1]<12,1024>>, %arg3: i1, %arg4: i1) -> !TFHE.glwe<sk[1]<12,1024>> {
+// CHECK-NEXT:   cf.cond_br %arg3, ^bb1, ^bb2
+// CHECK-NEXT: ^bb1:  // 2 preds: ^bb0, ^bb1
+// CHECK-NEXT:   cf.cond_br %arg4, ^bb1, ^bb2
+// CHECK-NEXT: ^bb2:  // 2 preds: ^bb0, ^bb1
+// CHECK-NEXT:   %0 = "TFHE.add_glwe"(%arg1, %arg0) : (!TFHE.glwe<sk[1]<12,1024>>, !TFHE.glwe<sk[1]<12,1024>>) -> !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT:   return %0 : !TFHE.glwe<sk[1]<12,1024>>
+// CHECK-NEXT: }
+func.func @multi_block(%arg0: !TFHE.glwe<sk[1]<12,1024>>, %arg1: !TFHE.glwe<sk[1]<12,1024>>, %arg2: !TFHE.glwe<sk[1]<12,1024>>, %cond: i1, %cond2: i1) -> !TFHE.glwe<sk?> {
+  %a0 = "TypeInference.propagate_downward"(%arg0) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  %a1 = "TypeInference.propagate_downward"(%arg1) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  %a2 = "TypeInference.propagate_downward"(%arg2) : (!TFHE.glwe<sk[1]<12,1024>>) -> (!TFHE.glwe<sk?>)
+  cf.cond_br %cond, ^bb0(%a0: !TFHE.glwe<sk?>), ^bb1(%a1: !TFHE.glwe<sk?>)
+^bb0(%bbarg0 : !TFHE.glwe<sk?>):
+  %0 = "TFHE.add_glwe"(%a1, %a1): (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  cf.cond_br %cond2, ^bb0(%0: !TFHE.glwe<sk?>), ^bb1(%0: !TFHE.glwe<sk?>)
+^bb1(%bbarg1 : !TFHE.glwe<sk?>):
+  %1 = "TFHE.add_glwe"(%a1, %a0): (!TFHE.glwe<sk?>, !TFHE.glwe<sk?>) -> (!TFHE.glwe<sk?>)
+  return %1: !TFHE.glwe<sk?>
+}
diff --git a/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc b/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc
index 7f9d80e1ee..719797e39b 100644
--- a/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc
+++ b/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc
@@ -68,6 +68,9 @@ class EndToEndTest : public ::testing::Test {
       ASSERT_OUTCOME_HAS_VALUE(maybeRes);
       auto result = maybeRes.value();
 
+      if (!mlir::concretelang::dfr::_dfr_is_root_node())
+        return;
+
       /* Check results */
       bool allgood = true;
       for (size_t i = 0; i < desc.outputs.size(); i++) {