/*
Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
*/
#ifndef NVIDIA_GXF_STD_TENSOR_COPIER_HPP_
#define NVIDIA_GXF_STD_TENSOR_COPIER_HPP_

#include "gxf/std/allocator.hpp"
#include "gxf/std/codelet.hpp"
#include "gxf/std/receiver.hpp"
#include "gxf/std/transmitter.hpp"

namespace nvidia {
namespace gxf {

// Codelet that copies tensors from host to device or from device to host
class TensorCopier : public Codelet {
 public:
  gxf_result_t registerInterface(Registrar* registrar) override;
  gxf_result_t initialize() override { return GXF_SUCCESS; }
  gxf_result_t deinitialize() override { return GXF_SUCCESS; }
  gxf_result_t tick() override;

 private:
  Parameter<Handle<Receiver>> receiver_;
  Parameter<Handle<Transmitter>> transmitter_;
  Parameter<Handle<Allocator>> allocator_;
  Parameter<int32_t> mode_;
};

}  // namespace gxf
}  // namespace nvidia

#endif  // NVIDIA_GXF_STD_TENSOR_COPIER_HPP_
