quark.torch.quantization.graph.optimization.opt_pass_manager

Module Contents

Classes

class quark.torch.quantization.graph.optimization.opt_pass_manager.OptPassBase

Base interface for implementing passes.

It is required to implement the call function so that we can directly pass instances of the Pass directly to the PassManager and call them as a function.

We can directly pass an instance of a class implementing this interface into the PassManager’s passes attribute.

abstract call(graph_module: torch.fx.graph_module.GraphModule) → torch.fx.graph_module.GraphModule

The pass that is run through the given graph module. To implement a pass, it is required to implement this function.

Args:

graph_module: The graph module we will run a pass on

requires(graph_module: torch.fx.graph_module.GraphModule) → None

This function will be called before the pass is run and will check that the given graph module contains the preconditions needed to run the pass. It is not required to implement this function.

Args:

graph_module: The graph module we will run checks on

ensures(graph_module: torch.fx.graph_module.GraphModule) → None

This function will be called after the pass is run and will check that the given graph module contains the postconditions needed to run the pass. It is not required to implement this function.

Args:

graph_module: The graph module we will run checks on

class quark.torch.quantization.graph.optimization.opt_pass_manager.OptPassManager(passes: List[Callable[[torch.fx.graph_module.GraphModule], torch.fx.graph_module.GraphModule]] | None = None)

Construct a OPTPassManager.

Collects passes. This defines the pass schedule

Args:

passes (Optional[List[Callable]]): List of passes. A pass is a

callable which modifies an object and returns a PassResu

add_pass(_pass: Callable[[torch.fx.graph_module.GraphModule], torch.fx.graph_module.GraphModule]) → None

Adds a pass into the current list of passes.

add_checks(check: Callable[[Any], Any]) → None

Adds a function which takes runs various checks on a given graph module. This function is run before and after each pass if the run_checks_after_each_pass flag is enabled.