mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-20 14:59:17 +00:00
Add tensor partition and generic copy for ck wrapper (#1108)
* Add tensor partition and generic copy for ck wrapper * Update changelog * Stylistic fixes * Change shape/strides logic to descriptor transforms * Fixes * Fix client example * Fix comments
This commit is contained in:
41
include/ck/wrapper/operations/copy.hpp
Normal file
41
include/ck/wrapper/operations/copy.hpp
Normal file
@@ -0,0 +1,41 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../utils/tensor_utils.hpp"
|
||||
|
||||
namespace ck {
|
||||
namespace wrapper {
|
||||
|
||||
/**
|
||||
* \brief Perform generic copy between two tensors. Tensors must have the
|
||||
* same size.
|
||||
*
|
||||
* \param src_tensor Source tensor.
|
||||
* \param dst_tensor Destination tensor.
|
||||
*/
|
||||
template <typename SrcTensorType, typename DstTensorType>
|
||||
__host__ __device__ void copy(const SrcTensorType& src_tensor, DstTensorType& dst_tensor)
|
||||
{
|
||||
if constexpr(!SrcTensorType::IsDynamicBuffer)
|
||||
{
|
||||
using SizeType = decltype(size(src_tensor));
|
||||
static_for<0, SizeType{}, 1>{}([&](auto i) { dst_tensor(i) = src_tensor(i); });
|
||||
}
|
||||
else if constexpr(!DstTensorType::IsDynamicBuffer)
|
||||
{
|
||||
using SizeType = decltype(size(dst_tensor));
|
||||
static_for<0, SizeType{}, 1>{}([&](auto i) { dst_tensor(i) = src_tensor(i); });
|
||||
}
|
||||
else
|
||||
{
|
||||
for(int i = 0; i < size(src_tensor); i++)
|
||||
{
|
||||
dst_tensor(i) = src_tensor(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace wrapper
|
||||
} // namespace ck
|
||||
Reference in New Issue
Block a user