set task_topology_max_group to cpu_count

by default. Can be overriden with a flag of the same str

shark/iree_utils/cpu_utils.py

@@ -16,6 +16,7 @@
 
 import subprocess
 import platform
+from shark.parser import shark_args
 
 
 def get_cpu_count():
@@ -44,4 +45,12 @@ def get_iree_cpu_args():
         error_message = f"OS Type f{os_name} not supported and triple can't be determined, open issue to dSHARK team please :)"
         raise Exception(error_message)
     print(f"Target triple found:{target_triple}")
-    return [f"--iree-llvmcpu-target-triple={target_triple}"]
+    cpu_count = (
+        get_cpu_count()
+        if shark_args.task_topology_max_group_count is None
+        else shark_args.task_topology_max_group_count
+    )
+    return [
+        f"--iree-llvmcpu-target-triple={target_triple}",
+        f"--task_topology_max_group_count={cpu_count}",
+    ]

shark/parser.py

@@ -119,5 +119,11 @@
     "to augment the base device allocator",
     choices=["debug", "caching"],
 )
+parser.add_argument(
+    "--task_topology_max_group_count",
+    type=str,
+    default=None,
+    help="passthrough flag for the iree flag of the same name. If None, defaults to cpu-count",
+)
 
 shark_args, unknown = parser.parse_known_args()