8000 [caffe2] Remove IValue include from operator.h · pytorch/pytorch@7e3bbbc · GitHub
[go: up one dir, main page]
Skip to content

Commit 7e3bbbc

Browse files
peterbell10peterbell10
committed
[caffe2] Remove IValue include from operator.h
ivalue.h includes Tensor.h, so creating a compilation barrier between operator.h and ivalue.h means non-exported caffe2 ops don't need to be rebuilt as often when developing PyTorch. [ghstack-poisoned]
1 parent de0edc0 commit 7e3bbbc

File tree

8 files changed

+299
-152
lines changed

8 files changed

+299
-152
lines changed

‎caffe2/contrib/aten/aten_op_template.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
#include <unordered_map>
33
#include <string>
44
#include <ATen/Functions.h>
5+
#include <ATen/core/List.h>
56
#include <c10/macros/Macros.h>
67
#include <caffe2/core/context.h>
78
#include <caffe2/core/operator.h>
@@ -21,7 +22,7 @@ using at::Half; // for AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, ...)
2122
namespace internal {
2223
TORCH_API at::Tensor index_with_uint8_handling(
2324
const at::Tensor& self,
24-
const torch::List<c10::optional<at::Tensor>>& indices);
25+
const c10::List<c10::optional<at::Tensor>>& indices);
2526
}
2627

2728
template <class Context>

