[jit] Move ModuleIndex operator to selective build. (#67483)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67483

Move ModuleIndex operator to selective build candidates.
ghstack-source-id: 141953898

Test Plan: eyes

Reviewed By: qihqi

Differential Revision: D32003895

fbshipit-source-id: 635c2bc37cd30a98f4a1e182fd6534eb9f1c4a69
This commit is contained in:
Zhengxu Chen 2021-10-29 13:29:54 -07:00 committed by Facebook GitHub Bot
parent 12ede84dbb
commit 60a80c5bbd
4 changed files with 19 additions and 20 deletions

View File

@ -15,11 +15,6 @@ Object::Object(
c10::StrongTypePtr(std::move(cu), type),
type->numAttributes())) {}
ObjectPtr Object::_ivalue() const {
TORCH_INTERNAL_ASSERT(_ivalue_);
return _ivalue_;
}
c10::optional<Method> Object::find_method(const std::string& basename) const {
for (Function* fn : type()->methods()) {
if (fn->name() == basename) {

View File

@ -31,7 +31,10 @@ struct TORCH_API Object {
std::shared_ptr<CompilationUnit> cu,
bool shouldMangle = false);
ObjectPtr _ivalue() const;
ObjectPtr _ivalue() const {
TORCH_INTERNAL_ASSERT(_ivalue_);
return _ivalue_;
}
c10::ClassTypePtr type() const {
return _ivalue()->type();

View File

@ -2936,6 +2936,21 @@ static const OperatorGeneratorArgs opGenArgs2[] = {
}
},
aliasAnalysisFromSchema()),
// This operator is generated inside the compiler for indexing into
// ModuleList without a statically determinable key. Accordingly,
// self must be a ModuleType and the output must be an InterfaceType.
OperatorGeneratorArgs(
TORCH_SELECTIVE_SCHEMA(
"prim::ModuleContainerIndex.list(Any self, int ind) -> Any"),
[](Stack& stack) {
IValue ind = pop(stack);
IValue module_dict = pop(stack);
std::stringstream ss;
ss << ind.toInt();
push(
stack, torch::jit::Object(module_dict.toObject()).attr(ss.str()));
},
aliasAnalysisFromSchema()),
#define DEFINE_DIVMOD_MIXED_OP(type_a, type_b) \
OperatorGeneratorArgs( \

View File

@ -170,20 +170,6 @@ RegisterOperators reg(
},
aliasAnalysisFromSchema()),
// This operator is generated inside the compiler for indexing into
// ModuleList without a statically determinable key. Accordingly,
// self must be a ModuleType and the output must be an InterfaceType.
OperatorGenerator(
TORCH_SELECTIVE_SCHEMA(
"prim::ModuleContainerIndex.list(Any self, int ind) -> Any"),
[](Stack& stack) {
IValue ind = pop(stack);
IValue module_dict = pop(stack);
std::stringstream ss;
ss << ind.toInt();
push(stack, module_dict.toModule().attr(ss.str()));
},
aliasAnalysisFromSchema()),
// This operator is generated inside the compiler for indexing into
// ModuleDict without a statically determinable key. Accordingly,
// self must be a ModuleType and the output must be an InterfaceType.
OperatorGenerator(