Skip to content

Commit 24ae9b5

Browse files
smessmerfacebook-github-bot
smessmer
authored and
facebook-github-bot
committed
Fix binary size of OpsAlreadyMovedToC10.cpp (#26237)
Summary: Pull Request resolved: #26237 Calling a lot of `std::string` constructors is horrible for binary size, see t53997334. Using `const char*` instead should make the binary size much smaller. ghstack-source-id: 90145501 Test Plan: size checks on the diff Differential Revision: D17386002 fbshipit-source-id: c5420adf225e535396e806a0df92419a7e2ad3e8
1 parent 976cefd commit 24ae9b5

File tree

5 files changed

+64
-21
lines changed

5 files changed

+64
-21
lines changed

‎aten/src/ATen/core/OpsAlreadyMovedToC10.cpp‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,34 @@
11
#include <ATen/core/OpsAlreadyMovedToC10.h>
22

3+
#include <string>
4+
#include <cstring>
5+
#include <utility>
6+
#include <unordered_set>
7+
#include <ATen/core/operator_name.h>
8+
39
// @generated by aten/src/ATen/gen.py
410

11+
// TODO Once all ATen ops are moved to c10, this file should be removed
12+
513
namespace at {
614

7-
// TODO Once all ATen ops are moved to c10, this file should be removed
15+
namespace {
16+
struct OpNameEquals final {
17+
bool operator()(const std::pair<const char*, const char*>& lhs, const std::pair<const char*, const char*>& rhs) const {
18+
return 0 == strcmp(lhs.first, rhs.first) && 0 == strcmp(lhs.second, rhs.second);
19+
}
20+
};
21+
22+
struct OpNameHash final {
23+
size_t operator()(const std::pair<const char*, const char*>& p) const {
24+
// use std::hash<std::string> because std::hash<const char*> would hash pointers and not pointed-to strings
25+
return std::hash<std::string>()(p.first) ^ (~ std::hash<std::string>()(p.second));
26+
}
27+
};
28+
}
829

9-
const std::unordered_set<c10::OperatorName>& aten_ops_already_moved_to_c10() {
10-
static std::unordered_set<c10::OperatorName> result {
30+
bool aten_op_is_already_moved_to_c10(const c10::OperatorName& opName) {
31+
static std::unordered_set<std::pair<const char*, const char*>, OpNameHash, OpNameEquals> ops {
1132
{"aten::_cast_Byte", ""},
1233
{"aten::_cast_Char", ""},
1334
{"aten::_cast_Double", ""},
@@ -582,11 +603,11 @@ const std::unordered_set<c10::OperatorName>& aten_ops_already_moved_to_c10() {
582603
{"aten::tanh_backward", ""},
583604
{"", ""}
584605
};
585-
return result;
606+
return ops.count(std::make_pair(opName.name.c_str(), opName.overload_name.c_str())) != 0;
586607
}
587608

588-
const std::unordered_set<c10::OperatorName>& aten_ops_not_moved_to_c10_yet() {
589-
static std::unordered_set<c10::OperatorName> result {
609+
bool aten_op_is_not_moved_to_c10_yet(const c10::OperatorName& opName) {
610+
static std::unordered_set<std::pair<const char*, const char*>, OpNameHash, OpNameEquals> ops {
590611
{"aten::backward", ""},
591612
{"aten::set_data", ""},
592613
{"aten::data", ""},
@@ -1366,7 +1387,7 @@ const std::unordered_set<c10::OperatorName>& aten_ops_not_moved_to_c10_yet() {
13661387
{"aten::im2col_backward", ""},
13671388
{"", ""}
13681389
};
1369-
return result;
1390+
return ops.count(std::make_pair(opName.name.c_str(), opName.overload_name.c_str())) != 0;
13701391
}
13711392

13721393
}

‎aten/src/ATen/core/OpsAlreadyMovedToC10.h‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
#pragma once
22

3-
#include <ATen/core/operator_name.h>
43
#include <c10/macros/Export.h>
5-
#include <unordered_set>
4+
5+
namespace c10 {
6+
struct OperatorName;
7+
}
68

79
namespace at {
810

@@ -21,9 +23,9 @@ register_c10_ops.cpp needs. We need to store two sets to be able to make both de
2123
*/
2224

2325
// list of ATen ops that got already moved to the c10 dispatcher
24-
CAFFE2_API const std::unordered_set<c10::OperatorName>& aten_ops_already_moved_to_c10();
26+
CAFFE2_API bool aten_op_is_already_moved_to_c10(const c10::OperatorName& opName);
2527

2628
// list of ATen ops that are still on the globalATenDispatch dispatcher.
27-
CAFFE2_API const std::unordered_set<c10::OperatorName>& aten_ops_not_moved_to_c10_yet();
29+
CAFFE2_API bool aten_op_is_not_moved_to_c10_yet(const c10::OperatorName& opName);
2830

2931
}

‎aten/src/ATen/core/op_registration/op_registration.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -378,7 +378,7 @@ class CAFFE2_API RegisterOperators final {
378378
static bool op_is_still_on_aten_dispatcher_(const char* schema_string) {
379379
// TODO Remove this function once all aten ops are on c10
380380
const auto op_name = parse_operator_name_(schema_string);
381-
return at::aten_ops_not_moved_to_c10_yet().count(op_name) != 0;
381+
return at::aten_op_is_not_moved_to_c10_yet(op_name);
382382
}
383383

384384
Options&& kernel(c10::optional<TensorTypeId>&& dispatch_key, KernelFunction* kernel_func, KernelCacheCreatorFunction&& cache_creator, void* unboxed_kernel_func, std::unique_ptr<FunctionSchema>&& inferred_function_schema) && {

‎aten/src/ATen/templates/OpsAlreadyMovedToC10.cpp‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,46 @@
11
#include <ATen/core/OpsAlreadyMovedToC10.h>
22

3+
#include <string>
4+
#include <cstring>
5+
#include <utility>
6+
#include <unordered_set>
7+
#include <ATen/core/operator_name.h>
8+
39
// ${generated_comment}
410

11+
// TODO Once all ATen ops are moved to c10, this file should be removed
12+
513
namespace at {
614

7-
// TODO Once all ATen ops are moved to c10, this file should be removed
15+
namespace {
16+
struct OpNameEquals final {
17+
bool operator()(const std::pair<const char*, const char*>& lhs, const std::pair<const char*, const char*>& rhs) const {
18+
return 0 == strcmp(lhs.first, rhs.first) && 0 == strcmp(lhs.second, rhs.second);
19+
}
20+
};
21+
22+
struct OpNameHash final {
23+
size_t operator()(const std::pair<const char*, const char*>& p) const {
24+
// use std::hash<std::string> because std::hash<const char*> would hash pointers and not pointed-to strings
25+
return std::hash<std::string>()(p.first) ^ (~ std::hash<std::string>()(p.second));
26+
}
27+
};
28+
}
829

9-
const std::unordered_set<c10::OperatorName>& aten_ops_already_moved_to_c10() {
10-
static std::unordered_set<c10::OperatorName> result {
30+
bool aten_op_is_already_moved_to_c10(const c10::OperatorName& opName) {
31+
static std::unordered_set<std::pair<const char*, const char*>, OpNameHash, OpNameEquals> ops {
1132
${c10_ops_already_moved_from_aten_to_c10}
1233
{"", ""}
1334
};
14-
return result;
35+
return ops.count(std::make_pair(opName.name.c_str(), opName.overload_name.c_str())) != 0;
1536
}
1637

17-
const std::unordered_set<c10::OperatorName>& aten_ops_not_moved_to_c10_yet() {
18-
static std::unordered_set<c10::OperatorName> result {
38+
bool aten_op_is_not_moved_to_c10_yet(const c10::OperatorName& opName) {
39+
static std::unordered_set<std::pair<const char*, const char*>, OpNameHash, OpNameEquals> ops {
1940
${c10_ops_not_moved_from_aten_to_c10_yet}
2041
{"", ""}
2142
};
22-
return result;
43+
return ops.count(std::make_pair(opName.name.c_str(), opName.overload_name.c_str())) != 0;
2344
}
2445

2546
}

‎torch/csrc/jit/register_c10_ops.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -177,8 +177,7 @@ Operator createOperatorFromC10(const c10::OperatorHandle& op) {
177177
class RegistrationListener final : public c10::OpRegistrationListener {
178178
public:
179179
void onOperatorRegistered(const c10::OperatorHandle& op) override {
180-
const auto& blacklist = at::aten_ops_already_moved_to_c10();
181-
if (0 != blacklist.count(op.schema().operator_name())) {
180+
if(at::aten_op_is_already_moved_to_c10(op.schema().operator_name())) {
182181
// Ignore ATen ops for now because they have their own code
183182
// to expose them to JIT in register_aten_ops.cpp
184183
// TODO Remove register_aten_ops.cpp and also use this registration here

0 commit comments

Comments
 (0)