host_tensor.cpp 1.47 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

Chao Liu's avatar
Chao Liu committed
4
#include <cassert>
Chao Liu's avatar
Chao Liu committed
5

6
#include "ck/library/utility/host_tensor.hpp"
Chao Liu's avatar
Chao Liu committed
7

Chao Liu's avatar
Chao Liu committed
8
void HostTensorDescriptor::CalculateStrides()
Chao Liu's avatar
Chao Liu committed
9
10
11
{
    mStrides.clear();
    mStrides.resize(mLens.size(), 0);
Chao Liu's avatar
Chao Liu committed
12
    if(mStrides.empty())
Chao Liu's avatar
Chao Liu committed
13
14
15
16
17
18
19
        return;

    mStrides.back() = 1;
    std::partial_sum(
        mLens.rbegin(), mLens.rend() - 1, mStrides.rbegin() + 1, std::multiplies<std::size_t>());
}

Chao Liu's avatar
Chao Liu committed
20
std::size_t HostTensorDescriptor::GetNumOfDimension() const { return mLens.size(); }
Chao Liu's avatar
Chao Liu committed
21

Chao Liu's avatar
Chao Liu committed
22
std::size_t HostTensorDescriptor::GetElementSize() const
Chao Liu's avatar
Chao Liu committed
23
24
25
26
27
28
{
    assert(mLens.size() == mStrides.size());
    return std::accumulate(
        mLens.begin(), mLens.end(), std::size_t{1}, std::multiplies<std::size_t>());
}

29
std::size_t HostTensorDescriptor::GetElementSpaceSize() const
Chao Liu's avatar
Chao Liu committed
30
{
Chao Liu's avatar
Chao Liu committed
31
    std::size_t space = 1;
32
    for(std::size_t i = 0; i < mLens.size(); ++i)
Chao Liu's avatar
Chao Liu committed
33
34
35
36
    {
        space += (mLens[i] - 1) * mStrides[i];
    }
    return space;
Chao Liu's avatar
Chao Liu committed
37
}
Chao Liu's avatar
Chao Liu committed
38

Chao Liu's avatar
Chao Liu committed
39
const std::vector<std::size_t>& HostTensorDescriptor::GetLengths() const { return mLens; }
Chao Liu's avatar
Chao Liu committed
40

Chao Liu's avatar
Chao Liu committed
41
const std::vector<std::size_t>& HostTensorDescriptor::GetStrides() const { return mStrides; }
42

43
44
45
46
47
48
49
50
51
52
53
54
55
56
std::ostream& operator<<(std::ostream& os, const HostTensorDescriptor& desc)
{
    os << "dim " << desc.GetNumOfDimension() << ", ";

    os << "lengths {";
    LogRange(os, desc.GetLengths(), ", ");
    os << "}, ";

    os << "strides {";
    LogRange(os, desc.GetStrides(), ", ");
    os << "}";

    return os;
}