Hetu/python/hetu/gpu_ops/Split.py

from __future__ import absolute_import
import numpy as np
from .Node import Op
from ..gpu_links import matrix_slice_simple
from ..gpu_links import matrix_slice_gradient_simple
from .. import ndarray


class SplitOp(Op):
    def __init__(self, node_A, axes, indices, splits, ctx=None):
        super().__init__(SplitOp, [node_A], ctx)
        self.axes = axes
        self.indices = indices
        self.splits = splits
        assert len(self.axes) == len(self.splits)
        assert all([x >= 0 for x in axes])
        assert all([x >= 1 for x in splits])
        assert all([x >= 0 and x < splits[i] for i, x in enumerate(indices)])

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            index = tuple([slice(i, i+j)
                           for i, j in zip(self.begin_pos, self.output_shape)])
            output_val[:] = input_vals[0].asnumpy()[index]
        else:
            # matrix_slice(input_vals[0], output_val, self.begin_pos, stream_handle)
            matrix_slice_simple(
                input_vals[0], output_val, self.gpu_buffer, stream_handle)

    def gradient(self, output_grad):
        self.grad_node = split_gradient_op(
            output_grad, self.axes, self.indices, self.splits, ctx=self.raw_ctx)
        return [self.grad_node]

    def infer_shape(self, input_shapes):
        assert len(input_shapes) == 1
        ori_shape = list(input_shapes[0])
        self.begin_pos = [0 for _ in ori_shape]
        self.output_shape = [x for x in ori_shape]
        for axe, ind, spl in zip(self.axes, self.indices, self.splits):
            part_size = ori_shape[axe] // spl
            self.begin_pos[axe] = ind * part_size
            self.output_shape[axe] = part_size if ind != spl - \
                1 else ori_shape[axe] - self.begin_pos[axe]

        if hasattr(self, 'grad_node'):
            self.grad_node.begin_pos = self.begin_pos
            self.grad_node.output_shape = ori_shape

        # here we save the information on device for GPU computation
        if self.on_gpu:
            ndim = len(ori_shape)
            gpu_buf = [0 for _ in range(3 * ndim)]
            for i in range(ndim):
                gpu_buf[i] = self.begin_pos[i]
                gpu_buf[ndim + i] = ori_shape[i]
                gpu_buf[2 * ndim + i] = self.output_shape[i]
            self.gpu_buffer = ndarray.array(
                gpu_buf, self.ctx, data_type=np.uintc)
        return self.output_shape


class SplitGradientOp(Op):
    def __init__(self, node_A, axes, indices, splits, ctx=None):
        super().__init__(SplitGradientOp, [node_A], ctx)
        self.axes = axes
        self.indices = indices
        self.splits = splits
        self.begin_pos = None
        self.output_shape = None
        assert len(self.axes) == len(self.splits)
        assert all([x >= 0 for x in axes])
        assert all([x >= 1 for x in splits])
        assert all([x >= 0 and x < splits[i] for i, x in enumerate(indices)])

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            output_val[:] = np.zeros(self.output_shape, dtype=np.float32)
            index = tuple([slice(i, i+j)
                           for i, j in zip(self.begin_pos, self.ori_shape)])
            output_val[index] = input_vals[0]
        else:
            # matrix_slice_gradient(input_vals[0], output_val, self.begin_pos, stream_handle)
            matrix_slice_gradient_simple(
                input_vals[0], output_val, self.gpu_buffer, stream_handle)

    def gradient(self, output_grad):
        raise NotImplementedError

    def infer_shape(self, input_shapes):
        assert self.output_shape != None and self.begin_pos != None
        assert len(input_shapes) == 1
        ori_shape = list(input_shapes[0])
        for i in range(len(ori_shape)):
            assert self.begin_pos[i] + ori_shape[i] <= self.output_shape[i]
        self.ori_shape = tuple(ori_shape)

        # here we save the information on device for GPU computation
        if self.on_gpu:
            ndim = len(ori_shape)
            gpu_buf = [0 for _ in range(3 * ndim)]
            for i in range(ndim):
                gpu_buf[i] = self.begin_pos[i]
                gpu_buf[ndim + i] = ori_shape[i]
                gpu_buf[2 * ndim + i] = self.output_shape[i]
            self.gpu_buffer = ndarray.array(
                gpu_buf, self.ctx, data_type=np.uintc)
        return self.output_shape


def split_op(node, axes, indices, splits, ctx=None):
    return SplitOp(node, axes, indices, splits, ctx=ctx)


def split_gradient_op(node, axes, indices, splits, ctx=None):
    return SplitGradientOp(node, axes, indices, splits, ctx=ctx)