in blocksparse/conv.py [0:0]
def spatial_grid(self, DHW, MPQ, mpqLut, mpq, trs):
# Find the most efficient super-block using a tile of size 32
# For ties then pick the larger tile in the W dim (more contiguous memory access)
# TODO: allow a mixture of superblock shapes, or maybe odd shapes to get better ulilization
ulilization = list()
# xxxxx yxxxx yyxxx zyyxx
for sb in ((1,1,32),(1,2,16),(1,4,8),(2,4,4)):
util = float(mpq) / reduce_mul( [ ceil_div(*dims) for dims in zip(MPQ, sb) ], 32)
ulilization.append((1.0 - util, 32 - sb[2], sb))
sb = sorted(ulilization)[0][2]
# Map the 32 positions in the superblock to MPQ coordinates
# superblock mask: zyyxx : (1,3,3), yxxxx : (0,1,15)
sb_mask = [ x - 1 for x in sb ]
# superblock cumulative right-shift: zyyxx : (4,2,0), yxxxx : (5,4,0)
shifts = [ len(bin(x)) - 3 for x in sb ]
sb_shift = [ shifts[1]+shifts[2], shifts[2], 0 ]
HW = DHW[1] * DHW[2]
W = DHW[2]
PQ = MPQ[1] * MPQ[2]
Q = MPQ[2]
# Get the dimension in super blocks
mpqDim = [ ceil_div(MPQ[i], sb[i]) for i in range(3) ]
mpq_lut = list()
# Iterate over superblocks to build the lut
for order, sb_mpq in sorted([ (z_order_3d(*mpq), mpq) for mpq in np.ndindex(*mpqDim) ]):
lut32 = [ list() for i in range(trs+1) ]
for i32 in range(32):
# get the mpq coord for each of the 32 positions in the superblock
m = sb_mpq[0] * sb[0] + ((i32 >> sb_shift[0]) & sb_mask[0])
p = sb_mpq[1] * sb[1] + ((i32 >> sb_shift[1]) & sb_mask[1])
q = sb_mpq[2] * sb[2] + ((i32 >> sb_shift[2]) & sb_mask[2])
# make sure we didn't fall off the edge
if all(lt(*mM) for mM in zip((m,p,q), MPQ)):
# add in all the input image offsets for each filter position
lut = [ d*HW + h*W + w if all(x >= 0 for x in (d,h,w)) else -1
for d in mpqLut[0][m]
for h in mpqLut[1][p]
for w in mpqLut[2][q] ]
# add the output image offset
lut.append( m*PQ + p*Q + q )
else:
# -1 offsets get zero padded
lut = [-1] * (trs+1)
# transpose lut data so contiguous rows are for 32 mpq coords of the same trs value
for i in range(trs+1):
lut32[i].append(lut[i])
mpq_lut.append(lut32)
return np.array(mpq_lut, dtype=np.int32)