Traverse the expression graph and apply fusion
(expr)
| 3229 | |
| 3230 | |
| 3231 | def optimize_blockwise_fusion(expr): |
| 3232 | """Traverse the expression graph and apply fusion""" |
| 3233 | |
| 3234 | def _fusion_pass(expr): |
| 3235 | # Full pass to find global dependencies |
| 3236 | seen = set() |
| 3237 | stack = [expr] |
| 3238 | dependents = defaultdict(set) |
| 3239 | dependencies = {} |
| 3240 | expr_mapping = {} |
| 3241 | |
| 3242 | while stack: |
| 3243 | next = stack.pop() |
| 3244 | |
| 3245 | if next._name in seen: |
| 3246 | continue |
| 3247 | seen.add(next._name) |
| 3248 | |
| 3249 | if is_valid_blockwise_op(next): |
| 3250 | dependencies[next._name] = set() |
| 3251 | if next._name not in dependents: |
| 3252 | dependents[next._name] = set() |
| 3253 | expr_mapping[next._name] = next |
| 3254 | |
| 3255 | for operand in next.dependencies(): |
| 3256 | stack.append(operand) |
| 3257 | if is_valid_blockwise_op(operand): |
| 3258 | if next._name in dependencies: |
| 3259 | dependencies[next._name].add(operand._name) |
| 3260 | dependents[operand._name].add(next._name) |
| 3261 | expr_mapping[operand._name] = operand |
| 3262 | expr_mapping[next._name] = next |
| 3263 | |
| 3264 | # Traverse each "root" until we find a fusable sub-group. |
| 3265 | # Here we use root to refer to a Blockwise Expr node that |
| 3266 | # has no Blockwise dependents |
| 3267 | roots = [ |
| 3268 | expr_mapping[k] |
| 3269 | for k, v in dependents.items() |
| 3270 | if v == set() |
| 3271 | or all(not is_valid_blockwise_op(expr_mapping[_expr]) for _expr in v) |
| 3272 | ] |
| 3273 | while roots: |
| 3274 | root = roots.pop() |
| 3275 | seen = set() |
| 3276 | stack = [root] |
| 3277 | group = [] |
| 3278 | while stack: |
| 3279 | next = stack.pop() |
| 3280 | |
| 3281 | if next._name in seen: |
| 3282 | continue |
| 3283 | seen.add(next._name) |
| 3284 | |
| 3285 | group.append(next) |
| 3286 | for dep_name in dependencies[next._name]: |
| 3287 | dep = expr_mapping[dep_name] |
| 3288 |