-
Notifications
You must be signed in to change notification settings - Fork 413
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
2024-10-30 nightly release (db38bcc)
- Loading branch information
pytorchbot
committed
Oct 30, 2024
1 parent
23327e9
commit c464a71
Showing
35 changed files
with
1,208 additions
and
262 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -28,6 +28,7 @@ | |
| op_relu, | ||
| op_repeat, | ||
| op_rsqrt, | ||
| op_select, | ||
| op_sigmoid, | ||
| op_slice, | ||
| op_squeeze, | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| # Copyright 2024 Arm Limited and/or its affiliates. | ||
| # | ||
| # This source code is licensed under the BSD-style license found in the | ||
| # LICENSE file in the root directory of this source tree. | ||
|
|
||
| from typing import List | ||
|
|
||
| import serializer.tosa_serializer as ts | ||
| from executorch.backends.arm.operators.node_visitor import ( | ||
| NodeVisitor, | ||
| register_node_visitor, | ||
| ) | ||
|
|
||
| from executorch.backends.arm.tosa_mapping import TosaArg | ||
|
|
||
| from executorch.backends.arm.tosa_utils import build_reshape, tosa_shape | ||
| from serializer.tosa_serializer import TosaOp | ||
| from torch.fx import Node | ||
|
|
||
|
|
||
| @register_node_visitor | ||
| class SelectVisitor(NodeVisitor): | ||
| target = "aten.select_copy.int" | ||
|
|
||
| def __init__(self, *args): | ||
| super().__init__(*args) | ||
|
|
||
| def define_node( | ||
| self, | ||
| node: Node, | ||
| tosa_graph: ts.TosaSerializer, | ||
| inputs: List[TosaArg], | ||
| output: TosaArg, | ||
| is_quant_node: bool, | ||
| ) -> None: | ||
|
|
||
| assert len(inputs) == 3 | ||
| input_node, dim, index = inputs | ||
| shape = input_node.shape | ||
| rank = len(shape) | ||
|
|
||
| dim = dim.number % rank if dim.number < 0 else dim.number | ||
| index = index.number % rank if index.number < 0 else index.number | ||
|
|
||
| # For aten.select_copy, the output will be rank[input_shape - 1] | ||
| # For TOSA rank(in) == rank(out). | ||
| # Add an intermediate with the same rank | ||
| expanded_shape = tuple(1 if i == dim else shape[i] for i in range(rank)) | ||
| expanded_shape = tosa_shape(expanded_shape, input_node.dim_order) | ||
|
|
||
| output_reshaped = tosa_graph.addIntermediate( | ||
| expanded_shape, ts.DType.INT8 if is_quant_node else output.dtype | ||
| ) | ||
|
|
||
| attr_slice = ts.TosaSerializerAttribute() | ||
|
|
||
| start_attr = [index if i == dim else 0 for i in input_node.dim_order] | ||
| size_attr = [ | ||
| 1 if i == dim else input_node.shape[i] for i in input_node.dim_order | ||
| ] | ||
|
|
||
| attr_slice.SliceAttribute(start_attr, size_attr) | ||
|
|
||
| tosa_graph.addOperator( | ||
| TosaOp.Op().SLICE, [input_node.name], [output_reshaped.name], attr_slice | ||
| ) | ||
|
|
||
| # Reshape back to original rank of output. | ||
| build_reshape(tosa_graph, output_reshaped.name, output.shape, output.name) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.