Zero-sized tensors abort the process on GPU/ANE: a 0-length split section fails MPSCommonRuntimeCanonicalization, a width-0 output fails in MPSNDArray
- Dominant language
- Python
- Stars
- 152
- Forks
- 45
- Avg merge
- 1d 7m
- Merged PRs (30d)
- 12
Description
## What happens
A graph that contains a zero-sized tensor converts cleanly, then aborts the process when the asset is loaded or run on the GPU or the Neural Engine. A CPU-only specialization loads and runs the same asset. Two shapes of it, with two different assertions:
**1. A `split` section of size 0** — fails a canonicalization pass at load. The empty output is never used.
```python
a, b = x.split([x.shape[1], 0], dim=1) # b is width-0 and dead
return a + 1.0
```
```
MPSGraphExecutable.mm:4419: failed assertion `Error: Optimize Original Module MLIR pass manager failed
Pass failed: MPSCommonRuntimeCanonicalization
Pass failed: mlir::detail::OpToOpPassAdaptor'
```
**2. A zero-sized model output** — fails at run.
```python
return x[:, :0]
```
```
MPSNDArray.mm:893: failed assertion `[MPSNDArray, initWithBufferImpl:...] Error: buffer is not large enough.
Must be 128 bytes'
```
Neither is catchable. `AIModel.load` and the inference call abort rather than returning an error.
| | CPU | GPU | Neural Engine |
|---|---|---|---|
| `split` with a 0-sized section | runs | abort (`MPSGraphExecutable.mm:4419`) | abort (`MPSGraphExecutable.mm:4419`) |
| width-0 output | runs | abort (`MPSNDArray.mm:893`) | abort (`MPSNDArray.mm:893`) |
Neighbouring constructs are fine on all three: `torch.cat` with a width-0 operand, `new_zeros(n, 0)` into a `cat`, and indexing a width-0 tensor all convert and run.
## Environment
M4 Max, macOS 27.0 (26A5416b), `coreai-torch` 0.4.2, `coreai-core` 1.0.0b2, torch 2.13.0, Python 3.11.
## Reproducer
```python
import asyncio, inspect, shutil, numpy as np, torch, coreai_torch
from pathlib import Path
from coreai_torch import TorchConverter
from coreai.runtime import AIModel, NDArray, SpecializationOptions, ComputeUnitKind
class SplitZero(torch.nn.Module):
def forward(self, x):
a, b = x.split([x.shape[1], 0], dim=1) # b is width-0 and never used
return a + 1.0
class ZeroOutput(torch.nn.Module):
def forward(self, x):
return x[:, :0]
async def _aw(v):
return await v if inspect.isawaitable(v) else v
def build_and_run(mod, path, unit):
x = torch.rand(8, 4)
with torch.no_grad():
ep = torch.export.export(mod.eval(), (x,)).run_decompositions(coreai_torch.get_decomp_table())
c = TorchConverter()
c.add_exported_program(ep, entrypoint_name="main", input_names=["x"], output_names=["o0"])
p = c.to_coreai()
p.optimize()
shutil.rmtree(path, ignore_errors=True)
p.save_asset(Path(path))
opts = (
SpecializationOptions.cpu_only()
if unit == "cpu"
else SpecializationOptions.from_preferred_compute_unit_kind(getattr(ComputeUnitKind, unit)())
)
loop = asyncio.new_event_loop()
m = loop.run_until_complete(_aw(AIModel.load(Path(path), opts)))
f = loop.run_until_complete(_aw(m.load_function("main")))
out = loop.run_until_complete(_aw(f({"x": NDArray(x.numpy())})))
print("ok", path, unit)
import sys
build_and_run(SplitZero(), "split_zero.aimodel", sys.argv[1]) # cpu -> ok, gpu / neural_engine -> abort
# build_and_run(ZeroOutput(), "zero_out.aimodel", sys.argv[1]) # cpu -> ok, gpu / neural_engine -> abort
```
Run it once per compute unit; each abort ends the process.
## What I am asking for
An error instead of an abort, as in #67 — a caller can fall back or report, but only if the call returns.
Beyond that, the zero-sized `split` section looks like it should simply work: the empty result is dead in the graph, so nothing needs to read from a zero-byte buffer.
## Where it came from
Adding a `format="coreai"` export to Ultralytics ([ultralytics/ultralytics#25926](https://github.com/ultralytics/ultralytics/pull/25926)). Their detection postprocess splits the prediction tensor into `[4, num_classes, extras]`, where `extras` is 0 for plain detection models — so the width-0 section arrives without anyone writing anything unusual. The same graph with that section omitted runs on the Neural Engine and matches PyTorch to 3.1e-4 px on box coordinates.
Contributor guide
Research direction
Start with the supplied SplitZero and ZeroOutput reproducer, following TorchConverter conversion through AIModel.load and the inference call for CPU, GPU, and Neural Engine specializations. Done means zero-sized tensors no longer abort the process: callers receive an error they can handle, while the dead zero-sized split section works as requested.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100