From 261e59d2b4a119f0bc4b1397c6a1892fc69bae92 Mon Sep 17 00:00:00 2001
From: Antoniu Pop <antoniu.pop@zama.ai>
Date: Tue, 28 Jun 2022 15:52:37 +0100
Subject: [PATCH] feat(dfr): add more operations to be converted to tasks.

---
 .../RT/Analysis/BuildDataflowTaskGraph.cpp       | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/compiler/lib/Dialect/RT/Analysis/BuildDataflowTaskGraph.cpp b/compiler/lib/Dialect/RT/Analysis/BuildDataflowTaskGraph.cpp
index 1a3cda21e..b8a890b57 100644
--- a/compiler/lib/Dialect/RT/Analysis/BuildDataflowTaskGraph.cpp
+++ b/compiler/lib/Dialect/RT/Analysis/BuildDataflowTaskGraph.cpp
@@ -41,13 +41,15 @@ namespace {
 
 // TODO: adjust these two functions based on cost model
 static bool isCandidateForTask(Operation *op) {
-  return isa<FHE::AddEintIntOp, FHE::AddEintOp, FHE::SubIntEintOp,
-             FHE::MulEintIntOp, FHE::ApplyLookupTableEintOp,
-             FHELinalg::MatMulIntEintOp, FHELinalg::MatMulEintIntOp,
-             FHELinalg::AddEintIntOp, FHELinalg::AddEintOp,
-             FHELinalg::SubIntEintOp, FHELinalg::NegEintOp,
-             FHELinalg::MulEintIntOp, FHELinalg::ApplyLookupTableEintOp,
-             FHELinalg::ApplyMultiLookupTableEintOp, FHELinalg::Dot>(op);
+  return isa<
+      FHE::AddEintIntOp, FHE::AddEintOp, FHE::SubIntEintOp, FHE::MulEintIntOp,
+      FHE::ApplyLookupTableEintOp, FHELinalg::MatMulIntEintOp,
+      FHELinalg::MatMulEintIntOp, FHELinalg::AddEintIntOp, FHELinalg::AddEintOp,
+      FHELinalg::SubIntEintOp, FHELinalg::NegEintOp, FHELinalg::MulEintIntOp,
+      FHELinalg::ApplyLookupTableEintOp, FHELinalg::ApplyMultiLookupTableEintOp,
+      FHELinalg::ApplyMappedLookupTableEintOp, FHELinalg::Dot,
+      FHELinalg::MatMulEintIntOp, FHELinalg::MatMulIntEintOp, FHELinalg::SumOp,
+      FHELinalg::ConcatOp, FHELinalg::FhelinalgConv2DNchwFchwOp>(op);
 }
 
 // Identify operations that are beneficial to sink into tasks.  These