tensor.cpp 9.0 KB
Newer Older
1 2 3 4
#include "lite/tensor.h"
#include <set>
#include <string>
#include <unordered_map>
M
Megvii Engine Team 已提交
5 6 7
#include "../../src/tensor_impl_base.h"
#include "common.h"
#include "lite-c/tensor_c.h"
8

M
Megvii Engine Team 已提交
9 10
const LiteLayout default_layout = {
        .shapes = {0, 0, 0, 0, 0}, .ndim = 0, .data_type = LiteDataType::LITE_FLOAT};
11

M
Megvii Engine Team 已提交
12 13 14 15 16
const LiteTensorDesc default_desc = {
        .is_pinned_host = false,
        .layout = default_layout,
        .device_type = LiteDeviceType::LITE_CPU,
        .device_id = 0};
17
namespace {
18 19

static LITE_MUTEX mtx_tensor;
M
Megvii Engine Team 已提交
20
std::unordered_map<void*, std::shared_ptr<lite::Tensor>>& get_global_tensor_holder() {
21
    static std::unordered_map<void*, std::shared_ptr<lite::Tensor>> global_holder;
22 23
    return global_holder;
}
24 25

static LITE_MUTEX mtx_attr;
M
Megvii Engine Team 已提交
26
std::unordered_map<std::string, lite::LiteAny>& get_global_tensor_attr_holder() {
27
    static std::unordered_map<std::string, lite::LiteAny> global_holder;
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
    return global_holder;
}
}  // namespace

//! convert the lite::Layout to Layout
LiteLayout convert_to_clayout(const lite::Layout& layout) {
    LiteLayout clayout;
    clayout.ndim = layout.ndim;
    LITE_ASSERT(layout.ndim < LAYOUT_MAX_DIM, "layout ndim is to large");
    for (size_t i = 0; i < layout.ndim; i++) {
        clayout.shapes[i] = layout.shapes[i];
    }
    clayout.data_type = layout.data_type;
    return clayout;
}

//! convert the C Layout to lite::Layout
lite::Layout convert_to_layout(const LiteLayout& clayout) {
    lite::Layout layout;
    layout.ndim = clayout.ndim;
    LITE_ASSERT(layout.ndim < LAYOUT_MAX_DIM, "clayout ndim is to large");
    for (size_t i = 0; i < layout.ndim; i++) {
        layout.shapes[i] = clayout.shapes[i];
    }
    layout.data_type = clayout.data_type;
    return layout;
}

int LITE_make_tensor(const LiteTensorDesc tensor_describe, LiteTensor* tensor) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE_make_tensor is null");
    lite::Layout layout = convert_to_layout(tensor_describe.layout);
    auto lite_tensor = std::make_shared<lite::Tensor>(
            tensor_describe.device_id, tensor_describe.device_type, layout,
            tensor_describe.is_pinned_host);
63 64 65 66
    {
        LITE_LOCK_GUARD(mtx_tensor);
        get_global_tensor_holder()[lite_tensor.get()] = lite_tensor;
    }
67 68 69 70 71 72 73
    *tensor = lite_tensor.get();
    LITE_CAPI_END();
}

int LITE_destroy_tensor(LiteTensor tensor) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
74
    LITE_LOCK_GUARD(mtx_tensor);
75 76 77 78
    auto& global_holder = get_global_tensor_holder();
    if (global_holder.find(tensor) != global_holder.end()) {
        global_holder.erase(tensor);
    }
79 80 81 82 83 84 85 86 87 88 89
    LITE_CAPI_END();
}

int LITE_set_tensor_layout(LiteTensor tensor, const LiteLayout layout) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    auto tensor_ptr = static_cast<lite::Tensor*>(tensor);
    tensor_ptr->set_layout(convert_to_layout(layout));
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
90 91
int LITE_reset_tensor_memory(
        LiteTensor tensor, void* prepared_data, size_t data_length_in_byte) {
92 93 94
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(prepared_data, "The prepared_data pass to LITE c_api is null");
M
Megvii Engine Team 已提交
95
    static_cast<lite::Tensor*>(tensor)->reset(prepared_data, data_length_in_byte);
96 97 98
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
99
int LITE_reset_tensor(LiteTensor tensor, const LiteLayout layout, void* prepared_data) {
100 101 102
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(prepared_data, "The prepared_data pass to LITE c_api is null");
M
Megvii Engine Team 已提交
103
    static_cast<lite::Tensor*>(tensor)->reset(prepared_data, convert_to_layout(layout));
104 105 106 107 108 109 110 111 112 113 114 115 116 117
    LITE_CAPI_END();
}

int LITE_tensor_reshape(LiteTensor tensor, const int* shape, int size) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor && shape, "The tensor pass to LITE c_api is null");
    std::vector<int> shapes;
    for (int i = 0; i < size; i++) {
        shapes.push_back(shape[i]);
    }
    static_cast<lite::Tensor*>(tensor)->reshape(shapes);
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
118 119 120
int LITE_tensor_slice(
        const LiteTensor tensor, const size_t* start, const size_t* end,
        const size_t* step, size_t size, LiteTensor* slice_tensor) {
121
    LITE_CAPI_BEGIN();
M
Megvii Engine Team 已提交
122 123 124
    LITE_ASSERT(
            tensor && start && end && slice_tensor,
            "The tensor pass to LITE c_api is null");
125 126 127 128 129 130 131 132
    std::vector<size_t> starts, ends, steps;
    for (size_t i = 0; i < size; i++) {
        starts.push_back(start[i]);
        ends.push_back(end[i]);
        if (step) {
            steps.push_back(step[i]);
        }
    }
M
Megvii Engine Team 已提交
133
    auto ret_tensor = static_cast<lite::Tensor*>(tensor)->slice(starts, ends, steps);
134 135 136 137
    {
        LITE_LOCK_GUARD(mtx_tensor);
        get_global_tensor_holder()[ret_tensor.get()] = ret_tensor;
    }
138 139 140 141 142 143 144 145 146 147 148 149 150
    *slice_tensor = ret_tensor.get();
    LITE_CAPI_END();
}

