stgraph.compiler package

Subpackages

• stgraph.compiler.backend package
  • Subpackages
    • stgraph.compiler.backend.mxnet package
      • Module contents
    • stgraph.compiler.backend.pytorch package
      • Submodules
      • stgraph.compiler.backend.pytorch.torch_callback module
      • stgraph.compiler.backend.pytorch.torch_kernel_wrapper module
      • Module contents
    • stgraph.compiler.backend.tensorflow package
      • Module contents
  • Submodules
  • stgraph.compiler.backend.callback module
    • STGraphBackend
      • STGraphBackend.backend_cb()
      • STGraphBackend.new_zeros_call_back()
      • STGraphBackend.tensor_raw_ptr()
  • stgraph.compiler.backend.kernel_wrapper module
    • KernelWrapper
      • KernelWrapper.backward()
      • KernelWrapper.forward()
      • KernelWrapper.setup_context()
  • Module contents
• stgraph.compiler.code_gen package
  • Subpackages
    • stgraph.compiler.code_gen.templates package
      • Subpackages
      • Module contents
  • Submodules
  • stgraph.compiler.code_gen.code_gen module
    • AggInfo
      • AggInfo.compute
      • AggInfo.init
      • AggInfo.inner_write
      • AggInfo.outter_write
    • ArgInfo
      • ArgInfo.is_ptr
      • ArgInfo.name
      • ArgInfo.type
    • EdgeInfo
      • EdgeInfo.compute
      • EdgeInfo.inner_write
      • EdgeInfo.load
    • NodeInfo
      • NodeInfo.compute
      • NodeInfo.inner_write
      • NodeInfo.load
    • gen_agg_info()
    • gen_arg_info()
    • gen_code()
    • gen_cuda()
    • gen_edge_info()
    • gen_node_info()
    • render_template()
  • stgraph.compiler.code_gen.compiler module
    • compile_cuda()
    • compile_with_nvcc()
    • compile_with_nvrtc()
  • stgraph.compiler.code_gen.cuda_check module
    • GPUDevice
      • GPUDevice.log()
  • stgraph.compiler.code_gen.cuda_driver module
  • stgraph.compiler.code_gen.cuda_error module
    • ASSERT_DRV()
  • stgraph.compiler.code_gen.device_info module
    • DeviceInfo
      • DeviceInfo.log()
  • stgraph.compiler.code_gen.kernel_context module
    • KernelContext
      • KernelContext.eq_dim()
      • KernelContext.get_offset_key()
      • KernelContext.init_offset_cache()
      • KernelContext.kernel_argument_used_in_stmt()
      • KernelContext.matrix_var_offset()
      • KernelContext.query_offset()
      • KernelContext.scalar_var_offset()
      • KernelContext.set_stmt_ctx()
      • KernelContext.vector_var_offset()
      • KernelContext.write_inner()
    • LinearizedKernelContext
      • LinearizedKernelContext.graph_type_key()
      • LinearizedKernelContext.query_offset()
    • StmtGenCtx
      • StmtGenCtx.kernel_arguments
      • StmtGenCtx.write_location
      • StmtGenCtx.write_type
  • Module contents
• stgraph.compiler.debugging package
  • Submodules
  • stgraph.compiler.debugging.print_variables module
  • stgraph.compiler.debugging.stgraph_logger module
    • print_log()
  • Module contents
• stgraph.compiler.op package
  • Subpackages
    • stgraph.compiler.op.agg package
      • Submodules
      • stgraph.compiler.op.agg.agg_op module
      • Module contents
    • stgraph.compiler.op.mxnet package
      • Module contents
    • stgraph.compiler.op.pytorch package
      • Submodules
      • stgraph.compiler.op.pytorch.torch_op module
      • Module contents
    • stgraph.compiler.op.tensorflow package
      • Module contents
  • Submodules
  • stgraph.compiler.op.op module
    • Op
      • Op.to_schema()
  • stgraph.compiler.op.op_factory module
    • OpFactory
      • OpFactory.create()
  • Module contents