‎caffe2/core/IValueInterface.cc

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
#define TORCH_ASSERT_NO_OPERATORS
2+
#include <caffe2/core/IValueInterface.h>
3+
#undef TORCH_ASSERT_NO_OPERATORS
4+
#include <ATen/core/ivalue.h>
5+
#include <c10/util/irange.h>
6+
7+
#include <memory>
8+
9+
namespace caffe2 {
10+
namespace detail {
11+
12+
caffe2::Tensor contiguous(caffe2::Tensor t) {
13+
return caffe2::Tensor(at::Tensor(std::move(t)).contiguous());
14+
}
15+
16+
17+
IValueInterface::IValueInterface(ArrayRef<c10::IValue> values):
18+
size_(values.size()) {
19+
auto values_tmp = std::make_unique<c10::IValue[]>(size_);
20+
std::copy_n(values.data(), size_, values_tmp.get());
21+
values_ = values_tmp.release();
22+
}
23+
24+
IValueInterface::~IValueInterface() {
25+
delete[] values_;
26+
}
27+
28+
bool IValueInterface::isTensorList(size_t idx) const {
29+
return at(idx).isTensorList();
30+
}
31+
32+
bool IValueInterface::isTensor(size_t idx) const {
33+
return at(idx).isTensor();
34+
}
35+
36+
const c10::IValue& IValueInterface::at(size_t idx) const {
37+
TORCH_INTERNAL_ASSERT(idx < size_, "caffe2: Input index ", idx,
38+
" is out of range for call with ", size_, " arguments");
39+
return values_[idx];
40+
}
41+
42+
std::vector<caffe2::Tensor> IValueInterface::toTensorVector(size_t idx) const {
43+
auto list = at(idx).toTensorList();
44+
std::vector<caffe2::Tensor> ret(list.size());
45+
for (auto i : c10::irange(list.size())) {
46+
ret[i] = caffe2::Tensor(list.get(i));
47+
}
48+
return ret;
49+
}
50+
51+
caffe2::Tensor IValueInterface::toTensor(size_t idx) const {
52+
return caffe2::Tensor(at(idx).toTensor());
53+
}
54+
55+
int IValueInterface::compute_input_size() const {
56+
if (empty()) {
57+
return 0;
58+
}
59+
if (values_[0].isTensorList()) {
60+
// if the first input is a tensor list, we get input tensors by indexing
61+
// into that list. currently, this means that only tensors from that list
62+
// are accessible as inputs. any hypothetical input tensors that come after
63+
// the list are not accessible.
64+
return values_[0].toTensorList().size();
65+
}
66+
// it's not a tensor list. Count the number of tensor inputs and return them.
67+
int num_tensor_inputs = 0;
68+
bool found_nontensor = false;
69+
for (auto i: c10::irange(size())) {
70+
if (values_[i].isTensor()) {
71+
TORCH_INTERNAL_ASSERT(
72+
!found_nontensor,
73+
"All tensor arguments must come before non-tensor arguments");
74+
++num_tensor_inputs;
75+
} else {
76+
found_nontensor = true;
77+
}
78+
}
79+
return num_tensor_inputs;
80+
}
81+
82+
template <typename T>
83+
typename c10::detail::ivalue_to_const_ref_overload_return<T>::type
84+
IValueInterface::to(size_t idx) const {
85+
return at(idx).template to<T>();
86+
}
87+
88+
#define INSTANTIATE_TO(Type) \
89+
template typename c10::detail::ivalue_to_const_ref_overload_return<Type>::type \
90+
IValueInterface::to<Type>(size_t idx) const
91+
92+
INSTANTIATE_TO(at::Tensor);
93+
INSTANTIATE_TO(at::Storage);
94+
INSTANTIATE_TO(c10::Stream);
95+
INSTANTIATE_TO(float);
96+
INSTANTIATE_TO(double);
97+
INSTANTIATE_TO(c10::complex<double>);
98+
INSTANTIATE_TO(unsigned char);
99+
INSTANTIATE_TO(signed char);
100+
INSTANTIATE_TO(unsigned short);
101+
INSTANTIATE_TO(short);
102+
INSTANTIATE_TO(int);
103+
INSTANTIATE_TO(uint32_t);
104+
INSTANTIATE_TO(uint64_t);
105+
INSTANTIATE_TO(int64_t);
106+
INSTANTIATE_TO(bool);
107+
INSTANTIATE_TO(c10::intrusive_ptr<caffe2::Blob>);;
108+
INSTANTIATE_TO(c10::intrusive_ptr<ivalue::ConstantString>);
109+
INSTANTIATE_TO(c10::intrusive_ptr<ivalue::Object>);
110+
INSTANTIATE_TO(at::Scalar);
111+
INSTANTIATE_TO(c10::List<int64_t>);
112+
INSTANTIATE_TO(c10::List<double>);
113+
INSTANTIATE_TO(c10::List<c10::complex<double>>);
114+
INSTANTIATE_TO(c10::List<bool>);
115+
INSTANTIATE_TO(c10::List<at::Tensor>);
116+
INSTANTIATE_TO(c10::impl::GenericList);
117+
INSTANTIATE_TO(c10::impl::GenericDict);
118+
INSTANTIATE_TO(c10::intrusive_ptr<ivalue::Tuple>);
119+
INSTANTIATE_TO(std::string);
120+
INSTANTIATE_TO(c10::string_view);
121+
INSTANTIATE_TO(c10::intrusive_ptr<ivalue::Future>);
122+
INSTANTIATE_TO(c10::intrusive_ptr<c10::RRefInterface>);
123+
INSTANTIATE_TO(c10::intrusive_ptr<at::Quantizer>);
124+
INSTANTIATE_TO(IValue);
125+
INSTANTIATE_TO(c10::Device);
126+
INSTANTIATE_TO(at::ScalarType);
127+
INSTANTIATE_TO(at::Layout);
128+
INSTANTIATE_TO(at::MemoryFormat);
129+
INSTANTIATE_TO(at::QScheme);
130+
INSTANTIATE_TO(at::Dimname);
131+
INSTANTIATE_TO(at::Generator);
132+
133+
template <typename T>
134+
struct list_value_type {
135+
using type = T;
136+
};
137+
138+
template <>
139+
struct list_value_type<int> {
140+
using type = int64_t;
141+
};
142+
143+
template <>
144+
struct list_value_type<int16_t> {
145+
using type = int64_t;
146+
};
147+
148+
template <>
149+
struct list_value_type<float> {
150+
using type = double;
151+
};
152+
153+
template <typename T, typename U>
154+
static std::vector<T> to_vector(const c10::List<U> &list) {
155+
std::vector<T> ret;
156+
for (auto i: c10::irange(list.size())) {
157+
ret.push_back(list.get(i));
158+
}
159+
return ret;
160+
}
161+
162+
template <typename T>
163+
std::vector<T> IValueInterface::toVec(size_t idx) const {
164+
using list_value_t = typename list_value_type<T>::type;
165+
return to_vector<T>(to<c10::List<list_value_t>>(idx));
166+
}
167+
168+
#define INSTANTIATE_TO_VEC(Type) \
169+
template std::vector<Type> IValueInterface::toVec<Type>(size_t idx) const;
170+
171+
INSTANTIATE_TO_VEC(int64_t);
172+
INSTANTIATE_TO_VEC(int32_t);
173+
INSTANTIATE_TO_VEC(int16_t);
174+
INSTANTIATE_TO_VEC(float);
175+
INSTANTIATE_TO_VEC(double);
176+
INSTANTIATE_TO_VEC(bool);
177+
178+
}} // namespace caffe2::detail

