llama.cpp/kompute/src/OpTensorCopy.cpp

91 lines
2.9 KiB
C++

// SPDX-License-Identifier: Apache-2.0
/**
* Copyright (c) 2023 Nomic, Inc. All rights reserved.
*
* This software is licensed under the terms of the Software for Open Models License (SOM),
* version 1.0, as detailed in the LICENSE_SOM.txt file. A copy of this license should accompany
* this software. Except as expressly granted in the SOM license, all rights are reserved by Nomic, Inc.
*/
#include "kompute/operations/OpTensorCopy.hpp"
#include "kompute/Tensor.hpp"
namespace kp {
OpTensorCopy::OpTensorCopy(const std::vector<std::shared_ptr<Tensor>>& tensors)
{
KP_LOG_DEBUG("Kompute OpTensorCopy constructor with params");
this->mTensors = tensors;
if (this->mTensors.size() < 2) {
throw std::runtime_error(
"Kompute OpTensorCopy called with less than 2 tensor");
}
kp::Tensor::TensorDataTypes dataType = this->mTensors[0]->dataType();
uint32_t size = this->mTensors[0]->size();
for (const std::shared_ptr<Tensor>& tensor : tensors) {
if (tensor->dataType() != dataType) {
throw std::runtime_error(fmt::format(
"Attempting to copy tensors of different types from {} to {}",
Tensor::toString(dataType),
Tensor::toString(tensor->dataType())));
}
if (tensor->size() != size) {
throw std::runtime_error(fmt::format(
"Attempting to copy tensors of different sizes from {} to {}",
size,
tensor->size()));
}
}
}
OpTensorCopy::~OpTensorCopy()
{
KP_LOG_DEBUG("Kompute OpTensorCopy destructor started");
}
void
OpTensorCopy::record(const vk::CommandBuffer& commandBuffer)
{
KP_LOG_DEBUG("Kompute OpTensorCopy record called");
// We iterate from the second tensor onwards and record a copy to all
for (size_t i = 1; i < this->mTensors.size(); i++) {
this->mTensors[i]->recordCopyFrom(commandBuffer, this->mTensors[0]);
}
}
void
OpTensorCopy::preEval(const vk::CommandBuffer& /*commandBuffer*/)
{
KP_LOG_DEBUG("Kompute OpTensorCopy preEval called");
}
void
OpTensorCopy::postEval(const vk::CommandBuffer& /*commandBuffer*/)
{
KP_LOG_DEBUG("Kompute OpTensorCopy postEval called");
// Do not copy on CPU side if source is storage tensor
if (this->mTensors[0]->tensorType() == kp::Tensor::TensorTypes::eStorage)
{
KP_LOG_DEBUG("Kompute OpTensorCopy not copying tensor source given it's of eStorage type");
return;
}
void* data = this->mTensors[0]->rawData();
// Copy the data from the first tensor into all the tensors
for (size_t i = 1; i < this->mTensors.size(); i++) {
if (this->mTensors[i]->tensorType() == kp::Tensor::TensorTypes::eStorage) {
KP_LOG_DEBUG("Kompute OpTensorCopy not copying to tensor dest given it's of eStorage type");
continue;
}
this->mTensors[i]->setRawData(data);
}
}
}