src/modules/data/dataloader.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// 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
//
//      https://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 "tenzr/modules/data/dataloader.h"

#include <cstring>
#include <future>
#include <utility>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/frontend/memory_pool.h"

namespace tenzr::data {

  DataItem DefaultCollate(const std::vector<DataItem>& batch_items) {
    TENZR_CHECK(!batch_items.empty(), "Cannot collate empty batch");
    size_t batch_size = batch_items.size();
    size_t num_tensors = batch_items[0].size();

    DataItem collated;
    collated.reserve(num_tensors);

    for (size_t col = 0; col < num_tensors; ++col) {
      const auto& sample = batch_items[0][col];
      std::vector<int64_t> batch_dims;
      batch_dims.push_back(static_cast<int64_t>(batch_size));

      for (size_t d = 0; d < sample.shape().Rank(); ++d) {
        batch_dims.push_back(sample.shape().Dim(d));
      }
      Shape batch_shape(batch_dims);

      size_t item_bytes = static_cast<size_t>(sample.shape().NumElements()) * SizeOf(sample.Dtype());
      size_t total_bytes = batch_size * item_bytes;

      auto buf = MemoryPool::Instance().Allocate(total_bytes);
      uint8_t* dst_ptr = static_cast<uint8_t*>(buf->Data());

      for (size_t b = 0; b < batch_size; ++b) {
        const auto& item_tensor = batch_items[b][col];
        TENZR_CHECK(item_tensor.HasData(), "Cannot collate tensor without bound data");
        TENZR_CHECK(item_tensor.shape() == sample.shape(), "Shape mismatch across batch items in column " + std::to_string(col));
        TENZR_CHECK(item_tensor.Dtype() == sample.Dtype(), "DType mismatch across batch items in column " + std::to_string(col));

        const uint8_t* src_ptr = static_cast<const uint8_t*>(item_tensor.Data()->Data());
        std::memcpy(dst_ptr + (b * item_bytes), src_ptr, item_bytes);
      }

      auto batch_tensor = Tensor::Parameter(batch_shape, sample.Dtype(), "batch_col_" + std::to_string(col));
      batch_tensor.BindData(buf);
      collated.push_back(std::move(batch_tensor));
    }

    return collated;
  }

  DataLoader::DataLoader(std::shared_ptr<Dataset> dataset, DataLoaderOptions options, CollateFn collate_fn)
      : dataset_(std::move(dataset)), options_(options), collate_fn_(std::move(collate_fn)) {
    TENZR_CHECK(dataset_ != nullptr, "Dataset cannot be null in DataLoader");
    TENZR_CHECK(options_.batch_size > 0, "batch_size must be positive");

    if (!collate_fn_) {
      collate_fn_ = DefaultCollate;
    }

    PrepareBatches();
  }

  DataLoader::~DataLoader() = default;

  DataLoader::DataLoader(DataLoader&&) noexcept = default;
  DataLoader& DataLoader::operator=(DataLoader&&) noexcept = default;

  void DataLoader::PrepareBatches() {
    std::shared_ptr<Sampler> sampler;
    if (options_.shuffle) {
      sampler = std::make_shared<RandomSampler>(dataset_->Size(), options_.seed);
    } else {
      sampler = std::make_shared<SequentialSampler>(dataset_->Size());
    }

    BatchSampler batch_sampler(sampler, options_.batch_size, options_.drop_last);
    batches_ = batch_sampler.Batches();
  }

  size_t DataLoader::Size() const { return batches_.size(); }

  size_t DataLoader::DatasetSize() const { return dataset_->Size(); }

  DataItem DataLoader::FetchBatch(size_t batch_idx) {
    TENZR_CHECK(batch_idx < batches_.size(), "Batch index out of bounds");
    const auto& indices = batches_[batch_idx];

    std::vector<DataItem> items(indices.size());
    if (options_.num_workers > 1 && indices.size() > 1) {
      // Multi-worker parallel fetching
      std::vector<std::future<void>> futures;
      futures.reserve(indices.size());

      for (size_t i = 0; i < indices.size(); ++i) {
        futures.push_back(std::async(std::launch::async, [&items, this, i, idx = indices[i]]() { items[i] = dataset_->Get(idx); }));
      }

      for (auto& fut : futures) {
        fut.get();
      }
    } else {
      // Sequential fetching
      for (size_t i = 0; i < indices.size(); ++i) {
        items[i] = dataset_->Get(indices[i]);
      }
    }

    return collate_fn_(items);
  }

  DataLoader::Iterator DataLoader::begin() { return Iterator(this, 0); }

  DataLoader::Iterator DataLoader::end() { return Iterator(this, batches_.size()); }

  // Iterator implementation
  DataLoader::Iterator::Iterator(DataLoader* loader, size_t batch_idx) : loader_(loader), batch_idx_(batch_idx) {}

  const DataItem& DataLoader::Iterator::operator*() const {
    TENZR_CHECK(loader_ != nullptr && batch_idx_ < loader_->batches_.size(), "Cannot dereference out-of-bounds iterator");
    current_batch_ = loader_->FetchBatch(batch_idx_);
    return current_batch_;
  }

  const DataItem* DataLoader::Iterator::operator->() const {
    operator*();
    return &current_batch_;
  }

  DataLoader::Iterator& DataLoader::Iterator::operator++() {
    batch_idx_++;
    return *this;
  }

  bool DataLoader::Iterator::operator==(const Iterator& other) const { return loader_ == other.loader_ && batch_idx_ == other.batch_idx_; }

  bool DataLoader::Iterator::operator!=(const Iterator& other) const { return !(*this == other); }

}  // namespace tenzr::data