add new fields

msccl/language/mscclpp/__init__.py

@@ -36,6 +36,8 @@ def __init__(
         replication_policy: ReplicationPolicy = ReplicationPolicy.duplicated,
         num_threads_per_block: int = 1024,
         use_double_scratch_buffer: bool = False,
+        min_message_size: int = -1,
+        max_message_size: int = -1 ,
     ):
         self.name = name
         self.topo = topo
@@ -47,6 +49,8 @@ def __init__(
         self.replication_policy = replication_policy
         self.num_threads_per_block = num_threads_per_block
         self.use_double_scratch_buffer = use_double_scratch_buffer
+        self.min_message_size = min_message_size
+        self.max_message_size = max_message_size
         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
@@ -147,6 +151,8 @@ def lower(self):
             self.collective.num_chunk_groups * self.instances,
             self.num_threads_per_block,
             self.use_double_scratch_buffer,
+            self.min_message_size,
+            self.max_message_size,
         )
         for gpu in program.gpus:
             gpu.input_chunks = len(self.buffers[gpu.rank][Buffer.input]) * self.instances

msccl/language/mscclpp/ir.py

@@ -377,5 +377,8 @@ def remove_empty_fields(d):
         "gpus": gpus,
         "num_threads_per_block": program.num_threads_per_block,
         "use_double_scratch_buffer": program.use_double_scratch_buffer,
+        "min_message_size": program.min_message_size,
+        "max_message_size": program.max_message_size,
+        "in_place": program.inplace,
     }
     return json.dumps(obj, indent=2)

msccl/language/types.py

@@ -18,7 +18,9 @@ class Program:
     num_chunk_groups: int = 1
     num_threads_per_block: int = 1024
     use_double_scratch_buffer: bool = False
-
+    min_message_size: int = -1
+    max_message_size: int = -1
+
 
 @dataclass
 class Gpu: