WIP

msccl/language/mscclpp/__init__.py

@@ -33,6 +33,7 @@ def __init__(
         protocol: str = "Simple",
         instr_fusion: bool = True,
         replication_policy: ReplicationPolicy = ReplicationPolicy.duplicated,
+        num_threads_per_block: int = 1024,
     ):
         self.name = name
         self.topo = topo
@@ -42,6 +43,7 @@ def __init__(
         self.protocol = protocol
         self.instr_fusion = instr_fusion
         self.replication_policy = replication_policy
+        self.num_threads_per_block = num_threads_per_block
         assert protocol == "Simple" or protocol == "LL", f"Given protocol: {protocol}. Must be either Simple, LL"
         self.run_opt = True  # Runs optimization passes
         # Initialize the input buffers
@@ -134,6 +136,7 @@ def lower(self):
             self.protocol,
             gpu_prgms,
             self.collective.num_chunk_groups * self.instances,
+            self.num_threads_per_block,
         )
 
     def generate_json(self):

msccl/language/mscclpp/ir.py

@@ -310,5 +310,6 @@ def remove_empty_fields(d):
         "protocol": program.protocol,
         "inplace": program.inplace,
         "gpus": gpus,
+        "num_threads_per_block": program.num_threads_per_block,
     }
     return json.dumps(obj, indent=2)

msccl/language/types.py

@@ -16,6 +16,7 @@ class Program:
     protocol: str
     gpus: list = field(default_factory=list)
     num_chunk_groups: int = 1
+    num_threads_per_block: int = 1024
 
 
 @dataclass