MmadWithSparse

Applicability

Function Usage

Performs matrix multiplication and addition operations. The passed left matrix A is a sparse matrix, and the passed right matrix B is a dense matrix. For matrix A, densification is completed during MmadWithSparse computation. For matrix B, densification is automatically completed during input data preparation before computation execution (densification is performed according to the densification algorithm described below). Therefore, matrix B passed to this API is a dense matrix. The dense matrix B needs to be loaded by calling LoadDataWithSparse, and the index matrix needs to be loaded at the same time. The index matrix is generated during the densification of matrix B and then used for the densification of matrix A.

Prototype

1
2

template <typename T = int32_t, typename U = int8_t, typename Std::enable_if<Std::is_same<PrimT<T>, int32_t>::value, bool>::type = true, typename Std::enable_if<Std::is_same<PrimT<U>, int8_t>::value, bool>::type = true>
__aicore__ inline void MmadWithSparse(const LocalTensor<T>& dst, const LocalTensor<U>& fm, const LocalTensor<U>& filter, const MmadParams& mmadParams)

Parameters

Restrictions

• In the original sparse matrix B, there should be a maximum of two non-zero elements in every four elements. If there are three or more non-zero elements, only the first two non-zero elements are used.
• If any of M, K, and N is 0, the instruction is not executed.

• For details about the operand address alignment requirements, see General Address Alignment Restrictions.

Dense Algorithm Description

It is assumed that there are at least two zeros in every four elements of the original sparse matrix B, and the matrix B after densification is a dense matrix in which two zeros are filtered out from every four elements. An index matrix is generated in a densification process of the matrix B. The process is as follows: For every four elements in the sparse matrix B, two 2-bit indexes are generated in the index matrix, and encoding is performed according to the following rule. The index must be in the range of {0, 1, 2}.

• The first index is used to indicate the relative location of the first non-zero element in the first three elements.
• The second index is used to indicate the relative location of the second non-zero element in the last three elements.

For details, see the following table. - indicates that the algorithm does not care about the value at the position because the value will be filtered.

The index matrix is used for densification of matrix A. Based on the index matrix, two elements are selected from the four elements in matrix A for computation, as shown in the following figure.

Examples

 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

int srcOffset = 0;
int dstOffset = 0;
AscendC::LocalTensor<int8_t> a1Local = inQueueA1.DeQue<int8_t>();
AscendC::LocalTensor<int8_t> a2Local = inQueueA2.AllocTensor<int8_t>();

// The a1Local matrix is not sparse.
AscendC::LoadData2DParams loadDataParams;
loadDataParams.repeatTimes = kBlocks * mBlocks;
loadDataParams.srcStride = 1;
loadDataParams.ifTranspose = false;

AscendC::LoadData(a2Local, a1Local, loadDataParams);

inQueueA2.EnQue<int8_t>(a2Local);
inQueueA1.FreeTensor(a1Local);

AscendC::LocalTensor<int8_t> b2Local = inQueueB2.AllocTensor<int8_t>();

// Transform NZ to ZN. Loading a sparse matrix requires using sparse indices together with the loadDataWithSparse instruction to properly load sparse data.
AscendC::LoadData2DParams loadDataParams;
loadDataParams.repeatTimes = kBlocks * nBlocks / 2;
loadDataParams.srcStride = 0;
loadDataParams.ifTranspose = false;

AscendC::LoadDataWithSparse(b2Local, b1Local, idxb1Local, loadDataParams);

inQueueB2.EnQue<int8_t>(b2Local);

AscendC::LocalTensor<int8_t> b2Local = inQueueB2.DeQue<int8_t>();
AscendC::LocalTensor<int32_t> c1Local = outQueueCO1.AllocTensor<int32_t>();

// mmad requires the matrix dimensions to be specified for computation.
uint32 m = 16;
uint32 k = 64;
uint32 n = 16;
AscendC::MmadWithSparse(c1Local, a2Local, b2Local, { m, n, k, false, 0, false, false, false });