Free srcBuffer and dstBuffer in Distributed::allToAllImpl (#20)

pbelevich · web-flow · commit 9d87d51a1998 · 2025-05-14T11:46:08.000+01:00
diff --git a/csrc/core/distributed.cpp b/csrc/core/distributed.cpp
@@ -27,4 +27,7 @@ void DistributedNVSHMEM::allToAllImpl(const void *input, void *output, size_t si
   nvshmem_quiet();
 
   CUDACHECK(cudaMemcpy(output, dstBuffer, size * count, cudaMemcpyDeviceToHost));
+
+  nvshmem_free(dstBuffer);
+  nvshmem_free(srcBuffer);
 }

Original file line number	Diff line number	Diff line change
`@@ -27,4 +27,7 @@ void DistributedNVSHMEM::allToAllImpl(const void input, void output, size_t si`
`27`	`27`	`nvshmem_quiet();`
`28`	`28`
`29`	`29`	`CUDACHECK(cudaMemcpy(output, dstBuffer, size * count, cudaMemcpyDeviceToHost));`
	`30`	`+`
	`31`	`+ nvshmem_free(dstBuffer);`
	`32`	`+ nvshmem_free(srcBuffer);`
`30`	`33`	`}`