Program Listing for File TensorBatch.hpp
↰ Return to documentation for file (nvcv_types/include/nvcv/TensorBatch.hpp
)
/*
* SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* 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
*
* 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.
*/
#ifndef NVCV_TENSORBATCH_HPP
#define NVCV_TENSORBATCH_HPP
#include "CoreResource.hpp"
#include "TensorBatch.h"
#include "TensorBatchData.hpp"
#include "alloc/Allocator.hpp"
#include <nvcv/Tensor.hpp>
#include <vector>
namespace nvcv {
NVCV_IMPL_SHARED_HANDLE(TensorBatch);
class TensorBatch : public CoreResource<NVCVTensorBatchHandle, TensorBatch>
{
public:
using Base = CoreResource<NVCVTensorBatchHandle, TensorBatch>;
using Requirements = NVCVTensorBatchRequirements;
using HandleType = NVCVTensorBatchHandle;
static Requirements CalcRequirements(int32_t capacity);
NVCV_IMPLEMENT_SHARED_RESOURCE(TensorBatch, Base);
TensorBatch(const Requirements &reqs, const Allocator &alloc = nullptr);
TensorBatch(int32_t capacity, const Allocator &alloc = nullptr);
int32_t capacity() const;
int32_t rank() const;
int32_t numTensors() const;
DataType dtype() const;
TensorLayout layout() const;
NVCVTensorBufferType type() const;
Allocator alloc() const;
template<typename It>
void pushBack(It begin, It end);
void pushBack(const Tensor &tensor);
void popTensors(int32_t numTensors);
void popTensor();
TensorBatchData exportData(CUstream stream);
void clear();
void setUserPointer(void *ptr);
void *getUserPointer() const;
Tensor operator[](int32_t idx) const;
void setTensor(int32_t index, const Tensor &tensor);
class Iterator;
Iterator begin() const;
Iterator end() const;
};
class TensorBatch::Iterator
{
public:
using value_type = Tensor;
using reference = const Tensor &;
using pointer = const Tensor *;
using iterator_category = std::random_access_iterator_tag;
using difference_type = int32_t;
reference operator*() const;
pointer operator->() const;
Iterator operator++(int);
Iterator &operator++();
Iterator operator--(int);
Iterator &operator--();
Iterator operator+(difference_type diff) const;
Iterator operator-(difference_type diff) const;
difference_type operator-(const Iterator &rhs) const;
bool operator==(const Iterator &rhs) const;
bool operator!=(const Iterator &rhs) const;
bool operator<(const Iterator &rhs) const;
bool operator>(const Iterator &rhs) const;
bool operator<=(const Iterator &rhs) const;
bool operator>=(const Iterator &rhs) const;
Iterator(Iterator &other)
: Iterator()
{
*this = other;
}
Iterator(Iterator &&other)
: Iterator()
{
*this = std::move(other);
}
Iterator &operator=(Iterator &other)
{
m_tensorBatch = other.m_tensorBatch;
m_idx = other.m_idx;
m_currentTensor = other.m_currentTensor;
return *this;
}
Iterator &operator=(Iterator &&other)
{
m_tensorBatch = other.m_tensorBatch;
m_idx = other.m_idx;
m_currentTensor = std::move(other.m_currentTensor);
return *this;
}
private:
friend class TensorBatch;
Iterator() = default;
Iterator(const TensorBatch *tensorBatch, int32_t idx)
: m_tensorBatch(tensorBatch)
, m_idx(idx)
, m_currentTensor{}
{
UpdateCurrentTensor();
}
void UpdateCurrentTensor();
const TensorBatch *m_tensorBatch = nullptr;
int32_t m_idx = 0;
mutable Tensor m_currentTensor = {};
};
using TensorBatchWrapHandle = NonOwningResource<TensorBatch>;
} // namespace nvcv
#include "detail/TensorBatchImpl.hpp"
#endif // NVCV_TENSORBATCH_HPP