int LITE_tensor_fill_zero(LiteTensor tensor) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    static_cast<lite::Tensor*>(tensor)->fill_zero();
    LITE_CAPI_END();
}

int LITE_tensor_copy(LiteTensor dst_tensor, const LiteTensor src_tensor) {
    LITE_CAPI_BEGIN();
M
Megvii Engine Team 已提交
151
    LITE_ASSERT(dst_tensor && src_tensor, "The tensor pass to LITE c_api is null");
152 153 154 155 156
    static_cast<lite::Tensor*>(dst_tensor)
            ->copy_from(*static_cast<lite::Tensor*>(src_tensor));
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
157
int LITE_tensor_share_memory_with(LiteTensor dst_tensor, const LiteTensor src_tensor) {
158
    LITE_CAPI_BEGIN();
M
Megvii Engine Team 已提交
159
    LITE_ASSERT(dst_tensor && src_tensor, "The tensor pass to LITE c_api is null");
160 161 162 163 164 165 166 167 168 169 170 171 172
    static_cast<lite::Tensor*>(dst_tensor)
            ->share_memory_with(*static_cast<lite::Tensor*>(src_tensor));
    LITE_CAPI_END();
}

int LITE_get_tensor_memory(const LiteTensor tensor, void** data) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(data, "The data ptr pass to LITE c_api is null");
    *data = static_cast<lite::Tensor*>(tensor)->get_memory_ptr();
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
173 174
int LITE_get_tensor_memory_with_index(
        const LiteTensor tensor, const size_t* index, size_t size, void** data) {
175
    LITE_CAPI_BEGIN();
M
Megvii Engine Team 已提交
176
    LITE_ASSERT(tensor && index && data, "The tensor pass to LITE c_api is null");
177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196
    std::vector<size_t> index_v;
    for (size_t i = 0; i < size; i++) {
        index_v.push_back(index[i]);
    }
    *data = static_cast<lite::Tensor*>(tensor)->get_memory_ptr(index_v);
    LITE_CAPI_END();
}

int LITE_get_tensor_total_size_in_byte(const LiteTensor tensor, size_t* size) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(size, "The size ptr pass to LITE c_api is null");
    *size = static_cast<lite::Tensor*>(tensor)->get_tensor_total_size_in_byte();
    LITE_CAPI_END();
}

int LITE_get_tensor_layout(const LiteTensor tensor, LiteLayout* layout) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(layout, "The layout ptr pass to LITE c_api is null");
M
Megvii Engine Team 已提交
197
    *layout = convert_to_clayout(static_cast<lite::Tensor*>(tensor)->get_layout());
198 199 200
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
201
int LITE_get_tensor_device_type(const LiteTensor tensor, LiteDeviceType* device_type) {
202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(device_type, "The device ptr pass to LITE c_api is null");
    *device_type = static_cast<lite::Tensor*>(tensor)->get_device_type();
    LITE_CAPI_END();
}

int LITE_get_tensor_device_id(const LiteTensor tensor, int* device_id) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor && device_id, "The tensor pass to LITE c_api is null");
    *device_id = static_cast<lite::Tensor*>(tensor)->get_device_id();
    LITE_CAPI_END();
}

int LITE_is_pinned_host(const LiteTensor tensor, int* is_pinned_host) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
M
Megvii Engine Team 已提交
219
    LITE_ASSERT(is_pinned_host, "The is_pinned_host ptr pass to LITE c_api is null");
220 221 222 223 224 225 226 227 228 229 230 231
    *is_pinned_host = static_cast<lite::Tensor*>(tensor)->is_pinned_host();
    LITE_CAPI_END();
}

int LITE_is_memory_continue(const LiteTensor tensor, int* is_continue) {
    LITE_CAPI_BEGIN();
    LITE_ASSERT(tensor, "The tensor pass to LITE c_api is null");
    LITE_ASSERT(is_continue, "The is_continue ptr pass to LITE c_api is null");
    *is_continue = static_cast<lite::Tensor*>(tensor)->is_continue_memory();
    LITE_CAPI_END();
}

M
Megvii Engine Team 已提交
232 233 234
int LITE_tensor_concat(
        LiteTensor* tensors, int nr_tensor, int dim, LiteDeviceType dst_device,
        int device_id, LiteTensor* result_tensor) {
235
    LITE_CAPI_BEGIN();
236
    LITE_ASSERT(result_tensor, "The tensor pass to LITE c_api is null");
237 238 239 240
    std::vector<lite::Tensor> v_tensors;
    for (int i = 0; i < nr_tensor; i++) {
        v_tensors.push_back(*static_cast<lite::Tensor*>(tensors[i]));
    }
M
Megvii Engine Team 已提交
241
    auto tensor = lite::TensorUtils::concat(v_tensors, dim, dst_device, device_id);
242 243 244 245
    {
        LITE_LOCK_GUARD(mtx_tensor);
        get_global_tensor_holder()[tensor.get()] = tensor;
    }
246 247 248 249 250
    *result_tensor = tensor.get();
    LITE_CAPI_END()
}

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}