Repository navigation
Expand file tree
/
Copy pathcompiler.py
More file actions
195 lines (156 loc) · 7.38 KB
/
Copy pathcompiler.py
File metadata and controls
195 lines (156 loc) · 7.38 KB
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
# *****************************************************************************
# Copyright (c) 2019-2020, Intel Corporation All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# *****************************************************************************
from llvmlite import binding
import sdc
import sdc.hiframes
import sdc.hiframes.hiframes_untyped
import sdc.hiframes.hiframes_typed
from sdc.hiframes.hiframes_untyped import HiFramesPass
from sdc.hiframes.hiframes_typed import HiFramesTypedPass
from sdc.hiframes.dataframe_pass import DataFramePass
import numba
import numba.core.compiler
from numba.core.compiler import DefaultPassBuilder
from numba import ir_utils, ir, postproc
from numba.core.registry import CPUDispatcher
from numba.core.ir_utils import guard, get_definition
from numba.core.inline_closurecall import inline_closure_call
from numba.typed_passes import (NopythonTypeInference, AnnotateTypes, ParforPass, IRLegalization)
from numba.untyped_passes import (DeadBranchPrune, InlineInlinables, InlineClosureLikes)
from sdc import config
from sdc.distributed import DistributedPass
from numba.core.compiler_machinery import FunctionPass, register_pass
# workaround for Numba #3876 issue with large labels in mortgage benchmark
binding.set_option("tmp", "-non-global-value-max-name-size=2048")
def inline_calls(func_ir, _locals):
work_list = list(func_ir.blocks.items())
while work_list:
label, block = work_list.pop()
for i, instr in enumerate(block.body):
if isinstance(instr, ir.Assign):
lhs = instr.target
expr = instr.value
if isinstance(expr, ir.Expr) and expr.op == 'call':
func_def = guard(get_definition, func_ir, expr.func)
if (isinstance(func_def, (ir.Global, ir.FreeVar))
and isinstance(func_def.value, CPUDispatcher)):
py_func = func_def.value.py_func
inline_out = inline_closure_call(
func_ir, py_func.__globals__, block, i, py_func,
work_list=work_list)
# TODO remove if when inline_closure_call() output fix
# is merged in Numba
if isinstance(inline_out, tuple):
var_dict = inline_out[1]
# TODO: update '##distributed' and '##threaded' in _locals
_locals.update((var_dict[k].name, v)
for k, v in func_def.value.locals.items()
if k in var_dict)
# for block in new_blocks:
# work_list.append(block)
# current block is modified, skip the rest
# (included in new blocks)
break
# sometimes type inference fails after inlining since blocks are inserted
# at the end and there are agg constraints (categorical_split case)
# CFG simplification fixes this case
func_ir.blocks = ir_utils.simplify_CFG(func_ir.blocks)
# TODO: remove these helper functions when Numba provide appropriate way to manipulate passes
def pass_position(pm, location):
assert pm.passes
pm._validate_pass(location)
for idx, (x, _) in enumerate(pm.passes):
if x == location:
return idx
raise ValueError("Could not find pass %s" % location)
def add_pass_before(pm, pass_cls, location):
assert pm.passes
pm._validate_pass(pass_cls)
position = pass_position(pm, location)
pm.passes.insert(position, (pass_cls, str(pass_cls)))
# if a pass has been added, it's not finalized
pm._finalized = False
def replace_pass(pm, pass_cls, location):
assert pm.passes
pm._validate_pass(pass_cls)
position = pass_position(pm, location)
pm[position] = (pass_cls, str(pass_cls))
# if a pass has been added, it's not finalized
pm._finalized = False
@register_pass(mutates_CFG=True, analysis_only=False)
class InlinePass(FunctionPass):
_name = "sdc_extention_inline_pass"
def __init__(self):
pass
def run_pass(self, state):
inline_calls(state.func_ir, state.locals)
return True
@register_pass(mutates_CFG=True, analysis_only=False)
class PostprocessorPass(FunctionPass):
_name = "sdc_extention_postprocessor_pass"
def __init__(self):
pass
def run_pass(self, state):
post_proc = postproc.PostProcessor(state.func_ir)
post_proc.run()
return True
class SDCPipeline(numba.core.compiler.CompilerBase):
"""SDC compiler pipeline
"""
def define_pipelines(self):
name = 'sdc_extention_pipeline_distributed'
pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
add_pass_before(pm, InlinePass, InlineClosureLikes)
pm.add_pass_after(HiFramesPass, InlinePass)
pm.add_pass_after(DataFramePass, AnnotateTypes)
pm.add_pass_after(PostprocessorPass, AnnotateTypes)
pm.add_pass_after(HiFramesTypedPass, DataFramePass)
pm.add_pass_after(DistributedPass, ParforPass)
pm.finalize()
return [pm]
@register_pass(mutates_CFG=True, analysis_only=False)
class ParforSeqPass(FunctionPass):
_name = "sdc_extention_parfor_seq_pass"
def __init__(self):
pass
def run_pass(self, state):
numba.parfors.parfor.lower_parfor_sequential(
state.typingctx, state.func_ir, state.typemap, state.calltypes)
return True
class SDCPipelineSeq(SDCPipeline):
"""SDC pipeline without the distributed pass (used in rolling kernels)
"""
def define_pipelines(self):
name = 'sdc_extention_pipeline_seq'
pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
add_pass_before(pm, InlinePass, InlineClosureLikes)
pm.add_pass_after(HiFramesPass, InlinePass)
pm.add_pass_after(DataFramePass, AnnotateTypes)
pm.add_pass_after(PostprocessorPass, AnnotateTypes)
pm.add_pass_after(HiFramesTypedPass, DataFramePass)
add_pass_before(pm, ParforSeqPass, IRLegalization)
pm.finalize()
return [pm]