Skip to content

Commit ea487f4

Browse files
yipjustinyipjustin
committed
[ET-VK][14/n] Add operators to Partitioner
Failure for Partial models: P1225514296 Differential Revision: [D56695929](https://our.internmc.facebook.com/intern/diff/D56695929/) ghstack-source-id: 224267707 Pull Request resolved: #3402
1 parent 23e04e2 commit ea487f4

File tree

4 files changed

+136
-3
lines changed

4 files changed

+136
-3
lines changed

‎backends/vulkan/partitioner/vulkan_partitioner.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727

2828
class VulkanSupportedOperators(OperatorSupportBase):
2929
def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
30+
if node.op == "call_function":
31+
print("AAA=>", node.target)
3032
supported = node.op == "call_function" and node.target in [
3133
# Binary arithmetic operators
3234
exir_ops.edge.aten.add.Tensor,
@@ -56,6 +58,12 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
5658
exir_ops.edge.aten.select_copy.int,
5759
exir_ops.edge.aten.unsqueeze_copy.default,
5860
exir_ops.edge.aten.view_copy.default,
61+
# Copy-releated operators
62+
exir_ops.edge.aten.clone.default,
63+
exir_ops.edge.aten.cat.default,
64+
exir_ops.edge.aten.split_with_sizes_copy.default,
65+
exir_ops.edge.aten.split.Tensor,
66+
exir_ops.edge.aten.slice_copy.Tensor,
5967
# Other
6068
operator.getitem,
6169
exir_ops.edge.aten.full.default,

‎backends/vulkan/runtime/graph/ops/impl/Split.cpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ void add_split_with_sizes_default_node(
106106
add_split_with_sizes_default_node(graph, in, split_sizes, dim, out);
107107
}
108108

109-
void split_with_sizes_default(
109+
void split_with_sizes_copy_default(
110110
ComputeGraph& graph,
111111
const std::vector<ValueRef>& args) {
112112
add_split_with_sizes_default_node(graph, args[0], args[1], args[2], args[3]);
@@ -134,7 +134,8 @@ void split_tensor(ComputeGraph& graph, const std::vector<ValueRef>& args) {
134134
}
135135

136136
REGISTER_OPERATORS {
137-
VK_REGISTER_OP(aten.split_with_sizes.default, split_with_sizes_default);
137+
VK_REGISTER_OP(
138+
aten.split_with_sizes_copy.default, split_with_sizes_copy_default);
138139
VK_REGISTER_OP(aten.split.Tensor, split_tensor);
139140
}
140141

‎backends/vulkan/serialization/vulkan_graph_builder.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ def create_node_value(self, node: Node) -> int:
133133
new_id = self.create_tensor_value(spec, constant_id)
134134
self.node_to_value_ids[node] = new_id
135135
return new_id
136-
elif isinstance(spec, tuple):
136+
elif isinstance(spec, tuple) or isinstance(spec, list):
137137
# Create a Value for each element in the tuple, wrap Values in a
138138
# ValueList, and map the Node to the ValueList id.
139139
new_id = self.create_value_list_value(spec)

‎backends/vulkan/test/test_vulkan_delegate.py‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ def run_test(memory_layout):
8080
compile_options = {
8181
"memory_layout_override": memory_layout,
8282
}
83+
84+
# At least model should run in eager mode.
85+
eager_output = model(*sample_inputs)
86+
8387
program: ExportedProgram = export(
8488
model, sample_inputs, dynamic_shapes=dynamic_shapes
8589
)
@@ -798,3 +802,123 @@ def forward(self, x):
798802
sample_inputs,
799803
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
800804
)
805+
806+
def DISABLED_test_vulkan_backend_permute_copy(self):
807+
# aten.permute_copy.default is not enabled yet in partitioner
808+
class PermuteModule(torch.nn.Module):
809+
def __init__(self):
810+
super().__init__()
811+
812+
def forward(self, x):
813+
return torch.permute(x, [3, 0, 2, 1])
814+
815+
sample_inputs = (torch.randn(size=(3, 6, 2, 7), dtype=torch.float32),)
816+
817+
self.lower_module_and_test_output(
818+
PermuteModule(),
819+
sample_inputs,
820+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
821+
)
822+
823+
def test_vulkan_backend_cat(self):
824+
class TestModule(torch.nn.Module):
825+
def __init__(self):
826+
super().__init__()
827+
828+
def forward(self, x, y, z, w):
829+
return torch.cat([x, y, z, w], dim=1)
830+
831+
sample_inputs = (
832+
torch.randn(size=(3, 6, 2, 7), dtype=torch.float32),
833+
torch.randn(size=(3, 1, 2, 7), dtype=torch.float32),
834+
torch.randn(size=(3, 9, 2, 7), dtype=torch.float32),
835+
torch.randn(size=(3, 3, 2, 7), dtype=torch.float32),
836+
)
837+
838+
self.lower_module_and_test_output(
839+
TestModule(),
840+
sample_inputs,
841+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
842+
)
843+
844+
def test_vulkan_backend_slice(self):
845+
class TestModule(torch.nn.Module):
846+
def __init__(self):
847+
super().__init__()
848+
849+
def forward(self, x):
850+
return x[:, 2:9:2, :]
851+
852+
sample_inputs = (torch.randn(size=(3, 13, 7, 3), dtype=torch.float32),)
853+
854+
self.lower_module_and_test_output(
855+
TestModule(),
856+
sample_inputs,
857+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
858+
)
859+
860+
def test_vulkan_backend_split_with_sizes(self):
861+
class TestModule(torch.nn.Module):
862+
def __init__(self):
863+
super().__init__()
864+
865+
def forward(self, x):
866+
return torch.split(x, (3, 6, 1, 3), dim=1)
867+
868+
sample_inputs = (torch.randn(size=(3, 13, 7, 3), dtype=torch.float32),)
869+
870+
self.lower_module_and_test_output(
871+
TestModule(),
872+
sample_inputs,
873+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
874+
)
875+
876+
def test_vulkan_backend_split_tensor(self):
877+
class TestModule(torch.nn.Module):
878+
def __init__(self):
879+
super().__init__()
880+
881+
def forward(self, x):
882+
return torch.tensor_split(x, 2, dim=1)
883+
884+
sample_inputs = (torch.randn(size=(3, 14, 7, 3), dtype=torch.float32),)
885+
886+
self.lower_module_and_test_output(
887+
TestModule(),
888+
sample_inputs,
889+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
890+
)
891+
892+
def test_vulkan_backend_clone(self):
893+
class TestModule(torch.nn.Module):
894+
def __init__(self):
895+
super().__init__()
896+
897+
def forward(self, x):
898+
return torch.clone(x)
899+
900+
sample_inputs = (torch.randn(size=(3, 14, 7, 3), dtype=torch.float32),)
901+
902+
self.lower_module_and_test_output(
903+
TestModule(),
904+
sample_inputs,
905+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
906+
)
907+
908+
def DISABLED_test_vulkan_backend_t_default(self):
909+
# aten.permute_copy.default is not enabled yet in partitioner
910+
class TestModule(torch.nn.Module):
911+
def __init__(self):
912+
super().__init__()
913+
914+
def forward(self, x):
915+
# torch.t is actually exported as aten::permute.
916+
return torch.t(x)
917+
918+
sample_inputs = (torch.randn(size=(3, 14), dtype=torch.float32),)
919+
920+
self.lower_module_and_test_output(
921+
TestModule(),
922+
sample_inputs,
923+
memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
924+
)

0 commit comments

Comments
 (0)