Function replay_dir

tools/Polygraphy/tests/tools/conftest.py:66–109

(request)

@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, tensor_format=trt.TensorFormat.LINEAR, dtype=trt.float32, strides=(1, 2)
    ):
        io_info = [
            TensorInfo(
                tensor_format=tensor_format, 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(TensorFormat.LINEAR, DataType.FLOAT, (1, 2), -1, 1),) | Outputs: (TensorInfo(TensorFormat.LINEAR, DataType.FLOAT, (1, 2), -1, 1),)']
            """
        )
        yield dir, EXPECTED_OUTPUT