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. The details are as follows.

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.

Restrictions

• For the following products, the TopK operator does not support the attribute sorted=false. When the input k is a non-constant tensor, TopK/TopKV2 is replaced with TopKV2. When the input k is a constant tensor, TopK/TopKV2 is replaced with TopKD/TopKV2D.
  • Atlas A3 training product/Atlas A3 inference product
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas 200I/500 A2 inference product
  • Atlas inference product
  • Atlas training product
• For Atlas 350 Accelerator Cards, this fusion pattern cannot be disabled.
• For Atlas 350 Accelerator Cards, the TopK/TopKV2 operator is replaced with the TopKV2 operator.

Availability

The pass effectiveness depends on whether the target product supports the TopK/TopKV2 operator. For details, see section .