TopKFusionPass

Description

Replaces the TopK/TopKV2 operator with the TopKV2 or TopKD/TopKV2D operator based on the platform, or directly splits the TopK/TopKV2 node into the combination of SegmentSort and MultiMerge. Specifically:

Scenario 1: The TopK/TopKV2 operator is replaced with the TopKV2 operator.

Scenario 2: The TopK/TopKV2 operator is replaced with the TopKD/TopKV2D operator.

Scenario 3: The TopK/TopKV2 operator is split into the combination of SegmentSort and MultiMerge.

Constraints

• For the following products, the TopK operator does not support the attribute sorted set to false. When the input k is a non-const tensor, TopK/TopKV2 is replaced with TopKV2. When the input k is a const tensor, TopK/TopKV2 is replaced with TopKD/TopKV2D.
  • Atlas A3 training products/Atlas A3 inference products
  • Atlas A2 training products/Atlas A2 inference products
  • Atlas 200I/500 A2 inference products
  • Atlas inference products
  • Atlas training products
• For Ascend 950PR/Ascend 950DT, this fusion pattern cannot be disabled.
• For Ascend 950PR/Ascend 950DT, the TopK/TopKV2 operator is replaced with the TopKV2 operator.

Applicable Products

The pass effectiveness depends on whether the target product supports the TopK/TopKV2 operator. For details, see Ascend IR Operator Specifications.