wehub-resource-sync
209 lines
6.5 KiB
Plaintext
209 lines
6.5 KiB
Plaintext
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
/*
|
|
DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a
|
|
https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py
|
|
*/
|
|
|
|
__global__ void NAME (TYPE* A __readonly __noalias __aligned(16),
|
|
TYPE* B __readonly __noalias __aligned(16),
|
|
TYPE* C __noalias __aligned(16),
|
|
int lda __multipleof(8),
|
|
int ldb __multipleof(8),
|
|
int ldc __multipleof(8),
|
|
long stride_za __multipleof(8),
|
|
long stride_zb __multipleof(8),
|
|
long stride_zc __multipleof(8),
|
|
long stride_ha __multipleof(8),
|
|
long stride_hb __multipleof(8),
|
|
long stride_hc __multipleof(8),
|
|
int DS0, int DS1,
|
|
int SDD_K __multipleof(16),
|
|
int SDD_off_width,
|
|
int* lut, int* locks, int nlocks) {
|
|
/* ---------------- */
|
|
/* Prologue */
|
|
/* ---------------- */
|
|
// program ids
|
|
int pid0 = get_program_id(0);
|
|
int pid1 = get_program_id(1);
|
|
int pidz = get_program_id(2);
|
|
#ifdef SDD
|
|
// load LUT header
|
|
pid1 = pid1 + SDD_off_width;
|
|
int blockidm[TM] = (0 ... TM) / BLOCK;
|
|
int blockidn[TN] = (0 ... TN) / BLOCK;
|
|
int offlutm[TM] = blockidm*(TN/BLOCK)*4;
|
|
int offlutn[TN] = blockidn*4;
|
|
int *header = lut + pid1 * (TM/BLOCK) * (TN/BLOCK) * 4;
|
|
int z = *(header + 0);
|
|
int i[TM] = *(header + 1 + offlutm);
|
|
int j[TN] = *(header + 2 + offlutn);
|
|
int AS1 = SDD_K / TZ;
|
|
int lockid = select(TZ > 1, 1, 0);
|
|
int offka = pid0 * AS1;
|
|
int offkb = pid0 * AS1;
|
|
int offmc = 0;
|
|
int offnc = 0;
|
|
int offpa = 0;
|
|
int offpb = 0;
|
|
int maxid = TZ;
|
|
int offhc = 0;
|
|
int offha = z;
|
|
int offhb = z;
|
|
int ram[TM] = i*BLOCK + ((0 ... TM) % BLOCK);
|
|
int rbn[TN] = j*BLOCK + ((0 ... TN) % BLOCK);
|
|
#else
|
|
// load LUT header
|
|
int *header = lut + pid0 * 6;
|
|
int offset = *(header + 0);
|
|
int AS1 = *(header + 1);
|
|
int column = *(header + 2);
|
|
int depth = *(header + 3);
|
|
int lockid = *(header + 4);
|
|
int maxid = *(header + 5);
|
|
int *pinc = lut + offset;
|
|
int offhc = depth;
|
|
#ifdef DSD
|
|
// output offset
|
|
int offnc = pid1 * TN;
|
|
int offmc = column * TM;
|
|
int offpc = 0;
|
|
// dense input offset
|
|
int offnb = pid1 * TN;
|
|
int offkb __multipleof(8) = *pinc;
|
|
int offpb = 0;
|
|
// sparse input offset
|
|
int offma = 0;
|
|
int offka = 0;
|
|
long offpa __multipleof(8) = *(pinc + 1);
|
|
offpa = offpa * BLOCK * BLOCK;
|
|
int offha = 0;
|
|
int offhb = depth;
|
|
#endif
|
|
#ifdef DDS
|
|
// output offset
|
|
int offmc = pid1 * TM;
|
|
int offnc = column * TN;
|
|
int offpc = 0;
|
|
// dense input offset
|
|
int offma = pid1 * TM;
|
|
int offka __multipleof(8) = *pinc;
|
|
int offpa = 0;
|
|
// sparse input offset
|
|
int offnb = 0;
|
|
int offkb = 0;
|
|
long offpb __multipleof(8) = *(pinc + 1);
|
|
offpb = offpb * BLOCK * BLOCK;
|
|
int offha = depth;
|
|
int offhb = 0;
|
|
#endif
|
|
int ram[TM] = offma + 0 ... TM;
|
|
int rbn[TN] = offnb + 0 ... TN;
|
|
#endif
|
|
// initialize a, b pointers
|
|
int rka[TK] = offka + 0 ... TK;
|
|
int rkb[TK] = offkb + 0 ... TK;
|
|
TYPE* pa[TM, TK] = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, newaxis] * STRIDE_AM + rka[newaxis, :] * STRIDE_AK;
|
|
TYPE* pb[TK, TN] = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[newaxis, :] * STRIDE_BN + rkb[:, newaxis] * STRIDE_BK;
|
|
// pre-fetch
|
|
#ifdef DDS
|
|
bool checkam[TM, TK] = ram[:, newaxis] < DS0;
|
|
#else
|
|
bool checkam[TM, TK] = AS1 > 0;
|
|
#endif
|
|
#ifdef DSD
|
|
bool checkbn[TK, TN] = rbn[newaxis, :] < DS0;
|
|
#else
|
|
bool checkbn[TK, TN] = AS1 > 0;
|
|
#endif
|
|
TYPE a[TM, TK] = checkam ? *pa : 0;
|
|
TYPE b[TK, TN] = checkbn ? *pb : 0;
|
|
|
|
/* ---------------- */
|
|
/* Inner Loop */
|
|
/* ---------------- */
|
|
// create result tile
|
|
float acc[TM, TN] = 0;
|
|
int step = TK;
|
|
for(int k = AS1; k > 0; k -= step) {
|
|
acc += a @ b;
|
|
// update pointers
|
|
#ifdef SDD
|
|
int inc_a = TK * STRIDE_AK;
|
|
int inc_b = TK * STRIDE_BK;
|
|
#else
|
|
pinc += 2;
|
|
#ifdef DSD
|
|
int inc_b __multipleof(8) = *pinc;
|
|
int inc_a __multipleof(8) = *(pinc + 1);
|
|
inc_b = inc_b * STRIDE_BK;
|
|
#endif
|
|
#ifdef DDS
|
|
int inc_a __multipleof(8) = *pinc;
|
|
int inc_b __multipleof(8) = *(pinc + 1);
|
|
inc_a = inc_a * STRIDE_AK;
|
|
#endif
|
|
#endif
|
|
pa += inc_a;
|
|
pb += inc_b;
|
|
// pre-fetch
|
|
bool checkak[TM, TK] = k > TK;
|
|
bool checkbk[TK, TN] = k > TK;
|
|
bool checka[TM, TK] = checkam && checkak;
|
|
bool checkb[TK, TN] = checkbk && checkbn;
|
|
a = *?(checka)pa;
|
|
b = *?(checkb)pb;
|
|
}
|
|
TYPE c[TM, TN] = acc;
|
|
|
|
/* ---------------- */
|
|
/* Epilogue */
|
|
/* ---------------- */
|
|
// initialize c pointers
|
|
#ifdef SDD
|
|
bool checkc[TM, TN] = 1;
|
|
// rematerialize
|
|
int rr_blockidm[TM] = (0 ... TM) / BLOCK;
|
|
int rr_blockidn[TN] = (0 ... TN) / BLOCK;
|
|
int rr_offlutm[TM] = rr_blockidm*(TN/BLOCK)*4;
|
|
int rr_offlutn[TN] = rr_blockidn*4;
|
|
int off_bkid[TM, TN] = 3 + rr_offlutm[:, newaxis] + rr_offlutn[newaxis, :];
|
|
int bkid[TM, TN] = *(header + off_bkid);
|
|
long offpc[TM, TN] = bkid * BLOCK * BLOCK;
|
|
// range within blocks
|
|
int rcm[TM] = (0 ... TM) % BLOCK;
|
|
int rcn[TN] = (0 ... TN) % BLOCK;
|
|
#else
|
|
int rcm[TM] = offmc + 0 ... TM;
|
|
int rcn[TN] = offnc + 0 ... TN;
|
|
#ifdef DSD
|
|
bool checkc[TM, TN] = rcn[newaxis, :] < DS0;
|
|
#endif
|
|
#ifdef DDS
|
|
bool checkc[TM, TN] = rcm[:, newaxis] < DS0;
|
|
#endif
|
|
#endif
|
|
TYPE* pc[TM, TN] = C + offpc + offhc*stride_hc + pidz*stride_zc + rcm[:, newaxis]*STRIDE_CM + rcn[newaxis, :]*STRIDE_CN;
|
|
// write-back directly
|
|
if(lockid == 0) {
|
|
*?(checkc) pc = c;
|
|
}
|
|
// accumulate partial result using spin-locks
|
|
else {
|
|
int *plock = locks + get_program_id(2)*nlocks*get_num_programs(1) + get_program_id(1)*nlocks + lockid - 1;
|
|
int *pcount = plock + get_num_programs(2)*get_num_programs(1)*nlocks;
|
|
for(int repeat = 1; repeat == 1; repeat = atomic_cas(plock, 0, 1));
|
|
int count = *pcount;
|
|
if(count == 0)
|
|
*?(checkc) pc = c;
|
|
else
|
|
*?(checkc) pc = c + *?(checkc)pc;
|
|
atomic_xchg(pcount, (count + 1) % maxid);
|
|
atomic_xchg(plock, 0);
|
|
}
|
|
}
|