‎caffe2/core/IValueInterface.h

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#pragma once
2+
#include <ATen/core/ivalue_to.h>
3+
#include <c10/util/ArrayRef.h>
4+
#include <caffe2/core/tensor.h>
5+
6+
7+
namespace c10 {
8+
class IValue;
9+
}
10+
11+
namespace caffe2 {
12+
namespace detail {
13+
14+
TORCH_API caffe2::Tensor contiguous(caffe2::Tensor);
15+
16+
class TORCH_API IValueInterface {
17+
public:
18+
IValueInterface(): values_(nullptr), size_(0) {}
19+
IValueInterface(ArrayRef<c10::IValue> values);
20+
~IValueInterface();
21+
22+
bool isTensorList(size_t idx) const;
23+
bool isTensor(size_t idx) const;
24+
25+
size_t size() const {
26+
return size_;
27+
}
28+
bool empty() const {
29+
return size() == 0;
30+
}
31+
32+
const c10::IValue& at(size_t idx) const;
33+
std::vector<caffe2::Tensor> toTensorVector(size_t idx) const;
34+
caffe2::Tensor toTensor(size_t idx) const;
35+
36+
int compute_input_size() const;
37+
38+
// These template functions are not defined in the header,
39+
// but are explicitly instantiated in IValueInterface.cc
40+
template <typename T>
41+
typename c10::detail::ivalue_to_const_ref_overload_return<T>::type to(size_t idx) const;
42+
43+
template <typename T>
44+
std::vector<T> toVec(size_t idx) const;
45+
46+
private:
47+
c10::IValue *values_;
48+
size_t size_;
49+
};
50+
51+
}} // namespace caffe2::detail

‎caffe2/core/export_caffe2_op_to_c10.h

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <ATen/core/grad_mode.h>
99
#include <ATen/core/op_registration/op_registration.h>
1010
#include <torch/csrc/jit/frontend/function_schema_parser.h>
11+
#include <caffe2/core/tensor.h>
1112
#include <c10/core/CompileTimeFunctionPointer.h>
1213
#include <torch/library.h>
1314
#include <vector>
@@ -20,16 +21,22 @@ constexpr const char* PREALLOCATED_OUTPUT_ARGNAME =
2021

2122
using _CallCaffe2OpFunc = c10::List<at::Tensor>(
2223
const c10::FunctionSchema& schema,
23-
std::vector<c10::IValue>&& inputs,
24-
c10::List<at::Tensor>&& outputs);
24+
const std::vector<c10::IValue> &inputs,
25+
c10::List<at::Tensor> &&outputs);
2526

