stgraph.compiler.passes package

Submodules

stgraph.compiler.passes.cf module

stgraph.compiler.passes.cf.CF(prog)

Constant folding. Modify prog in place. Currently, we replace var2 = mul(var1, 1) with var1

we replace the var3 in var2(valType.NodeType) = aggsum(var1(valType.E))

var3(valType.E) = mul(var2(valType.NodeType), 1)

with var1 to bypass the aggsum and mul

stgraph.compiler.passes.cse module

stgraph.compiler.passes.cse.CSE(prog)

Common subexpression elimination Modify prog in place

stgraph.compiler.passes.dce module

stgraph.compiler.passes.dce.DCE(prog, output_vars)

DCE, modify prog in place. A statement is marked dead if there is no use of its return value and its return value is not used as outputs

stgraph.compiler.passes.dependency_analysis module

stgraph.compiler.passes.dependency_analysis.dep_program(target_var, stopping_vars)

Step1: Analyze the part of the program that computes a “target_var”. Recursively 1. find the stmt that creates the current var 2. find the dependant stmts of its inputs until encounter “stopping_vars”.

Step2: Set the dependent stmt of stopping var to be None

Returns a copy of the dependency program of target_var

stgraph.compiler.passes.fusion module

class stgraph.compiler.passes.fusion.FusionStateMachine(init_stmt=None)

Bases: object

accept(stmt)

advance(stmt)

current_fusion_type()

state_trans = {0: {'a2d': 1, 'a2s': 1, 'd': 3, 'e': 1, 's': 5}, 1: {'d': 2, 'e': 1, 's': 2}, 2: {'d': 2, 's': 2}, 3: {'a2d': 1, 'd': 3, 's': 4}, 4: {'d': 4, 's': 4}, 5: {'a2s': 1, 'd': 4, 's': 5}}

static stmt_to_trans(stmt)

stgraph.compiler.passes.fusion.find_var(var, prog_list)

stgraph.compiler.passes.fusion.fusable(downstream_s, upstream_s, stmt2state_machine, new_fsm)

stgraph.compiler.passes.fusion.fuse(progs, outputs)

Generate one/multiple execution units from one or more programs, which are used for code generation. Parallel mode of execution unit is determined by the ValType of ret var.

stgraph.compiler.passes.fusion.mergable(prog1, prog2)

stgraph.compiler.passes.fusion.merge_independent(exec_units)

stgraph.compiler.passes.fusion.merge_program(prog_list)

stgraph.compiler.passes.fusion.merge_stmt(cur_stmt, p_stmt, stmt2fused_prog, stmt2state_machine, stmt_stack, var_stack, prog_list, new_fsm)

stgraph.compiler.passes.fusion.unit_independent(u1, u2)

stgraph.compiler.passes.mem_planning module

stgraph.compiler.passes.mem_planning.mem_planning(funits, BProg, grad_map, grads)

Conduct memory planning by considering both FProg and BProg Create annotations for the graph so that code-generation can be done TODO: Add cost-model to choose a different materialized vars

stgraph.compiler.passes.peephole module

stgraph.compiler.passes.peephole.PH(BProg, known_vars, output_vars)

Peephole optimization pass. Mofiy program in place. It replaces unfusable expressions(‘sum’ and ‘aggsum’) in forward and backward programs with known_vars by applying various mathmatically equivelent tansformations.

stgraph.compiler.passes.peephole.execute_sym_program(prog, sym_table, rmv_list)

stgraph.compiler.passes.peephole.generate_stmts_from_expr(expr, var_table)

stgraph.compiler.passes.peephole.shape_propogation(s)

stgraph.compiler.passes.peephole.sum_propogation(sum_stmt)

stgraph.compiler.passes.visualize module

stgraph.compiler.passes.visualize.plog_program(prog)

stgraph.compiler.passes.visualize.plot_exec_units(units, filename='egl-fused-dag')

stgraph.compiler.passes.visualize.plot_programs(progs, filename='egl-dag')

Module contents

stgraph.compiler.passes.joint_optimize(F, B)

stgraph.compiler.passes.optimize(prog)