Incorrect output failure of TensorRT 11.1 and 11.2 when running dynamic If → Transpose → Resize on GPU GTX 1660 SUPER and A10G

"""Standalone synthetic TensorRT If -> Transpose -> Resize reproducer."""

import argparse
import ctypes
import ctypes.util
from pathlib import Path

import numpy as np
import onnx
from onnx import helper as h
from onnx import numpy_helper as nh


def model(batch: int, mode: str, transpose: bool = True, before: bool = False) -> onnx.ModelProto:
    vi = h.make_tensor_value_info
    branch_shape = [batch, 1, "W", "H"] if transpose and before else [batch, 1, "H", "W"]
    branch_input = "ORIENTED" if transpose and before else "X"
    yes = h.make_graph([h.make_node("Add", [branch_input, "TWO"], ["YES"])], "then", [], [vi("YES", 1, branch_shape)])
    no = h.make_graph([h.make_node("Add", [branch_input, "ONE"], ["NO"])], "else", [], [vi("NO", 1, branch_shape)])
    nodes = [h.make_node("If", ["COND"], ["SELECTED"], then_branch=yes, else_branch=no)]
    if transpose and before:
        nodes.insert(0, h.make_node("Transpose", ["X"], ["ORIENTED"], perm=[0, 1, 3, 2]))
    if transpose and not before:
        nodes.append(h.make_node("Transpose", ["SELECTED"], ["ORIENTED"], perm=[0, 1, 3, 2]))
    nodes.extend(
        [
            h.make_node("Cast", ["SIZE"], ["SIZE64"], to=7),
            h.make_node("Concat", ["NC", "SIZE64"], ["TARGET"], axis=0),
            h.make_node(
                "Resize",
                ["ORIENTED" if transpose and not before else "SELECTED", "", "", "TARGET"],
                ["Y"],
                mode=mode,
                coordinate_transformation_mode="align_corners",
                cubic_coeff_a=-0.75,
            ),
        ]
    )
    weights = [
        nh.from_array(np.full((1, 1, 1, 1), 2.0, np.float32), "TWO"),
        nh.from_array(np.ones((1, 1, 1, 1), np.float32), "ONE"),
        nh.from_array(np.array([batch, 1], np.int64), "NC"),
    ]
    result = h.make_model(
        h.make_graph(
            nodes,
            "conditional-transpose",
            [vi("X", 1, [batch, 1, "H", "W"]), vi("COND", 9, []), vi("SIZE", 6, [2])],
            [vi("Y", 1, [batch, 1, "OH", "OW"])],
            initializer=weights,
        ),
        opset_imports=[h.make_opsetid("", 18)],
        ir_version=10,
    )
    onnx.checker.check_model(result, full_check=True)
    return result


def build(args: argparse.Namespace) -> None:
    import tensorrt as trt  # noqa: PLC0415 - optional GPU dependency

    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(0)
    if args.api:
        x = network.add_input("X", trt.float32, (args.batch, 1, -1, -1))
        cond = network.add_input("COND", trt.bool, ())
        size = network.add_input("SIZE", trt.int32, (2,))
        if not args.no_transpose and args.before:
            orient = network.add_shuffle(x)
            orient.second_transpose = (0, 1, 3, 2)
            x = orient.get_output(0)
        branch = network.add_if_conditional()
        branch.set_condition(cond)
        inside = branch.add_input(x).get_output(0)
        one = np.ones((1, 1, 1, 1), np.float32)
        two = one * 2
        yes = network.add_elementwise(
            inside, network.add_constant(two.shape, two).get_output(0), trt.ElementWiseOperation.SUM
        )
        no = network.add_elementwise(
            inside, network.add_constant(one.shape, one).get_output(0), trt.ElementWiseOperation.SUM
        )
        value = branch.add_output(yes.get_output(0), no.get_output(0)).get_output(0)
        if not args.no_transpose and not args.before:
            orient = network.add_shuffle(value)
            orient.second_transpose = (0, 1, 3, 2)
            value = orient.get_output(0)
        nc = np.array([args.batch, 1], np.int32)
        sizes = network.add_concatenation([network.add_constant(nc.shape, nc).get_output(0), size])
        sizes.axis = 0
        resize = network.add_resize(value)
        resize.resize_mode = getattr(trt.InterpolationMode, args.interpolation.upper())
        resize.coordinate_transformation = trt.ResizeCoordinateTransformation.ALIGN_CORNERS
        resize.cubic_coeff = -0.75
        resize.set_input(1, sizes.get_output(0))
        result = resize.get_output(0)
        result.name = "Y"
        network.mark_output(result)
    else:
        parser = trt.OnnxParser(network, logger)
        assert parser.parse(args.onnx.read_bytes())
    config = builder.create_builder_config()
    config.clear_flag(trt.BuilderFlag.TF32)
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 512 << 20)
    config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
    profile = builder.create_optimization_profile()
    profile.set_shape(
        "X",
        (args.batch, 1, 1, 1),
        (args.batch, 1, args.opt_height, args.opt_width),
        (args.batch, 1, 1536, 3072),
    )
    profile.set_shape_input("SIZE", (1, 1), (512, 512), (3072, 3072))
    config.add_optimization_profile(profile)
    plan = builder.build_serialized_network(network, config)
    assert plan is not None
    args.plan.write_bytes(bytes(plan))


