guest2180 opened a new issue, #19925:
URL: https://github.com/apache/tvm/issues/19925
When importing some ONNX models into Relax and compiling for CUDA, TVM may
lower a shape/index helper path into a host `call_pure_packed` tensor producer
followed by a CUDA `call_tir` consumer, causing a device mismatch at runtime.
This appears to be triggered by ONNX frontend lowering patterns such as
`Gather` on a singleton dimension, where the imported Relax graph contains:
- `shape_of`
- `shape_to_tensor`
- `take` on the shape tensor
- then a later `take` on CUDA data
## Expected Behavior
The imported ONNX graph should run successfully on CUDA without device
mismatch errors.
## Actual Behavior
Runtime failure similar to:
```text
ValueError: Mismatched lv16.device_type on argument #0 when calling:
take(lv16: Tensor([3], int64), B: Tensor([], int64), T_take: Tensor([],
int64)),
expected cuda
```
## Root Cause Analysis
The issue is not in ONNX semantics themselves, but in the Relax lowering
path after ONNX import.
`relax.shape_to_tensor` is registered with
`FCallPacked("relax.run.shape_to_tensor")`, so after `LegalizeOps` it becomes a
packed runtime call instead of `call_tir`.
Then `relax.run.shape_to_tensor` constructs its return value using:
```python
tvm.runtime.tensor([int(v) for v in shape_tuple])
```
Since `tvm.runtime.tensor(..., device=None)` defaults to CPU, this produces
a host tensor.
Later, `relax.take` is legalized to CUDA `call_tir`, so CUDA TIR receives a
CPU tensor as input and fails device checking.
## Minimal Relax Pattern
```python
shape = R.shape_of(x)
shape_tensor = R.shape_to_tensor(shape)
idx = R.take(shape_tensor, R.const(1, "int64"), axis=0)
y = R.take(x, idx, axis=1)
```
After legalization, this becomes conceptually:
```python
shape_tensor = R.call_pure_packed("relax.run.shape_to_tensor", ...)
idx = R.call_tir(take, (shape_tensor, ...), ...)
y = R.call_tir(take, (x, idx), ...)
```
## ONNX behavior
The triggering graph comes from ONNX import, for example a `ReduceMax +
Constant + Gather` tail that is semantically close to `Squeeze(axis=1)`.
ONNX Runtime executes the original ONNX graph correctly.
The failure happens because the ONNX frontend imports the graph into a Relax
form that introduces `shape_to_tensor`, and that path later crosses host/device
boundaries incorrectly.
## Source Pointers
Relevant upstream code in current `main`:
- `relax.shape_to_tensor` still uses
`FCallPacked("relax.run.shape_to_tensor")`
- `LegalizeOps` still prefers `FLegalize`, otherwise falls back to
`FCallPacked`
- `relax.run.shape_to_tensor` still returns `tvm.runtime.tensor(...)`
- similar host-tensor risk may also exist for `relax.nonzero` and
`relax.unique`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
