*: small updates for resource scoping

Signed-off-by: Rohan Yadav <[email protected]>
rohany · Jan 20, 2023 · c097032 · c097032
1 parent 46ce39a
commit c097032
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/cunumeric/runtime.py b/cunumeric/runtime.py
@@ -138,13 +138,7 @@ def __init__(self, legate_context: LegateContext) -> None:
                 ty.int32,
             )
         )
-        self.num_procs = int(
-            self.legate_context.get_tunable(
-                CuNumericTunable.NUM_PROCS,
-                ty.int32,
-            )
-        )
-        self.num_gpus = int(
+        num_gpus = int(
             self.legate_context.get_tunable(
                 CuNumericTunable.NUM_GPUS,
                 ty.int32,
@@ -161,9 +155,17 @@ def __init__(self, legate_context: LegateContext) -> None:
         self.args = parse_library_command_args("cunumeric", ARGS)
         self.args.warning = self.args.warning or self.args.test_mode
 
-        if self.num_gpus > 0 and self.args.preload_cudalibs:
+        if num_gpus > 0 and self.args.preload_cudalibs:
             self._load_cudalibs()
 
+    @property
+    def num_procs(self) -> int:
+        return self.legate_runtime.machine.num_procs
+
+    @property
+    def num_gpus(self) -> int:
+        return self.legate_runtime.num_gpus
+
     def _register_dtypes(self) -> None:
         type_system = self.legate_context.type_system
         for numpy_type, core_type in SUPPORTED_DTYPES.items():

diff --git a/src/cunumeric/random/bitgenerator.h b/src/cunumeric/random/bitgenerator.h
@@ -84,6 +84,9 @@ class BitGeneratorTask : public CuNumericTask<BitGeneratorTask> {
 
  public:
   static void cpu_variant(legate::TaskContext& context);
+#ifdef LEGATE_USE_OPENMP
+  static void omp_variant(legate::TaskContext& context) { BitGeneratorTask::cpu_variant(context); }
+#endif
 #ifdef LEGATE_USE_CUDA
   static void gpu_variant(legate::TaskContext& context);
 #endif

diff --git a/src/cunumeric/set/unique_reduce.h b/src/cunumeric/set/unique_reduce.h
@@ -26,6 +26,9 @@ class UniqueReduceTask : public CuNumericTask<UniqueReduceTask> {
 
  public:
   static void cpu_variant(legate::TaskContext& context);
+#ifdef LEGATE_USE_OPENMP
+  static void omp_variant(legate::TaskContext& context) { UniqueReduceTask::cpu_variant(context); }
+#endif
 };
 
 }  // namespace cunumeric