in blocksparse/conv.py [0:0]
def __init__(self, BCK, TRS, DHW, MPQ=None, strides=(1,1,1), dilates=(1,1,1), padding="SAME", debug=False, deconv=False):
# save this so we know the users perfered number of dims (before we pad 1's out to 3 dims)
self.userTRS = list(TRS)
# support 1-3 dims (additional dimensions are possible by purely extending this python code)
for a in (TRS, DHW, MPQ, strides, dilates, padding):
if type(a) in (tuple, list):
assert 1 <= len(a) <= 3
assert len(TRS) == len(DHW)
# Process the spatial dimensions
# pad sizes and strides out to 3 dimensions
TRS = expand_dims(TRS)
DHW = expand_dims(DHW)
strides = expand_dims(strides)
dilates = expand_dims(dilates)
padding = get_padding(padding, TRS, dilates)
if MPQ is None:
MPQ = [ out_dim(*dims) for dims in zip(TRS, DHW, padding, strides, dilates) ]
else:
MPQ = expand_dims(MPQ)
trs = reduce_mul(TRS)
dhw = reduce_mul(DHW)
mpq = reduce_mul(MPQ)
# contruct feature portion of the grid data loaded to each cuda block
cMax = kMax = sizeF = 0
overlapC = overlapK = False
cSet, kSet = set(), set()
ckLut = list()
fpropGridF = list()
bpropGridF = list()
updatGridF = list()
normList = list()
blkSizes = set()
for listC, listK in BCK:
offset_C = list()
for c in listC:
offset_C.append(c * dhw)
if c in cSet:
overlapC = True
else:
cSet.add(c)
offset_K = list()
for k in listK:
offset_K.append(k * mpq)
if k in kSet:
overlapK = True
else:
kSet.add(k)
block_C = len(listC)
block_K = len(listK)
offset_CK = len(ckLut)
cMax = max(cMax, block_C)
kMax = max(kMax, block_K)
CTRS = block_C*trs
KTRS = block_K*trs
blkSizes.add((block_K, block_C))
# fprop: K is the outer product dim
fpropGridF.append( [ ceil_div(block_K, 32), block_C, block_K, offset_CK, sizeF ] )
# bprop: C is the outer product dim
bpropGridF.append( [ ceil_div(block_C, 32), block_C, block_K, offset_CK, sizeF ] )
# update: K and CTRS are the outer dims (KCRS = KPQ x CHW.T)
updatGridF.append( [ ceil_div(CTRS, 32), ceil_div(block_K, 32), block_C, block_K, offset_CK, sizeF ] )
# setup luts for weight norm
if deconv:
# for deconv, C and K were swapped coming in, so we need to unswap them
for c in range(block_C):
normList.append((c, KTRS, CTRS, sizeF))
else:
for k in range(block_K):
normList.append((sizeF + k * CTRS, CTRS))
# total filter size (and current filter block offset)
sizeF += block_K * block_C * trs
ckLut.extend(offset_C)
ckLut.extend(offset_K)
ckLut = np.array(ckLut, dtype=np.int32)
# Assume no missing mappings.
self.C = len(cSet)
self.K = len(kSet)
self.fixed_block_size = len(blkSizes) == 1
# Process the spatial component of the grid
self.mpqLut = list()
self.dhwLut = list()
self.mpqSlice = None
fdata = list(zip(TRS, padding, strides, dilates))
for i in range(3):
self.mpqLut.append( [ fprop_lut( x, DHW[i], *fdata[i]) for x in range(MPQ[i]) ] )
self.dhwLut.append( [ bprop_lut( x, MPQ[i], *fdata[i]) for x in range(DHW[i]) ] )
mpq_lut = self.spatial_grid(DHW, MPQ, self.mpqLut, mpq, trs)
dhw_lut = self.spatial_grid(MPQ, DHW, self.dhwLut, dhw, trs)
# get the super block dimension
dim_O = mpq_lut.shape[0]
dim_I = dhw_lut.shape[0]
# merge the spatial and feature outer product grid info
fpropGrid = list()
for dim_K, block_C, block_K, offset_CK, offset_F in fpropGridF:
for order, idx_MPQ, idx_K in sorted([ (z_order_3d(0,o,k), o,k) for o,k in np.ndindex(dim_O, dim_K) ]):
# idx_K/idx_MPQ, block_K/block_C, offset_CK, offset_F
fpropGrid.append( [
idx_MPQ + (idx_K << 16),
block_C + (block_K << 16),
offset_CK, offset_F ] )
bpropGrid = list()
for dim_C, block_C, block_K, offset_CK, offset_F in bpropGridF:
for order, idx_DHW, idx_C in sorted([ (z_order_3d(0,i,c), i,c) for i,c in np.ndindex(dim_I, dim_C) ]):
# idx_C/idx_DHW, block_K/block_C, offset_CK, offset_F
bpropGrid.append( [
idx_DHW + (idx_C << 16),
block_C + (block_K << 16),
offset_CK, offset_F ] )
updatGrid = list()
for dim_CTRS, dim_K, block_C, block_K, offset_CK, offset_F in updatGridF:
for order, idx_MPQ, idx_K, idx_CTRS in sorted([ (z_order_3d(o,k,c), o,k,c) for o,k,c in np.ndindex(dim_O, dim_K, dim_CTRS) ]):
# idx_MPQ, idx_CTRS/idx_K, block_C, block_K, offset_CK, offset_F
updatGrid.append( [
idx_MPQ, idx_CTRS + (idx_K << 16),
block_C, block_K,
offset_CK, offset_F ] )
fpropGrid = np.array(fpropGrid, dtype=np.int32)
bpropGrid = np.array(bpropGrid, dtype=np.int32)
updatGrid = np.array(updatGrid, dtype=np.int32)
normLut = np.array(normList, dtype=np.int32)
self.fshared = (trs*32 + 32 + ceil_div(cMax,4)*4 + min(kMax,32)) * 4
self.bshared = (trs*32 + 32 + ceil_div(kMax,4)*4 + min(cMax,32)) * 4
# flops per image of minibatch
self.flops = sizeF * mpq * 2
self.blocks = len(BCK)
self.debug = bool(debug)
self.BCK = BCK
self.TRS = TRS
self.DHW = DHW
self.MPQ = MPQ
self.sizeF = sizeF
self.strides = strides
self.dilates = dilates
self.padding = padding
# For integer division we'd like to do this in a single XMAD sass instruction (plus shift).
# We need to be inside of 16 bits for this to work.
# An additional XMAD could be added at a slight performance loss to support larger dimensions.
# But I'm not sure these larger dimensions are needed in practice.
cktrsMax = ceil_div(max(cMax, kMax)*trs, 32) * 32
cktrsMagic = magic32u(cktrsMax, trs)
assert cktrsMax < 2**16 and cktrsMagic[0] < 2**16, \
"Use cuDNN for large single blocks, but email me if you think there is a use case for this: scott@openai.com"
# kernel params
self.trs = trs
self.magic_trs = cktrsMagic
self.overlapC = overlapC
self.overlapK = overlapK
self.normSize = len(normList)
self.ck_lut = tf.constant(ckLut, name="ck_lut")
self.mpq_lut = tf.constant(mpq_lut, name="mpq_lut")
self.dhw_lut = tf.constant(dhw_lut, name="dhw_lut")
self.fprop_grid = tf.constant(fpropGrid, name="fprop_grid")
self.bprop_grid = tf.constant(bpropGrid, name="bprop_grid")
self.updat_grid = tf.constant(updatGrid, name="updat_grid")
self.norm_lut = tf.constant(normLut, name="norm_lut")