Program Listing for File rmm_tensor.hpp

Return to documentation for file (morpheus/_lib/include/morpheus/objects/rmm_tensor.hpp)


/* * SPDX-FileCopyrightText: Copyright (c) 2021-2024, 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 * * * * 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 "morpheus/objects/dtype.hpp" // for DType #include "morpheus/objects/memory_descriptor.hpp" #include "morpheus/objects/tensor_object.hpp" #include "morpheus/types.hpp" // for RankType, ShapeType, TensorIndex, TensorSize #include <rmm/device_buffer.hpp> #include <cstdint> #include <memory> #include <vector> namespace morpheus { #pragma GCC visibility push(default) /****** Component public implementations *******************/ /****** RMMTensor****************************************/ class RMMTensor : public ITensor { public: RMMTensor(std::shared_ptr<rmm::device_buffer> device_buffer, TensorSize offset, DType dtype, ShapeType shape, ShapeType stride = {}); ~RMMTensor() override = default; bool is_compact() const final; DType dtype() const override; RankType rank() const final; std::shared_ptr<ITensor> deep_copy() const override; std::shared_ptr<ITensor> reshape(const ShapeType& dims) const override; std::shared_ptr<ITensor> slice(const ShapeType& min_dims, const ShapeType& max_dims) const override; std::shared_ptr<ITensor> copy_rows(const std::vector<RangeType>& selected_rows, TensorIndex num_rows) const override; std::shared_ptr<MemoryDescriptor> get_memory() const override; TensorSize bytes() const final; TensorSize count() const final; TensorIndex shape(TensorIndex idx) const final; TensorIndex stride(TensorIndex idx) const final; void* data() const override; void get_shape(ShapeType& s) const; void get_stride(ShapeType& s) const; intptr_t stream() const override; // Tensor reshape(ShapeType shape) // { // CHECK(is_compact()); // return Tensor(descriptor_shared(), dtype_size(), shape); // } std::shared_ptr<ITensor> as_type(DType dtype) const override; protected: private: TensorSize offset_bytes() const; // Memory info std::shared_ptr<MemoryDescriptor> m_mem_descriptor; std::shared_ptr<rmm::device_buffer> m_md; TensorSize m_offset; // Type info DType m_dtype; // Shape info ShapeType m_shape; ShapeType m_stride; }; #pragma GCC visibility pop// end of group } // namespace morpheus

© Copyright 2024, NVIDIA. Last updated on Apr 25, 2024.