Cube–Vector software pipelining

This document describes the CV pipelining pass in HIVM. It optimizes CV-style kernels. Before reading this document, we recommend CV Optimization for CV terminology.

Hardware Background

Each Ascend Core contains Cube and Vector units, performing matrix-multiplication and other vectorized computations respectively. These cores are capable of running in parallel to each other asynchronously whenever there are no data dependencies. In order to achieve better performance, it is essential to leverage this capability for better hardware utilization.

CV-Pipelining optimizes loops in Mix kernels where multiple Cube and Vector instructions depend on each other (e.g. FlashAttention). By running Vector and Cube in parallel, it improves hardware utilization (ILP) and performance.

The pass uses multi-buffering, which can increase UB usage; tune the number of pipeline stages for your workload.

Algorithm Principle

The pass finds suitable scf.for loops, splits Cube and Vector into separate work items, sets up data dependencies, expands tensors that need multi-buffering, unrolls the loop, and places each work item in its own loop.

Before:

scf.for 0 to N step S {
    %c = Cube() : tensor<16x16xf32>
    %v = Vector(%c) : tensor<16x16xf32>
    %c1 = Cube(%v) : tensor<16x16xf32>
    %v1 = Vector(%c1) : tensor<16x16xf32>
}

After:

scf.for 0 to N step 3*S {
    %c = scf.for 0 to 3 -> tensor<3x16x16xf32> {
        Cube();
        tensor.insert_slice
    } {cube_loop}
    %v = scf.for 0 to 3 -> tensor<3x16x16xf32> {
        %c_slice = extract_slice %c
        Vector(%c_slice) : tensor<16x16xf32>
        tensor.insert_slice
    } {vector_loop}
    %c1 = scf.for 0 to 3 -> tensor<3x16x16xf32> {
        %v_slice = extract_slice %v
        Cube(%v_slice) : tensor<16x16xf32>
        tensor.insert_slice
    } {cube_loop}
    // When no other Work Item needs the result, no buffer expansion needed
    %v1 = scf.for 0 to 3 -> tensor<16x16xf32> {
        %c_slice = extract_slice %c1
        Vector(%c_slice) : tensor<16x16xf32>
    } {vector_loop}
}

API Description

Compile Options

Option	Default Value	Description
set-workspace-multibuffer	2	Number of software pipeline stages and multi-buffering count

Constraints and Capabilities

1. Only scf.for and scf.if in the pipelined loop may have regions/blocks, and each region must contain only Cube or only Vector instructions.

2. Cross-iteration dependencies must be separated into distinct work items.

   • CV-pipelining is not applied when v0 and v1 cannot be in the same work item (due to a Cube between them) but arg0 is defined in v1 and used by v0.

   • If Cube does not use v0, v0 is placed in the same work item as v1 and CV-pipelining applies.

scf.for iter_args(%arg0 = %init) {
    %v0 = Vector(%arg0)
    %c = Cube(%v0)
    %v1 = Vector(%c)
    yield %v1
}