From ee2f743a787379dea3f506d224a83310d2c94ede Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Thu, 1 Dec 2022 14:00:46 +0100
Subject: [PATCH] feat(compiler): Add new action dump-sdfg

Add a new action `dump-sdfg` to `concretecompiler` that causes the IR
to be dumpred right after the extraction of SDFG operations.
---
 compiler/include/concretelang/Support/CompilerEngine.h | 4 ++++
 compiler/lib/Support/CompilerEngine.cpp                | 4 ++++
 compiler/src/main.cpp                                  | 6 ++++++
 3 files changed, 14 insertions(+)

diff --git a/compiler/include/concretelang/Support/CompilerEngine.h b/compiler/include/concretelang/Support/CompilerEngine.h
index 5e7c77e2f..04038b1e6 100644
--- a/compiler/include/concretelang/Support/CompilerEngine.h
+++ b/compiler/include/concretelang/Support/CompilerEngine.h
@@ -184,6 +184,10 @@ public:
     /// BConcrete operations
     BCONCRETE,
 
+    /// Read sources and lower all FHE, TFHE and Concrete operations to
+    /// BConcrete, then extract SDFG operations
+    SDFG,
+
     /// Read sources and lower all FHE, TFHE and Concrete
     /// operations to canonical MLIR dialects. Cryptographic operations
     /// are lowered to invocations of the concrete library.
diff --git a/compiler/lib/Support/CompilerEngine.cpp b/compiler/lib/Support/CompilerEngine.cpp
index f5b2a6776..3af3dd0a7 100644
--- a/compiler/lib/Support/CompilerEngine.cpp
+++ b/compiler/lib/Support/CompilerEngine.cpp
@@ -405,6 +405,10 @@ CompilerEngine::compile(llvm::SourceMgr &sm, Target target, OptionalLib lib) {
     }
   }
 
+  if (target == Target::SDFG) {
+    return std::move(res);
+  }
+
   // BConcrete -> Canonical dialects
   if (mlir::concretelang::pipeline::lowerBConcreteToStd(mlirContext, module,
                                                         enablePass)
diff --git a/compiler/src/main.cpp b/compiler/src/main.cpp
index bcb726c4a..5b4919db3 100644
--- a/compiler/src/main.cpp
+++ b/compiler/src/main.cpp
@@ -50,6 +50,7 @@ enum Action {
   DUMP_CONCRETE,
   DUMP_CONCRETEWITHLOOPS,
   DUMP_BCONCRETE,
+  DUMP_SDFG,
   DUMP_STD,
   DUMP_LLVM_DIALECT,
   DUMP_LLVM_IR,
@@ -128,6 +129,8 @@ static llvm::cl::opt<enum Action> action(
     llvm::cl::values(
         clEnumValN(Action::DUMP_BCONCRETE, "dump-bconcrete",
                    "Lower to Bufferized Concrete and dump result")),
+    llvm::cl::values(clEnumValN(Action::DUMP_SDFG, "dump-sdfg",
+                                "Lower to SDFG operations annd dump result")),
     llvm::cl::values(clEnumValN(Action::DUMP_STD, "dump-std",
                                 "Lower to std and dump result")),
     llvm::cl::values(clEnumValN(Action::DUMP_LLVM_DIALECT, "dump-llvm-dialect",
@@ -489,6 +492,9 @@ mlir::LogicalResult processInputBuffer(
     case Action::DUMP_BCONCRETE:
       target = mlir::concretelang::CompilerEngine::Target::BCONCRETE;
       break;
+    case Action::DUMP_SDFG:
+      target = mlir::concretelang::CompilerEngine::Target::SDFG;
+      break;
     case Action::DUMP_STD:
       target = mlir::concretelang::CompilerEngine::Target::STD;
       break;