Implementing Branch Judgment Using SIMT

This performance optimization suggestion applies to the following models:

  • Atlas 350 Accelerator Card

[Priority] High

[Description] Batch data computation based on the SIMD programming model delivers high performance. However, when branch judgment is involved in the operator implementation logic, SIMD-based computation operations become relatively complex, resulting in performance degradation. In this case, you can consider using the SIMT mode because SIMT programming is more flexible and more suitable for scenarios involving branch judgment.

[Example] The floor_mod operator is used as an example. The operator divides each element of the input self by the corresponding element of the input other to obtain the remainder. The remainder must have the same sign as the divisor other, and its absolute value must be less than the absolute value of other. During the calculation, you need to determine the sign of each element in other and the relationship between the remainder and the absolute value of the element.

[Negative Example]

Implementation of the SIMD-based floor_mod operator: SIMD cannot directly implement the branch judgment logic. Therefore, multiple Reg vector calculation APIs are required to complete branch judgment. The related code is as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
template <typename T>
__simd_vf__ inline void floor_mod_int_simd(__ubuf__ T* dstAddr, __ubuf__ T* input1Addr, __ubuf__ T* input2Addr,
    __ubuf__ T* divAddr, const uint32_t count)
{
    uint32_t vecLen = VECTOR_LENGTH / sizeof(T);
    uint16_t loopTimes = (count + vecLen - 1) / vecLen;
    AscendC::Reg::RegTensor<T> zeroValue;
    AscendC::Reg::RegTensor<T> defaultValue;
    AscendC::Reg::RegTensor<T> signValue;
    AscendC::Reg::RegTensor<T> input1Value;
    AscendC::Reg::RegTensor<T> input2Value;
    AscendC::Reg::RegTensor<T> divValue;
    AscendC::Reg::RegTensor<T> mulValue;
    AscendC::Reg::RegTensor<T> subValue;
    AscendC::Reg::RegTensor<T> modValue;
    AscendC::Reg::RegTensor<T> modSignValue;
    AscendC::Reg::RegTensor<T> addValue;
    AscendC::Reg::RegTensor<T> input2SignValue;
    AscendC::Reg::RegTensor<T> resValue;
    AscendC::Reg::MaskReg preg;
    AscendC::Reg::MaskReg cmpValue;
    AscendC::Reg::MaskReg negValue;
    AscendC::Reg::MaskReg signNegValue;
    AscendC::Reg::MaskReg resMaskValue;
    uint32_t sregMask = count;
    AscendC::Reg::Duplicate(zeroValue, T(0));
    AscendC::Reg::Duplicate(defaultValue, T(-1));
    AscendC::Reg::Duplicate(signValue, FMOD_B32_SIGN);
    for (uint16_t j = 0; j < loopTimes; j++) {
        // handel -1
        preg = AscendC::Reg::UpdateMask<T>(sregMask);
        AscendC::Reg::DataCopy<T, AscendC::Reg::LoadDist::DIST_NORM>(input2Value, input2Addr + vecLen * j);
        AscendC::Reg::DataCopy<T, AscendC::Reg::LoadDist::DIST_NORM>(divValue, divAddr + vecLen * j);
        AscendC::Reg::Mul(mulValue, input2Value, divValue, preg);
        AscendC::Reg::DataCopy<T, AscendC::Reg::LoadDist::DIST_NORM>(input1Value, input1Addr + vecLen * j);
        AscendC::Reg::Sub(subValue, input1Value, mulValue, preg);
        AscendC::Reg::Compare<T, AscendC::CMPMODE::NE>(cmpValue, input2Value, zeroValue, preg);
        AscendC::Reg::Select(modValue, subValue, defaultValue, cmpValue);
        // post handel
        AscendC::Reg::Add(addValue, modValue, input2Value, preg);
        AscendC::Reg::Compare<T, AscendC::CMPMODE::NE>(negValue, modValue, zeroValue, preg);
        AscendC::Reg::And(input2SignValue, input2Value, signValue, preg);
        AscendC::Reg::And(modSignValue, modValue, signValue, preg);
        AscendC::Reg::Compare<T, AscendC::CMPMODE::NE>(signNegValue, modSignValue, input2SignValue, preg);
        AscendC::Reg::MaskAnd(resMaskValue, signNegValue, negValue, preg);
        AscendC::Reg::Select(resValue, addValue, modValue, resMaskValue);
        AscendC::Reg::DataCopy<T, AscendC::Reg::StoreDist::DIST_NORM>(dstAddr + vecLen * j, resValue, preg);
    }
}

[Positive Example]

SIMT-based floor_mod operator implementation: The SIMT programming mode is used to implement the calculation process, and the if else statement is used to complete branch judgment. The code is as follows, which is concise and easy to implement. For details about the complete operator implementation code, see floor_mod operator sample.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
template <typename T>
__simt_vf__ inline void floor_mod_simt(
    __ubuf__ T* self,
    __ubuf__ T* other,
    __ubuf__ T* out,
    uint32_t input_total_length)
{
    uint32_t index = threadIdx.x;
    auto rem = self[index] % other[index];
    bool signs_differ = ((rem < 0) != (other[index] < 0));
    if (signs_differ && (rem != 0)) {
        out[index] = rem + other[index];
    } else {
        out[index] = rem;
    }
}

[Performance Comparison]

As shown in the following figure, the kernel execution time of the floor_mod operator implemented based on SIMD is 4.03 μs.
Figure 1 Time required for implementing floor_mod using SIMD
As shown in the following figure, the kernel execution time of the floor_mod operator implemented based on SIMT is 3.444 μs.
Figure 2 Time required for implementing floor_mod using SIMT

When the number of cores remains unchanged, the amount of data processed by each core is the same, and data is transferred to the Unified Buffer for computation, the performance of branch judgment implemented using SIMT is 14.6% higher than that implemented using SIMD.