1# Copyright (c) Meta Platforms, Inc. and affiliates. 2# All rights reserved. 3# 4# This source code is licensed under the BSD-style license found in the 5# LICENSE file in the root directory of this source tree. 6 7from typing import Dict 8 9import torch 10from executorch.backends.xnnpack.operators.node_visitor import ( 11 NodeVisitor, 12 register_node_visitor, 13) 14from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( 15 XNNBatchMatrixMultiply, 16 XNNGraph, 17 XNode, 18) 19from executorch.backends.xnnpack.utils.utils import get_input_node 20 21 22@register_node_visitor 23class BMMVisitor(NodeVisitor): 24 target = "aten.bmm.default" 25 26 def __init__(self, *args) -> None: 27 super().__init__(*args) 28 29 def define_node( 30 self, 31 node: torch.fx.Node, 32 xnn_graph: XNNGraph, 33 vals_to_ids: Dict[torch.fx.Node, int], 34 debug_handle: int, 35 ) -> None: 36 37 self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids) 38 39 # input1 40 input1_id = vals_to_ids[get_input_node(node, 0)] 41 42 # input2 43 input2_id = vals_to_ids[get_input_node(node, 1)] 44 45 # output 46 output_id = vals_to_ids[node] 47 48 ser_node = XNode( 49 xnode_union=XNNBatchMatrixMultiply( 50 input1_id=input1_id, input2_id=input2_id, output_id=output_id, flags=0 51 ), 52 debug_handle=debug_handle, 53 ) 54 xnn_graph.xnodes.append(ser_node) 55