[PYTHON][EXAMPLES] Removed BlockSparse examples; see
https://github.com/ptillet/torch-blocksparse.git
This commit is contained in:
Philippe Tillet
committed by
Philippe Tillet
Philippe Tillet
parent
268894a5ce
commit
9fda39f64c
@@ -1,364 +0,0 @@
|
||||
import triton
|
||||
import torch
|
||||
|
||||
class _linear(torch.autograd.Function):
|
||||
|
||||
src = '''
|
||||
__global__ void main (TYPE* A __readonly __noalias __aligned(16),
|
||||
TYPE* B __readonly __noalias __aligned(16),
|
||||
TYPE* C __writeonly __noalias __aligned(16),
|
||||
int lda, int ldb, int ldc,
|
||||
int M, int Kmax,
|
||||
int* lut,
|
||||
int* locks, int nlocks) {
|
||||
/* ---------------- */
|
||||
/* Prologue */
|
||||
/* ---------------- */
|
||||
// program ids
|
||||
int pid0 = get_program_id(0);
|
||||
int pid1 = get_program_id(1);
|
||||
#ifdef DW
|
||||
// load LUT header
|
||||
int *header = lut + pid0 * 2;
|
||||
int i = *(header + 0);
|
||||
int j = *(header + 1);
|
||||
int K = Kmax / TZ;
|
||||
int lockid = select(TZ > 1, 1, 0);
|
||||
int offk = pid1 * K;
|
||||
int offm = i * TM;
|
||||
int offn = j * TN;
|
||||
int maxid = get_num_programs(1);
|
||||
#else
|
||||
// load LUT header
|
||||
int *header = lut + pid1 * 5;
|
||||
int offset = *(header + 0);
|
||||
int K = *(header + 1);
|
||||
int column = *(header + 2);
|
||||
int lockid = *(header + 3);
|
||||
int maxid = *(header + 4);
|
||||
int *pinc = lut + offset;
|
||||
int offk = (*pinc) * TK;
|
||||
int offm = pid0 * TM;
|
||||
int offn = column * TN;
|
||||
#endif
|
||||
// initialize a, b pointers
|
||||
int rka[TK] = offk + 0 ... TK;
|
||||
int rkb[TK] = offk + 0 ... TK;
|
||||
int ram[TM] = offm + (0 ... TM);
|
||||
int rbn[TN] = offn + (0 ... TN);
|
||||
TYPE* pa[TM, TK] = A + ram[:, newaxis] * STRIDE_AM + rka[newaxis, :] * STRIDE_AK;
|
||||
TYPE* pb[TK, TN] = B + rbn[newaxis, :] * STRIDE_BN + rkb[:, newaxis] * STRIDE_BK;
|
||||
// pre-fetch
|
||||
bool checka[TM, TK] = ram[:, newaxis] < M;
|
||||
bool checkb[TK, TN] = 1;
|
||||
TYPE a[TM, TK] = checka ? *pa : 0;
|
||||
TYPE b[TK, TN] = checkb ? *pb : 0;
|
||||
|
||||
/* ---------------- */
|
||||
/* Inner Loop */
|
||||
/* ---------------- */
|
||||
// create result tile
|
||||
float acc[TM, TN] = 0;
|
||||
#ifdef DW
|
||||
int step = TK;
|
||||
#else
|
||||
int step = 1;
|
||||
#endif
|
||||
for(int k = K; k > 0; k -= step) {
|
||||
acc += a @ b;
|
||||
// update pointers
|
||||
#ifdef DW
|
||||
int inc_a = TK * STRIDE_AK;
|
||||
int inc_b = TK * STRIDE_BK;
|
||||
#else
|
||||
pinc += 1;
|
||||
int inc_a = (*pinc) * TK * STRIDE_AK;
|
||||
int inc_b = (*pinc) * TK * STRIDE_BK;
|
||||
#endif
|
||||
pa += inc_a;
|
||||
pb += inc_b;
|
||||
// pre-fetch
|
||||
bool checka[TM, TK] = k > 1;
|
||||
bool checkb[TK, TN] = k > 1;
|
||||
a = *?(checka)pa;
|
||||
b = *?(checkb)pb;
|
||||
}
|
||||
TYPE c[TM, TN] = acc;
|
||||
|
||||
/* ---------------- */
|
||||
/* Epilogue */
|
||||
/* ---------------- */
|
||||
// initialize c pointers
|
||||
int rcm[TM] = offm + (0 ... TM);
|
||||
int rcn[TN] = offn + (0 ... TN);
|
||||
TYPE* pc[TM, TN] = C + rcm[:, newaxis]*ldc + rcn[newaxis, :];
|
||||
bool checkc[TM, TN] = rcm[:, newaxis] < M;
|
||||
// write-back directly
|
||||
if(lockid == 0) {
|
||||
*?(checkc) pc = c;
|
||||
}
|
||||
// accumulate partial result using spin-locks
|
||||
else {
|
||||
int *plock = locks + get_program_id(0)*nlocks + lockid - 1;
|
||||
int *pcount = plock + get_num_programs(0)*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);
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
# dictionaries for cached triton kernels
|
||||
y_kernel = dict()
|
||||
dx_kernel = dict()
|
||||
dw_kernel = dict()
|
||||
|
||||
# Given an array sizes representing reduction size for each
|
||||
# column of a block-sparse matrix multiplication,
|
||||
# performs load-balancing to achieve more smaller reductions
|
||||
# of size seg_size
|
||||
@staticmethod
|
||||
def load_balance(sizes, seg_size=8):
|
||||
div = sizes // seg_size
|
||||
rem = sizes % seg_size
|
||||
packs = div + (rem != 0).long()
|
||||
width = packs.sum()
|
||||
# split reduction into segments
|
||||
segments = torch.empty(width, dtype=sizes.dtype)
|
||||
column = torch.empty_like(segments)
|
||||
lockid = torch.zeros_like(segments)
|
||||
maxid = torch.zeros_like(segments)
|
||||
nlocks = 0
|
||||
current = 0
|
||||
col_idx = 0
|
||||
for i in range(len(sizes)):
|
||||
d, r = div[i], rem[i]
|
||||
last = current + d + (r > 0)
|
||||
# column id
|
||||
column[current:last] = col_idx
|
||||
# lock id
|
||||
if d > 1 or (d == 1 and r > 0):
|
||||
nlocks += 1
|
||||
lockid[current:last] = nlocks
|
||||
maxid[current:last] = last - current
|
||||
# segment size
|
||||
segments[current:current+d] = seg_size
|
||||
if r > 0:
|
||||
segments[current+d] = r
|
||||
current = last
|
||||
col_idx += 1
|
||||
offsets = torch.zeros_like(segments)
|
||||
offsets[1:] = torch.cumsum(segments[:-1], dim=0)
|
||||
return segments, column, lockid, maxid, offsets
|
||||
|
||||
# Given a binary mask of 0s and 1s,
|
||||
# Construct look-up table for efficient execution on GPUs
|
||||
@staticmethod
|
||||
def make_ydx_lut(mask, block_size):
|
||||
# offsets in lookup table
|
||||
sizes = torch.sum(mask, 0)
|
||||
offsets = torch.zeros_like(sizes)
|
||||
offsets[1:] = torch.cumsum(sizes[:-1], dim=0)
|
||||
# load-balancing
|
||||
segments, column, lockid, maxid, offsets = dot.load_balance(sizes)
|
||||
# pointer increments
|
||||
nnz = torch.nonzero(mask.T)
|
||||
idx = nnz[:, 1]
|
||||
incs = idx.clone()
|
||||
incs[1:] -= idx[:-1]
|
||||
incs[offsets] = idx[offsets]
|
||||
# create header
|
||||
width = column.size(0)
|
||||
offsets += 5*width
|
||||
header = torch.stack((offsets, segments, column, lockid, maxid), dim=1).view(-1).contiguous()
|
||||
# create lut
|
||||
lut = torch.cat((header, incs)).type(torch.int32).cuda()
|
||||
# create locks
|
||||
num_locks = max(1, lockid.max())
|
||||
locks = torch.zeros((2*mask.size(0), num_locks), dtype=torch.int32).cuda()
|
||||
return lut, locks, width
|
||||
|
||||
@staticmethod
|
||||
def make_dw_lut(mask, depth, block_size):
|
||||
nnz = torch.nonzero(mask)
|
||||
# create lut
|
||||
width = nnz.size(0)
|
||||
i = nnz[:, 0]
|
||||
j = nnz[:, 1]
|
||||
lut = torch.stack((i, j), dim=1).view(-1).contiguous()
|
||||
lut = lut.type(torch.int32).cuda()
|
||||
# create locks
|
||||
num_locks = 1
|
||||
locks = torch.zeros((2*width, num_locks), dtype=torch.int32).cuda()
|
||||
return lut, locks, width
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, w, block_size,
|
||||
y_lut, y_locks, y_width,
|
||||
dx_lut, dx_locks, dx_width,
|
||||
dw_lut, dw_locks, dw_width):
|
||||
M, Kx = x.size()
|
||||
Kw, N = w.size()
|
||||
dtype = x.dtype
|
||||
# memory strides
|
||||
lda = Kx
|
||||
ldb = N
|
||||
ldc = N
|
||||
# create kernel
|
||||
key = (dtype, block_size)
|
||||
if key not in dot.y_kernel:
|
||||
defines = {'TM': 64, 'TN': block_size, 'TK': block_size, 'TYPE': dtype,
|
||||
'STRIDE_AM': 'lda', 'STRIDE_AK': '1',
|
||||
'STRIDE_BN': '1', 'STRIDE_BK': 'ldb'}
|
||||
dot.y_kernel[key] = triton.kernel(dot.src, defines=defines)
|
||||
kernel = dot.y_kernel[key]
|
||||
# allocate output
|
||||
y = torch.empty((M, N), dtype=dtype, device=x.device)
|
||||
# launch kernel
|
||||
grid = lambda opt: [triton.cdiv(M, opt.d('TM')), y_width]
|
||||
kernel(x, w, y, lda, ldb, ldc, M, K, y_lut, y_locks, y_locks.size(1), grid=grid)
|
||||
# save information in context
|
||||
ctx.dx_width = dx_width
|
||||
ctx.dw_width = dw_width
|
||||
ctx.kernel = kernel
|
||||
ctx.block_size = block_size
|
||||
ctx.save_for_backward(x, w, dx_lut, dx_locks, dw_lut, dw_locks)
|
||||
return y
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy):
|
||||
# retrieve information in context
|
||||
x, w, dx_lut, dx_locks, dw_lut, dw_locks = ctx.saved_tensors
|
||||
dx_width = ctx.dx_width
|
||||
dw_width = ctx.dw_width
|
||||
block_size = ctx.block_size
|
||||
kernel = ctx.kernel
|
||||
# shapes
|
||||
M, N = dy.size()
|
||||
_, K = x.size()
|
||||
dtype = x.dtype
|
||||
################
|
||||
# input gradient
|
||||
################
|
||||
dx = None
|
||||
if ctx.needs_input_grad[0]:
|
||||
# create kernel
|
||||
key = (dtype, block_size)
|
||||
if key not in dot.dx_kernel:
|
||||
defines = {'TM': 64, 'TN': block_size, 'TK': block_size, 'TYPE': dtype,
|
||||
'STRIDE_AM': 'lda', 'STRIDE_AK': '1',
|
||||
'STRIDE_BN': 'ldb', 'STRIDE_BK': '1'}
|
||||
dot.dx_kernel[key] = triton.kernel(dot.src, defines=defines)
|
||||
kernel = dot.dx_kernel[key]
|
||||
# allocate output
|
||||
dx = torch.empty_like(x)
|
||||
# launch kernel
|
||||
grid = lambda opt: [triton.cdiv(M, opt.d('TM')), dx_width]
|
||||
kernel(dy, w, dx, N, N, K, M, N, dx_lut, dx_locks, dx_locks.size(1), grid=grid)
|
||||
#################
|
||||
# weight gradient
|
||||
#################
|
||||
dw = None
|
||||
if ctx.needs_input_grad[1]:
|
||||
# create kernel
|
||||
key = (dtype, block_size)
|
||||
if key not in dot.dw_kernel:
|
||||
defines = {'TM': block_size, 'TN': block_size, 'TK': 8, 'TYPE': dtype,
|
||||
'STRIDE_AM': '1', 'STRIDE_AK': 'lda',
|
||||
'STRIDE_BN': '1', 'STRIDE_BK': 'ldb',
|
||||
'DW': True, 'TZ': 2}
|
||||
dot.dw_kernel[key] = triton.kernel(dot.src, defines=defines)
|
||||
kernel = dot.dw_kernel[key]
|
||||
# allocate output
|
||||
dw = torch.zeros_like(w)
|
||||
# launch kernel
|
||||
grid = lambda opt: [dw_width, opt.d('TZ')]
|
||||
kernel(x, dy, dw, K, N, N, K, M, dw_lut, dw_locks, dw_locks.size(1), grid=grid)
|
||||
# done
|
||||
return dx, dw, None,\
|
||||
None, None, None,\
|
||||
None, None, None,\
|
||||
None, None, None
|
||||
linear = _linear.apply
|
||||
|
||||
class Linear(torch.nn.Module):
|
||||
|
||||
def __init__(self, in_features, out_features, block_size, mask):
|
||||
super(Linear, self).__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.weight = torch.nn.Parameter(torch.Tensor(out_features, in_features))
|
||||
self.reset_parameter()
|
||||
# create look-up tables
|
||||
self.y_lut, self.y_locks, self.y_width = _linear.make_ydx_lut(mask, block_size)
|
||||
self.dx_lut, self.dx_locks, self.dx_width = _linear.make_ydx_lut(mask.T, block_size)
|
||||
self.dw_lut, self.dw_locks, self.dw_width = _linear.make_dw_lut(mask, M, block_size)
|
||||
|
||||
def reset_parameters(self):
|
||||
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def forward(self, input):
|
||||
return linear(input, self.weight, self.block_size,
|
||||
self.y_lut, self.y_locks, self.y_width,
|
||||
self.dx_lut, self.dx_locks, self.dx_width,
|
||||
self.dw_lut, self.dw_locks, self.dw_width)
|
||||
|
||||
|
||||
def reference_dot(x, w, mask):
|
||||
WS0, WS1 = w.size()
|
||||
MS0, MS1 = mask.size()
|
||||
assert WS0 % MS0 == 0
|
||||
assert WS1 % MS1 == 0
|
||||
block_size_0 = WS0 // MS0
|
||||
block_size_1 = WS1 // MS1
|
||||
assert block_size_0 == block_size_1
|
||||
maskedw = w.clone()
|
||||
for bi, wi in enumerate(range(0, WS0, block_size_0)):
|
||||
for bj, wj in enumerate(range(0, WS1, block_size_1)):
|
||||
maskedw[wi : wi+block_size_0,
|
||||
wj : wj+block_size_1] *= mask[bi, bj]
|
||||
return torch.matmul(x, maskedw)
|
||||
|
||||
torch.manual_seed(0)
|
||||
# parameters
|
||||
M, N, K = 256, 256, 256
|
||||
BS = 16
|
||||
# initialize inputs
|
||||
mask = torch.randint(0, 2, (K//BS, N//BS))
|
||||
x = torch.rand((M, K), dtype=torch.float32, requires_grad=True).cuda()
|
||||
w = torch.rand((K, N), dtype=torch.float32, requires_grad=True).cuda()
|
||||
x.retain_grad()
|
||||
w.retain_grad()
|
||||
# reference result
|
||||
ry = reference_dot(x, w, mask)
|
||||
dy = torch.rand_like(ry)
|
||||
ry.backward(dy)
|
||||
rdx = x.grad.clone()
|
||||
rdw = w.grad.clone()
|
||||
# reset gradients
|
||||
x.grad.zero_()
|
||||
w.grad.zero_()
|
||||
# triton result
|
||||
y_lut, y_locks, y_width = _linear.make_ydx_lut(mask, BS)
|
||||
dx_lut, dx_locks, dx_width = _linear.make_ydx_lut(mask.T, BS)
|
||||
dw_lut, dw_locks, dw_width = _linear.make_dw_lut(mask, M, BS)
|
||||
ty = _linear.apply(x, w, BS,
|
||||
y_lut, y_locks, y_width,
|
||||
dx_lut, dx_locks, dx_width,
|
||||
dw_lut, dw_locks, dw_width)
|
||||
ty.backward(dy)
|
||||
tdx = x.grad.clone()
|
||||
tdw = w.grad.clone()
|
||||
# test
|
||||
print((ty - ry).abs().max())
|
||||
print((tdx - rdx).abs().max())
|
||||
print((tdw - rdw).abs().max())
|
||||
Reference in New Issue
Block a user