def run(a: argparse.Namespace) -> None:  # noqa: PLR0915 - keep raw CUDA execution explicit
    import onnxruntime as ort  # noqa: PLC0415 - only required by the runner
    import tensorrt as trt  # noqa: PLC0415 - optional GPU dependency

    cuda = ctypes.CDLL(a.cudart or ctypes.util.find_library("cudart") or "libcudart.so")
    for name, args in [
        ("cudaSetDevice", [ctypes.c_int]),
        ("cudaMalloc", [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]),
        ("cudaFree", [ctypes.c_void_p]),
        ("cudaMemcpy", [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]),
        ("cudaStreamCreate", [ctypes.POINTER(ctypes.c_void_p)]),
        ("cudaStreamSynchronize", [ctypes.c_void_p]),
        ("cudaDeviceSynchronize", []),
        ("cudaStreamDestroy", [ctypes.c_void_p]),
    ]:
        fn = getattr(cuda, name)
        fn.argtypes = args
        fn.restype = ctypes.c_int

    def call(name: str, *args: object) -> None:
        code = getattr(cuda, name)(*args)
        if code:
            raise RuntimeError(f"{name}: CUDA {code}")

    def allocate(array: np.ndarray) -> ctypes.c_void_p:
        ptr = ctypes.c_void_p()
        call("cudaMalloc", ctypes.byref(ptr), array.nbytes)
        call("cudaMemcpy", ptr, ctypes.c_void_p(array.ctypes.data), array.nbytes, 1)
        return ptr

    call("cudaSetDevice", 0)
    runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
    engine = runtime.deserialize_cuda_engine(a.plan.read_bytes())
    context = engine.create_execution_context()
    stream = ctypes.c_void_p()
    if not a.default_stream:
        call("cudaStreamCreate", ctypes.byref(stream))
    handle = stream.value or 0
    assert context.set_optimization_profile_async(0, handle)
    call("cudaStreamSynchronize", stream)
    opt = ort.SessionOptions()
    opt.intra_op_num_threads = 1
    ref = ort.InferenceSession(str(a.onnx), sess_options=opt, providers=["CPUExecutionProvider"])
    x = (np.arange(a.batch * a.height * a.width, dtype=np.float32).reshape(a.batch, 1, a.height, a.width) % 113) * 0.1
    for condition in (True, False, True):
        feed = {"X": x, "COND": np.array(condition, np.bool_), "SIZE": np.array([19, 33], np.int32)}
        expected = ref.run(None, feed)[0]
        allocs = {}
        for name, value in feed.items():
            ptr = ctypes.c_void_p(value.ctypes.data)
            if engine.get_tensor_location(name) == trt.TensorLocation.DEVICE:
                ptr = allocate(value)
                allocs[name] = ptr
            assert context.set_input_shape(name, tuple(value.shape))
            assert context.set_tensor_address(name, ptr.value)
        assert not context.infer_shapes()
        actual = np.full(tuple(context.get_tensor_shape("Y")), np.nan, np.float32)
        out = allocate(actual)
        allocs["Y"] = out
        assert context.set_tensor_address("Y", out.value)
        call("cudaDeviceSynchronize")
        print("execute raw", a.batch, a.api, condition, flush=True)
        assert context.execute_async_v3(handle)
        call("cudaStreamSynchronize", stream)
        call("cudaMemcpy", ctypes.c_void_p(actual.ctypes.data), out, actual.nbytes, 2)
        print(
            "actual",
            actual.min(),
            actual.max(),
            "expected",
            expected.min(),
            expected.max(),
            "max_error",
            np.max(abs(actual - expected)),
            flush=True,
        )
        for ptr in allocs.values():
            call("cudaFree", ptr)
        np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-4)
    print("PASS raw", flush=True)

    if not a.default_stream:
        call("cudaStreamDestroy", stream)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=("build", "run"))
    parser.add_argument("--batch", type=int, default=4)
    parser.add_argument("--height", type=int, default=19)
    parser.add_argument("--width", type=int, default=31)
    parser.add_argument("--opt-height", type=int, default=1024)
    parser.add_argument("--opt-width", type=int, default=1024)
    parser.add_argument("--api", action="store_true", help="Bypass the ONNX parser when building TensorRT")
    parser.add_argument("--no-transpose", action="store_true")
    parser.add_argument("--before", action="store_true", help="Move Transpose before If (control case)")
    parser.add_argument("--default-stream", action="store_true")
    parser.add_argument("--cudart", help="Path to libcudart.so when unavailable on the library search path")
    parser.add_argument("--interpolation", choices=("cubic", "linear"), default="cubic")
    parser.add_argument("--directory", type=Path, default=Path("dist/source-prepost/conditional-investigation/repro"))
    args = parser.parse_args()
    args.directory.mkdir(parents=True, exist_ok=True)
    order = "plain" if args.no_transpose else ("transpose-before" if args.before else "transpose")
    stem = f"b{args.batch}-{args.interpolation}-{order}"
    args.onnx = args.directory / f"{stem}.onnx"
    args.plan = args.directory / f"{stem}-{'api' if args.api else 'onnx'}.plan"
    if args.action == "build":
        onnx.save(model(args.batch, args.interpolation, not args.no_transpose, args.before), args.onnx)
        build(args)
    else:
        run(args)


if __name__ == "__main__":
    main()