2627
template <class Caffe2Operator>
2728
inline c10::List<at::Tensor> _call_caffe2_op(
2829
const c10::FunctionSchema& schema,
29-
std::vector<c10::IValue>&& inputs,
30-
c10::List<at::Tensor>&& outputs) {
31-
Caffe2Operator op(schema, std::move(inputs), std::move(outputs), -1);
30+
const std::vector<c10::IValue> &inputs,
31+
c10::List<at::Tensor> &&outputs) {
32+
c10::SmallVector<caffe2::Tensor, 6> outputs_caffe2(outputs.size());
33+
for (auto i : c10::irange(outputs.size())) {
34+
outputs_caffe2[i] = caffe2::Tensor(outputs.get(i));
35+
}
36+
37+
Caffe2Operator op(schema, inputs, outputs_caffe2, -1);
3238
op.Run(-1);
39+
3340
auto op_outputs = std::move(op).move_output_tensors();
3441
TORCH_INTERNAL_ASSERT(outputs.size() == op_outputs.size());
3542
for (auto i : c10::irange(op_outputs.size())) {

‎caffe2/core/operator.cc

Lines changed: 13 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#define TORCH_ASSERT_NO_OPERATORS
12
#include "caffe2/core/operator.h"
23

34
#include <algorithm>
@@ -15,8 +16,12 @@
1516
#include "caffe2/proto/caffe2_pb.h"
1617
#include "caffe2/utils/proto_utils.h"
1718
#include "caffe2/utils/string_utils.h"
18-
#if !defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
19-
#include <ATen/core/List.h>
19+
20+
#undef TORCH_ASSERT_NO_OPERATORS
21+
22+
#if defined(EXPOSE_C2_OPS) || \
23+
!defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
24+
#include <ATen/core/function_schema.h>
2025
#endif
2126

2227
#include "caffe2/core/export_c10_op_to_caffe2.h"
@@ -86,51 +91,18 @@ OperatorBase::OperatorBase(const OperatorDef& operator_def, Workspace* ws)
8691

8792
#if defined(EXPOSE_C2_OPS) || \
8893
!defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
89-
namespace {
90-
int C10_UNUSED // Suppress unused function warning on mobile.
91-
compute_input_size_(const std::vector<c10::IValue>& inputs) {
92-
if (inputs.empty()) {
93-
return 0;
94-
}
95-
if (inputs[0].isTensorList()) {
96-
// if the first input is a tensor list, we get input tensors by indexing
97-
// into that list. currently, this means that only tensors from that list
98-
// are accessible as inputs. any hypothetical input tensors that come after
99-
// the list are not accessible.
100-
return inputs[0].toTensorVector().size();
101-
}
102-
// it's not a tensor list. Count the number of tensor inputs and return them.
103-
size_t num_tensor_inputs = 0;
104-
bool found_nontensor = false;
105-
for (const auto& input : inputs) {
106-
if (input.isTensor()) {
107-
AT_ASSERTM(
108-
!found_nontensor,
109-
"All tensor arguments must come before non-tensor arguments");
110-
++num_tensor_inputs;
111-
} else {
112-
found_nontensor = true;
113-
}
114-
}
115-
return num_tensor_inputs;
116-
}
117-
} // namespace
11894

11995
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
12096
OperatorBase::OperatorBase(
12197
const c10::FunctionSchema& fn_schema,
122-
std::vector<c10::IValue> inputs,
123-
const c10::List<at::Tensor> &outputs)
98+
c10::ArrayRef<c10::IValue> inputs,
99+
c10::ArrayRef<caffe2::Tensor> outputs)
124100
// NOLINTNEXTLINE(performance-move-const-arg)
125-
: fn_schema_(make_unique<c10::FunctionSchema>(std::move(fn_schema))),
126-
newstyle_inputs_(std::move(inputs)),
127-
input_size_(compute_input_size_(newstyle_inputs_)) {
101+
: fn_schema_(new c10::FunctionSchema(std::move(fn_schema))),
102+
newstyle_inputs_(inputs),
103+
output_tensors_(outputs.vec()),
104+
input_size_(newstyle_inputs_.compute_input_size()) {
128105
input_tensors_.resize(input_size_);
129-
130-
output_tensors_.reserve(outputs_.size());
131-
for (auto i : c10::irange(outputs.size())) {
132-
output_tensors_.emplace_back(outputs.extract(i));
133-
}
134106
}
135107
#endif
136108

0 commit comments

Comments
 (0)
0