guest2180 opened a new issue, #19887:
URL: https://github.com/apache/tvm/issues/19887
TensorRT BYOC offload gaps on YOLO-style graphs: operators run on CUDA
fallback and/or direct TensorRT, but fail to offload cleanly through TVM BYOC
## Summary
We are using **TVM Relax BYOC** with **TensorRT** as the external codegen
backend.
This came up while trying to run a **YOLO-style model** through TVM TensorRT
BYOC. After reducing the problem, we found operator-level gaps where an op can
execute in CUDA fallback and/or direct TensorRT, but does not offload cleanly
through the TVM TensorRT BYOC path.
We found a case where an operator:
1. runs correctly with normal TVM CUDA fallback
2. runs correctly with direct TensorRT
3. but **fails when offloaded through TVM TensorRT BYOC**
The concrete operator we can prove with a self-contained reproducer is:
- `relax.permute_dims`
- `relax.image.resize2d`
In addition, while working on YOLO-style graphs, we also observed related
offload/partition problems around:
- `relax.strided_slice`
- `relax.concat`
- `relax.split`
- many `relax.reshape` glue nodes
This suggests the issue is not that:
- CUDA cannot execute the operator
- TensorRT cannot execute the operator
Instead, the problem appears to be in the **TVM TensorRT BYOC offload /
operator mapping path**.
## Environment
- TVM: local TensorRT BYOC build under `/home/perception/tvm025`, using the
current `#19810` TensorRT BYOC file set
- TensorRT: `10.16.1.11`
- Python: `3.11`
- Target: `cuda`
## Expected behavior
If an operator can:
1. run on CUDA fallback in TVM
2. be implemented and executed directly in TensorRT
then the same operator should be offloadable through TVM TensorRT BYOC, or
at least fail with a clearer explanation about missing BYOC support.
## Actual behavior
For `permute_dims`:
- TVM CUDA fallback works
- direct TensorRT works
- TVM TensorRT BYOC fails with:
```text
InternalError: Check failed: (it != map.end()) is false:
tensorrt.permute_dims: Unsupported operator
```
This means the operator reaches the TensorRT BYOC path, but the TensorRT
runtime/builder side cannot handle the generated BYOC operator name.
For `resize2d`:
- in our current single-op BYOC reproducer, it forms a TensorRT region
- but execution then fails with:
```text
InternalError: Check failed: (it != map.end()) is false:
tensorrt.image.resize2d: Unsupported operator
```
- so `resize2d` has the same class of BYOC lower/runtime-side support gap as
`permute_dims`
For other operators seen in YOLO-style graphs, the observed problems are:
1. `permute_dims`
- explicit BYOC runtime failure
- reaches TensorRT BYOC region
- fails with `tensorrt.permute_dims: Unsupported operator`
2. `resize2d`
- single-op test produces `trt_regions: 1`
- but runtime fails with `tensorrt.image.resize2d: Unsupported operator`
3. `strided_slice`
- does not form a TensorRT region in our single-op offload test
- so it currently appears to stay on the TVM side instead of being
offloaded
4. `concat`
- can be offloaded as a small standalone region in some experiments
- but becomes problematic when trying to form larger TensorRT regions on
YOLO-style graphs
5. `split`
- similarly acts as a partition boundary in YOLO-style graphs
- it prevents larger continuous TensorRT regions from forming cleanly
6. `reshape`
- not necessarily individually unsupported
- but YOLO-style graphs contain many reshape glue nodes, and they are
part of the fragmentation problem when trying to build larger TensorRT regions
## Minimal reproduction
The attached script is self-contained and demonstrates current BYOC behavior
for single operators, including:
1. `permute_dims`
2. `strided_slice`
3. `resize2d`
Script:
```python
from __future__ import annotations
import traceback
from dataclasses import dataclass
import numpy as np
import tvm
from tvm import relax
from tvm.relax import transform
from tvm.relax.dpl import is_op, wildcard
@dataclass(slots=True)
class Case:
name: str
pattern_name: str
mod: tvm.IRModule
input_shape: tuple[int, ...]
def build_permute_dims_case() -> Case:
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorStructInfo((1, 3, 4, 5), "float32"))
with bb.function("main", [data]):
with bb.dataflow():
out = bb.emit(relax.op.permute_dims(data, axes=[0, 2, 3, 1]))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return Case(
name="permute_dims",
pattern_name="tensorrt.permute_dims",
mod=bb.finalize(),
input_shape=(1, 3, 4, 5),
)
def build_strided_slice_case() -> Case:
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorStructInfo((1, 3, 8, 8), "float32"))
with bb.function("main", [data]):
with bb.dataflow():
out = bb.emit(
relax.op.strided_slice(
data,
axes=[2, 3],
begin=[1, 1],
end=[7, 7],
strides=[2, 2],
)
)
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return Case(
name="strided_slice",
pattern_name="tensorrt.strided_slice",
mod=bb.finalize(),
input_shape=(1, 3, 8, 8),
)
def build_resize2d_case() -> Case:
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorStructInfo((1, 3, 4, 4), "float32"))
with bb.function("main", [data]):
with bb.dataflow():
out = bb.emit(
relax.op.image.resize2d(
data,
size=(8, 8),
layout="NCHW",
method="nearest_neighbor",
coordinate_transformation_mode="asymmetric",
rounding_method="floor",
)
)
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return Case(
name="resize2d",
pattern_name="tensorrt.image.resize2d",
mod=bb.finalize(),
input_shape=(1, 3, 4, 4),
)
def build_concat_case() -> Case:
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorStructInfo((1, 3, 4, 4), "float32"))
with bb.function("main", [data]):
with bb.dataflow():
out = bb.emit(relax.op.concat((data, data), axis=1))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return Case(
name="concat",
pattern_name="tensorrt.concatenate",
mod=bb.finalize(),
input_shape=(1, 3, 4, 4),
)
def partition_single_op(case: Case) -> tvm.IRModule:
if case.name == "concat":
pattern = is_op("relax.concat")(wildcard())
else:
pattern = is_op(f"relax.{case.name}")(wildcard()) if case.name !=
"resize2d" else is_op("relax.image.resize2d")(wildcard(), wildcard())
# Build patterns explicitly per case to avoid relying on local helper
files.
if case.name == "permute_dims":
patterns = [(case.pattern_name,
is_op("relax.permute_dims")(wildcard()))]
elif case.name == "strided_slice":
patterns = [(case.pattern_name,
is_op("relax.strided_slice")(wildcard()))]
elif case.name == "resize2d":
patterns = [(case.pattern_name,
is_op("relax.image.resize2d")(wildcard(), wildcard()))]
elif case.name == "concat":
patterns = [(case.pattern_name, is_op("relax.concat")(wildcard()))]
else:
raise ValueError(case.name)
mod = transform.FuseOpsByPattern(
patterns,
bind_constants=True,
annotate_codegen=False,
)(case.mod)
mod = transform.MergeCompositeFunctions()(mod)
mod = transform.RunCodegen()(mod)
return mod
def count_trt_regions(mod: tvm.IRModule) -> int:
return mod.script().count('R.call_dps_packed("')
def run_case(case: Case) -> None:
print(f"\n=== case: {case.name} ===")
try:
lowered = partition_single_op(case)
print("run_codegen_ok: true")
print("trt_regions:", count_trt_regions(lowered))
with tvm.transform.PassContext(opt_level=0):
ex = tvm.compile(lowered, "cuda")
dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)
inp = np.random.rand(*case.input_shape).astype("float32")
out = vm["main"](tvm.runtime.tensor(inp, dev))
print("runtime_ok: true")
if hasattr(out, "shape"):
print("output_shape:", out.shape)
else:
print("output_type:", type(out).__name__)
except Exception as err: # noqa: BLE001
print("runtime_ok: false")
print("error_type:", type(err).__name__)
print("error_message:", err)
print("[traceback_begin]")
print(traceback.format_exc().rstrip())
print("[traceback_end]")
def main() -> None:
cases = [
build_permute_dims_case(),
build_strided_slice_case(),
build_resize2d_case(),
]
for case in cases:
run_case(case)
if __name__ == "__main__":
main()
```
## Repro command
```bash
export TVM_HOME=/path/to/tvm
export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH
python repro_trt_single_op_offload.py
```
## Observed output
```text
=== case: permute_dims ===
run_codegen_ok: true
trt_regions: 1
runtime_ok: false
error_type: InternalError
error_message: Check failed: (it != map.end()) is false:
tensorrt.permute_dims: Unsupported operator
=== case: resize2d ===
run_codegen_ok: true
trt_regions: 1
runtime_ok: false
error_type: InternalError
error_message: Check failed: (it != map.end()) is false:
tensorrt.image.resize2d: Unsupported operator
```
## Why this demonstrates a BYOC offload gap
These reproducers isolate two distinct offload-gap modes:
- `permute_dims`: the op reaches the TensorRT BYOC runtime path, but fails
there as unsupported
- `resize2d`: the op also reaches the TensorRT BYOC runtime path, but fails
there as unsupported
- `strided_slice`: the op does not form a TensorRT region and therefore
stays on the TVM side
So this is not just a generic runtime limitation of CUDA or TensorRT.
It is specifically a **TVM TensorRT BYOC offload support gap** in the
partitioning / lowering path.
## Additional YOLO context
This issue was found while trying to run a YOLO-style model with TensorRT
BYOC.
On that graph, the main difficulty is not the backbone convolution
operators, but the neck/tail glue operators. In practice, these operators
fragment the partition or fail to lower cleanly when we try to build larger
TensorRT regions.
Examples observed on the YOLO-style graph include:
- `permute_dims`
- `image.resize2d`
- `strided_slice`
- `concat`
- `split`
- many `reshape`
So while the minimal reproducer here uses `permute_dims`, the broader
real-world motivation is that these gaps make YOLO-style TensorRT BYOC offload
much harder than expected.
## Notes
One of the failures is currently triggered during the TensorRT BYOC
execution path with:
```text
tensorrt.permute_dims: Unsupported operator
```
This may indicate a mismatch between:
1. the BYOC pattern/operator name emitted by Relax partitioning
2. the operator names registered in the TensorRT runtime converter table
If helpful, I can also provide similar minimal reproducers for other
YOLO-related operators, especially `concat` and `split`, or the full model and
script I used to test yolo.
--
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]
