MCPcopy Create free account
hub / github.com/ml-explore/mlx-data / next

Method next

mlx/data/stream/Prefetch.cpp:29–56  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

27}
28
29Sample Prefetch::next() const {
30 std::unique_lock lock(mutex_);
31
32 // First time we are called so enqueue all the fetching
33 if (prefetchCache_.size() < prefetchSize_) {
34 for (int i = 0; i < prefetchSize_; i++) {
35 prefetchCache_.emplace(
36 pool_->enqueue([s = stream_] { return s->next(); }));
37 }
38 }
39
40 // We are looping prefetchSize_ times. If all we get is empty then the
41 // underlying stream is indeed exhausted.
42 Sample res;
43 for (int i = 0; i < prefetchSize_; i++) {
44 std::future<Sample> fsample;
45 fsample = std::move(prefetchCache_.front());
46 prefetchCache_.pop();
47 prefetchCache_.emplace(pool_->enqueue([s = stream_] { return s->next(); }));
48 res = fsample.get();
49
50 if (!res.empty()) {
51 break;
52 }
53 }
54
55 return res;
56}
57
58void Prefetch::reset() {
59 std::unique_lock lock(mutex_);

Callers

nothing calls this directly

Calls 3

enqueueMethod · 0.80
sizeMethod · 0.45
getMethod · 0.45

Tested by

no test coverage detected