MCPcopy Index your code
hub / github.com/pytorch/pytorch / compile

Method compile

caffe2/python/session.py:85–131  ·  view source on GitHub ↗
(cls, runnable, workspace_type=None, setup_net_list=None)

Source from the content-addressed store, hash-verified

83
84 @classmethod
85 def compile(cls, runnable, workspace_type=None, setup_net_list=None):
86 if isinstance(runnable, CompiledRunnable):
87 assert cls == runnable.session_class, (
88 'Runnable was compiled for different session type. ' +
89 'Need: %s, got: %s' % (
90 cls.__name__, runnable.session_class.__name__))
91 return runnable
92
93 if runnable in cls._compiled_cache:
94 return cls._compiled_cache[runnable]
95
96 if isinstance(runnable, TaskGroup):
97 if workspace_type:
98 if runnable.workspace_type():
99 assert runnable.workspace_type() == workspace_type, \
100 "Require {} but already have {}".format(
101 workspace_type, runnable.workspace_type())
102 else:
103 runnable._workspace_type = workspace_type
104 tg = runnable
105 else:
106 if workspace_type is None:
107 workspace_type = WorkspaceType.GLOBAL
108 tg = TaskGroup(workspace_type=workspace_type)
109 if isinstance(runnable, Task):
110 tg.add(runnable)
111 elif isinstance(runnable, core.ExecutionStep):
112 tg.add(Task(step=runnable))
113 elif isinstance(runnable, core.Plan):
114 # ExecutionSteps in Plan() object is supposed to run sequentially, while
115 # tasks in TaskGroup run in parallel. So if we have multiple
116 # ExecutionSteps in Plan() object, we choose to have a root
117 # ExecutionStep to wrap all ExecutionSteps.
118 assert len(runnable.Steps()) > 0
119 if len(runnable.Steps()) == 1:
120 tg.add(Task(step=runnable.Steps()[0]))
121 else:
122 # Task takes a list of ExecutionSteps and automatically wrap into
123 # a root ExecutionStep
124 tg.add(Task(step=runnable.Steps()))
125 else:
126 step = core.execution_step('runnable', runnable)
127 tg.add(Task(step=step))
128 compiled = CompiledRunnable(
129 cls._compile_task_group(tg, setup_net_list), session_class=cls)
130 cls._compiled_cache[runnable] = compiled
131 return compiled
132
133 def run(self, runnable, workspace_type=None, setup_net_list=None):
134 """Run the given runnable.

Callers 1

runMethod · 0.95

Calls 10

addMethod · 0.95
TaskGroupClass · 0.90
TaskClass · 0.90
isinstanceFunction · 0.85
CompiledRunnableClass · 0.85
StepsMethod · 0.80
execution_stepMethod · 0.80
workspace_typeMethod · 0.45
formatMethod · 0.45
_compile_task_groupMethod · 0.45

Tested by

no test coverage detected