include/tenzr/modules/data/dataloader.h¶

// 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.

#ifndef TENZR_MODULES_DATA_DATALOADER_H_
#define TENZR_MODULES_DATA_DATALOADER_H_

#include <cstddef>
#include <cstdint>
#include <deque>
#include <functional>
#include <future>
#include <memory>
#include <mutex>
#include <optional>
#include <vector>

#include "tenzr/frontend/tensor.h"
#include "tenzr/modules/data/dataset.h"
#include "tenzr/modules/data/sampler.h"

namespace tenzr::data {

  struct DataLoaderOptions {
    size_t batch_size{32};
    bool shuffle{false};
    std::optional<uint64_t> seed{std::nullopt};
    size_t num_workers{0};
    bool drop_last{false};
    bool pin_memory{true};
    size_t prefetch_factor{2};
  };

  using CollateFn = std::function<DataItem(const std::vector<DataItem>&)>;

  DataItem DefaultCollate(const std::vector<DataItem>& batch_items);

  class DataLoader {
   public:
    DataLoader(std::shared_ptr<Dataset> dataset, DataLoaderOptions options, CollateFn collate_fn = nullptr);
    ~DataLoader();

    DataLoader(const DataLoader&) = delete;
    DataLoader& operator=(const DataLoader&) = delete;
    DataLoader(DataLoader&&) noexcept;
    DataLoader& operator=(DataLoader&&) noexcept;

    class Iterator {
     public:
      using iterator_category = std::forward_iterator_tag;
      using value_type = DataItem;
      using difference_type = std::ptrdiff_t;
      using pointer = const DataItem*;
      using reference = const DataItem&;

      Iterator(DataLoader* loader, size_t batch_idx);

      reference operator*() const;
      pointer operator->() const;
      Iterator& operator++();
      bool operator==(const Iterator& other) const;
      bool operator!=(const Iterator& other) const;

     private:
      DataLoader* loader_{nullptr};
      size_t batch_idx_{0};
      mutable DataItem current_batch_;
    };

    Iterator begin();
    Iterator end();

    size_t Size() const;
    size_t DatasetSize() const;

    const DataLoaderOptions& Options() const { return options_; }

   private:
    friend class Iterator;

    DataItem FetchBatch(size_t batch_idx);
    void PrepareBatches();

    std::shared_ptr<Dataset> dataset_;
    DataLoaderOptions options_;
    CollateFn collate_fn_;
    std::vector<std::vector<size_t>> batches_;
  };

}  // namespace tenzr::data

#endif  // TENZR_MODULES_DATA_DATALOADER_H_