torch.max(input) fails at XNNPACK runtime

[kinghchan]

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