"...composable_kernel_rocm.git" did not exist on "62996211076a70e7c2fecd3ce0a0ef2f49201236"
tensor.cpp 1.42 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
5
#include <boost/range/adaptor/transformed.hpp>
#include <cassert>

#include "tensor.hpp"

Chao Liu's avatar
Chao Liu committed
6
HostTensorDescriptor::HostTensorDescriptor(std::initializer_list<std::size_t> lens) : mLens(lens)
Chao Liu's avatar
Chao Liu committed
7
8
9
10
{
    this->CalculateStrides();
}

Chao Liu's avatar
Chao Liu committed
11
12
HostTensorDescriptor::HostTensorDescriptor(std::vector<std::size_t> lens,
                                           std::vector<std::size_t> strides)
Chao Liu's avatar
Chao Liu committed
13
    : mLens(lens), mStrides(strides)
Chao Liu's avatar
Chao Liu committed
14
15
16
{
}

Chao Liu's avatar
Chao Liu committed
17
void HostTensorDescriptor::CalculateStrides()
Chao Liu's avatar
Chao Liu committed
18
19
20
{
    mStrides.clear();
    mStrides.resize(mLens.size(), 0);
Chao Liu's avatar
Chao Liu committed
21
    if(mStrides.empty())
Chao Liu's avatar
Chao Liu committed
22
23
24
25
26
27
28
        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
29
std::size_t HostTensorDescriptor::GetNumOfDimension() const { return mLens.size(); }
Chao Liu's avatar
Chao Liu committed
30

Chao Liu's avatar
Chao Liu committed
31
std::size_t HostTensorDescriptor::GetElementSize() const
Chao Liu's avatar
Chao Liu committed
32
33
34
35
36
37
{
    assert(mLens.size() == mStrides.size());
    return std::accumulate(
        mLens.begin(), mLens.end(), std::size_t{1}, std::multiplies<std::size_t>());
}

Chao Liu's avatar
Chao Liu committed
38
std::size_t HostTensorDescriptor::GetElementSpace() const
Chao Liu's avatar
Chao Liu committed
39
{
Chao Liu's avatar
Chao Liu committed
40
    auto ls = mLens | boost::adaptors::transformed([](std::size_t v) { return v - 1; });
Chao Liu's avatar
Chao Liu committed
41
42
    return std::inner_product(ls.begin(), ls.end(), mStrides.begin(), std::size_t{0}) + 1;
}
Chao Liu's avatar
Chao Liu committed
43

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

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