global daphne context

daphne-eu · Jul 22, 2024 · e9d977d · e9d977d
1 parent 233e1a5
commit e9d977d
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 25 deletions.
diff --git a/src/api/cli/DaphneUserConfig.h b/src/api/cli/DaphneUserConfig.h
@@ -19,10 +19,12 @@
 
 #include <api/daphnelib/DaphneLibResult.h>
 #include <compiler/catalog/KernelCatalog.h>
-#include <runtime/local/vectorized/LoadPartitioningDefs.h>
+#include <runtime/local/context/IContext.h>
 #include <runtime/local/datastructures/IAllocationDescriptor.h>
+#include <runtime/local/vectorized/LoadPartitioningDefs.h>
 #include <util/LogConfig.h>
 #include <util/DaphneLogger.h>
+
 class DaphneLogger;
 
 #include <vector>
@@ -117,6 +119,8 @@ struct DaphneUserConfig {
 
     KernelCatalog kernelCatalog;
 
+    IContext* dctx_ptr{};
+
     /**
      * @brief Replaces the prefix `"{exedir}/"` in the field `libdir` by the path
      * of the directory in which the currently running executable resides.

diff --git a/src/compiler/lowering/InsertDaphneContextPass.cpp b/src/compiler/lowering/InsertDaphneContextPass.cpp
@@ -37,7 +37,11 @@ using namespace mlir;
 struct InsertDaphneContextPass : public PassWrapper<InsertDaphneContextPass, OperationPass<func::FuncOp>>
 {
     const DaphneUserConfig& user_config;
-    explicit InsertDaphneContextPass(const DaphneUserConfig& cfg) : user_config(cfg) {}
+    std::shared_ptr<spdlog::logger> logger;
+
+    explicit InsertDaphneContextPass(const DaphneUserConfig& cfg) : user_config(cfg) {
+        logger = spdlog::get("compiler");
+    }
     void runOnOperation() final;
 };
 
@@ -51,21 +55,18 @@ void InsertDaphneContextPass::runOnOperation()
 
     // Insert a CreateDaphneContextOp as the first operation in the block.
     builder.create<daphne::CreateDaphneContextOp>(
-        loc, daphne::DaphneContextType::get(&getContext()),
-        builder.create<daphne::ConstantOp>(
-            loc, reinterpret_cast<uint64_t>(&user_config)),
-        builder.create<daphne::ConstantOp>(
-            loc,
-            reinterpret_cast<uint64_t>(&KernelDispatchMapping::instance())),
+            loc, daphne::DaphneContextType::get(&getContext()),
+            builder.create<daphne::ConstantOp>(loc, reinterpret_cast<uint64_t>(&user_config)),
+            builder.create<daphne::ConstantOp>(loc, reinterpret_cast<uint64_t>(&KernelDispatchMapping::instance())),
         builder.create<daphne::ConstantOp>(
             loc,
             reinterpret_cast<uint64_t>(&Statistics::instance())));
 #ifdef USE_CUDA
-    if(user_config.use_cuda) {
+    if (user_config.use_cuda) {
         builder.create<daphne::CreateCUDAContextOp>(loc);
     }
 #endif
-    if (user_config.use_distributed){
+    if (user_config.use_distributed) {
         builder.create<daphne::CreateDistributedContextOp>(loc);
     }
 #ifdef USE_FPGAOPENCL
@@ -74,13 +75,16 @@ void InsertDaphneContextPass::runOnOperation()
     }
 #endif
 
-
-    // Insert a DestroyDaphneContextOp as the last operation in the block, but
-    // before the block's terminator.
-    builder.setInsertionPoint(b.getTerminator());
-    builder.create<daphne::DestroyDaphneContextOp>(loc);
+    // only destroy the DAPHNE context at the end of the main function
+    if(f.getName().str() == "main") {
+        // Insert a DestroyDaphneContextOp as the last operation in the block, but
+        // before the block's terminator.
+        builder.setInsertionPoint(b.getTerminator());
+        builder.create<daphne::DestroyDaphneContextOp>(loc);
+    }
 }
 
+
 std::unique_ptr<Pass> daphne::createInsertDaphneContextPass(const DaphneUserConfig& cfg)
 {
     return std::make_unique<InsertDaphneContextPass>(cfg);

diff --git a/src/runtime/local/kernels/CUDA/CreateCUDAContext.cpp b/src/runtime/local/kernels/CUDA/CreateCUDAContext.cpp
@@ -4,8 +4,11 @@
 namespace CUDA {
     void createCUDAContext(DCTX(ctx)) {
         // ToDo: one context per device
-        if(ctx->getUserConfig().log_ptr)
-            ctx->getUserConfig().log_ptr->registerLoggers();
-        ctx->cuda_contexts.emplace_back(CUDAContext::createCudaContext(0));
+     auto dctx = reinterpret_cast<DaphneContext*>(ctx->getUserConfig().dctx_ptr);
+        if(dctx && dctx->cuda_contexts.empty()) {
+            if (ctx->getUserConfig().log_ptr)
+                ctx->getUserConfig().log_ptr->registerLoggers();
+            ctx->cuda_contexts.emplace_back(CUDAContext::createCudaContext(0));
+        }
     }
 }
diff --git a/src/runtime/local/kernels/CreateDaphneContext.cpp b/src/runtime/local/kernels/CreateDaphneContext.cpp
@@ -20,10 +20,15 @@
 void createDaphneContext(DaphneContext *&res, uint64_t configPtr,
                          uint64_t dispatchMappingPtr, uint64_t statisticsPtr) {
     auto config = reinterpret_cast<DaphneUserConfig *>(configPtr);
-    auto dispatchMapping =
-        reinterpret_cast<KernelDispatchMapping *>(dispatchMappingPtr);
-    auto statistics = reinterpret_cast<Statistics *>(statisticsPtr);
-    if (config->log_ptr != nullptr)
-        config->log_ptr->registerLoggers();
-    res = new DaphneContext(*config, *dispatchMapping, *statistics);
+        if(config->dctx_ptr) {
+        res = reinterpret_cast<DaphneContext*>(config->dctx_ptr);
+    }
+    else {
+        auto dispatchMapping =
+            reinterpret_cast<KernelDispatchMapping *>(dispatchMappingPtr);
+        auto statistics = reinterpret_cast<Statistics *>(statisticsPtr);
+        if (config->log_ptr != nullptr)
+            config->log_ptr->registerLoggers();
+        res = new DaphneContext(*config, *dispatchMapping, *statistics);
+    }
 }
diff --git a/src/runtime/local/kernels/CreateDaphneContext.h b/src/runtime/local/kernels/CreateDaphneContext.h
@@ -16,8 +16,8 @@
 
 #pragma once
 
-#include <api/cli/DaphneUserConfig.h>
 #include <runtime/local/context/DaphneContext.h>
+#include <api/cli/DaphneUserConfig.h>
 
 #include <cstdint>