Function check_runtime

tests/python/codegen/test_target_codegen_cuda_fastmath.py:267–282

Check the runtime for the given dtype and case.

(dtype: str, case: MathCase, executable: Executable)

def check_runtime(dtype: str, case: MathCase, executable: Executable):
    """Check the runtime for the given dtype and case."""
    dev = tvm.cuda(0)

    np_inputs = make_numpy_inputs(dtype, case)
    expected = case.np_ref(*[arr.astype(dtype) for arr in np_inputs]).astype(dtype)

    tvm_inputs = [tvm.runtime.tensor(arr, device=dev) for arr in np_inputs]
    output = tvm.runtime.empty((VECTOR_N_INPUTS,), dtype, dev)

    executable(*tvm_inputs, output)
    dev.sync()

    actual = output.numpy()

    np.testing.assert_allclose(actual, expected, rtol=case.rtol, atol=case.atol)


@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])