llvm / llvm/llvm-project

[Bug][ONNX] DepthToSpace mode="CRD" produces wrong results for blocksize >= 3

Open
#192,916 0 comments 0 reactions 0 assignees View on GitHub
new issue
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

## Description

`DepthToSpace` with `mode="CRD"` produces **incorrect output** when `blocksize >= 3`. The `DCR` mode works correctly for all blocksizes, and `CRD` mode works for `blocksize=2`.

## Reproduction

```python
import numpy as np
import onnx
from onnx import helper, TensorProto
import onnxruntime as ort
import subprocess
from iree.compiler import compile_str
from iree import runtime as ireert

blocksize = 3
C = blocksize * blocksize * 2 # 18
x = np.arange(1 * C * 3 * 3, dtype=np.float32).reshape(1, C, 3, 3)

X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, C, 3, 3])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)
node = helper.make_node("DepthToSpace", ["X"], ["Y"], blocksize=blocksize, mode="CRD")
graph = helper.make_graph([node], "test", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
model = onnx.shape_inference.infer_shapes(model)

# ORT reference
sess = ort.InferenceSession(model.SerializeToString())
ort_out = sess.run(None, {"X": x})[0]

# IREE
onnx.save(model, "/tmp/d2s.onnx")
r = subprocess.run(["iree-import-onnx", "/tmp/d2s.onnx"], capture_output=True, text=True)
compiled = compile_str(r.stdout, input_type="onnx", target_backends=["llvm-cpu"])
config = ireert.Config("local-task")
ctx = ireert.SystemContext(config=config)
vm_module = ireert.VmModule.copy_buffer(ctx.instance, compiled)
ctx.add_vm_module(vm_module)
iree_out = ctx.modules.module.test(x).to_host()

print(f"ORT (first 8): {ort_out.flatten()[:8]}")
# [ 0. 9. 18. 1. 10. 19. 2. 11.]
print(f"IREE (first 8): {iree_out.flatten()[:8]}")
# [ 0. 9. 1. 10. 2. 11. 18. 27.]
print(f"Max diff: {np.max(np.abs(ort_out - iree_out))}")
# 68.0
```

## Test matrix

| blocksize | DCR mode | CRD mode |
|-----------|----------|----------|
| 2 | OK | OK |
| 3 | OK | **WRONG** (max_diff=68) |
| 4 | OK | **WRONG** (max_diff=86) |
| 5 | OK | **WRONG** (max_diff=208) |

## Analysis

For CRD mode, the ONNX spec defines the reshape/transpose as:
- Reshape to `(B, C', blocksize, blocksize, H, W)`
- Transpose to `(B, C', H, blocksize, W, blocksize)`
- Reshape to `(B, C', H*blocksize, W*blocksize)`

The CRD mode lowering appears to use an incorrect permutation when `blocksize >= 3`. The fact that `blocksize=2` works but `blocksize=3+` fails suggests a dimension ordering bug in the transpose indices, possibly an off-by-one in the permutation generation that only manifests when the blocksize dimension is > 2.

## Environment

- iree-base-compiler: 3.11.0 (IREE compiler version 3.11.0rc20260316)
- iree-base-runtime: 3.11.0
- Backend: llvm-cpu
- Python: 3.11
- OS: Linux

Contributor guide

Open the contributing guide

Research direction

Start by running the provided Python reproduction with blocksize 3 and compare the ONNX Runtime output against the IREE output. Trace the DepthToSpace CRD lowering reached through iree-import-onnx, focusing on the reshape and transpose permutation. Done means CRD matches ONNX Runtime for blocksizes 3, 4, and 5 while DCR and blocksize 2 remain correct.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.