• stgraph.compiler.passes package
  • Submodules
  • stgraph.compiler.passes.cf module
    • CF()
  • stgraph.compiler.passes.cse module
    • CSE()
  • stgraph.compiler.passes.dce module
    • DCE()
  • stgraph.compiler.passes.dependency_analysis module
    • dep_program()
  • stgraph.compiler.passes.fusion module
    • FusionStateMachine
      • FusionStateMachine.accept()
      • FusionStateMachine.advance()
      • FusionStateMachine.current_fusion_type()
      • FusionStateMachine.state_trans
      • FusionStateMachine.stmt_to_trans()
    • find_var()
    • fusable()
    • fuse()
    • mergable()
    • merge_independent()
    • merge_program()
    • merge_stmt()
    • unit_independent()
  • stgraph.compiler.passes.mem_planning module
    • mem_planning()
  • stgraph.compiler.passes.peephole module
    • PH()
    • execute_sym_program()
    • generate_stmts_from_expr()
    • shape_propogation()
    • sum_propogation()
  • stgraph.compiler.passes.visualize module
    • plog_program()
    • plot_exec_units()
    • plot_programs()
  • Module contents
    • joint_optimize()
    • optimize()
• stgraph.compiler.val package
  • Subpackages
    • stgraph.compiler.val.mxnet package
      • Module contents
    • stgraph.compiler.val.pytorch package
      • Submodules
      • stgraph.compiler.val.pytorch.torch_val module
      • Module contents
    • stgraph.compiler.val.tensorflow package
      • Module contents
  • Submodules
  • stgraph.compiler.val.val module
    • Val
      • Val.backend()
      • Val.backend_key()
      • Val.device()
      • Val.dtype()
      • Val.id
      • Val.layout()
      • Val.requires_grad()
      • Val.size()
      • Val.sum()
      • Val.v
      • Val.val_type()
      • Val.view()
  • stgraph.compiler.val.val_factory module
    • ValFactory
      • ValFactory.create()
      • ValFactory.get_val_backend()
  • Module contents

Submodules

stgraph.compiler.autodiff module

The auto-differentiation module

stgraph.compiler.autodiff.diff(vars, grads, forward_units, fprog)

The forward graph differentiator

For each var we find the statment that computes it, then we use itself as well as its grad to get the statements to calculate or accumulate the gradient for each of its imputs.

We need to synchrounize the graidents for a same var in order to keep the statements ordered by return var id. In order to do that, we use processed_count == forward_num_users_of_var as condition to determine whether we can push it to queue. If processed_count < forward* it means that we need to wait for computing more of its gradients before propogate back through that variable. If the var is in stopping var, there is no need to propogate further as the task will be delegated to backend system.

Parameters:

• vars (Var) – The var that has gradient to propogate back, determined by zoomOut

• grads – The coresponding gradient for each var

• forward_units – Forward execution units

Returns:

The differentiated program to compute the gradients

Return type:

BProg

stgraph.compiler.execution_unit module

The fundamental execution unit of STGraph

class stgraph.compiler.execution_unit.ExecutionUnit(args, tmps, prog, compiled=False)

Bases: object

add_parent_unit(parent)

add_ret_val(ret_val)

all_rets()

calculate_kernel_launch_params(num_nodes)

calculate_kernel_params(num_nodes)

calculate_kernel_params_fa(num_nodes)

property compiled

compute_tile_sizes(blockDimx, blockDimy)

create_context(index_type)

depends_on(other)

feature_size()

first_pow2_less_than_n(n, upper_bound)

get_all_args()

return the set of all vars used/returned in the program of this exec unit

get_all_vars()

return the set of all vars used/returned in the program of this exec unit

has_parent()

is_child_of(other)

kernel_args()

property kernel_name

kernel_run(tensor_list)

materilized_vars()

max_dims()

max_ret_id()

merge_with_independent_unit(other)

nearest_pow2(targ)

parallel_mode()

prepare_compiled_kernel(graph, compiled_module)

property program

reset_graph_info(graph)

set_parallel_mode(mode)

property tmps

unit_args()

unit_count = 0

unit_rets()

use_fa_tmpl()

class stgraph.compiler.execution_unit.FeatureAdaptiveKernel(num_nodes, row_offsets_ptr, col_indices_ptr, eids_ptr, node_ids_ptr, max_dims, kernel_name, compiled_module, launch_config)

Bases: Kernel

class stgraph.compiler.execution_unit.Kernel

Bases: object

reset_graph_info(num_nodes, row_offsets_ptr, col_indices_ptr, eids_ptr, node_ids_ptr)

run(tensor_list)

class stgraph.compiler.execution_unit.V2Kernel(num_nodes, row_offsets_ptr, col_indices_ptr, eids_ptr, max_dims, kernel_name, compiled_module, launch_config, tile_sizes)

Bases: Kernel

The Version 2 Kernel

This class contains the parameters for the second version of the kernel written for STGraph

Parameters:

• num_nodes (int) – Number of nodes present in the graph

• row_offsets_ptr (c_type) – Pointer to the row offset array

• col_indices_ptr (c_type) – Pointer to the column indicies array

scalar_args

List of the scalar arguments passed to the kernel

Type:

list[c_types]

launch_config

List of the kernel launch configurations

Type:

list[int]

stgraph.compiler.executor module

class stgraph.compiler.executor.ExeState

Bases: object

all_bu_executed()

clear_current_tensor_state()

is_executed_bu(bu)

reset(input_map, f_merged_units, bunits)

track_executed_bu(bu)

track_tensor(key, val)

class stgraph.compiler.executor.Executor(graph, forward_exec_units, backward_exec_units, compiled_module, rets)

Bases: object

backward_cb(kid, grad_list)

FuncWrapper will call this function in backward pass

construct_backward_mappping(funits, bunits)

create_tensor_for_grad_vars(var_list, tensor_map)

create_tensor_for_vars(var_list)

execute(FuncWrapper)

Execute forward pass

execute_compiled(uid, FuncWrapper)

execute_prog(units)

execute_unit(unit, tensor_list)

forward_cb(uid, kernel_args, rets, tensor_list)

FuncWrapper will call this function in forward pass

merge_units(exec_units)

restart(input_map, graph=None)

set_new_zeros_cb(cb)

set_raw_ptr_cb(cb)

class stgraph.compiler.executor.MergedUnit(units)

Bases: object

append(unit)

compiled()

joint_args()

joint_inputs()

joint_rets()

kernel_arg_list()

last()

union_of_rets()

class stgraph.compiler.executor.Stack(val=None)

Bases: object

pop()

print()

push(val)

top()

stgraph.compiler.node module

class stgraph.compiler.node.CentralNode

Bases: object

update_allnode(feat_map)

class stgraph.compiler.node.NbEdge(center, direction, nbnodes)

Bases: object

class stgraph.compiler.node.NbNode(center, direction)

Bases: object

stgraph.compiler.program module

class stgraph.compiler.program.Program

Bases: object

append_stmt(stmt)

begin()

clear_stmts()

copy_append_prog(other_prog)

copy_append_stmt(stmt)

copy_append_stmts(stmts)

copy_prepend_prog(other_prog)

copy_prepend_stmt(stmt)

deepcopy()

empty()

end()

find_ret_var_by_id(targ_id)

find_var_by_id(targ_id)

has_stmt(stmt)

input_vars()

insert_stmts_before(stmt, stmts_list)

last_stmt()

prepend_stmt(stmt)

resort_vars()

set_input(var_name)

class stgraph.compiler.program.Stmt(op_schema, args, ret, op_type, callback)

Bases: object

class StmtInfo(op_schema, args)

Bases: tuple

args

Alias for field number 1

op_schema

Alias for field number 0

copy()

classmethod create_add_stmt(args)

classmethod create_binary_bcast_stmt(op_schema, args, callback=None)

classmethod create_mul_stmt(args)

classmethod create_stmt(op_schema=None, args=None, ret=None, callback=None)

depends_on(stmt)

execute(args, **kargs)

gen_code(ctx)

grad(y, grad_y)

insert_after(new_stmt)

insert_before(new_stmt)

is_agg()

is_dst()

is_edgewise()

is_element_wise_fusable()

is_nodewise()

is_src()

is_supported()

property op_name

print_stmt_args()

remove_cur()

replace_arg_with(old, new, propogate_shape)

shape_propogation()

stmt_info()

type_eq(var1, var2)

class stgraph.compiler.program.Var(var_id, val_type, var_shape, var_dtype, device, requires_grad)

Bases: object

add_user(stmt)

classmethod copy(other)

classmethod create_var(var_shape=None, var_dtype=None, val_type=None, var_id=None, device=None, requires_grad=True)

detach_from_stmt()

property device

property id

property int_id

is_dstvar()

is_edgevar()

is_nodevar()

is_srcvar()

replace_all_uses_with(other_var, propogate_shape)

replace_grad(new_grad)

property requires_grad

rmv_user(stmt)

set_grad(other_var)

set_to_be_grad_of(other_var)

property stmt

used_by(stmt)

property users

property val_type

property var_dtype

property var_shape

stgraph.compiler.registry module

class stgraph.compiler.registry.AddOp(fstmt, create_var_cb, create_stmt_cb)

Bases: BinaryOpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

class stgraph.compiler.registry.AggMaxOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

gen_init(var)

grad_impl(pos, x, y, grad_y)

y = AggMax(x) => dydx = (x)

class stgraph.compiler.registry.AggSumOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

gen_init(var)

grad_impl(pos, x, y, grad_y)

y = AggSum(x) => dydx = Bcast(x)

class stgraph.compiler.registry.BackwardAMaxOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

class stgraph.compiler.registry.BackwardLeakyReluOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

return a map with keys: args, grad_x and op_schema

class stgraph.compiler.registry.BackwardReluOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

return a map with keys: args, grad_x and op_schema

class stgraph.compiler.registry.BinaryOpImpl(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

grad_impl(pos, x, y, grad_y)

return a map with keys: args, grad_x and op_schema

class stgraph.compiler.registry.ExpOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

y = exp(x) => dydx = exp(x) = y

class stgraph.compiler.registry.GTypeCastOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

return a map with keys: args, grad_x and op_schema

class stgraph.compiler.registry.GradInfo(targ, args, grad_x, op_schema)

Bases: tuple

args

Alias for field number 1

grad_x

Alias for field number 2

op_schema

Alias for field number 3

targ

Alias for field number 0

class stgraph.compiler.registry.LeakyReluOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

y = leaky_relu(x) => dydx = backward_leaky_relu(x)

class stgraph.compiler.registry.MulOp(fstmt, create_var_cb, create_stmt_cb)

Bases: BinaryOpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

class stgraph.compiler.registry.OpImpl(fstmt, create_var_cb, create_stmt_cb)

Bases: ABC

New ops need to inherit from this class with name “XXXOp”

property args

create_var_like(x)

gen_agg_info_map(ctx)

abstract gen_code(ctx)

return the cuda code that corresponding to this op

gen_edge_info_map(ctx)

gen_load(ctx)

gen_var(var, ctx)

gen_write(ctx)

grad(y, grad_y)

abstract grad_impl(pos, x, y, grad_y)

return a map with keys: args, grad_x and op_schema

multiply_grad(dzdy, dydx, x)

property op_schema

property ret

class stgraph.compiler.registry.ReluOp(fstmt, create_var_cb, create_stmt_cb)

Bases: OpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

grad_impl(pos, x, y, grad_y)

return a map with keys: args, grad_x and op_schema

class stgraph.compiler.registry.SubOp(fstmt, create_var_cb, create_stmt_cb)

Bases: BinaryOpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

class stgraph.compiler.registry.TrueDivOp(fstmt, create_var_cb, create_stmt_cb)

Bases: BinaryOpImpl

gen_code(ctx)

return the cuda code that corresponding to this op

stgraph.compiler.registry.look_up_registry(stmt)

stgraph.compiler.registry.register_ops()

stgraph.compiler.registry.register_or_look_up_backend_cb(stmt, cb)

stgraph.compiler.schema module

class stgraph.compiler.schema.Schema(op_name, **kargs)

Bases: object

Schema of an op

stgraph.compiler.stgraph module

class stgraph.compiler.stgraph.Context(func, nspace, run_cb)

Bases: object

class stgraph.compiler.stgraph.STGraph(backend_framework: STGraphBackend)

Bases: object

compile(gnn_module, hetero_graph=False)

stgraph.compiler.utils module

class stgraph.compiler.utils.EdgeDirection(value)

Bases: Enum

An enumeration.

IN = 0

OUT = 1

class stgraph.compiler.utils.FusionType(value)

Bases: Enum

An enumeration.

NEAN = 0

NN = 1

NOT_FUSIBLE = 2

class stgraph.compiler.utils.OpType(value)

Bases: Enum

An enumeration.

A = 2

D = 3

E = 1

S = 0

class stgraph.compiler.utils.ParallelMode(value)

Bases: Enum

An enumeration.

DstParallel = 1

SrcParallel = 0

class stgraph.compiler.utils.ValType(value)

Bases: Enum

An enumeration.

DEST = 1

EDGE = 2

PARAM = 3

SRC = 0

class stgraph.compiler.utils.WriteLocation(value)

Bases: Enum

An enumeration.

INNER = 0

NONE = 2

OUTER = 1

class stgraph.compiler.utils.WriteType(value)

Bases: Enum

An enumeration.

ADD = 0

ASSIGN = 2

ATOMIC = 1

NONE = 3

stgraph.compiler.utils.any_var(var_list)

stgraph.compiler.utils.bcast_dim(var_list)

stgraph.compiler.utils.infer_op_type_from_args(op_schema, args)

stgraph.compiler.utils.infer_val_type(vals)

When type of vals are different, we return edge type. P type is considered same with every one.

stgraph.compiler.utils.is_const_scalar(val)

Module contents

The backbone of STGraph