CustomPassStage
Custom pass execution phase. Header file: #include <register/register_custom_pass.h>
enum class CustomPassStage : uint32_t {
kBeforeInferShape = 0,
kAfterInferShape = 1,
kAfterAssignLogicStream = 2, // only support CustomAllocateStreamPassFunc in this stage
kAfterBuiltinFusionPass = 3,
kAfterOriginGraphOptimize=4,
kCompatibleInherited=5,
kInvalid
};
- kBeforeInferShape (default value): The custom pass is executed before InferShape at the framework entry.
- kAfterInferShape: The custom pass is executed after InferShape.
If the custom pass is executed after InferShape, the pass must ensure the shape continuity after graph modification. You can use the InferShapeAndType API to ensure the continuity.
1 2 3 4 5 6 7
// 1. Obtain the output description of node1, an input node. TensorDesc output_desc; node1.GetOutputDesc(0, output_desc); // 2. Update the input description of node2. node2.UpdateInputDesc(0, output_desc); // 3. Perform InferShape for node2. operator2.InferShapeAndType();
When the InferShape function is called, the original shape of the input is flushed into the operator's input shape before InferShape runs, and the operator's output shape is flushed into its original output shape after InferShape completes. When you set InputDesc for an operator, the original shape must be explicitly set.
- kAfterAssignLogicStream: The custom pass is executed after the logical stream allocation phase. This phase only takes passes for logical stream allocation. (For details about how to register the execution function of a custom logical stream allocation pass, see CustomAllocateStreamPassFn.) Since graph modification is not allowed in this phase, any graph modification pass intended for other scenarios will fail verification and trigger an error.
- kAfterBuiltinFusionPass: The custom pass is executed after the built-in original graph fusion pass.
- kAfterOriginGraphOptimize: The custom pass is executed after the original image optimization phase.
- kCompatibleInherited: Applicable only to built-in passes. You do not need to register this field. It does not take effect if it is registered for a custom pass.
- kInvalid: It does not take effect.
Parent topic: Basic Data Structures and APIs of Graphs