1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
|
#include <caffe2/ideep/ideep_utils.h>
#include <caffe2/ideep/operators/operator_fallback_ideep.h>
#include "caffe2/operators/utility_ops.h"
#include "caffe2/operators/elementwise_add_op.h"
using namespace caffe2;
namespace {
class IDEEPSumOp final : public IDEEPOperator {
public:
USE_IDEEP_DEF_ALIASES();
USE_IDEEP_OPERATOR_FUNCTIONS();
using FALLBACK_SUM = IDEEPFallbackOp<SumOp<CPUContext>, SkipIndices<0>>;
using FALLBACK_ADD = IDEEPFallbackOp<BinaryElementwiseOp<
NumericTypes, CPUContext, AddFunctor<CPUContext>>, SkipIndices<0>>;
IDEEPSumOp(const OperatorDef& operator_def, Workspace* ws)
: IDEEPOperator(operator_def, ws),
fallback_sum_(operator_def, ws),
fallback_add_(operator_def, ws) {}
// NOLINTNEXTLINE(modernize-use-equals-default)
~IDEEPSumOp() override {}
bool RunOnDevice() override {
itensor::dims input_dims;
bool fallback_to_cpu = false;
vector<itensor> inputs_itensor;
// We only support element-wise sum for ideep tensors here.
// If a CPU tensor is detected in input list, we have to fallback
// to corresponding CPU operator.
for (int i = 0; i < InputSize(); ++i) {
if (OperatorBase::InputBlob(i).template IsType<itensor>()) {
auto& tensor_ideep = Input(i);
if (input_dims.empty()) {
input_dims = tensor_ideep.get_dims();
} else if (input_dims != tensor_ideep.get_dims()) {
fallback_to_cpu = true;
break;
}
inputs_itensor.emplace_back(tensor_ideep);
} else {
CAFFE_ENFORCE(
BlobIsTensorType(OperatorBase::InputBlob(i), CPU),
"Expect cpu tensor if not itensor");
fallback_to_cpu = true;
break;
}
}
if (!fallback_to_cpu) {
auto* Y = Output(OUTPUT);
if (InputSize() == 1) {
const auto& X = Input(INPUT0);
ideep::direct_copy::compute(X, *Y);
} else {
const vector<float> scales(InputSize(), 1.0);
ideep::sum::compute(scales, inputs_itensor, *Y);
}
return true;
}
if (InputSize() == 2) {
return fallback_add_.Run(0);
}
return fallback_sum_.Run(0);
}
private:
FALLBACK_SUM fallback_sum_;
FALLBACK_ADD fallback_add_;
INPUT_TAGS(INPUT0);
OUTPUT_TAGS(OUTPUT);
};
REGISTER_IDEEP_OPERATOR(Sum, IDEEPSumOp);
REGISTER_IDEEP_OPERATOR(Add, IDEEPSumOp);
} // namespace
|
