chore: import upstream snapshot with attribution

paddle/phi/core/operators/reader/CMakeLists.txt

@@ -0,0 +1,5 @@
+file(
+  GLOB reader_cc
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "*.cc")
+collect_srcs(core_srcs SRCS ${reader_cc})

paddle/phi/core/operators/reader/blocking_queue.h

@@ -0,0 +1,195 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <condition_variable>  // NOLINT
+#include <deque>
+#include <utility>
+
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle {
+namespace operators {
+namespace reader {
+
+template <typename T>
+class BlockingQueue {
+  // BlockingQueue is for buffered reading and is supposed to use only the
+  // reader package. It is true that we could and we should have been using
+  // framework::Channel, but which has currently a deadlock bug. BlockingQueue
+  // is a workaround and a simplified version of framework::Channel as it
+  // doesn't support GPU and it implements on buffered blocking queue.
+ public:
+  explicit BlockingQueue(size_t capacity, bool speed_test_mode = false)
+      : capacity_(capacity), speed_test_mode_(speed_test_mode) {
+    PADDLE_ENFORCE_GT(capacity_,
+                      static_cast<size_t>(0),
+                      common::errors::InvalidArgument(
+                          "The capacity of a reader::BlockingQueue must be "
+                          "greater than 0, but received capacity is %d.",
+                          capacity_));
+  }
+
+  bool Send(const T& elem) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    send_cv_.wait(
+        lock, [&] { return queue_.size() < capacity_ || closed_ || killed_; });
+    if (killed_) {
+      // VLOG(3)
+      //     << "WARNING:: Sending an element to a killed
+      //     reader::BlockingQueue";
+      return false;
+    }
+    if (closed_) {
+      // VLOG(5)
+      //     << "WARNING: Sending an element to a closed
+      //     reader::BlockingQueue.";
+      return false;
+    }
+    PADDLE_ENFORCE_LT(
+        queue_.size(),
+        capacity_,
+        common::errors::PermissionDenied(
+            "The queue size cannot exceed the set queue capacity. Expected "
+            "queue size is less than %d. But received %d",
+            capacity_,
+            queue_.size()));
+    queue_.push_back(elem);
+    receive_cv_.notify_one();
+    return true;
+  }
+
+  bool Send(T&& elem) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    send_cv_.wait(
+        lock, [&] { return queue_.size() < capacity_ || closed_ || killed_; });
+    if (killed_) {
+      // VLOG(3)
+      //     << "WARNING:: Sending an element to a killed
+      //     reader::BlockingQueue";
+      return false;
+    }
+    if (closed_) {
+      // VLOG(5)
+      //     << "WARNING: Sending an element to a closed
+      //     reader::BlockingQueue.";
+      return false;
+    }
+    PADDLE_ENFORCE_LT(
+        queue_.size(),
+        capacity_,
+        common::errors::PermissionDenied(
+            "The queue size cannot exceed the set queue capacity. Expected "
+            "queue size is less than %d. But received %d",
+            capacity_,
+            queue_.size()));
+    queue_.emplace_back(std::move(elem));
+    receive_cv_.notify_one();
+    return true;
+  }
+
+  bool Receive(T* elem) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    receive_cv_.wait(lock,
+                     [&] { return !queue_.empty() || closed_ || killed_; });
+    EnforceNotKilled();
+    if (!queue_.empty()) {
+      PADDLE_ENFORCE_NOT_NULL(
+          elem,
+          common::errors::InvalidArgument(
+              "The holder to receive queue data is null pointer."));
+      *elem = queue_.front();
+      if (LIKELY(!speed_test_mode_)) {
+        queue_.pop_front();
+      }
+      send_cv_.notify_one();
+      return true;
+    } else {
+      PADDLE_ENFORCE_EQ(closed_,
+                        true,
+                        common::errors::PermissionDenied(
+                            "Blocking queue status error, if queue is empty "
+                            "when pop data, it should be closed."));
+      // VLOG(3) << "queue is closed! return nothing.";
+      return false;
+    }
+  }
+
+  void ReOpen() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    EnforceNotKilled();
+    // VLOG(1) << "reopen queue";
+    closed_ = false;
+    std::deque<T> new_deque;
+    queue_.swap(new_deque);
+    send_cv_.notify_all();
+    receive_cv_.notify_all();
+  }
+
+  void Close() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    // VLOG(1) << "close queue";
+    closed_ = true;
+    send_cv_.notify_all();
+    receive_cv_.notify_all();
+  }
+
+  bool IsClosed() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return closed_;
+  }
+
+  size_t Cap() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return capacity_;
+  }
+
+  size_t Size() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return queue_.size();
+  }
+
+  void Kill() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    // VLOG(1) << "kill queue";
+    closed_ = true;
+    killed_ = true;
+    send_cv_.notify_all();
+    receive_cv_.notify_all();
+  }
+
+ private:
+  inline void EnforceNotKilled() {
+    PADDLE_ENFORCE_NE(
+        killed_,
+        true,
+        common::errors::Fatal("Blocking queue is killed because the "
+                              "data reader raises an exception."));
+  }
+
+ private:
+  size_t capacity_;
+  bool speed_test_mode_;
+  bool closed_{false};
+  bool killed_{false};  // the queue is broken since exception raises
+  std::deque<T> queue_;
+
+  mutable std::mutex mutex_;
+  mutable std::condition_variable receive_cv_;
+  mutable std::condition_variable send_cv_;
+};
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle

