<commit_before>from __future__ import print_function
import framework
from framework import Program, default_main_program, Parameter, Variable
import optimizer
from layer_helper import LayerHelper
from distributed_spliter import *
import math


class VarBlock:
    def __init__(self, varname, offset, size):
        self.varname = varname
        # NOTE: real offset is offset * size
        self.offset = offset
        self.size = size

    def __str__(self):
        return "%s:%d:%d" % (self.varname, self.offset, self.size)


def split_dense_variable(var_list,
                         pserver_count,
                         min_block_size=1024,
                         max_block_size=1048576):
    """
        We may need to split dense tensor to one or several blocks and put
        them equally onto parameter server. One block is a sub-tensor
        aligned by dim[0] of the tensor.
        
        We need to have a minimal block size so that the calculations in
        the parameter server side can gain better performance. By default
        mininum block size is 1024. The max block size is used to prevent
        too large block that may causing send error.
    """
    blocks = []
    for var in var_list:
        split_count = pserver_count
        var_numel = reduce(lambda x, y: x * y, var.shape)
        max_pserver_count = int(math.floor(var_numel / float(min_block_size)))
        if max_pserver_count == 0:
            max_pserver_count = 1
        if max_pserver_count < pserver_count:
            split_count = max_pserver_count
        block_size = int(math.ceil(var_numel / float(split_count)))

        if len(var.shape) >= 2:
            # align by dim1(width)
            dim1 = reduce(lambda x, y: x * y, var.shape[1:])
            remains = block_size % dim1
            if remains != 0:
                block_size += dim1 - remains
        # update split_count after align
        split_count = int(math.ceil(var_numel / float(block_size)))
        for block_id in xrange(split_count):
            curr_block_size = min(block_size, var_numel - (
                (block_id) * block_size))
            block = VarBlock(var.name, block_id, curr_block_size)
            blocks.append(str(block))
    return blocks


class DistributeTranspiler:
    def transpile(self,
                  optimize_ops,
                  params_grads,
                  program=None,
                  pservers="127.0.0.1:6174",
                  trainers=1,
                  split_method=round_robin):
        """<commit_msg>dist train support split selectedrows<commit_after>from __future__ import print_function
import framework
from framework import Program, default_main_program, Parameter, Variable
import optimizer
from layer_helper import LayerHelper
from distributed_spliter import *
import math


class VarBlock:
    def __init__(self, varname, offset, size):
        self.varname = varname
        # NOTE: real offset is offset * size
        self.offset = offset
        self.size = size

    def __str__(self):
        return "%s:%d:%d" % (self.varname, self.offset, self.size)


def split_dense_variable(var_list,
                         pserver_count,
                         min_block_size=1024,
                         max_block_size=1048576):
    """
        We may need to split dense tensor to one or several blocks and put
        them equally onto parameter server. One block is a sub-tensor
        aligned by dim[0] of the tensor.
        
        We need to have a minimal block size so that the calculations in
        the parameter server side can gain better performance. By default
        mininum block size is 1024. The max block size is used to prevent
        too large block that may causing send error.
    """
    blocks = []
    for var in var_list:
        split_count = pserver_count
        var_numel = reduce(lambda x, y: x * y, var.shape)
        max_pserver_count = int(math.floor(var_numel / float(min_block_size)))
        if max_pserver_count == 0:
            max_pserver_count = 1
        if max_pserver_count < pserver_count:
            split_count = max_pserver_count
        block_size = int(math.ceil(var_numel / float(split_count)))

        if len(var.shape) >= 2:
            # align by dim1(width)
            dim1 = reduce(lambda x, y: x * y, var.shape[1:])
            remains = block_size % dim1
            if remains != 0:
                block_size += dim1 - remains
        # update split_count after align
        split_count = int(math.ceil(var_numel / float(block_size)))
        for block_id in xrange(split_count):
            curr_block_size = min(block_size, var_numel - (
                (block_id) * block_size))
            block = VarBlock(var.name, block_id, curr_block_size)
            blocks.append(str(block))
    return blocks


def split_selected_rows(var,
                        pserver_count,
                        min_block_size=1024,
                        max_block_size=1048576):
    assert ((len(var.shape)) <= 1)

    split_count = pserver_count
    indices = var.desc.selected_rows().dims()
    var_width = reduce(lambda x, y: x * y, var.shape[1:])
    row_count = len(indices)
    rows_per_block = 1
    if var_width < min_block_size:
        rows_per_block = 1
        split_count = row_count
    else:
        rows_per_block = row_count / pserver_count
        if not rows_per_block % pserver_count:
            rows_per_block += 1
        split_count = row_count / rows_per_block
        if not row_count % rows_per_block:
            split_count += 1
    blocks = []
    for block_id in xrange(split_count):
        curr_block_rows = min(rows_per_block,
                              row_count - (block_id * rows_per_block))
        block = VarBlock(var.name, block_id, curr_block_rows)
        blocks.append(block)
    return blocks


def split_variable(var_list,
                   pserver_count,
                   min_block_size=1024,
                   max_block_size=1048576):
    for var in var_list:
        if var.type == core.VarDesc.VarType.LOD_TENSOR:
            split_dense_variable(var_list, pserver_count, min_block_size,
                                 max_block_size)
        elif var.type == core.VarDesc.VarType.SELECTED_ROWS:
            split_selected_rows(var_list, pserver_count, min_block_size,
                                max_block_size)
        else:
            raise TypeError("variable must be lodtensor or selected rows")


class DistributeTranspiler:
    def transpile(self,
                  optimize_ops,
                  params_grads,
                  program=None,
                  pservers="127.0.0.1:6174",
                  trainers=1,
                  split_method=round_robin):
        """
{'fraction_non_alphanumeric': 0.07265463154779521, 'fraction_numerical': 0.01479727730097662, 'mean_word_length': 1.9986690328305234, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '406974', 'n_tokens_mistral': 2006, 'n_tokens_neox': 1905, 'n_words': 530}