Performs actions for the call_function node of a graph module. Dispatches based on 'target' and emits the corresponding function. 'call_function' is a powerful node that contains many operations ranging from control_flow, to memory management, to delegate and operator calls.
( # pyre-fixme[14]
self, target: _Target, args: Tuple[_Argument, ...], kwargs: Dict[str, _Argument]
)
| 1761 | self.chain.instructions.append(instruction) |
| 1762 | |
| 1763 | def call_function( # pyre-fixme[14] |
| 1764 | self, target: _Target, args: Tuple[_Argument, ...], kwargs: Dict[str, _Argument] |
| 1765 | ) -> _EmitterValue: |
| 1766 | """Performs actions for the call_function node of a graph module. |
| 1767 | |
| 1768 | Dispatches based on 'target' and emits the corresponding function. 'call_function' is a |
| 1769 | powerful node that contains many operations ranging from control_flow, to memory management, |
| 1770 | to delegate and operator calls. |
| 1771 | """ |
| 1772 | |
| 1773 | # Delegate and operator calls are the only functions that should have a debug handle |
| 1774 | # associated with them. All the others such as memory.alloc, getitem should be ignored. |
| 1775 | # Default to none and let delegates and ops override. |
| 1776 | if target == operator.getitem: |
| 1777 | assert len(args) == 2 |
| 1778 | head = typing.cast(Mapping[int, _EmitterValue], args[0]) |
| 1779 | index = typing.cast(int, args[1]) |
| 1780 | return head[index] |
| 1781 | |
| 1782 | elif target == memory.alloc: |
| 1783 | assert len(args) == 1 |
| 1784 | return self._emit_spec(self.node.meta["spec"]) |
| 1785 | |
| 1786 | elif target == memory.view: |
| 1787 | return self._emit_view(args) |
| 1788 | |
| 1789 | elif target == memory.free: |
| 1790 | assert len(args) == 1 |
| 1791 | # pyre-ignore |
| 1792 | return self._emit_free(args[0]) |
| 1793 | |
| 1794 | elif target in ( |
| 1795 | torch.ops.higher_order.cond, |
| 1796 | torch.ops.higher_order.map_impl, |
| 1797 | torch.ops.higher_order.while_loop, |
| 1798 | torch.ops.higher_order.scan, |
| 1799 | ): |
| 1800 | return self._emit_control_flow(target, args, kwargs) |
| 1801 | |
| 1802 | elif target == executorch_call_delegate: |
| 1803 | lowered_module = args[0] |
| 1804 | assert is_lowered_module(lowered_module) |
| 1805 | v = self._emit_delegate(lowered_module, args[1:], kwargs) |
| 1806 | delegate_instruction_id = len(self.chain.instructions) - 1 |
| 1807 | self._add_debug_handle(delegate_instruction_id, target, lowered_module) |
| 1808 | self._add_delegate_map(lowered_module, delegate_instruction_id) |
| 1809 | return v |
| 1810 | |
| 1811 | elif isinstance( |
| 1812 | target, (torch._ops.OpOverload, EdgeOpOverload, BackendOpOverload) |
| 1813 | ): |
| 1814 | return self._emit_operator(target, args, kwargs) |
| 1815 | |
| 1816 | else: |
| 1817 | raise InternalError( |
| 1818 | self._emit_node_specific_error( |
| 1819 | self.node, f"invalid target for call_function {target}" |
| 1820 | ) |