Skip to content

Commit 338118f

Browse files
committed
Update on "[ET-VK] Introduce AOT operator registry"
## Changes Move the following files to the root directory of Vulkan backend: * `backends/vulkan/partitioner/supported_ops.py` -> `backends/vulkan/op_registry.py` * `backends/vulkan/_passes/custom_ops_defs.py` -> `backends/vulkan/custom_ops_lib.py` In the new `op_registry.py` file, the way operator features are specified is reworked to provide much more detail about the features of the operator implementation in Vulkan. See the new `OpFeatures` class for more details. An example of registering a new operator to the export flow is ``` update_features( [ exir_ops.edge.aten._log_softmax.default, exir_ops.edge.aten._softmax.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.sum.dim_IntList, exir_ops.edge.aten.amax.default, exir_ops.edge.aten.amin.default, ] ) def register_reduce_op(features: OpFeatures): features.texture_impl = TextureImplFeatures( uses_packed_dim=True, ) features.resize_fn = True def check_reduce_node(node: torch.fx.Node) -> bool: dim_list = node.args[1] assert isinstance(dim_list, list) if len(dim_list) != 1: return False keepdim = node.args[2] assert isinstance(keepdim, bool) if not keepdim: return False return True features.check_node_fn = check_reduce_node return features ``` ## Rationale The purpose of these changes is to centralize operator definitions so that there is a common source of truth about the capabilities of operator implementation in Vulkan. This way, the partitioner does not have to implement ad-hoc functions for specific operators (i.e. `is_valid_to_copy`) and graph transforms do not have to maintain their own operator metadata (`USES_WEIGHTS` in `insert_prepack_nodes`). Differential Revision: [D64915640](https://our.internmc.facebook.com/intern/diff/D64915640/) [ghstack-poisoned]
2 parents cad7edf + 5e48c1a commit 338118f

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

backends/vulkan/runtime/VulkanBackend.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
*/
88

99
#include <executorch/backends/vulkan/runtime/VulkanDelegateHeader.h>
10-
#include <executorch/backends/vulkan/schema_generated.h>
10+
#include <executorch/backends/vulkan/serialization/schema_generated.h>
1111

1212
#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
1313

backends/vulkan/targets.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,9 +189,9 @@ def define_common_targets(is_fbcode = False):
189189
],
190190
deps = [
191191
":vulkan_graph_runtime",
192+
"//executorch/backends/vulkan/serialization:vk_delegate_schema",
192193
"//executorch/runtime/core:event_tracer",
193194
"//executorch/runtime/backend:interface",
194-
"//executorch/runtime/backends/vulkan/serialization:vk_delegate_schema",
195195
"//executorch/runtime/core/exec_aten/util:tensor_util",
196196
],
197197
define_static_target = False,

0 commit comments

Comments
 (0)