nvidia--tensorrt
wehub-resource-sync
129 行
4.7 KiB
Python
129 行
4.7 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import os
|
|
import tempfile
|
|
from collections import namedtuple
|
|
from textwrap import dedent
|
|
from typing import List
|
|
|
|
import pytest
|
|
import tensorrt as trt
|
|
from polygraphy.backend.trt import Algorithm, TacticReplayData, TensorInfo
|
|
from polygraphy.json import save_json
|
|
|
|
|
|
def make_poly_fixture(subtool: List[str]):
|
|
@pytest.fixture()
|
|
def poly_fixture(script_runner):
|
|
def poly_fixture_impl(
|
|
additional_opts: List[str] = [], expect_error: bool = False, *args, **kwargs
|
|
):
|
|
cmd = ["polygraphy"] + subtool + ["-v"] + additional_opts
|
|
# NOTE: script_runner does not work very well in `in-process`` mode if you need to inspect stdout/stderr.
|
|
# Occasionally, the output comes out empty - not clear why. Cave emptor!
|
|
# Decorate your tests with `@pytest.mark.script_launch_mode("subprocess")` to use `subprocess` to avoid this issue.
|
|
status = script_runner.run(*cmd, *args, **kwargs)
|
|
assert status.success == (not expect_error)
|
|
return status
|
|
|
|
return poly_fixture_impl
|
|
|
|
return poly_fixture
|
|
|
|
|
|
poly = make_poly_fixture([])
|
|
poly_run = make_poly_fixture(["run"])
|
|
poly_convert = make_poly_fixture(["convert"])
|
|
poly_inspect = make_poly_fixture(["inspect"])
|
|
poly_check = make_poly_fixture(["check"])
|
|
poly_multi_device_shard = make_poly_fixture(["multi-device", "shard"])
|
|
poly_surgeon = make_poly_fixture(["surgeon"])
|
|
poly_surgeon_extract = make_poly_fixture(["surgeon", "extract"])
|
|
poly_template = make_poly_fixture(["template"])
|
|
poly_template_shard = make_poly_fixture(["template", "shard-hints"])
|
|
poly_debug = make_poly_fixture(["debug"])
|
|
poly_data = make_poly_fixture(["data"])
|
|
poly_plugin_match = make_poly_fixture(["plugin", "match"])
|
|
poly_plugin_list_plugins = make_poly_fixture(["plugin", "list"])
|
|
poly_plugin_replace = make_poly_fixture(["plugin", "replace"])
|
|
|
|
FakeAlgorithmContext = namedtuple(
|
|
"FakeAlgorithmContext", ["name", "num_inputs", "num_outputs"]
|
|
)
|
|
FakeAlgorithm = namedtuple("FakeAlgorithm", ["algorithm_variant", "io_info"])
|
|
FakeAlgorithm.get_algorithm_io_info = lambda this, index: this.io_info[index]
|
|
|
|
FakeAlgorithmVariant = namedtuple("FakeAlgorithmVariant", ["implementation", "tactic"])
|
|
|
|
|
|
@pytest.fixture(scope="session", params=["", "subdir"])
|
|
def replay_dir(request):
|
|
def fake_context(name, num_inputs=1, num_outputs=1):
|
|
return FakeAlgorithmContext(
|
|
name=name, num_inputs=num_inputs, num_outputs=num_outputs
|
|
)
|
|
|
|
def fake_algo(
|
|
implementation=6, tactic=0, num_io=2, dtype=trt.float32, strides=(1, 2)
|
|
):
|
|
io_info = [
|
|
TensorInfo(
|
|
dtype=dtype,
|
|
strides=strides,
|
|
vectorized_dim=-1,
|
|
components_per_element=1,
|
|
)
|
|
] * num_io
|
|
return FakeAlgorithm(
|
|
algorithm_variant=FakeAlgorithmVariant(implementation, tactic),
|
|
io_info=io_info,
|
|
)
|
|
|
|
def make_replay(tactic):
|
|
return TacticReplayData().add(
|
|
"layer0", Algorithm.from_trt(fake_context("layer0"), fake_algo(0, tactic))
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as dir:
|
|
|
|
def make_path(prefix, *args):
|
|
path = os.path.join(dir, prefix)
|
|
if request.param:
|
|
path = os.path.join(path, request.param)
|
|
path = os.path.join(path, *args)
|
|
return path
|
|
|
|
# Good tactics
|
|
save_json(make_replay(0), make_path("good", "0.json"))
|
|
save_json(make_replay(1), make_path("good", "1.json"))
|
|
|
|
# Bad tactics
|
|
save_json(make_replay(1), make_path("bad", "0.json"))
|
|
save_json(make_replay(2), make_path("bad", "1.json"))
|
|
|
|
EXPECTED_OUTPUT = dedent(
|
|
"""
|
|
[I] Loaded {num} good tactic replays.
|
|
[I] Loaded {num} bad tactic replays.
|
|
[I] Found potentially bad tactics:
|
|
[I] Layer: layer0
|
|
Algorithms: ['(Implementation: 0, Tactic: 2) | Inputs: (TensorInfo(DataType.FLOAT, (1, 2), -1, 1),) | Outputs: (TensorInfo(DataType.FLOAT, (1, 2), -1, 1),)']
|
|
"""
|
|
)
|
|
yield dir, EXPECTED_OUTPUT
|