WIp

microsoft · Dec 7, 2024 · 8cb4269 · 8cb4269
1 parent dcda43b
commit 8cb4269
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/executor/execution_plan.cc b/src/executor/execution_plan.cc
@@ -166,7 +166,7 @@ std::vector<BufferType> ExecutionPlan::Impl::getConnectedBufferTypes(int rank) c
   return std::vector<BufferType>(bufferTypes.begin(), bufferTypes.end());
 }
 
-size_t ExecutionPlan::Impl::getScratchBufferSize(int rank, size_t inputSize, size_t outputSize) const {
+size_t ExecutionPlan::Impl::getScratchBufferSize(int rank, size_t inputSize) const {
   size_t sizePerChunk = 0;
   if (this->inputChunks.at(rank) != 0)
     sizePerChunk = inputSize / this->inputChunks.at(rank);
@@ -187,7 +187,7 @@ size_t ExecutionPlan::Impl::getMaxScratchBufferSize(int rank) const {
   else
     throw mscclpp::Error("Input chunks must be greater than 0", mscclpp::ErrorCode::ExecutorError);
 
-  return this->getScratchBufferSize(rank, this->maxMessageSize, sizePerChunk * this->outputChunks.at(rank));
+  return this->getScratchBufferSize(rank, sizePerChunk * this->inputChunks.at(rank));
 }
 
 std::vector<Operation> ExecutionPlan::Impl::getOperations(int rank, int threadblock) const {

diff --git a/src/executor/executor.cc b/src/executor/executor.cc
@@ -169,7 +169,7 @@ struct Executor::Impl {
     ExecutionContext context;
     size_t maxScratchBufferSize = plan.impl_->getMaxScratchBufferSize(rank);
     size_t scratchBufferSize =
-        std::min(plan.impl_->getScratchBufferSize(rank, sendMemRange, recvMemRange), maxScratchBufferSize);
+        std::min(plan.impl_->getScratchBufferSize(rank, sendMemRange), maxScratchBufferSize);
     std::shared_ptr<char> scratchBuffer;
     if (isNvlsSupported()) {
       scratchBuffer = allocSharedPhysicalCuda<char>(scratchBufferSize);