Skip to content

Commit

Permalink
readd skip wait if callback is not present
Browse files Browse the repository at this point in the history
  • Loading branch information
MS-Renan committed May 3, 2024
1 parent 9c8e88d commit 90dc2a6
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 22 deletions.
1 change: 1 addition & 0 deletions AnyBuildLogs/latest.txt
@@ -0,0 +1 @@
20231019-111207-d314f8bf
55 changes: 33 additions & 22 deletions src/pq_flash_index.cpp
Expand Up @@ -1253,37 +1253,48 @@ int PQFlashIndex<T, LabelT>::load_from_separate_paths(uint32_t num_threads, cons
}

#ifdef USE_BING_INFRA
bool getNextCompletedRequest(const IOContext &ctx, size_t size, int &completedIndex)
bool getNextCompletedRequest(std::shared_ptr<AlignedFileReader> &reader, IOContext &ctx, size_t size,
int &completedIndex)
{
bool waitsRemaining = false;
long completeCount = ctx.m_completeCount;
do
if ((*ctx.m_pRequests)[0].m_callback)
{
for (int i = 0; i < size; i++)
bool waitsRemaining = false;
long completeCount = ctx.m_completeCount;
do
{
auto ithStatus = (*ctx.m_pRequestsStatus)[i];
if (ithStatus == IOContext::Status::READ_SUCCESS)
for (int i = 0; i < size; i++)
{
completedIndex = i;
return true;
auto ithStatus = (*ctx.m_pRequestsStatus)[i];
if (ithStatus == IOContext::Status::READ_SUCCESS)
{
completedIndex = i;
return true;
}
else if (ithStatus == IOContext::Status::READ_WAIT)
{
waitsRemaining = true;
}
}
else if (ithStatus == IOContext::Status::READ_WAIT)

// if we didn't find one in READ_SUCCESS, wait for one to complete.
if (waitsRemaining)
{
waitsRemaining = true;
WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), 100);
// this assumes the knowledge of the reader behavior (implicit
// contract). need better factoring?
}
}
} while (waitsRemaining);

// if we didn't find one in READ_SUCCESS, wait for one to complete.
if (waitsRemaining)
{
WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), 100);
// this assumes the knowledge of the reader behavior (implicit
// contract). need better factoring?
}
} while (waitsRemaining);

completedIndex = -1;
return false;
completedIndex = -1;
return false;
}
else
{
reader->wait(ctx, completedIndex);
return completedIndex != -1;
}
}
#endif

Expand Down Expand Up @@ -1671,7 +1682,7 @@ void PQFlashIndex<T, LabelT>::cached_beam_search(const T *query1, const uint64_t
long requestCount = static_cast<long>(frontier_read_reqs.size());
// If we issued read requests and if a read is complete or there are
// reads in wait state, then enter the while loop.
while (requestCount > 0 && getNextCompletedRequest(ctx, requestCount, completedIndex))
while (requestCount > 0 && getNextCompletedRequest(reader, ctx, requestCount, completedIndex))
{
assert(completedIndex >= 0);
auto &frontier_nhood = frontier_nhoods[completedIndex];
Expand Down

0 comments on commit 90dc2a6

Please sign in to comment.