Program Listing for File TensorCuda.hpp

↰ Return to documentation for file (pennylane_lightning/core/src/simulators/lightning_tensor/tensor/tncuda/TensorCuda.hpp)

// Copyright 2024 Xanadu Quantum Technologies Inc.

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

#pragma once

#include <complex>
#include <memory>

#include "DataBuffer.hpp"
#include "Error.hpp"
#include "TensorBase.hpp"
#include "cuda_helpers.hpp"

namespace {
namespace cuUtil = Pennylane::LightningGPU::Util;
using namespace Pennylane::LightningGPU;
} // namespace

namespace Pennylane::LightningTensor::TNCuda {

template <class PrecisionT>
class TensorCuda final : public TensorBase<PrecisionT, TensorCuda<PrecisionT>> {
  public:
    using BaseType = TensorBase<PrecisionT, TensorCuda>;
    using CFP_t = decltype(cuUtil::getCudaType(PrecisionT{}));

    explicit TensorCuda(const std::size_t rank,
                        const std::vector<std::size_t> &modes,
                        const std::vector<std::size_t> &extents,
                        const DevTag<int> &dev_tag, bool device_alloc = true)
        : TensorBase<PrecisionT, TensorCuda<PrecisionT>>(rank, modes, extents),
          data_buffer_{std::make_shared<DataBuffer<CFP_t>>(
              BaseType::getLength(), dev_tag, device_alloc)} {}

    explicit TensorCuda(const std::vector<std::size_t> &extents,
                        const std::vector<CFP_t> &host_tensor,
                        const DevTag<int> &dev_tag, bool device_alloc = true)
        : TensorBase<PrecisionT, TensorCuda<PrecisionT>>(extents),
          data_buffer_{std::make_shared<DataBuffer<CFP_t>>(
              BaseType::getLength(), dev_tag, device_alloc)} {
        data_buffer_->CopyHostDataToGpu(host_tensor.data(),
                                        BaseType::getLength());
    }

    TensorCuda() = delete;

    ~TensorCuda() = default;

    inline void CopyGpuDataToHost(std::complex<PrecisionT> *host_tensor,
                                  std::size_t length,
                                  bool async = false) const {
        PL_ABORT_IF_NOT(BaseType::getLength() == length,
                        "Sizes do not match for Host and GPU data");
        data_buffer_->CopyGpuDataToHost(host_tensor, length, async);
    }

    DataBuffer<CFP_t> &getDataBuffer() { return *data_buffer_; }

  private:
    std::shared_ptr<DataBuffer<CFP_t>> data_buffer_;
};
} // namespace Pennylane::LightningTensor::TNCuda

api/program_listing_file_pennylane_lightning_core_src_simulators_lightning_tensor_tensor_tncuda_TensorCuda.hpp

 

Download Python script

 

Download Notebook

 

View on GitHub