From a76023aa55f34602d3d1677bd5228d6b6f0332b0 Mon Sep 17 00:00:00 2001
From: Antoniu Pop <antoniu.pop@zama.ai>
Date: Wed, 24 Apr 2024 15:39:12 +0100
Subject: [PATCH] fix(compiler): [GPU backend] add SDFG_MAX_BATCH_SIZE
 environment variable to control batching limit.

---
 .../include/concretelang/Support/CompilerEngine.h        | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Support/CompilerEngine.h b/compilers/concrete-compiler/compiler/include/concretelang/Support/CompilerEngine.h
index e1fd0772ee..4100a43141 100644
--- a/compilers/concrete-compiler/compiler/include/concretelang/Support/CompilerEngine.h
+++ b/compilers/concrete-compiler/compiler/include/concretelang/Support/CompilerEngine.h
@@ -135,11 +135,16 @@ struct CompilationOptions {
       loopParallelize = true;
       break;
     case Backend::GPU:
+      loopParallelize = true;
       batchTFHEOps = true;
       emitGPUOps = true;
       emitSDFGOps = true;
-      if (maxBatchSize == std::numeric_limits<int64_t>::max())
-        maxBatchSize = 1 << 16;
+      char *env = getenv("SDFG_MAX_BATCH_SIZE");
+      if (env != nullptr) {
+        int64_t targetMax = strtoul(env, NULL, 10);
+        if (targetMax > 0)
+          maxBatchSize = targetMax;
+      }
       break;
     }
   }