def __init__()

in blocksparse/matmul.py [0:0]

• 55 lines of code
• 16 McCabe index (conditional complexity)

    def __init__(self, layout, block_size=32, feature_axis=0, z_order=True, name=None):

        if (feature_axis == 0 and block_size in (8,16,32)) or \
           (feature_axis == 1 and block_size in (32,64)):
           self.axis   = feature_axis
           self.bsize  = block_size
        else:
            raise ValueError("Unsupported block size with this feature axis")

        assert len(layout.shape) == 2
        CB, KB = layout.shape

        group_sizes = layout.sum(axis=0) # assume symetrical transpose
        max_group = group_sizes.max()
        min_group = group_sizes[np.nonzero(group_sizes)].min()
        if max_group / min_group > 2.0:
            segment_size = max(ceil_div(max_group,4), min_group*2)
        else:
            segment_size = SEG_MAX # not worth segmenting
        #print(max_group, min_group, segment_size, KB)
        #segment_size = SEG_MAX

        # don't creat any segments smaller than this
        seg_min = max(ceil_div(segment_size, 4), 4)

        # segment_size = seg_min = 2

        if layout.dtype != np.int32:
            layout = layout.astype(np.int32)

        # convert to csr for vastly more efficient python iteration on large matrices
        csr = sparse.csr_matrix(layout)
        cs, ks, vs = sparse.find(csr) # ks is in sorted order by default
        blocks = len(vs)
        idx  = list(range(blocks))
        idxT = sorted(idx, key=lambda i: cs[i]) # transpose view

        # morton order (z-order) the blocks for efficient L2 cache utilization across all 3 ops
        updat_list = list()
        if z_order:
            blk = 0
            for _, i in sorted( [ (z_order_2d(cs[i], ks[i]), i) for i in range(blocks) ] ):
                vs[i] = blk
                updat_list.append((cs[i], ks[i]))
                blk += 1
        else:
            # row contiguous
            updat_list = list( zip(cs, ks) )
            vs = list(range(blocks))
            # cs = [b[0] for b in updat_list]
            # ks = [b[1] for b in updat_list]

        self.updat_list = updat_list
        self.updat_lut  = np.array(updat_list, dtype=np.int32)

        fsetup = self.xprop_lut(KB, cs, ks, vs, idx,  segment_size, seg_min)
        bsetup = self.xprop_lut(CB, ks, cs, vs, idxT, segment_size, seg_min)

        self.fprop_list, self.fprop_lut, self.l2_lut, self.fprop_shared, self.l2_shared, self.fprop_segments, self.fprop_locks = fsetup
        self.bprop_list, self.bprop_lut,           _, self.bprop_shared,              _, self.bprop_segments, self.bprop_locks = bsetup

        if name is None:
            name = "BlocksparseMatMul"

        self.z_order = z_order
        self.name    = name
        self.flops   = blocks * block_size * block_size * 2
        self.blocks  = blocks
        self.w_shape = (blocks, block_size, block_size)
        self.g_shape = (blocks,)
        self.count   = 0

        self.CB = CB
        self.KB = KB
        self.C  = CB * block_size
        self.K  = KB * block_size

        self.sparsity = round(float(blocks) / float(CB * KB), 3)

        # save boolean version for serialization purposes, TODO save csr version
        self.layout = layout > 0