From 2cd06580ee42c746d7acd7202058634b73a40c06 Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Wed, 5 Apr 2023 17:07:48 +0200
Subject: [PATCH] feat(compiler): Add support for `scf.forall` and associated
 ops in TFHE parametrization

---
 .../TFHEGlobalParametrization.cpp                     | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/compilers/concrete-compiler/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp b/compilers/concrete-compiler/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp
index 5679fd04a..074075b8f 100644
--- a/compilers/concrete-compiler/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp
@@ -353,6 +353,9 @@ void TFHEGlobalParametrizationPass::runOnOperation() {
     mlir::concretelang::addDynamicallyLegalTypeOp<mlir::tensor::EmptyOp>(
         target, converter);
 
+    patterns.add<RegionOpTypeConverterPattern<
+        mlir::scf::InParallelOp, TFHEGlobalParametrizationTypeConverter>>(
+        &getContext(), converter);
     patterns.add<RegionOpTypeConverterPattern<
         mlir::linalg::GenericOp, TFHEGlobalParametrizationTypeConverter>>(
         &getContext(), converter);
@@ -362,6 +365,9 @@ void TFHEGlobalParametrizationPass::runOnOperation() {
     patterns.add<RegionOpTypeConverterPattern<
         mlir::scf::ForOp, TFHEGlobalParametrizationTypeConverter>>(
         &getContext(), converter);
+    patterns.add<RegionOpTypeConverterPattern<
+        mlir::scf::ForallOp, TFHEGlobalParametrizationTypeConverter>>(
+        &getContext(), converter);
     patterns.add<RegionOpTypeConverterPattern<
         mlir::func::ReturnOp, TFHEGlobalParametrizationTypeConverter>>(
         &getContext(), converter);
@@ -372,6 +378,8 @@ void TFHEGlobalParametrizationPass::runOnOperation() {
         &getContext(), converter);
     mlir::concretelang::addDynamicallyLegalTypeOp<mlir::linalg::YieldOp>(
         target, converter);
+    mlir::concretelang::addDynamicallyLegalTypeOp<
+        mlir::tensor::ParallelInsertSliceOp>(target, converter);
 
     mlir::concretelang::populateWithTensorTypeConverterPatterns(
         patterns, target, converter);
@@ -389,6 +397,9 @@ void TFHEGlobalParametrizationPass::runOnOperation() {
     mlir::concretelang::populateWithRTTypeConverterPatterns(patterns, target,
                                                             converter);
 
+    mlir::concretelang::GenericTypeConverterPattern<
+        mlir::tensor::ParallelInsertSliceOp>(&getContext(), converter);
+
     // Apply conversion
     if (mlir::applyPartialConversion(op, target, std::move(patterns))
             .failed()) {