Open
Description
🐛 Describe the bug
Vulkan models containing torch.linalg.norm can fail to partition with an index out of range error.
import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower, EdgeCompileConfig, to_edge
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
from typing import Callable, List, Optional, Tuple, Union
class LinalgNormModel(torch.nn.Module):
def __init__(
self,
ord: Optional[Union[int, float, str]] = None,
dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None,
keepdim: bool = False,
dtype: Optional[torch.dtype] = None
):
super().__init__()
self.ord = ord
self.dim = dim
self.keepdim = keepdim
self.dtype = dtype
def forward(self, x):
return torch.linalg.norm(x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype)
model = LinalgNormModel(ord=1)
inputs = (
torch.randn(5, 10),
)
eager_outputs = model(*inputs)
ep = torch.export.export(model.eval(), inputs)
print(ep)
lowered = to_edge_transform_and_lower(
ep,
partitioner=[VulkanPartitioner()],
compile_config=EdgeCompileConfig(_check_ir_validity=False)
).to_executorch()
print(lowered.exported_program())
et_model = _load_for_executorch_from_buffer(lowered.buffer)
et_outputs = et_model([*inputs])[0]
print(f"Inputs: {inputs}")
print(f"Eager: {eager_outputs}")
print(f"ET: {et_outputs}")
Outputs:
...
File /data/users/gjcomer/fbsource/buck-out/v2/gen/fbcode/fdcb6705e87e1def/bento_kernels/cria/__bento_kernel_cria_binary__/bento_kernel_cria_binary#link-tree/executorch/backends/vulkan/op_registry.py:457, in register_reduce_op.<locals>.check_reduce_node(node)
454 if isinstance(dim_list, list) and len(dim_list) != 1:
455 return False
--> 457 keepdim = node.args[2]
458 if isinstance(keepdim, bool) and not keepdim:
459 return False
IndexError: tuple index out of range
Versions
Run on Meta internal master, Jul 3, fbcode/SwiftShader