[Options] Add option for controlling parallel build with number of jobs or memory reserved for each job

python/hidet/driver.py

@@ -278,13 +278,14 @@ def build_ir_module_batch(
         ]
         build_results = []
         if parallel:
+            num_jobs, mem_for_worker = option.get_parallel_tune()
+            num_jobs = os.cpu_count() if num_jobs == -1 else num_jobs
+            mem_for_worker *= 1024**3
             # Set the affinity of current process. Some package such as numpy will change affinity of current process,
             # which might limit the parallelism of compilation.
-            os.sched_setaffinity(0, range(os.cpu_count()))
+            os.sched_setaffinity(0, range(num_jobs))
 
-            # the maximum number of processes is limited by the number of cores and memory
-            mem_for_worker = 1.5 * 1024 * 1024 * 1024  # 1.5 GiB
-            num_workers = min(max(int(psutil.virtual_memory().available // mem_for_worker), 1), psutil.cpu_count())
+            num_workers = min(max(int(psutil.virtual_memory().available // mem_for_worker), 1), num_jobs)
 
             _lazy_initialize_cuda()
             with multiprocessing.Pool(processes=num_workers) as pool:

python/hidet/option.py

@@ -99,6 +99,12 @@ def register_hidet_options():
         description='Whether to build operators in parallel.',
         choices=[True, False],
     )
+    register_option(
+        name='parallel_tune',
+        type_hint='int, float',
+        default_value=(-1, 1.5),
+        description='The number of parallel jobs and memory reserved for each job',
+    )
     register_option(
         name='save_lower_ir',
         type_hint='bool',
@@ -476,6 +482,33 @@ def get_parallel_build() -> bool:
     return OptionContext.current().get_option('parallel_build')
 
 
+def parallel_tune(num_jobs: int = -1, mem_gb_per_job: float = 1.5):
+    """
+    Whether to build operators in parallel.
+
+    Parameters
+    ----------
+    num_jobs: int
+        The maximum number of parallel jobs allowed, default -1 (decided by the system)
+    mem_gb_per_job: float
+        The minimum amount of memory (in GiB) reserved for each tuning job, default 1.5GiB.
+    """
+    OptionContext.current().set_option('parallel_tune', (num_jobs, mem_gb_per_job))
+
+
+def get_parallel_tune() -> Tuple[int, float]:
+    """
+    Get the option value of whether to build operators in parallel.
+
+    Returns
+    -------
+    ret: Tuple[int, float]
+        Get the maximum number of jobs and minumum amount of memory reserved for tuning.
+
+    """
+    return OptionContext.current().get_option('parallel_tune')
+
+
 def save_lower_ir(enabled: bool = True):
     """
     Whether to save the lower IR.