executorch
executorch copied to clipboard
`torch.max(input)` fails at XNNPACK runtime
Hi team,
Below is a minimal example script used to generate xnnpack_lowered_module.pte
import time
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from torch import nn
from torch._export import capture_pre_autograd_graph
import torch
from torch.export import ExportedProgram, export
from executorch.exir import EdgeProgramManager, to_edge
class Wrapper(nn.Module):
def __init__(self):
super(Wrapper, self).__init__()
def forward(self, im: torch.Tensor):
return torch.max(im)
wrapper = Wrapper()
wrapper.eval()
input = torch.randn((3,2))
result = wrapper(input)
print(result)
pre_autograd_aten_dialect = capture_pre_autograd_graph(wrapper, (input,))
print("Pre-Autograd ATen Dialect Graph")
aten_dialect: ExportedProgram = export(pre_autograd_aten_dialect, (input,))
print("ATen Dialect Graph")
edge_program: EdgeProgramManager = to_edge(aten_dialect)
print("Edge Dialect Graph")
to_be_lowered_module = edge_program.exported_program()
##### Lowering to XNNPACK
start = time.time()
xnnpack_lowered_module = edge_program.to_backend(XnnpackPartitioner())
end = time.time()
print(f"xnnpack_lowered_module elapsed: {end - start}")
exec_prog = xnnpack_lowered_module.to_executorch()
with open("xnnpack_lowered_module.pte", "wb") as f:
exec_prog.write_to_file(f)
The Edge Program graph is this:
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, arg0_1: "f32[3, 2]"):
aten_max_default: "f32[]" = executorch_exir_dialects_edge__ops_aten_max_default(arg0_1); arg0_1 = None
return (aten_max_default,)
Then, running with XNNPACK runner:
./cmake-out/backends/xnnpack/xnn_executor_runner --model_path=./xnnpack_lowered_module.pte
Gives error:
I 00:00:00.000653 executorch:executor_runner.cpp:82] Using method forward
I 00:00:00.000655 executorch:executor_runner.cpp:129] Setting up planned buffer 0, size 48.
E 00:00:00.000686 executorch:method.cpp:536] Missing operator: [0] aten::max.unary_out
E 00:00:00.000687 executorch:method.cpp:724] There are 1 instructions don't have corresponding operator registered. See logs for details
F 00:00:00.000688 executorch:executor_runner.cpp:151] In function main(), assert failed (method.ok()): Loading of method forward failed with status 0x14
No subgraphs were delegated to XNNPACK, yet this fails?
I'm surprised that the Executorch runtime can't run torch.max
This is just a minimal example, the real graph contains XNNPACK subgraphs, so I do need to run it with the XNNPACK runtime.
Thank you!
Package versions used
torchvision==0.18.0
executorch==0.2.0
torch==2.3.0