From c8ca318fde000983257957b8f2c9fa5d27dfd821 Mon Sep 17 00:00:00 2001
From: youben11 <ayoub.benaissa@zama.ai>
Date: Thu, 14 Apr 2022 07:42:13 +0100
Subject: [PATCH] feat: throw clear error if DF jit-exec is not supported

was previously getting an undefined symbol error as the runtime lib
didn't support dataflow execution
---
 compiler/include/concretelang/Support/Jit.h | 6 ++++++
 compiler/lib/Support/JITSupport.cpp         | 3 +++
 compiler/lib/Support/Jit.cpp                | 7 +++++++
 3 files changed, 16 insertions(+)

diff --git a/compiler/include/concretelang/Support/Jit.h b/compiler/include/concretelang/Support/Jit.h
index 76613c937..ec9477d90 100644
--- a/compiler/include/concretelang/Support/Jit.h
+++ b/compiler/include/concretelang/Support/Jit.h
@@ -38,6 +38,8 @@ public:
   llvm::Expected<std::unique_ptr<clientlib::PublicResult>>
   call(clientlib::PublicArguments &args);
 
+  void setUseDataflow(bool option) { this->useDataflow = option; }
+
 private:
   /// invokeRaw execute the jit lambda with a list of Argument, the last one is
   /// used to store the result of the computation.
@@ -52,6 +54,10 @@ private:
   mlir::LLVM::LLVMFunctionType type;
   std::string name;
   std::unique_ptr<mlir::ExecutionEngine> engine;
+  // Tell if the DF parallelization was on or during compilation. This will be
+  // useful to abort execution if the runtime doesn't support dataflow
+  // execution, instead of having undefined symbol issues
+  bool useDataflow = false;
 };
 
 } // namespace concretelang
diff --git a/compiler/lib/Support/JITSupport.cpp b/compiler/lib/Support/JITSupport.cpp
index 2bd0f1da6..d7a3c2fa7 100644
--- a/compiler/lib/Support/JITSupport.cpp
+++ b/compiler/lib/Support/JITSupport.cpp
@@ -45,6 +45,9 @@ JITSupport::compile(llvm::SourceMgr &program, CompilationOptions options) {
   }
   auto result = std::make_unique<JitCompilationResult>();
   result->lambda = std::shared_ptr<concretelang::JITLambda>(std::move(*lambda));
+  // Mark the lambda as compiled using DF parallelization
+  result->lambda->setUseDataflow(options.dataflowParallelize ||
+                                 options.autoParallelize);
   result->clientParameters =
       compilationResult.get().clientParameters.getValue();
   return std::move(result);
diff --git a/compiler/lib/Support/Jit.cpp b/compiler/lib/Support/Jit.cpp
index d5027f337..c455331dc 100644
--- a/compiler/lib/Support/Jit.cpp
+++ b/compiler/lib/Support/Jit.cpp
@@ -77,6 +77,13 @@ uint64_t numArgOfRankedMemrefCallingConvention(uint64_t rank) {
 
 llvm::Expected<std::unique_ptr<clientlib::PublicResult>>
 JITLambda::call(clientlib::PublicArguments &args) {
+#ifndef CONCRETELANG_PARALLEL_EXECUTION_ENABLED
+  if (this->useDataflow) {
+    return StreamStringError(
+        "call: current runtime doesn't support dataflow execution, while "
+        "compilation used dataflow parallelization");
+  }
+#endif
   // invokeRaw needs to have pointers on arguments and a pointers on the result
   // as last argument.
   // Prepare the outputs vector to store the output value of the lambda.