fix(compiler): [GPU backend] add SDFG_MAX_BATCH_SIZE environment vari…

…able to control batching limit.
zama-ai · Apr 25, 2024 · a76023a · a76023a
1 parent 0414c0f
commit a76023a
Showing 1 changed file with 7 additions and 2 deletions.
diff --git 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;
     }
   }