()
| 75 | |
| 76 | |
| 77 | def test_tensorrt_offload(): |
| 78 | @memoize("relax.tests.test_codegen_tensorrt.conv2d_residual") |
| 79 | def get_ref(): |
| 80 | data_np = np.random.randn(1, 64, 56, 56).astype("float32") |
| 81 | weight1_np = np.random.randn(64, 64, 3, 3).astype("float32") |
| 82 | weight2_np = np.random.randn(64, 64, 3, 3).astype("float32") |
| 83 | inputs = [data_np, weight1_np, weight2_np] |
| 84 | ref = build_and_run(Conv2dResidualBlock, inputs, "llvm", legalize=True) |
| 85 | return inputs, ref |
| 86 | |
| 87 | inputs, ref = get_ref() |
| 88 | |
| 89 | conv_pat = make_fused_bias_activation_pattern( |
| 90 | "relax.nn.conv2d", with_bias=False, activation=None |
| 91 | ) |
| 92 | relu_pat = is_op("relax.nn.relu")(wildcard()) |
| 93 | add_pat = is_op("relax.add")(wildcard(), wildcard()) |
| 94 | |
| 95 | patterns = [ |
| 96 | ("tensorrt.nn.conv2d", conv_pat), |
| 97 | ("tensorrt.nn.relu", relu_pat), |
| 98 | ("tensorrt.add", add_pat), |
| 99 | ] |
| 100 | |
| 101 | params_np = {"weight1": inputs[1], "weight2": inputs[2]} |
| 102 | |
| 103 | mod = tvm.transform.Sequential( |
| 104 | [ |
| 105 | relax.transform.BindParams("main", params_np), |
| 106 | relax.transform.FuseOpsByPattern(patterns), |
| 107 | relax.transform.MergeCompositeFunctions(), |
| 108 | relax.transform.RunCodegen(), |
| 109 | ] |
| 110 | )(Conv2dResidualBlock) |
| 111 | |
| 112 | out = build_and_run(mod, inputs[:1], "cuda") |
| 113 | |
| 114 | tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3) |
| 115 | |
| 116 | |
| 117 | def _offload_and_compare(mod, params_np, patterns, data_np, rtol=1e-2, atol=1e-2): |
nothing calls this directly
no test coverage detected
searching dependent graphs…