afe.backends.mla.afe_to_n2a_compiler.insert_nodes

Functions

insert_pre_mla_segment_nodes(net, mla_node, pack_params)

Insert nodes into the AwesomeNet in place prior to the AwesomeNode containing the

insert_post_mla_segment_nodes(net, mla_node, ...)

Insert nodes into the AwesomeNet in place after the AwesomeNode containing the

Module Contents

afe.backends.mla.afe_to_n2a_compiler.insert_nodes.insert_pre_mla_segment_nodes(net: afe.ir.net.AwesomeNet, mla_node: afe.ir.node.AwesomeNode, pack_params: afe.backends.mla.afe_to_n2a_compiler.defines.PackParameters)[source]

Insert nodes into the AwesomeNet in place prior to the AwesomeNode containing the BackendIR which is to be executed on MLA. Each input to the MLA node should be tessellated in accordance with provided TessellateParameters. If enable_mla flag is set to True, tessellation will be performed on the MLA, otherwise tessellation node will be inserted. In case of multiple inputs to the MLA node, the (tessellated) inputs may be packed into a single or multiple section buffers.

Parameters:

  • net – Top-level AwesomeNet.

  • mla_node – AwesomeNode containing the MLA sub-graph.

  • pack_params – A dictionary containing input tessellation parameters for each section buffer to the MLA subgraph. Dictionary keys are the section names. Values in the dictionary are tuples of input ids and corresponding tessellation params. The input ids are the indices in the MLA input list before the backend compilation. The order of dictionary items is the order of MLA arguments after the backend compilation.

Returns:

None. Mutates the AwesomeNet by inserting the tessellation and if needed, pack nodes. The AwesomeNet is in invalid state on the return. The topological_sort is needed to return the AwesomeNet to valid state.

afe.backends.mla.afe_to_n2a_compiler.insert_nodes.insert_post_mla_segment_nodes(net: afe.ir.net.AwesomeNet, mla_node: afe.ir.node.AwesomeNode, unpack_params: afe.backends.mla.afe_to_n2a_compiler.defines.PackParameters, uncompiled_nodes: list[afe.ir.defines.NodeName])[source]

Insert nodes into the AwesomeNet in place after the AwesomeNode containing the BackendIR which is to be executed on MLA.

In case of multiple output sections from the MLA node, new TupleGetItem AwesomeNodes are generated to collect the outputs from the MLA node.

In case of multiple outputs in an output section, the section outputs are unpacked into multiple buffers.

Finally, a detessellate node is inserted if the detessellation of an output is not done on MLA; otherwise, a slice node is inserted if there is padding to be removed.

Parameters:

  • net – Top-level AwesomeNet.

  • mla_node – AwesomeNode containing the MLA sub-graph.

  • unpack_params – A dictionary containing output detessellation parameters for each section buffer from the MLA subgraph. Dictionary keys are the section names. Values in the dictionary are tuples of output ids and corresponding tessellation params. The output ids are the indices in the MLA output list before the backend compilation. The order of dictionary items is the order of MLA outputs after the backend compilation.

  • uncompiled_nodes – The list of nodes that have not been visited during compilation. The list is mutated so that potential TupleGetItem nodes are removed from the list while unpack nodes are generated.

Returns:

None. Mutates the AwesomeNet by inserting the detessellation nodes and unpack node, if needed. The AwesomeNet is in invalid state on the return. The topological_sort is needed to return the AwesomeNet to valid state.