tico 0.1.0.dev250411__py3-none-any.whl

tico/serialize/operators/hashable_opcode.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from circle_schema import circle
+
+
+class OpCode(circle.OperatorCode.OperatorCodeT):
+    """
+    Wrapper class for operator code in circle schema
+    This implements __eq__ and __hash__ for use with dict()
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def __eq__(self, other):
+        if self.version != other.version:
+            return False
+
+        if self.builtinCode == circle.BuiltinOperator.BuiltinOperator.CUSTOM:
+            return self.customCode == other.customCode
+
+        return self.builtinCode == other.builtinCode
+
+    def __hash__(self):
+        val = (
+            self.deprecatedBuiltinCode,
+            self.customCode,
+            self.version,
+            self.builtinCode,
+        )
+        return hash(val)

tico/serialize/operators/node_visitor.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Type, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+
+
+class NodeVisitor:
+    """
+    Node visitor for lowering edge IR to circle
+    """
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        # For setting opcode index in circle model
+        # This is updated during serialization
+        self._op_codes = op_codes
+        self.graph = graph
+
+    # Define circle model operator
+    def define_node(
+        self,
+        node: torch.fx.node.Node,
+    ) -> circle.Operator.OperatorT:
+        raise NotImplementedError("NodeVisitor must be extended.")
+
+
+# container for all node visitors
+_node_visitor_dict: Dict[torch._ops.OpOverload, Type[NodeVisitor]] = {}
+
+
+# Decorator for each visitor
+def register_node_visitor(visitor):
+    for target in visitor.target:
+        _node_visitor_dict[target] = visitor
+    return visitor
+
+
+def get_node_visitor(target: torch._ops.OpOverload) -> Type[NodeVisitor]:
+    """
+    Get a single node visitor (for unittest purpose)
+    """
+    _visitor = _node_visitor_dict.get(target, None)
+
+    if not _visitor:
+        raise LookupError(f"NodeVisitor for {target} is not registered")
+
+    return _visitor
+
+
+# Get all node visitors
+def get_node_visitors(
+    op_codes: Dict[OpCode, int], graph: CircleSubgraph
+) -> Dict[torch._ops.OpOverload, NodeVisitor]:
+    node_visitors = {}
+    for target, visitor in _node_visitor_dict.items():
+        node_visitors[target] = visitor(op_codes, graph)
+
+    return node_visitors
+
+
+def get_support_targets():
+    return _node_visitor_dict.keys()

tico/serialize/operators/op_add.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import AddTensorArgs
+
+
+@register_node_visitor
+class AddVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.add.Tensor,
+        torch.ops.aten.add.Scalar,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = AddTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        other = args.other
+
+        inputs = [input, other]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.ADD, self._op_codes
+        )
+
+        inputs = [input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.AddOptions
+        option = circle.AddOptions.AddOptionsT()
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+        option.potScaleInt16 = False
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_alias_copy.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import AliasCopyArgs
+
+
+@register_node_visitor
+class AliasCopyVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.alias.default,
+        torch.ops.aten.alias_copy.default,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = AliasCopyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.TRANSPOSE, self._op_codes
+        )
+
+        permute = torch.IntTensor(list(range(len(input.meta["val"].shape))))
+
+        inputs = [input, permute]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.TransposeOptions
+        )
+        option = circle.TransposeOptions.TransposeOptionsT()
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_any.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.circle_mapping import (
+    circle_legalize_dtype_to,
+    extract_circle_dtype,
+    extract_shape,
+    extract_torch_dtype,
+)
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import AnyArgs
+
+
+@register_node_visitor
+class AnyVisitor(NodeVisitor):
+    """
+    Let's take NotEqual0 -> ReduceMax workaround for float, int
+    [RESTRICTION]
+        1. ReduceAny is not supported (luci-interpreter)
+    [CASE: BOOL]
+        (Bool tensors don't need 'Not Equal 0' at the first step.)
+        bool[d0..dN]      --- Reduce Max      ---> bool[]
+    [CASE: FLOAT, INT]
+        int/float[d0..dN] --- Not Equal 0     ---> bool[d0,...dN]
+                          --- Reduce Max      ---> bool[]
+        * [d0..dN] means a tensor with any shape
+        * [] means Scalar
+    """
+
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.any.default,
+        torch.ops.aten.any.dim,
+        torch.ops.aten.any.dims,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_max_node(
+        self, inputs: List, outputs: List, keepdims: bool
+    ) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.REDUCE_MAX, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.ReducerOptions
+        )
+        option = circle.ReducerOptions.ReducerOptionsT()
+        option.keepDims = keepdims
+
+        operator.builtinOptions = option
+
+        return operator
+
+    def define_ne_node(self, inputs: List, outputs: List) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.NOT_EQUAL, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.NotEqualOptions
+        )
+        option = circle.NotEqualOptions.NotEqualOptionsT()
+        operator.builtinOptions = option
+        return operator
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = AnyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        dim = args.dim
+        keepdim = args.keepdim
+
+        input_shape = list(extract_shape(input))
+        output_shape = list(extract_shape(node))
+
+        if dim is None:
+            dims = tuple(i for i in range(0, len(input_shape)))
+            dim_i32 = tuple(
+                circle_legalize_dtype_to(dim, dtype=torch.int32) for dim in dims
+            )
+        if isinstance(dim, int):
+            dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+        if isinstance(dim, tuple):
+            dim_i32 = tuple(circle_legalize_dtype_to(d, dtype=torch.int32) for d in dim)
+
+        inputs = [input, dim_i32]
+        outputs = [node]
+
+        dtype_torch = extract_torch_dtype(input)
+        input_tensor: torch.fx.node.Node | circle.Tensor.TensorT = input
+
+        if dtype_torch in [torch.int32, torch.int64, torch.float32, torch.float64]:
+            dst_dtype_circle = circle.TensorType.TensorType.BOOL
+            dst_dtype_torch = torch.bool
+            ne_tensor: circle.Tensor.TensorT = self.graph.add_tensor_from_scratch(
+                prefix=f"{input.name}_ne", shape=input_shape, dtype=dst_dtype_circle
+            )
+            ne_node = self.define_ne_node(
+                [input_tensor, torch.Tensor([0]).to(dtype_torch)], [ne_tensor]
+            )
+            self.graph.add_operator(ne_node)
+
+            dtype_torch = dst_dtype_torch
+            input_tensor = ne_tensor
+            inputs = [ne_tensor, dim_i32]
+
+        inputs = [input_tensor, dim_i32]
+
+        reduce_node: circle.Operator.OperatorT = self.define_max_node(
+            inputs, outputs, keepdim
+        )
+
+        return reduce_node

tico/serialize/operators/op_arange_start_step.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.utils.validate_args_kwargs import ArangeStartStepArgs
+
+
+@register_node_visitor
+class ArangeStartStepVisitor(NodeVisitor):
+    """
+    Fuse arange_start_step to const_tensor
+    """
+
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.arange.start_step]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = ArangeStartStepArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        start = args.start
+        end = args.end
+        step = args.step
+        delta = 1
+
+        if step is not None:
+            delta = step[0]  # type: ignore[index]
+        # assert False, "This pass must not be in use."
+
+        arange_dtype: torch.dtype = torch.float32
+        if isinstance(start, int) and isinstance(end, int):
+            arange_dtype = torch.int64
+
+        output_data = torch.arange(start=start, end=end, step=delta, dtype=arange_dtype)
+        self.graph.update_tensor_buffer(output_data, node.name)
+
+        return None  # type: ignore[return-value]

tico/serialize/operators/op_argmax.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.circle_mapping import circle_legalize_dtype_to
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import ArgMaxArgs
+
+
+@register_node_visitor
+class ArgMaxVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.argmax.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = ArgMaxArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        tensor = args.tensor
+        dim = args.dim
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.ARG_MAX, self._op_codes
+        )
+
+        dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+        inputs = [tensor, dim_i32]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.ArgMaxOptions
+        option = circle.ArgMaxOptions.ArgMaxOptionsT()
+        option.outputType = circle.TensorType.TensorType.INT64
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_avg_pool2d.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.define import define_pad_node
+from tico.utils.validate_args_kwargs import AvgPool2dArgs
+
+
+@register_node_visitor
+class AvgPool2DVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.circle_custom.avgpool2d]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = AvgPool2dArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        kernel_size = args.kernel_size
+        stride = args.stride
+        padding = args.padding
+
+        avgpool_input: torch.fx.Node | circle.Tensor.TensorT = input
+
+        def define_padding_node():
+            assert isinstance(padding, list), type(padding)
+            padding_vec = torch.tensor(
+                [
+                    [0, 0],
+                    [padding[0], padding[0]],
+                    [padding[1], padding[1]],
+                    [0, 0],
+                ],
+                dtype=torch.int32,
+            )
+            input_shape = list(extract_shape(input))
+            input_dtype: int = extract_circle_dtype(input)
+            padded_input_shape = [
+                input_shape[0],
+                input_shape[1],
+                input_shape[2],
+                input_shape[3],
+            ]
+            padded_input_shape[1] += padding[0] * 2
+            padded_input_shape[2] += padding[1] * 2
+            # create padded input tensor
+            padded_input_tensor = self.graph.add_tensor_from_scratch(
+                prefix=f"{input.name}_pad_output",
+                shape=padded_input_shape,
+                dtype=input_dtype,
+            )
+            pad_operator = define_pad_node(
+                self.graph, self._op_codes, [input, padding_vec], [padded_input_tensor]
+            )
+            self.graph.add_operator(pad_operator)
+            return padded_input_tensor
+
+        if padding is not None:
+            avgpool_input = define_padding_node()
+
+        inputs = [avgpool_input]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.AVERAGE_POOL_2D,
+            self._op_codes,
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.Pool2DOptions
+        option = circle.Pool2DOptions.Pool2DOptionsT()
+
+        SAME, VALID = 0, 1
+        option.padding = VALID
+        option.strideH = stride[0]
+        option.strideW = stride[1]
+        option.filterHeight = kernel_size[0]
+        option.filterWidth = kernel_size[1]
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_bmm.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import BmmArgs
+
+
+@register_node_visitor
+class BatchMatmulVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.bmm.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = BmmArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        mat2 = args.mat2
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
+        )
+
+        inputs = [input, mat2]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.BatchMatMulOptions
+        )
+        option = circle.BatchMatMulOptions.BatchMatMulOptionsT()
+        option.adjointLhs, option.adjointRhs = False, False
+        operator.builtinOptions = option
+
+        return operator