binary.h 1.71 KB
Newer Older
1
2
#ifndef __INFINIOP_BINARY_H__
#define __INFINIOP_BINARY_H__
3

4
#include "../tensor.h"
5
6
#include <numeric>

7
namespace op::binary {
8
9

// Stores metadata for binary operations on CPU
10
11
12
13
14
struct BinaryInfo {
private:
    BinaryInfo(infiniopTensorDescriptor_t c_desc,
               infiniopTensorDescriptor_t a_desc,
               infiniopTensorDescriptor_t b_desc)
15
16
17
18
19
20
21
22
23
24
25
        : ndim(c_desc->ndim()),
          c_shape(std::move(c_desc->shape())),
          a_shape(std::move(a_desc->shape())),
          b_shape(std::move(b_desc->shape())),
          c_strides(std::move(c_desc->strides())),
          a_strides(std::move(a_desc->strides())),
          b_strides(std::move(b_desc->strides())) {
        this->c_data_size = std::accumulate(c_shape.begin(), c_shape.end(), 1ULL, std::multiplies<size_t>());
        this->broadcasted = (a_strides != c_strides) || (b_strides != c_strides);
    }

26
27
28
29
30
31
32
33
34
35
public:
    size_t c_data_size;
    size_t ndim;
    bool broadcasted;
    std::vector<size_t> c_shape;
    std::vector<size_t> a_shape;
    std::vector<size_t> b_shape;
    std::vector<ptrdiff_t> c_strides;
    std::vector<ptrdiff_t> a_strides;
    std::vector<ptrdiff_t> b_strides;
36

37
38
39
40
41
42
43
    static infiniStatus_t create(
        BinaryInfo **instance,
        infiniopTensorDescriptor_t c_desc,
        infiniopTensorDescriptor_t a_desc,
        infiniopTensorDescriptor_t b_desc) {
        if (!c_desc || !a_desc || !b_desc) {
            return INFINI_STATUS_BAD_PARAM;
44
45
        }

46
47
48
49
50
        try {
            *instance = new BinaryInfo(c_desc, a_desc, b_desc);
            return INFINI_STATUS_SUCCESS;
        } catch (const std::exception &e) {
            return INFINI_STATUS_INTERNAL_ERROR;
51
52
        }
    }
53
54
};
} // namespace op::binary
55

56
#endif // __INFINIOP_BINARY_H__