array_arith.cc 7.31 KB
Newer Older
sangwzh's avatar
sangwzh committed
1
// !!! This is a file automatically generated by hipify!!!
2
/**
3
 *  Copyright (c) 2019 by Contributors
4
5
 * @file array/array_aritch.cc
 * @brief DGL array arithmetic operations
6
7
8
 */
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/container.h>
9
10
#include <dgl/runtime/ndarray.h>

11
#include "../c_api_common.h"
sangwzh's avatar
sangwzh committed
12
13
#include "arith.h"
#include "array_op.h"
14
15
16
17
18
19
20

using namespace dgl::runtime;

namespace dgl {
namespace aten {

// Generate operators with both operations being NDArrays.
21
22
23
24
25
26
27
28
29
30
31
#define BINARY_ELEMENT_OP(name, op)                                  \
  IdArray name(IdArray lhs, IdArray rhs) {                           \
    IdArray ret;                                                     \
    CHECK_SAME_DTYPE(lhs, rhs);                                      \
    CHECK_SAME_CONTEXT(lhs, rhs);                                    \
    ATEN_XPU_SWITCH_CUDA(lhs->ctx.device_type, XPU, #name, {         \
      ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {                      \
        ret = impl::BinaryElewise<XPU, IdType, arith::op>(lhs, rhs); \
      });                                                            \
    });                                                              \
    return ret;                                                      \
32
33
34
  }

// Generate operators with only lhs being NDArray.
35
36
37
38
39
40
41
42
43
#define BINARY_ELEMENT_OP_L(name, op)                                \
  IdArray name(IdArray lhs, int64_t rhs) {                           \
    IdArray ret;                                                     \
    ATEN_XPU_SWITCH_CUDA(lhs->ctx.device_type, XPU, #name, {         \
      ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {                      \
        ret = impl::BinaryElewise<XPU, IdType, arith::op>(lhs, rhs); \
      });                                                            \
    });                                                              \
    return ret;                                                      \
44
45
46
  }

// Generate operators with only lhs being NDArray.
47
48
49
50
51
52
53
54
55
#define BINARY_ELEMENT_OP_R(name, op)                                \
  IdArray name(int64_t lhs, IdArray rhs) {                           \
    IdArray ret;                                                     \
    ATEN_XPU_SWITCH_CUDA(rhs->ctx.device_type, XPU, #name, {         \
      ATEN_ID_TYPE_SWITCH(rhs->dtype, IdType, {                      \
        ret = impl::BinaryElewise<XPU, IdType, arith::op>(lhs, rhs); \
      });                                                            \
    });                                                              \
    return ret;                                                      \
56
57
58
  }

// Generate operators with only lhs being NDArray.
59
60
61
62
63
64
65
66
67
#define UNARY_ELEMENT_OP(name, op)                             \
  IdArray name(IdArray lhs) {                                  \
    IdArray ret;                                               \
    ATEN_XPU_SWITCH_CUDA(lhs->ctx.device_type, XPU, #name, {   \
      ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {                \
        ret = impl::UnaryElewise<XPU, IdType, arith::op>(lhs); \
      });                                                      \
    });                                                        \
    return ret;                                                \
68
69
70
71
72
73
  }

BINARY_ELEMENT_OP(Add, Add)
BINARY_ELEMENT_OP(Sub, Sub)
BINARY_ELEMENT_OP(Mul, Mul)
BINARY_ELEMENT_OP(Div, Div)
74
BINARY_ELEMENT_OP(Mod, Mod)
75
76
77
78
79
80
81
82
83
84
85
BINARY_ELEMENT_OP(GT, GT)
BINARY_ELEMENT_OP(LT, LT)
BINARY_ELEMENT_OP(GE, GE)
BINARY_ELEMENT_OP(LE, LE)
BINARY_ELEMENT_OP(EQ, EQ)
BINARY_ELEMENT_OP(NE, NE)

BINARY_ELEMENT_OP_L(Add, Add)
BINARY_ELEMENT_OP_L(Sub, Sub)
BINARY_ELEMENT_OP_L(Mul, Mul)
BINARY_ELEMENT_OP_L(Div, Div)
86
BINARY_ELEMENT_OP_L(Mod, Mod)
87
88
89
90
91
92
93
94
95
96
97
BINARY_ELEMENT_OP_L(GT, GT)
BINARY_ELEMENT_OP_L(LT, LT)
BINARY_ELEMENT_OP_L(GE, GE)
BINARY_ELEMENT_OP_L(LE, LE)
BINARY_ELEMENT_OP_L(EQ, EQ)
BINARY_ELEMENT_OP_L(NE, NE)

BINARY_ELEMENT_OP_R(Add, Add)
BINARY_ELEMENT_OP_R(Sub, Sub)
BINARY_ELEMENT_OP_R(Mul, Mul)
BINARY_ELEMENT_OP_R(Div, Div)
98
BINARY_ELEMENT_OP_R(Mod, Mod)
99
100
101
102
103
104
105
106
107
108
109
110
111
BINARY_ELEMENT_OP_R(GT, GT)
BINARY_ELEMENT_OP_R(LT, LT)
BINARY_ELEMENT_OP_R(GE, GE)
BINARY_ELEMENT_OP_R(LE, LE)
BINARY_ELEMENT_OP_R(EQ, EQ)
BINARY_ELEMENT_OP_R(NE, NE)

