AutoTuneCombination

class tensorrt.plugin.AutoTuneCombination(io_types: Optional[str] = None, layouts: Optional[str] = None, tactics: Optional[Iterable[int]] = None)

Construct a set of supported type/format combinations of a plugin’s I/O.

Any custom tactic s per each such type/format combination can also be advertised. A tactic is simply another way to calculate the output of a plugin for the same type/format combination of the I/O (e.g. if there are multiple kernels available).

Parameters

• io_types (str, optional) –

  A string representation of a type combination.

  Valid format is “type0,type1,…,type#io” where ‘type’ is of the form “TYPE0[sep]TYPE1[sep]…”.

  TYPE is a valid string representation of a trt.DataType. These include “FP32” for trt.float32, “FP16” for trt.float16. The string representation of other data types is the same as their name in the trt.DataType enum.

  [sep] is a valid separator, which is either ‘|’ or ‘*’. Only one of these separators can appear in a given io_types.

  (1). ‘|’ indicates a dependent combination: the dependence of the type of one I/O to another I/O. e.g. “FP32|FP16,FP32|FP16” indicates the IO can only be both FP32 or both FP16.

  (2). ‘*’ indicates an independent combination. e.g. “FP32*FP16,FP32|FP16,FP32|FP16” indicates that the first input is independently either FP32 or FP16 regardless of the rest of the IO.

• layouts (str, optional) –

  A string representation of a format combination.

  Valid format is “format0,format1,…,format#io” where ‘format’ is of the form “FORMAT0[sep]FORMAT1[sep]…”.

  FORMAT is a valid string representation of a trt.TensorFormat. These are string versions for the enum values of trt.TensorFormat. e.g. “LINEAR” for trt.TensorFormat.LINEAR.

  [sep] is a valid separator, which is either ‘|’ or ‘*’. The rules are the same as for io_types.

• tactics (Iterable[int], optional) – Custom tactics for this type/format combination. Each custom tactic must be a positive integer. Defaults to default tactic (0).

For a plugin with 3 I/Os, I/O indices 0 and 1 are dependently either FP32/FP16 and index 2 is independently FP32/FP16.

@trtp.autotune("my::plugin")
def autotune(inp0: trtp.TensorDesc, inp1: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
    # The following would result in the following type combinations:
    # [FP32, FP32, FP32], [FP16, FP16, FP32], [FP32, FP32, FP16], [FP16, FP16, FP16]
    return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16, FP32|FP16", "LINEAR", [1, 2])]

For a plugin with 2 I/Os, the input/output supports either LINEAR or HWC format for FP32 and LINEAR format for FP16.

@trtp.autotune("my::plugin")
def autotune(inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
    # Even though (FP16, HWC) is not a valid combination (see next example), TRT should intelligently reject those
    # and pass the following combinations to the impl function:
    # [{FP32, FP32}, {LINEAR, LINEAR}], [{FP32, FP32}, {HWC, LINEAR}], [{FP16, FP32}, {LINEAR, LINEAR}]
    return [trtp.AutoTuneCombination("FP32*FP16, FP32", "LINEAR*HWC, LINEAR", [1, 2])]

For a plugin with 2 I/Os, the input/output supports either LINEAR or HWC format for FP32 and LINEAR format for FP16 (second method).

@trtp.autotune("my::plugin")
def autotune(inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
    # We can use two AutoTuneCombination objects to avoid communicating illegal combinations
    return [trtp.AutoTuneCombination("FP32*FP16, FP32", "LINEAR, LINEAR", [1, 2]), trtp.AutoTuneCombination("FP32, FP32", "HWC, LINEAR", [1, 2])]

pos(pos: Iterable[int], io_types: str, layouts: str = 'LINEAR') → None

Specify I/O types and formats for a specified set of I/O indices.

Parameters

• pos (Iterable[int]) – I/O indices. Input indices are [0, 1, …, #inputs - 1] and output indices are [#inputs, #inputs + 1, …, #inputs + #outputs - 1].

• io_types (str) – Data types for these I/O indices.

• layouts (str, optional) – Tensor format(s) for these I/O indices. Defaults to “LINEAR”.

Raises

ValueError – If types or layouts for any of these I/O indices is already specified.

For a plugin with 3 I/Os, I/O indices 0 and 1 are dependently either FP32/FP16 and index 2 is independently FP32/FP16.

@trtp.autotune("my::plugin")
def autotune(inp0: trtp.TensorDesc, inp1: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
    c = trtp.AutoTuneCombination()
    c.pos([0, 1], "FP32|FP16", "LINEAR")
    c.pos(2, "FP32*FP16") # Omitting format is the same as declaring it to be LINEAR.
    c.tactics([1, 2])
    return [c]

tactics(tactics: Iterable[int]) → None

Specify custom tactics for this type/format combination

Parameters

tactics (Iterable[int]) – Custom tactics. These must be positive integers.