Function test_conv2d_offload

tests/python/relax/test_codegen_cudnn.py:179–210

(data_shape, weight_shape, dtype, with_bias, activation)

    ],
)
def test_conv2d_offload(data_shape, weight_shape, dtype, with_bias, activation):
    input = np.random.randn(*data_shape).astype(dtype)
    weight = np.random.randn(*weight_shape).astype(dtype)

    if with_bias:
        oc = weight_shape[0]
        bias = np.random.randn(1, 1, 1, oc).astype(dtype)
        args = (input, weight, bias)
    else:
        bias = None
        args = (input, weight)

    activation = _activation_table[activation]

    mod = get_relax_conv2d_module(
        data_shape,
        weight_shape,
        dtype,
        with_bias=with_bias,
        activation=activation,
    )

    out = get_result_with_relax_cudnn_offload(mod, args)
    ref = build_and_run(mod, args, "llvm", legalize=True)
    if dtype == "float16":
        # FIXME(lei): currently raise into 3e-1 to prevent flaky test
        # see https://github.com/apache/tvm/pull/18319
        tvm.testing.assert_allclose(out, ref, rtol=3e-1, atol=3e-1)
    else:
        # Increased tolerance to 2.5e-2 to prevent flaky test due to numerical
        # differences between cuDNN and LLVM implementations
        tvm.testing.assert_allclose(out, ref, rtol=2.5e-2, atol=2.5e-2)


@pytest.mark.skip(reason="flaky test")