paddle/phi/core/operators/reader/buffered_reader.cc

@@ -0,0 +1,417 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/operators/reader/buffered_reader.h"
+
+#include "paddle/phi/core/framework/convert_utils.h"
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+#include "paddle/phi/core/platform/profiler.h"
+#include "paddle/phi/core/platform/profiler/event_tracing.h"
+
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/backends/device_guard.h"
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+
+#include "glog/logging.h"
+
+namespace paddle::operators::reader {
+BufferedReader::~BufferedReader() {
+  VLOG(1) << "~BufferedReader";
+  reader_->Shutdown();
+  while (!position_.empty()) {
+    auto &front = position_.front();
+    if (front.valid()) {
+      front.wait();
+    }
+    position_.pop();
+  }
+}
+
+BufferedReader::BufferedReader(
+    const std::shared_ptr<framework::ReaderBase> &reader,
+    const Place &place,
+    size_t buffer_size,
+    bool pin_memory)
+    : framework::DecoratedReader(reader),
+      thread_pool_(1),
+      place_(place),
+      buffer_size_(buffer_size),
+      pin_memory_(pin_memory),
+      position_(),
+      cpu_buffer_(),
+      cuda_buffer_(),
+      xpu_buffer_(),
+      custom_device_buffer_() {
+  VLOG(1) << "BufferedReader";
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  if (place_.GetType() == AllocationType::GPU && !pin_memory) {
+    int dev_idx = place_.device;  // NOLINT
+    compute_stream_ =
+        ((phi::GPUContext *)(phi::DeviceContextPool::Instance().Get(place_)))
+            ->stream();
+    events_.resize(buffer_size);
+    for (auto &event : events_) {
+      event = platform::CudaEventResourcePool::Instance().New(dev_idx);
+    }
+    stream_ = platform::CudaStreamResourcePool::Instance().New(dev_idx);
+  }
+#endif
+
+#ifdef PADDLE_WITH_XPU
+  if (place_.GetType() == AllocationType::XPU) {
+    int dev_idx = place_.device;
+    compute_stream_ =
+        ((phi::XPUContext *)(phi::DeviceContextPool::Instance().Get(place_)))
+            ->stream();
+    events_.resize(buffer_size);
+    for (auto &event : events_) {
+      event = platform::XpuEventResourcePool::Instance().New(dev_idx);
+    }
+    stream_ = platform::XpuStreamResourcePool::Instance().New(dev_idx);
+  }
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (place_.GetType() == AllocationType::CUSTOM) {
+    auto stream =
+        ((phi::CustomContext *)(phi::DeviceContextPool::Instance().Get(place_)))
+            ->stream();
+    custom_device_compute_stream_ =
+        std::make_shared<phi::stream::Stream>(place_, stream);
+
+    custom_device_events_.resize(buffer_size);
+    for (auto &event : custom_device_events_) {
+      event = std::make_shared<phi::event::Event>();
+      event->Init(place_);
+    }
+    custom_device_stream_ = std::make_shared<phi::stream::Stream>();
+    custom_device_stream_->Init(place_);
+  }
+#endif
+
+  cpu_buffer_.resize(buffer_size);
+  cuda_buffer_.resize(buffer_size);
+  xpu_buffer_.resize(buffer_size);
+  custom_device_buffer_.resize(buffer_size);
+  ReadTillBufferFullAsync();
+}
+
+void BufferedReader::ReadTillBufferFullAsync() {
+  for (size_t i = 0; i < buffer_size_; ++i) {
+    ReadAsync(i);
+  }
+}
+
+void BufferedReader::ReadAsync(size_t i) {
+  position_.emplace(thread_pool_.enqueue([this, i]() -> size_t {
+    TensorVec &cpu = cpu_buffer_[i];
+    reader_->ReadNext(&cpu);
+
+    if (cpu.empty()) {
+      return -1UL;
+    }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)  // @{ Group GPU Place
+    if (place_.GetType() == AllocationType::GPU) {
+      auto dev_ctx_gpu =
+          (phi::GPUContext *)(phi::DeviceContextPool::Instance().Get(place_));
+      TensorVec &cuda = cuda_buffer_[i];
+      if (cuda.empty()) {
+        cuda.resize(cpu.size());
+      } else {
+        PADDLE_ENFORCE_EQ(
+            cuda.size(),
+            cpu.size(),
+            common::errors::InvalidArgument(
+                "Input tensor number on GPU and CPU devices are not matched."));
+      }
+      if (pin_memory_) {
+        // NOTE: [Copy processing of different input devices]
+        // We may accept input tensor in three different devices:
+        //   - CPUPlace
+        //   - CUDAPinnedPlace
+        //   - CUDAPlace
+        // CUDA Stream Synchronizing is slow, in order to avoid Synchronizing
+        // in BufferedReader thread, we do data copy as follows:
+        //   - If src Tensor on CPU memory, we copy it to CUDAPinned memory
+        //   - IF src Tensor on CUDAPinned memory, we use it directly
+        //   - IF src Tensor on CUDA memory, we use it directly
+        phi::GPUPinnedPlace cuda_pinned_place;
+        std::vector<void *> cuda_pinned_ptrs;
+        cuda_pinned_ptrs.reserve(cpu.size());
+        phi::RecordEvent record_event(
+            "BufferedReader:MemoryCopy", phi::TracerEventType::UserDefined, 1);
+        // NODE(chenweihang): When we use CUDAPinned Memory, we need call
+        // cudaHostAlloc, that is a CUDA API, calling CUDA API need load
+        // cuda lib into device, it will cost hundreds of MB of GPU memory.
+        // If we don't set Device here, which will use CUDAPlace(0) default.
+        platform::SetDeviceId(place_.device);
+        for (size_t i = 0; i < cpu.size(); ++i) {
+          if (cpu[i].place().GetType() == AllocationType::CPU) {
+            cuda[i].Resize(cpu[i].dims());
+            cuda[i].set_layout(cpu[i].layout());
+            cuda_pinned_ptrs[i] =
+                dev_ctx_gpu->Alloc(&cuda[i],
+                                   cpu[i].type(),
+                                   0,
+                                   true);  // pinned=true to use GPUPinnedPlace
+            auto size = cpu[i].numel() * phi::SizeOf(cpu[i].dtype());
+
+            phi::memory_utils::Copy(cuda_pinned_place,
+                                    cuda_pinned_ptrs[i],
+                                    cpu[i].place(),
+                                    cpu[i].data(),
+                                    size);
+
+            cuda[i].set_lod(cpu[i].lod());
+          } else {
+            // Here the cpu[i]'s place may be CUDAPlace, CUDAPinnedPlace, or
+            // others, we don't copy the memory of it to CUDAPinnedPlace, but
+            // we should share tensor data to cuda[i]
+            cuda[i].ShareDataWith(cpu[i]);
+          }
+        }
+      } else {
+        // NOTE(liangdun): using async copy instead of TensorCopySync
+        // TensorCopySync would block other stream, because TensorCopySync
+        // issues the copying command to the default stream, it will make two
+        // commands from different streams cannot run concurrently.
+        std::vector<void *> gpu_ptrs;
+        gpu_ptrs.reserve(cpu.size());
+        for (size_t i = 0; i < cpu.size(); ++i) {
+          cuda[i].Resize(cpu[i].dims());
+          cuda[i].set_layout(cpu[i].layout());
+          gpu_ptrs.emplace_back(dev_ctx_gpu->Alloc(&cuda[i], cpu[i].type()));
+        }
+
+        // NOTE(zjl): cudaStreamWaitEvent() must be called after all
+        // Alloc(cuda[i]) is called, since some ops release
+        // cuda memory immediately without waiting cuda kernel ends
+        platform::SetDeviceId(place_.device);
+#ifdef PADDLE_WITH_HIP
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            hipEventRecord(events_[i].get(), compute_stream_));
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            hipStreamWaitEvent(stream_.get(), events_[i].get(), 0));
+#else
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            cudaEventRecord(events_[i].get(), compute_stream_));
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            cudaStreamWaitEvent(stream_.get(), events_[i].get(), 0));
+#endif
+
+        phi::RecordEvent record_event(
+            "BufferedReader:MemoryCopy", phi::TracerEventType::UserDefined, 1);
+        for (size_t i = 0; i < cpu.size(); ++i) {
+          auto cpu_place = cpu[i].place();
+          auto cpu_ptr = cpu[i].data();
+          auto gpu_ptr = gpu_ptrs[i];
+          auto size = cpu[i].numel() * phi::SizeOf(cpu[i].dtype());
+          if (cpu_place.GetType() == AllocationType::GPUPINNED ||
+              cpu_place.GetType() == AllocationType::GPU) {
+            phi::memory_utils::Copy(
+                place_, gpu_ptr, cpu_place, cpu_ptr, size, stream_.get());
+          } else {
+            phi::GPUPinnedPlace cuda_pinned_place;
+            DenseTensor cuda_pinned_tensor;
+            cuda_pinned_tensor.Resize(cpu[i].dims());
+            auto cuda_pinned_ptr =
+                dev_ctx_gpu->Alloc(&cuda_pinned_tensor,
+                                   cpu[i].type(),
+                                   0,
+                                   true);  // pinned=true to use GPUPinnedPlace
+            phi::memory_utils::Copy(
+                cuda_pinned_place, cuda_pinned_ptr, cpu_place, cpu_ptr, size);
+            phi::memory_utils::Copy(place_,
+                                    gpu_ptr,
+                                    cuda_pinned_place,
+                                    cuda_pinned_ptr,
+                                    size,
+                                    stream_.get());
+
+            phi::backends::gpu::GpuStreamSync(stream_.get());
+          }
+          cuda[i].set_lod(cpu[i].lod());
+        }
+        phi::backends::gpu::GpuStreamSync(stream_.get());
+      }
+    }
+#endif
+
+#ifdef PADDLE_WITH_XPU
+    if (place_.GetType() == AllocationType::XPU) {
+      auto dev_ctx_xpu =
+          (phi::XPUContext *)(phi::DeviceContextPool::Instance().Get(place_));
+      TensorVec &xpu = xpu_buffer_[i];
+      if (xpu.empty()) {
+        xpu.resize(cpu.size());
+      } else {
+        PADDLE_ENFORCE_EQ(
+            xpu.size(),
+            cpu.size(),
+            common::errors::InvalidArgument(
+                "Input tensor number on XPU and CPU devices are not matched. "
+                "The number on XPU is %d, on CPU is %d",
+                xpu.size(),
+                cpu.size()));
+      }
+
+      std::vector<void *> xpu_ptrs;
+      xpu_ptrs.reserve(cpu.size());
+      for (size_t i = 0; i < cpu.size(); ++i) {
+        xpu[i].Resize(cpu[i].dims());
+        xpu[i].set_layout(cpu[i].layout());
+        xpu_ptrs.emplace_back(dev_ctx_xpu->Alloc(&xpu[i], cpu[i].type()));
+      }
+
+      phi::backends::xpu::XPUDeviceGuard guard(place_.device);
+      int r = xpu_event_record(events_[i].get(), compute_stream_);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_event_record");
+      r = xpu_stream_wait_event(stream_.get(), events_[i].get());
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_stream_wait_event");
+
+      phi::RecordEvent record_event(
+          "BufferedReader:MemoryCopy", phi::TracerEventType::UserDefined, 1);
+      for (size_t i = 0; i < cpu.size(); ++i) {
+        auto cpu_place = cpu[i].place();
+        auto cpu_ptr = cpu[i].data();
+        auto xpu_ptr = xpu_ptrs[i];
+        auto size = cpu[i].numel() * phi::SizeOf(cpu[i].dtype());
+        // TODO(zhanghuan) for now hardware not support xpu_memcpy_async, maybe
+        // KL3
+        if ((cpu_place.GetType() == AllocationType::XPU)) {
+          platform::XPUStreamSync(stream_.get());
+          char *tmp = new char[size];
+          PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(
+              tmp, cpu_ptr, size, XPUMemcpyKind::XPU_DEVICE_TO_HOST));
+          PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(
+              xpu_ptr, tmp, size, XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+          delete[] tmp;
+        } else {
+          phi::memory_utils::Copy(place_, xpu_ptr, cpu_place, cpu_ptr, size);
+        }
+        xpu[i].set_lod(cpu[i].lod());
+      }
+      platform::XPUStreamSync(stream_.get());
+    }
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+    if (place_.GetType() == AllocationType::CUSTOM) {
+      auto dev_ctx_custom =
+          (phi::CustomContext *)(phi::DeviceContextPool::Instance().Get(
+              place_));
+      phi::DeviceManager::SetDevice(place_);
+
+      TensorVec &custom_device = custom_device_buffer_[i];
+      if (custom_device.empty()) {
+        custom_device.resize(cpu.size());
+      } else {
+        PADDLE_ENFORCE_EQ(custom_device.size(),
+                          cpu.size(),
+                          common::errors::InvalidArgument(
+                              "Input tensor number on CustomDevice and CPU "
+                              "devices are not matched. "
+                              "The number on CustomDevice is %d, on CPU is %d",
+                              custom_device.size(),
+                              cpu.size()));
+      }
+
+      std::vector<void *> custom_device_ptrs;
+      custom_device_ptrs.reserve(cpu.size());
+      for (size_t i = 0; i < cpu.size(); ++i) {
+        custom_device[i].Resize(cpu[i].dims());
+        custom_device[i].set_layout(cpu[i].layout());
+        custom_device_ptrs.emplace_back(
+            dev_ctx_custom->Alloc(&custom_device[i], cpu[i].type()));
+      }
+
+      phi::DeviceManager::GetDeviceWithPlace(place_)->RecordEvent(
+          custom_device_events_[i].get(), custom_device_compute_stream_.get());
+      phi::DeviceManager::GetDeviceWithPlace(place_)->StreamWaitEvent(
+          custom_device_stream_.get(), custom_device_events_[i].get());
+
+      phi::RecordEvent record_event(
+          "BufferedReader:MemoryCopy", phi::TracerEventType::UserDefined, 1);
+      for (size_t i = 0; i < cpu.size(); ++i) {
+        auto cpu_place = cpu[i].place();
+        auto cpu_ptr = cpu[i].data();
+        auto custom_device_ptr = custom_device_ptrs[i];
+        auto size = cpu[i].numel() * phi::SizeOf(cpu[i].dtype());
+        if ((cpu_place.GetType() == AllocationType::CUSTOM)) {
+          phi::memory_utils::Copy(
+              place_, custom_device_ptr, cpu_place, cpu_ptr, size);
+          custom_device_stream_->Synchronize();
+        } else {
+          phi::memory_utils::Copy(
+              place_, custom_device_ptr, cpu_place, cpu_ptr, size);
+        }
+        custom_device[i].set_lod(cpu[i].lod());
+      }
+      custom_device_stream_->Synchronize();
+    }
+#endif
+    return i;
+  }));
+}
+
+void BufferedReader::ShutdownImpl() {
+  VLOG(1) << "ShutdownImpl";
+  reader_->Shutdown();
+  while (!position_.empty()) {
+    position_.pop();
+  }
+  prev_pos_ = -1UL;
+}
+
+void BufferedReader::StartImpl() {
+  reader_->Start();
+  ReadTillBufferFullAsync();
+}
+
+void BufferedReader::ReadNextImpl(phi::TensorArray *out) {
+  if (position_.empty()) {
+    out->clear();
+    return;
+  }
+  size_t i = position_.front().get();
+  position_.pop();
+
+  if (i == -1UL) {
+    ReadNextImpl(out);
+    return;
+  }
+
+  if (place_.GetType() == AllocationType::GPU) {  // NOLINT
+    *out = std::move(cuda_buffer_[i]);
+  } else if (place_.GetType() == AllocationType::XPU) {
+    *out = std::move(xpu_buffer_[i]);
+  } else if (place_.GetType() == AllocationType::CUSTOM) {
+    *out = std::move(custom_device_buffer_[i]);
+  } else {
+    *out = std::move(cpu_buffer_[i]);
+  }
+
+  // Do not push current position into ReadAsync. Push the previous position
+  // Since all computation in fluid are async, change the data of
+  // current position may cause data error.
+  if (prev_pos_ != -1Ul) {
+    ReadAsync(prev_pos_);
+  }
+  prev_pos_ = i;
+}
+
+}  // namespace paddle::operators::reader

