Done for nvls

Azure · Nov 5, 2024 · c569822 · c569822
1 parent 6ac918b
commit c569822
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/examples/mscclang/mscclpp/allreduce_nvls.py b/examples/mscclang/mscclpp/allreduce_nvls.py
@@ -37,7 +37,7 @@ def allreduce_allpairs(gpus, instances):
             c.group_store(ngbrs, sendtb=0)
 
         Json()
-        # Check()
+        Check()
 
 
 parser = argparse.ArgumentParser()

diff --git a/msccl/language/mscclpp/__init__.py b/msccl/language/mscclpp/__init__.py
@@ -132,7 +132,7 @@ def lower(self):
             self.instr_dag.optimize()
         self.instr_dag.lower_pt1(self.instances)
         gpu_prgms = self.instr_dag.lower_pt2(self.instances, self.replication_policy)
-        return Program(
+        program = Program(
             self.name,
             self.collective.name,
             self.collective.inplace,
@@ -142,6 +142,11 @@ def lower(self):
             self.num_threads_per_block,
             self.use_double_scratch_buffer,
         )
+        for gpu in program.gpus:
+            gpu.input_chunks = len(self.buffers[gpu.rank][Buffer.input]) * self.instances
+            if not self.collective.inplace:
+                gpu.output_chunks = len(self.buffers[gpu.rank][Buffer.output]) * self.instances
+        return program
 
     def generate_json(self):
         return ir_to_json(self.lower())

diff --git a/msccl/language/mscclpp/ir.py b/msccl/language/mscclpp/ir.py
@@ -201,7 +201,7 @@ def remove_empty_fields(d):
         gpu_instance["channels"] = list(filter(lambda x: x["type"] != "none", gpu_instance["channels"]))
         gpu_instance["channels"] = sorted(gpu_instance["channels"], key=lambda x: (x["srcbuff"], x["dstbuff"]))
 
-        # render GPU NVLS channels
+        # render for GPU NVLS channels
         for i, chan in enumerate(gpu_instance["channels"]):
             if chan["type"] == "nvls":
                 buff = chan["srcbuff"]