mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-14 11:07:44 +00:00
94 lines
3.0 KiB
C++
94 lines
3.0 KiB
C++
// SPDX-License-Identifier: MIT
|
|
// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
|
|
|
|
#include <numeric>
|
|
#include <cstdlib>
|
|
#include <iostream>
|
|
#include <initializer_list>
|
|
#include <vector>
|
|
#include <gtest/gtest.h>
|
|
|
|
#include "ck/utility/common_header.hpp"
|
|
|
|
#include "ck/wrapper/layout.hpp"
|
|
#include "ck/wrapper/tensor.hpp"
|
|
|
|
// Compare data in tensor with offset with layout.
|
|
// Data and offset should match if physical memory has been initialized with
|
|
// sequentially increasing values.
|
|
template <typename TensorType, typename LayoutType, typename Idxs>
|
|
__host__ __device__ bool TestTensorReadWriteCheckCustom(const TensorType& tensor,
|
|
const LayoutType& layout,
|
|
const std::vector<Idxs> idxs)
|
|
{
|
|
for(size_t i = 0; i < idxs.size(); i++)
|
|
{
|
|
if(tensor(idxs[i]) != layout(idxs[i]))
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
template <typename TensorType, typename LayoutType>
|
|
__host__ __device__ bool TestTensorReadWriteCheck1d(const TensorType& tensor,
|
|
const LayoutType& layout)
|
|
{
|
|
for(ck::index_t w = 0; w < ck::wrapper::size<0>(layout); w++)
|
|
{
|
|
if(tensor(w) != layout(ck::make_tuple(w)))
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
TEST(TestTensorReadWrite, HostMemory)
|
|
{
|
|
constexpr ck::index_t nelems = 8;
|
|
|
|
std::array<ck::index_t, nelems> data;
|
|
std::iota(data.begin(), data.end(), 0);
|
|
const auto layout = ck::wrapper::make_layout(ck::make_tuple(4, 2));
|
|
const auto tensor = ck::wrapper::make_tensor(&data[0], layout);
|
|
|
|
std::vector<ck::Tuple<ck::index_t, ck::index_t>> idxs;
|
|
for(ck::index_t h = 0; h < ck::wrapper::size<0>(layout); h++)
|
|
{
|
|
for(ck::index_t w = 0; w < ck::wrapper::size<1>(layout); w++)
|
|
{
|
|
idxs.emplace_back(h, w);
|
|
}
|
|
}
|
|
|
|
EXPECT_TRUE(TestTensorReadWriteCheck1d(tensor, layout));
|
|
EXPECT_TRUE(TestTensorReadWriteCheckCustom(tensor, layout, idxs));
|
|
}
|
|
|
|
TEST(TestTensorReadWrite, HostMemoryNested)
|
|
{
|
|
constexpr ck::index_t nelems = 8;
|
|
|
|
std::array<ck::index_t, nelems> data;
|
|
std::iota(data.begin(), data.end(), 0);
|
|
const auto layout = ck::wrapper::make_layout(ck::make_tuple(ck::make_tuple(2, 2), 2));
|
|
const auto tensor = ck::wrapper::make_tensor(&data[0], layout);
|
|
|
|
std::vector<ck::Tuple<ck::Tuple<ck::index_t, ck::index_t>, ck::index_t>> idxs;
|
|
for(ck::index_t d = 0; d < ck::wrapper::size<0>(ck::wrapper::get<0>(layout)); d++)
|
|
{
|
|
for(ck::index_t h = 0; h < ck::wrapper::size<1>(ck::wrapper::get<0>(layout)); h++)
|
|
{
|
|
for(ck::index_t w = 0; w < ck::wrapper::size<1>(layout); w++)
|
|
{
|
|
idxs.emplace_back(ck::make_tuple(d, h), w);
|
|
}
|
|
}
|
|
}
|
|
|
|
EXPECT_TRUE(TestTensorReadWriteCheck1d(tensor, layout));
|
|
EXPECT_TRUE(TestTensorReadWriteCheckCustom(tensor, layout, idxs));
|
|
}
|