paddle/phi/core/operators/reader/buffered_reader.h

@@ -0,0 +1,104 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <list>
+#include <memory>
+#include <queue>
+#include <vector>
+
+#include "ThreadPool.h"
+#include "paddle/phi/core/framework/reader.h"
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_resource_pool.h"
+#endif
+
+#ifdef PADDLE_WITH_XPU
+#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
+#include "paddle/phi/core/platform/device/xpu/xpu_resource_pool.h"
+#endif
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/backends/event.h"
+#include "paddle/phi/backends/stream.h"
+#endif
+namespace paddle {
+namespace operators {
+namespace reader {
+
+class PADDLE_API BufferedReader : public framework::DecoratedReader {
+  using TensorVec = phi::TensorArray;
+  using VecFuture = std::future<TensorVec>;
+
+ public:
+  BufferedReader(const std::shared_ptr<framework::ReaderBase>& reader,
+                 const Place& place,
+                 size_t buffer_size,
+                 bool pin_memory = false);
+
+  ~BufferedReader() override;
+
+  Place GetPlace() const { return place_; }
+
+ private:
+  void ReadTillBufferFullAsync();
+
+  void ReadAsync(size_t i);
+
+ protected:
+  void ShutdownImpl() override;
+  void StartImpl() override;
+  void ReadNextImpl(phi::TensorArray* out) override;
+
+ private:
+  ThreadPool thread_pool_;
+  Place place_;
+  const size_t buffer_size_;
+  bool pin_memory_;
+
+  std::queue<std::future<size_t>> position_;
+
+  // The buffer for reading data.
+  // NOTE: the simplest way to implement buffered reader is do not use any
+  // buffer, just read async and create futures as buffer size. However, to
+  // malloc tensors every time is extremely slow. Here we store all data in
+  // buffers and prevent alloc every time.
+  std::vector<TensorVec> cpu_buffer_;
+  std::vector<TensorVec> cuda_buffer_;
+  std::vector<TensorVec> xpu_buffer_;
+  std::vector<TensorVec> custom_device_buffer_;
+  size_t prev_pos_{-1UL};
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  gpuStream_t compute_stream_;
+  std::shared_ptr<platform::CudaStreamObject> stream_ = nullptr;
+  std::vector<std::shared_ptr<platform::CudaEventObject>> events_{};
+#endif
+
+#ifdef PADDLE_WITH_XPU
+  xpuStream compute_stream_;
+  std::shared_ptr<platform::XpuStreamObject> stream_ = nullptr;
+  std::vector<std::shared_ptr<platform::XpuEventObject>> events_{};
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  std::shared_ptr<phi::stream::Stream> custom_device_compute_stream_ = nullptr;
+  std::shared_ptr<phi::stream::Stream> custom_device_stream_ = nullptr;
+  std::vector<std::shared_ptr<phi::event::Event>> custom_device_events_{};
+#endif
+};
+
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle

paddle/phi/core/operators/reader/dense_tensor_blocking_queue.h

@@ -0,0 +1,253 @@
+//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/operators/reader/blocking_queue.h"
+#include "paddle/phi/core/tensor_array.h"
+
+namespace paddle {
+namespace operators {
+namespace reader {
+
+class DenseTensorBlockingQueue {
+ public:
+  explicit DenseTensorBlockingQueue(size_t capacity,
+                                    bool speed_test_mode = false)
+      : queue_(capacity, speed_test_mode) {}
+
+  ~DenseTensorBlockingQueue() {
+    // VLOG(10) << "Destruct DenseTensorBlockingQueue";
+  }
+
+  bool Push(const phi::TensorArray& lod_tensor_vec) {
+    return queue_.Send(lod_tensor_vec);
+  }
+
+  bool Push(phi::TensorArray&& lod_tensor_vec) {
+    return queue_.Send(std::move(lod_tensor_vec));
+  }
+
+  phi::TensorArray Pop(bool* ok = nullptr) {
+    phi::TensorArray lod_tensor_vec;
+    bool success = queue_.Receive(&lod_tensor_vec);
+    if (ok != nullptr) *ok = success;
+    return lod_tensor_vec;
+  }
+
+  inline size_t Cap() const { return queue_.Cap(); }
+
+  inline size_t Size() const { return queue_.Size(); }
+
+  inline void ReOpen() { queue_.ReOpen(); }
+
+  inline void Close() {
+    // VLOG(1) << "DenseTensorBlockingQueue close";
+    queue_.Close();
+  }
+
+  inline bool IsClosed() const { return queue_.IsClosed(); }
+
+  inline void Kill() { queue_.Kill(); }
+
+  inline bool WaitForInited(size_t) { return true; }
+
+ private:
+  BlockingQueue<phi::TensorArray> queue_;
+};
+
+class OrderedMultiDeviceDenseTensorBlockingQueue {
+ public:
+  OrderedMultiDeviceDenseTensorBlockingQueue(size_t capacity,
+                                             bool speed_test_mode = false)
+      : capacity_(capacity), speed_test_mode_(speed_test_mode) {}
+
+  ~OrderedMultiDeviceDenseTensorBlockingQueue() {
+    // VLOG(10) << "Destruct OrderedMultiDeviceDenseTensorBlockingQueue";
+  }
+
+  bool WaitForInited(size_t milliseconds) {
+    std::unique_lock<std::mutex> lock(init_mutex_);
+    return cv_.wait_for(lock, std::chrono::milliseconds(milliseconds), [this] {
+      return !queues_.empty();
+    });
+  }
+
+  void SetDeviceCount(size_t dev_cnt) {
+    {
+      std::lock_guard<std::mutex> lock(init_mutex_);
+      PADDLE_ENFORCE_GE(dev_cnt,
+                        1,
+                        common::errors::InvalidArgument(
+                            "Device count to init "
+                            "OrderedMultiDeviceDenseTensorBlockingQueue"
+                            " must be larger than 1"));
+      if (!queues_.empty()) {
+        PADDLE_ENFORCE_EQ(queues_.size(),
+                          dev_cnt,
+                          common::errors::InvalidArgument(
+                              "queues should be only inited once"));
+        return;
+      }
+
+      // VLOG(1) << "Init queue with size " << dev_cnt;
+      queues_.resize(dev_cnt);
+      for (auto& item : queues_) {
+        auto cap = (capacity_ + dev_cnt - 1) / dev_cnt;
+        item =
+            std::make_unique<DenseTensorBlockingQueue>(cap, speed_test_mode_);
+      }
+    }
+    cv_.notify_all();
+  }
+
+  const std::shared_ptr<DenseTensorBlockingQueue>& GetQueue(size_t idx) const {
+    EnforceIsInited();
+    PADDLE_ENFORCE_LT(
+        idx,
+        queues_.size(),
+        common::errors::OutOfRange("The queue index is out of range"));
+    return queues_[idx];
+  }
+
+  bool Push(const phi::TensorArray& lod_tensor_vec) {
+    return CurQueue()->Push(lod_tensor_vec);
+  }
+
+  inline size_t Size() const {
+    size_t size = 0;
+    for (auto& item : queues_) {
+      size += item->Size();
+    }
+    return size;
+  }
+
+  inline void Close() {
+    for (auto& item : queues_) {
+      item->Close();
+    }
+  }
+
+  inline void Kill() {
+    for (auto& item : queues_) {
+      item->Kill();
+    }
+  }
+
+  inline void Reset() {
+    {
+      std::lock_guard<std::mutex> reset_lock(reset_mutex_);
+      for (auto& method : reset_methods_) {
+        if (method) method();
+      }
+    }
+
+    auto dev_cnt = queues_.size();
+    for (auto& item : queues_) {
+      auto cap = (capacity_ + dev_cnt - 1) / dev_cnt;
+      item = std::make_unique<DenseTensorBlockingQueue>(cap, speed_test_mode_);
+    }
+    data_index_ = 0;
+  }
+
+  inline void SetResetMethod(size_t idx,
+                             const std::function<void()>& reset_method) {
+    std::lock_guard<std::mutex> reset_lock(reset_mutex_);
+    EnforceIsInited();
+    if (reset_methods_.size() <= idx) {
+      reset_methods_.resize(idx + 1);
+    }
+    reset_methods_[idx] = reset_method;
+  }
+
+  inline size_t Cap() const { return capacity_; }
+
+ private:
+  const std::shared_ptr<DenseTensorBlockingQueue>& CurQueue() {
+    return queues_[(data_index_++) % queues_.size()];
+  }
+
+ private:
+  void EnforceIsInited() const {
+    PADDLE_ENFORCE_EQ(queues_.empty(),
+                      false,
+                      common::errors::NotFound("queue has not been inited"));
+  }
+
+ private:
+  std::vector<std::shared_ptr<DenseTensorBlockingQueue>> queues_;
+  mutable uint64_t data_index_{0};
+
+  size_t dev_cnt_{0};
+  const size_t capacity_;
+  const bool speed_test_mode_;
+  bool is_closed_{false};
+
+  std::vector<std::function<void()>> reset_methods_;
+  mutable std::mutex reset_mutex_;
+
+  mutable std::mutex init_mutex_;
+  mutable std::condition_variable cv_;
+};
+
+class DenseTensorBlockingQueueHolder {
+ public:
+  void InitOnce(size_t capacity, bool speed_test_mode = false) {
+    PADDLE_ENFORCE_EQ(
+        queue_,
+        nullptr,
+        common::errors::AlreadyExists("DenseTensorBlockingQueueHolder::"
+                                      "InitOnce() can only be called once"));
+    queue_ =
+        std::make_unique<DenseTensorBlockingQueue>(capacity, speed_test_mode);
+  }
+
+  inline const std::shared_ptr<DenseTensorBlockingQueue>& GetQueue() const {
+    return queue_;
+  }
+
+ private:
+  std::shared_ptr<DenseTensorBlockingQueue> queue_;
+};
+
+class OrderedMultiDeviceDenseTensorBlockingQueueHolder {
+ public:
+  void InitOnce(size_t capacity, bool speed_test_mode = false) {
+    PADDLE_ENFORCE_EQ(queue_,
+                      nullptr,
+                      common::errors::AlreadyExists(
+                          "OrderedMultiDeviceDenseTensorBlockingQueueHolder::"
+                          "InitOnce() can only be called once"));
+    queue_ = std::make_unique<OrderedMultiDeviceDenseTensorBlockingQueue>(
+        capacity, speed_test_mode);
+  }
+
+  inline const std::shared_ptr<OrderedMultiDeviceDenseTensorBlockingQueue>&
+  GetQueue() const {
+    return queue_;
+  }
+
+ private:
+  std::shared_ptr<OrderedMultiDeviceDenseTensorBlockingQueue> queue_;
+};
+
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle

paddle/phi/core/operators/reader/py_reader.cc

@@ -0,0 +1,45 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/operators/reader/py_reader.h"
+
+namespace paddle::operators::reader {
+
+PyReader::PyReader(
+    const std::shared_ptr<DenseTensorBlockingQueue>& queue,
+    const std::vector<phi::DDim>& dims,
+    const std::vector<framework::proto::VarType::Type>& var_types,
+    const std::vector<bool>& need_check_feed)
+    : framework::FileReader(dims, var_types, need_check_feed) {
+  PADDLE_ENFORCE_NOT_NULL(queue,
+                          common::errors::PreconditionNotMet(
+                              "DenseTensorBlockingQueue must not be null."));
+  queue_ = queue;
+}
+
+void PyReader::ReadNext(phi::TensorArray* out) {
+  bool success = false;
+  *out = queue_->Pop(&success);
+  if (!success) out->clear();
+}
+
+PyReader::~PyReader() {  // NOLINT
+  queue_->Close();
+}
+
+void PyReader::Shutdown() { queue_->Close(); }
+
+void PyReader::Start() { queue_->ReOpen(); }
+
+}  // namespace paddle::operators::reader

paddle/phi/core/operators/reader/py_reader.h

@@ -0,0 +1,52 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+#include <memory>
+#include <vector>
+
+#include "paddle/phi/core/framework/reader.h"
+#include "paddle/phi/core/operators/reader/dense_tensor_blocking_queue.h"
+
+namespace paddle {
+namespace operators {
+namespace reader {
+
+class DenseTensorBlockingQueue;
+
+class PADDLE_API PyReader : public framework::FileReader {
+ public:
+  explicit PyReader(
+      const std::shared_ptr<DenseTensorBlockingQueue>& queue,
+      const std::vector<phi::DDim>& dims,
+      const std::vector<framework::proto::VarType::Type>& var_types,
+      const std::vector<bool>& need_check_feed);
+
+  void ReadNext(phi::TensorArray* out) override;
+
+  ~PyReader();
+
+  void Shutdown() override;
+
+  void Start() override;
+
+ private:
+  std::shared_ptr<DenseTensorBlockingQueue> queue_;
+};
+
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle