Class OperatorNodeBase

Inheritance Relationships

Base Types

Class Documentation

class mgb::cg::OperatorNodeBase : public GraphNodeBase, public Hashable, public mgb::cg::GraphExecutable

Node for an operator.

An operator is defined to be a node that could generate one or more VarNode as output.

Each operator node must be purely functional, i.e. the same node evaluated on the same input value must produce the same output value

Each operator has an owner, the computing graph that it belongs to

Public Types

enum AddInputSortType

used by add_input() to sort vars for deduplication


enumerator NONE
enumerator CUR_ADDED

sort newly added vars

enumerator ALL

sort all currently added vars

using NodeProp = OperatorNodeProp
using InputWaitingSpec = SmallVector<InputWaitingSpecElem, 1>

Public Functions

~OperatorNodeBase() noexcept
const std::string &name() const
const char *cname() const
void name(std::string name)
const VarNodeArray &input() const
const VarNodeArray &output() const
const VarNodeArray usable_output() const
VarNode *input(size_t idx) const
VarNode *output(size_t idx) const
size_t hash() const final override

hash that combines all inputs, m_config.comp_node() and all add_equivalence_component calls

const NodeProp &node_prop() const

get node prop, which is available and constant after node construction

Note that this function calls do_make_node_prop() on first call

void set_inserted_in_graph()

called by ComputingGraph to mark that this node has been inserted in graph; inputs and outputs could not be later changed

bool inserted_in_graph() const
const OperatorNodeConfig &config() const
void execute(ExecEnv &env) final override

Execute the operator by starting all kernels on device.

  1. wait on input as indicated by get_input_waiting_spec

  2. allocate memory for dynamic outputs

  3. call do_execute

  4. set_ready on output

const InputWaitingSpec &input_waiting_spec() const

get computing nodes that need to wait on other vars

This is only valid after the computing func has been compiled.

void input_waiting_spec(InputWaitingSpec &&spec)

set input waiting spec

This should only be called from SeqCompNodeOptimizerImpl::init_ready_event() or EagerEvalManager

void add_input_layout_constraint()

add layout constraint for input vars by calling VarNode::add_layout_constraint

Note that this method is always called exactly once for operators that are inserted into the computing sequence

void mem_plan_fwd_in2out_readonly()

called by graph compiler to setup readonly memory forwarding

This function would always be called unless input has dynamic storage but output has static storage

void mem_plan_fwd_in2out_writable()

called by graph compiler to setup writable memory forwarding

This function would always be called unless input has dynamic storage but output has static storage

OprEventCallback get_opr_event_callback()

get callbacks to be invoked on events related to this operator; default implementation returns empty event

void on_output_comp_node_stream_changed() = 0

called when stream of comp node of output vars is changed for graph optimization

void init_output_dtype()

initialize output dtype by calling VarNode::dtype

The default implementation requires all inputs to have the same dtype and set output dtype to it

This function is called once during operator insertion.

void init_output_format()

initialize output format by calling VarNode::format

The default implementation require all inputs to have the same non-default format and set all non-volatile outputs format to it.

This function is called once during operator insertion

void init_output_comp_node() = 0

inititialize output comp_node by calling VarNode::comp_node

This function is called once during operator insertion.

void init_rt_force_dynamic_mem_alloc_imply_chain()

call VarNode::add_rt_force_dynamic_mem_alloc_imply_chain on input and output vars

This function is called once during operator insertion.

void init_output_static_infer_desc() = 0

register static infer descriptors for output vars by calling methods on ComputingGraph::static_infer_manager()

This function is called once during operator insertion.

void init_output_mem_plan(bool dynamic)

initialize mem alloc plan for output nodes

Mem plans are used for memory optimization; the storage of var node’s device tensor should always come from mem plan

Default implmentation works by calling VarNode::init_mem_plan on vars that match dynamic param

output(…)->shape() is guaranteed to be valid before calling this function.

Remember to add Flag::IMPURE_OUTPUT_MEM_PLAN if needed.

  • dynamic: if true, initialize mem plans for vars that could not be statically inferred; otherwise for statically inferable vars

void add_input(std::initializer_list<VarNode*> list, AddInputSortType sort_type = AddInputSortType::NONE)

add input var to this operator

VarNode *add_output(const Maybe<std::string> &name)

allocate a new output VarNode; the name would be appended to this->name to form the final name

template<typename T, typename ...Args>
void add_equivalence_component(Args&&... args)

add extra component for equivalence check

This is only a helper function to make the default hash() and is_same() implementation consider other components in addition to all the input nodes; you can also override hash() and is_same() to implement deduplication.

Note that the order for calling add_equivalence_component matters. Also note that all input vars are used for deduplication by default.

NodeProp *do_make_node_prop() const

allocate a new node prop and initialize dep entry as all inputs

bool update_priority() const

Update operator priority.

This method would be invoked if and only if initializing m_node_prop or after the graph optimizer modified the opr’s priority.


whether the priority would be changed.

Protected Functions

OperatorNodeBase(ComputingGraph *owner, const OperatorNodeConfig &config, const std::string &default_name, const VarNodeArrayView &input_var_naming)

  • input_var_naming: used for generating default node name

OperatorNodeBase(const CtorParamPack &param)
void do_execute(ExecEnv &env) = 0

actually execute; all input and output have been checked, and the subclasses only need to perform the actual computing

struct CtorParamPack

pack of params in constructor, to ease inheritance

Public Members

ComputingGraph *owner
const OperatorNodeConfig &config
const std::string &default_name
const VarNodeArrayView &input_var_naming
struct InputWaitingSpecElem

specifies waiting strategy on one comp node for input vars

Public Members

CompNode comp_node

on which comp node to wait other inputs

VarNodeArray dev_ready

vars that must be ready on device

struct OprEventCallback

struct to specify the callback function pointers for various operator events.

This exists mainly for optimization: if a callback is not needed, related surronding code would not be inserted into execution queue.

Public Types

using cbptr_t = thin_function<void()>

Public Members

Maybe<cbptr_t> on_mem_status_changed

called when memory status changed

Memory status is defined by all layouts and addresses of DEV_VALUE deps and outputs; if any of it changes, this callback would be called before execution