UNARY_ELEMENT_OP(Neg, Neg)

}  // namespace aten
}  // namespace dgl

///////////////// Operator overloading for NDArray /////////////////
112
NDArray operator+(const NDArray& lhs, const NDArray& rhs) {
113
114
  return dgl::aten::Add(lhs, rhs);
}
115
NDArray operator-(const NDArray& lhs, const NDArray& rhs) {
116
117
  return dgl::aten::Sub(lhs, rhs);
}
118
NDArray operator*(const NDArray& lhs, const NDArray& rhs) {
119
120
  return dgl::aten::Mul(lhs, rhs);
}
121
NDArray operator/(const NDArray& lhs, const NDArray& rhs) {
122
123
  return dgl::aten::Div(lhs, rhs);
}
124
NDArray operator%(const NDArray& lhs, const NDArray& rhs) {
125
126
  return dgl::aten::Mod(lhs, rhs);
}
127
NDArray operator+(const NDArray& lhs, int64_t rhs) {
128
129
  return dgl::aten::Add(lhs, rhs);
}
130
NDArray operator-(const NDArray& lhs, int64_t rhs) {
131
132
  return dgl::aten::Sub(lhs, rhs);
}
133
NDArray operator*(const NDArray& lhs, int64_t rhs) {
134
135
  return dgl::aten::Mul(lhs, rhs);
}
136
NDArray operator/(const NDArray& lhs, int64_t rhs) {
137
138
  return dgl::aten::Div(lhs, rhs);
}
139
NDArray operator%(const NDArray& lhs, int64_t rhs) {
140
141
  return dgl::aten::Mod(lhs, rhs);
}
142
NDArray operator+(int64_t lhs, const NDArray& rhs) {
143
144
  return dgl::aten::Add(lhs, rhs);
}
145
NDArray operator-(int64_t lhs, const NDArray& rhs) {
146
147
  return dgl::aten::Sub(lhs, rhs);
}
148
NDArray operator*(int64_t lhs, const NDArray& rhs) {
149
150
  return dgl::aten::Mul(lhs, rhs);
}
151
NDArray operator/(int64_t lhs, const NDArray& rhs) {
152
153
  return dgl::aten::Div(lhs, rhs);
}
154
NDArray operator%(int64_t lhs, const NDArray& rhs) {
155
156
  return dgl::aten::Mod(lhs, rhs);
}
157
NDArray operator-(const NDArray& array) { return dgl::aten::Neg(array); }
158

159
NDArray operator>(const NDArray& lhs, const NDArray& rhs) {
160
161
  return dgl::aten::GT(lhs, rhs);
}
162
NDArray operator<(const NDArray& lhs, const NDArray& rhs) {
163
164
  return dgl::aten::LT(lhs, rhs);
}
165
NDArray operator>=(const NDArray& lhs, const NDArray& rhs) {
166
167
  return dgl::aten::GE(lhs, rhs);
}
168
NDArray operator<=(const NDArray& lhs, const NDArray& rhs) {
169
170
  return dgl::aten::LE(lhs, rhs);
}
171
NDArray operator==(const NDArray& lhs, const NDArray& rhs) {
172
173
  return dgl::aten::EQ(lhs, rhs);
}
174
NDArray operator!=(const NDArray& lhs, const NDArray& rhs) {
175
176
  return dgl::aten::NE(lhs, rhs);
}
177
NDArray operator>(const NDArray& lhs, int64_t rhs) {
178
179
  return dgl::aten::GT(lhs, rhs);
}
180
NDArray operator<(const NDArray& lhs, int64_t rhs) {
181
182
  return dgl::aten::LT(lhs, rhs);
}
183
NDArray operator>=(const NDArray& lhs, int64_t rhs) {
184
185
  return dgl::aten::GE(lhs, rhs);
}
186
NDArray operator<=(const NDArray& lhs, int64_t rhs) {
187
188
  return dgl::aten::LE(lhs, rhs);
}
189
NDArray operator==(const NDArray& lhs, int64_t rhs) {
190
191
  return dgl::aten::EQ(lhs, rhs);
}
192
NDArray operator!=(const NDArray& lhs, int64_t rhs) {
193
194
  return dgl::aten::NE(lhs, rhs);
}
195
NDArray operator>(int64_t lhs, const NDArray& rhs) {
196
197
  return dgl::aten::GT(lhs, rhs);
}
198
NDArray operator<(int64_t lhs, const NDArray& rhs) {
199
200
  return dgl::aten::LT(lhs, rhs);
}
201
NDArray operator>=(int64_t lhs, const NDArray& rhs) {
202
203
  return dgl::aten::GE(lhs, rhs);
}
204
NDArray operator<=(int64_t lhs, const NDArray& rhs) {
205
206
  return dgl::aten::LE(lhs, rhs);
}
207
NDArray operator==(int64_t lhs, const NDArray& rhs) {
208
209
  return dgl::aten::EQ(lhs, rhs);
}
210
NDArray operator!=(int64_t lhs, const NDArray& rhs) {
211
212
  return dgl::aten::NE(lhs, rhs);
}