Repository: chainer/chainer

Issue:
Refactoring of Multi Node Batch Normalization in ChainerMN
There are many parts where Multi Node Batch Normalization in ChainerMN overlaps with BatchNormalization in Chainer.
So, we should refactor multi node batch normalization.

ToDo:

- [ ] add args for comm and backend to BatchNormalization in functions of Chainer
- [ ] use BatchNormalization in functions of Chainer in Multi Node Batch Normalization.
 


Discussion:

Resolving patch:
diff --git a/chainer/functions/normalization/batch_normalization.py b/chainer/functions/normalization/batch_normalization.py
--- a/chainer/functions/normalization/batch_normalization.py
+++ b/chainer/functions/normalization/batch_normalization.py
@@ -15,6 +15,216 @@
 import chainerx
 
 
+class _BatchNormalizationImpl:
+
+    def forward(self, axis, gamma, x, xp, expander,
+                beta, eps, decay, running_mean, running_var):
+        raise NotImplementedError()
+
+    def backward(self, axis, gamma, gy, x, xp,
+                 expander, mean, inv_std, eps, var):
+        raise NotImplementedError()
+
+
+class GeneralBatchNormalizationImpl(_BatchNormalizationImpl):
+    def forward(self, axis, gamma, x, xp, expander,
+                beta, eps, decay, running_mean, running_var):
+        interm_dtype = numpy.promote_types(x.dtype, gamma.dtype)
+
+        gamma = gamma[expander].astype(interm_dtype, copy=False)
+        beta = beta[expander].astype(interm_dtype, copy=False)
+        mean, var = self.get_mean_and_var(axis, gamma,
+                                          x, xp, interm_dtype)
+        if xp is numpy:
+            inv_std = numpy.reciprocal(numpy.sqrt(
+                var + eps, dtype=interm_dtype))
+        else:
+            inv_std = cuda.cupyx.rsqrt(
+                var + eps, dtype=interm_dtype)
+
+        y = _apply_bn_fwd(xp, x, mean[expander],
+                          inv_std[expander], gamma, beta)
+
+        # Update running statistics if given
+        if running_mean is not None:
+            m = x.size // gamma.size
+            adjust = m / max(m - 1., 1.)  # unbiased estimation
+
+            xp = backend.get_array_module(
+                running_mean, running_var)
+
+            if xp is chainerx:
+                running_mean, running_var = backend.from_chx(
+                    (running_mean, running_var))
+
+            if xp is numpy:
+                running_mean *= decay
+                running_mean += (1 - decay) * mean
+                running_var *= decay
+                running_var += (1 - decay) * adjust * var
+            else:
+                cuda.elementwise(
+                    'U mean, U var, U decay, U adjust',
+                    'U r_mean, U r_var',
+                    '''
+                    r_mean = r_mean * decay + mean * (1 - decay);
+                    r_var = r_var * decay + var * (1 - decay) * adjust;
+                    ''',
+                    'update_mean_var')(mean, var, decay, adjust,
+                                       running_mean, running_var)
+
+            if xp is chainerx:
+                running_mean = backend.to_chx(running_mean)
+                running_var = backend.to_chx(running_var)
+
+        return y, running_mean, running_var, mean, var, inv_std
+
+    def get_mean_and_var(self, axis, gamma, x, xp, interm_dtype):
+        mean = x.mean(axis=axis, dtype=interm_dtype)
+        var = x.var(axis=axis, dtype=interm_dtype)
+        return mean, var
+
+    def get_ggamma_and_gbeta(self, axis, gamma, gy, x_hat, xp):
+        gbeta = gy.sum(axis=axis, dtype=gamma.dtype)
+        ggamma = (gy * x_hat).sum(axis=axis, dtype=gamma.dtype)
+
+        return gbeta, ggamma
+
+    def backward(self, axis, gamma, gy, x, xp,
+                 expander, mean, inv_std, eps, var):
+        if isinstance(gy, intel64.mdarray):
+            # intel64.mdarray does not support dtype option in sum, so we
+            # convert it to numpy here.
+            gy = numpy.asarray(gy)
+
+        x_hat = _x_hat(x, mean[expander], inv_std[expander])
+        gbeta, ggamma = self.get_ggamma_and_gbeta(axis, gamma, gy, x_hat, xp)
+
+        inv_m = gamma.dtype.type(1. / (x.size // gamma.size))
+        if xp is numpy:
+            gx = (gamma * inv_std)[expander] * (
+                gy - (x_hat * ggamma[expander] + gbeta[expander]) * inv_m)
+            gx = gx.astype(dtype=x.dtype, copy=False)
+        else:
+            gx = cuda.elementwise(
+                '''
+                T gy, U x_hat, U gamma, U inv_std, U ggamma, U gbeta,
+                U inv_m
+                ''',
+                'T gx',
+                '''
+                gx = (gamma * inv_std) * (
+                    gy - (x_hat * ggamma + gbeta) * inv_m)
+                ''', 'bn_bwd')(gy, x_hat, gamma[expander],
+                               inv_std[expander], ggamma[expander],
+                               gbeta[expander], inv_m)
+
+        return gx, ggamma, gbeta
+
+
+class _IDeepBatchNormalizationImpl(_BatchNormalizationImpl):
+    def forward(self, axis, gamma, x, xp, expander,
+                beta, eps, decay, running_mean, running_var):
+        expand_dim = False
+        if x.ndim == 2:
+            expand_dim = True
+            x = x[:, :, None, None]
+
+        y, mean, var, inv_std = (
+            intel64.ideep.batchNormalization.Forward(
+                intel64.ideep.array(x.astype(gamma.dtype, copy=False)),
+                intel64.ideep.array(gamma),
+                intel64.ideep.array(beta),
+                None,
+                None,
+                eps
+            ))
+        y = y.astype(x.dtype, copy=False)
+
+        if expand_dim:
+            y = numpy.squeeze(y, axis=(2, 3))
+
+        # Update running statistics if given
+        if running_mean is not None:
+            m = x.size // gamma.size
+            adjust = m / max(m - 1., 1.)
+
+            # Update running_mean
+            if isinstance(running_mean, intel64.ideep.mdarray):
+                running_mean.inplace_axpby(
+                    decay, (1 - decay), mean)
+            else:
+                running_mean *= decay
+                running_mean += mean * (1 - decay)
+
+            # Update running_var
+            if isinstance(running_var, intel64.ideep.mdarray):
+                running_var.inplace_axpby(
+                    decay, (1 - decay), var * adjust)
+            else:
+                running_var *= decay
+                running_var += var * adjust * (1 - decay)
+        return y, running_mean, running_var, mean, var, inv_std
+
+    def backward(self, axis, gamma, gy, x, xp,
+                 expander, mean, inv_std, eps, var):
+        expand_dim = False
+        if x.ndim == 2:
+            expand_dim = True
+            x = x[:, :, None, None]
+            gy = gy[:, :, None, None]
+
+        gx, gW = intel64.ideep.batchNormalization.Backward(
+            intel64.ideep.array(x.astype(gamma.dtype, copy=False)),
+            intel64.ideep.array(gy.astype(gamma.dtype, copy=False)),
+            mean,
+            var,
+            intel64.ideep.array(gamma),
+            eps)
+
+        gx = gx.astype(x.dtype, copy=False)
+        ggamma, gbeta = gW[:2]
+
+        if expand_dim:
+            gx = numpy.squeeze(gx, axis=(2, 3))
+
+        return gx, ggamma, gbeta
+
+
+class _CudnnBatchNormalizationImpl(_BatchNormalizationImpl):
+    def __init__(self, is_for_conv2d, cudnn_mode):
+        self.is_for_conv2d = is_for_conv2d
+        self.cudnn_mode = cudnn_mode
+
+    def forward(self, axis, gamma, x, xp, expander,
+                beta, eps, decay, running_mean, running_var):
+        if running_mean is not None:
+            mean = running_mean
+            var = running_var
+        else:
+            # Create dummies.
+            mean = xp.zeros_like(gamma, dtype=x.dtype)
+            var = xp.zeros_like(gamma, dtype=x.dtype)
+
+        # mean and inv_std are used as buffers to save
+        # intermediate results computed during forward pass. These buffers
+        # are used to speed-up backward pass.
+        y, mean, inv_std = (
+            cudnn.batch_normalization_forward_training(
+                x, gamma, beta, mean, var, None, None,
+                eps, decay, self.is_for_conv2d,
+                self.cudnn_mode, chainer.is_debug()))
+
+        return y, running_mean, running_var, mean, var, inv_std
+
+    def backward(self, axis, gamma, gy, x, xp,
+                 expander, mean, inv_std, eps, var):
+        return cudnn.batch_normalization_backward(
+            x, gamma, gy, mean, inv_std, eps,
+            self.is_for_conv2d, self.cudnn_mode,
+            chainer.is_debug())
+
+
 if cuda.cudnn_enabled:
     cudnn = cuda.cudnn
     libcudnn = cuda.libcudnn
@@ -34,12 +244,29 @@ def _compute_key_axis(x_ndim, gamma_ndim=1, axis=None):
     return key_axis
 
 
+def _impl_selector(batch_norm_func, inputs):
+    x, gamma, _ = inputs
+    xp = backend.get_array_module(x)
+    mode = _BNMode(x, gamma, batch_norm_func.key_axis)
+    use_cudnn = mode.can_use_cudnn(xp)
+    use_ideep = mode.can_use_ideep()
+
+    if use_ideep:
+        return _IDeepBatchNormalizationImpl()
+    elif use_cudnn:
+        return _CudnnBatchNormalizationImpl(mode.is_for_conv2d,
+                                            mode.get_cudnn_mode())
+    else:
+        return GeneralBatchNormalizationImpl()
+
+
 class BatchNormalization(function_node.FunctionNode):
 
     mean = None
     inv_std = None
 
-    def __init__(self, eps=2e-5, mean=None, var=None, decay=0.9, axis=None):
+    def __init__(self, eps=2e-5, mean=None, var=None, decay=0.9, axis=None,
+                 impl_selector=_impl_selector):
         self.running_mean = mean
         self.running_var = var
 
@@ -65,6 +292,8 @@ def __init__(self, eps=2e-5, mean=None, var=None, decay=0.9, axis=None):
             raise RuntimeError('axis must be int, tuple of int or None')
         self.axis = axis
 
+        self._impl_selector = impl_selector
+
     def check_type_forward(self, in_types):
         type_check.expect(in_types.size() == 3)
         x_type, gamma_type, beta_type = in_types
@@ -152,109 +381,18 @@ def forward(self, inputs):
         expander = tuple(expander)
         self.expander = expander
 
-        self.mode = _BNMode(x, gamma, self.key_axis)
         xp = backend.get_array_module(x)
-        self.use_cudnn = self.mode.can_use_cudnn(xp)
-        self.use_ideep = self.mode.can_use_ideep()
-
-        if self.use_ideep:
-            # TODO(niboshi): Refactor iDeep part into a separate method
-            expand_dim = False
-            if x.ndim == 2:
-                expand_dim = True
-                x = x[:, :, None, None]
-
-            y, self.mean, self.var, self.inv_std = (
-                intel64.ideep.batchNormalization.Forward(
-                    intel64.ideep.array(x.astype(gamma.dtype, copy=False)),
-                    intel64.ideep.array(gamma),
-                    intel64.ideep.array(beta),
-                    None,
-                    None,
-                    self.eps
-                ))
-            y = y.astype(x.dtype, copy=False)
-
-            if expand_dim:
-                y = numpy.squeeze(y, axis=(2, 3))
-
-            # Update running statistics if given
-            if self.running_mean is not None:
-                m = x.size // gamma.size
-                adjust = m / max(m - 1., 1.)
-
-                # Update running_mean
-                if isinstance(self.running_mean, intel64.ideep.mdarray):
-                    self.running_mean.inplace_axpby(
-                        self.decay, (1 - self.decay), self.mean)
-                else:
-                    self.running_mean *= self.decay
-                    self.running_mean += self.mean * (1 - self.decay)
-
-                # Update running_var
-                if isinstance(self.running_var, intel64.ideep.mdarray):
-                    self.running_var.inplace_axpby(
-                        self.decay, (1 - self.decay), self.var * adjust)
-                else:
-                    self.running_var *= self.decay
-                    self.running_var += self.var * adjust * (1 - self.decay)
-
-        elif self.use_cudnn:
-            if self.running_mean is not None:
-                mean = self.running_mean
-                var = self.running_var
-            else:
-                # Create dummies.
-                mean = xp.zeros_like(gamma, dtype=x.dtype)
-                var = xp.zeros_like(gamma, dtype=x.dtype)
-
-            # self.mean and self.inv_std are used as buffers to save
-            # intermediate results computed during forward pass. These buffers
-            # are used to speed-up backward pass.
-            y, self.mean, self.inv_std = (
-                cudnn.batch_normalization_forward_training(
-                    x, gamma, beta, mean, var, None, None,
-                    self.eps, self.decay, self.mode.is_for_conv2d,
-                    self.mode.get_cudnn_mode(), chainer.is_debug()))
-        else:
-            # Generic CPU and GPU implementation
-
-            interm_dtype = numpy.promote_types(x.dtype, gamma.dtype)
-
-            gamma = gamma[expander].astype(interm_dtype, copy=False)
-            beta = beta[expander].astype(interm_dtype, copy=False)
-            self.mean = x.mean(axis=self.axis, dtype=interm_dtype)
-            var = x.var(axis=self.axis, dtype=interm_dtype)
-            if xp is numpy:
-                self.inv_std = numpy.reciprocal(numpy.sqrt(
-                    var + self.eps, dtype=interm_dtype))
-            else:
-                self.inv_std = cuda.cupyx.rsqrt(
-                    var + self.eps, dtype=interm_dtype)
-
-            y = _apply_bn_fwd(xp, x, self.mean[expander],
-                              self.inv_std[expander], gamma, beta)
 
-            # Update running statistics if given
-            if self.running_mean is not None:
-                m = x.size // gamma.size
-                adjust = m / max(m - 1., 1.)  # unbiased estimation
+        self._impl = self._impl_selector(self, inputs)
 
-                xp = backend.get_array_module(
-                    self.running_mean, self.running_var)
-
-                if xp is chainerx:
-                    self.running_mean, self.running_var = backend.from_chx(
-                        (self.running_mean, self.running_var))
-
-                self.running_mean *= self.decay
-                self.running_mean += (1 - self.decay) * self.mean
-                self.running_var *= self.decay
-                self.running_var += (1 - self.decay) * adjust * var
-
-                if xp is chainerx:
-                    self.running_mean = backend.to_chx(self.running_mean)
-                    self.running_var = backend.to_chx(self.running_var)
+        y, self.running_mean, self.running_var,\
+            self.mean, self.var, self.inv_std = \
+            self._impl.forward(axis=self.axis, gamma=gamma, x=x,
+                               xp=xp, expander=expander,
+                               beta=beta, eps=self.eps,
+                               decay=self.decay,
+                               running_mean=self.running_mean,
+                               running_var=self.running_var)
 
         return y,
 
@@ -262,32 +400,34 @@ def backward(self, indexes, grad_outputs):
         x, gamma = self.get_retained_inputs()
         gy, = grad_outputs
 
-        if self.use_ideep:
+        if isinstance(self._impl, _IDeepBatchNormalizationImpl):
             assert self.var is not None
             var = self.var
         else:
             var = None
 
         f = BatchNormalizationGrad(
-            self.eps, self.use_cudnn, self.mode, self.expander, self.axis,
-            self.mean, var, self.inv_std, self.key_axis)
+            self.eps, self.expander,
+            self.axis,
+            self.mean, var,
+            self.inv_std, self.key_axis,
+            self._impl)
+
         return f.apply((x, gamma, gy))
 
 
 class BatchNormalizationGrad(function_node.FunctionNode):
 
-    def __init__(self, eps, use_cudnn, mode, expander, axis, mean, var,
-                 inv_std, key_axis):
+    def __init__(self, eps, expander, axis, mean, var,
+                 inv_std, key_axis, impl):
         self.eps = eps
-        self.use_cudnn = use_cudnn
-        self.use_ideep = mode.can_use_ideep()
-        self.mode = mode
         self.expander = expander
         self.axis = axis
         self.mean = mean
         self.var = var  # Only used in iDeep implementation
         self.inv_std = inv_std
         self.key_axis = key_axis
+        self._impl = impl
 
     def forward(self, inputs):
         self.retain_inputs((0, 1, 2))
@@ -295,62 +435,10 @@ def forward(self, inputs):
         expander = self.expander
         xp = backend.get_array_module(x)
 
-        if self.use_ideep:
-            # TODO(niboshi): Refactor iDeep part into a separate method
-            expand_dim = False
-            if x.ndim == 2:
-                expand_dim = True
-                x = x[:, :, None, None]
-                gy = gy[:, :, None, None]
-
-            gx, gW = intel64.ideep.batchNormalization.Backward(
-                intel64.ideep.array(x.astype(gamma.dtype, copy=False)),
-                intel64.ideep.array(gy.astype(gamma.dtype, copy=False)),
-                self.mean,
-                self.var,
-                intel64.ideep.array(gamma),
-                self.eps)
-
-            gx = gx.astype(x.dtype, copy=False)
-            ggamma, gbeta = gW[:2]
-
-            if expand_dim:
-                gx = numpy.squeeze(gx, axis=(2, 3))
-
-        elif self.use_cudnn:
-            gx, ggamma, gbeta = cudnn.batch_normalization_backward(
-                x, gamma, gy, self.mean, self.inv_std, self.eps,
-                self.mode.is_for_conv2d, self.mode.get_cudnn_mode(),
-                chainer.is_debug())
-        else:
-            # CPU and GPU implementation
-            if isinstance(gy, intel64.mdarray):
-                # intel64.mdarray does not support dtype option in sum, so we
-                # convert it to numpy here.
-                gy = numpy.asarray(gy)
-
-            gbeta = gy.sum(axis=self.axis, dtype=gamma.dtype)
-            x_hat = _x_hat(x, self.mean[expander], self.inv_std[expander])
-            ggamma = (gy * x_hat).sum(axis=self.axis, dtype=gamma.dtype)
-
-            inv_m = gamma.dtype.type(1. / (x.size // gamma.size))
-            if xp is numpy:
-                gx = (gamma * self.inv_std)[expander] * (
-                    gy - (x_hat * ggamma[expander] + gbeta[expander]) * inv_m)
-                gx = gx.astype(dtype=x.dtype, copy=False)
-            else:
-                gx = cuda.elementwise(
-                    '''
-                    T gy, U x_hat, U gamma, U inv_std, U ggamma, U gbeta,
-                    U inv_m
-                    ''',
-                    'T gx',
-                    '''
-                    gx = (gamma * inv_std) * (
-                        gy - (x_hat * ggamma + gbeta) * inv_m)
-                    ''', 'bn_bwd')(gy, x_hat, gamma[expander],
-                                   self.inv_std[expander], ggamma[expander],
-                                   gbeta[expander], inv_m)
+        gx, ggamma, gbeta = self._impl.backward(self.axis, gamma, gy,
+                                                x, xp, expander,
+                                                self.mean, self.inv_std,
+                                                self.eps, self.var)
         self.retain_inputs((0, 1, 2))
         self.retain_outputs((0, 1))
         return gx, ggamma, gbeta
diff --git a/chainermn/functions/batch_normalization.py b/chainermn/functions/batch_normalization.py
--- a/chainermn/functions/batch_normalization.py
+++ b/chainermn/functions/batch_normalization.py
@@ -1,41 +1,15 @@
-# This file is heavily based on Chainer's batch normalization implementation.
-# See: chainer/functions/normalization/batch_normalization.py (dbb650)
-
-from abc import ABCMeta
-from abc import abstractmethod
 import chainer
-from chainer import backend
 from chainer import cuda
-from chainer import function
 import chainer.utils
-from chainer.utils import type_check
-import numpy
-import six
-
-
-def _xhat(x, mean, std, expander):
-    x_mu = x - mean[expander]
-    x_mu /= std[expander]
-    return x_mu
-
-
-class _MultiNodeBatchNormalizationBackend(six.with_metaclass(ABCMeta)):
-
-    @abstractmethod
-    def forward(self, axis, gamma, x, xp):
-        raise NotImplementedError()
-
-    @abstractmethod
-    def backward(self, axis, gamma, gy, x_hat, x, xp):
-        raise NotImplementedError()
+from chainer.functions.normalization.batch_normalization \
+    import GeneralBatchNormalizationImpl
 
 
-class _MpiBackend(_MultiNodeBatchNormalizationBackend):
-
+class _MpiImpl(GeneralBatchNormalizationImpl):
     def __init__(self, comm):
         self.comm = comm
 
-    def forward(self, axis, gamma, x, xp):
+    def get_mean_and_var(self, axis, gamma, x, xp, interm_dtype):
         tmp = xp.empty(gamma.size * 2, dtype=gamma.dtype)
         x.mean(axis=axis, out=tmp[:gamma.size], dtype=gamma.dtype)
         xp.square(x).mean(axis=axis, out=tmp[gamma.size:], dtype=gamma.dtype)
@@ -47,7 +21,7 @@ def forward(self, axis, gamma, x, xp):
         var = sqmean - xp.square(mean)
         return mean, var
 
-    def backward(self, axis, gamma, gy, x_hat, x, xp):
+    def get_ggamma_and_gbeta(self, axis, gamma, gy, x_hat, xp):
         tmp = xp.empty(gamma.size * 2, dtype=gamma.dtype)
         gy.sum(axis=axis, out=tmp[:gamma.size], dtype=gamma.dtype)
         (gy * x_hat).sum(axis=axis, out=tmp[gamma.size:], dtype=gamma.dtype)
@@ -59,7 +33,7 @@ def backward(self, axis, gamma, gy, x_hat, x, xp):
         return gbeta, ggamma
 
 
-class _NcclBackend(_MultiNodeBatchNormalizationBackend):
+class _NcclImpl(GeneralBatchNormalizationImpl):
 
     def __init__(self, comm):
         self.comm = comm
@@ -68,7 +42,7 @@ def __init__(self, comm):
         import chainermn.communicators._memory_utility as memory_utility_module
         self.memory_utility_module = memory_utility_module
 
-    def forward(self, axis, gamma, x, xp):
+    def get_mean_and_var(self, axis, gamma, x, xp, interm_dtype):
         gpu_buffer_n_elems = gamma.size * 2
         gpu_buffer_size = gamma.dtype.itemsize * gpu_buffer_n_elems
         gpu_buffer_a = self.memory_utility_module.DeviceMemory()
@@ -94,7 +68,7 @@ def forward(self, axis, gamma, x, xp):
         var = sqmean - xp.square(mean)
         return mean, var
 
-    def backward(self, axis, gamma, gy, x_hat, x, xp):
+    def get_ggamma_and_gbeta(self, axis, gamma, gy, x_hat, xp):
         gpu_buffer_n_elems = gamma.size * 2
         gpu_buffer_size = gamma.dtype.itemsize * gpu_buffer_n_elems
         gpu_buffer_a = self.memory_utility_module.DeviceMemory()
@@ -142,192 +116,13 @@ def get_communication_backend(comm, communication_backend='auto'):
     return selected_communication_backend
 
 
-class MultiNodeBatchNormalizationFunction(function.Function):
-
-    def __init__(self, comm, eps=2e-5, mean=None, var=None, decay=0.9,
-                 communication_backend='auto'):
-        chainer.utils.experimental(
-            'chainermn.functions.MultiNodeBatchNormalizationFunction')
-
+class MultiNodeBNImplSelector:
+    def __init__(self, comm, communication_backend_name):
         self.comm = comm
-        self.running_mean = mean
-        self.running_var = var
-
-        # Note: cuDNN v5 requires that eps be greater than 1e-5. Otherwise, an
-        # error will occur.
-        # See CUDNN_BN_MIN_EPSILON value in cudnn.h to verify minimum allowable
-        # value.
-        self.eps = eps
-        if chainer.should_use_cudnn('>=auto'):
-            if eps < 1e-5:
-                msg = 'cuDNN does not allow an eps value less than 1e-5.'
-                raise RuntimeError(msg)
-        self.mean_cache = None
-        self.decay = decay
+        self.communication_backend_name = communication_backend_name
 
-        selected_communication_backend = \
-            get_communication_backend(comm, communication_backend)
-
-        if selected_communication_backend == 'nccl':
-            self._backend = _NcclBackend(comm)
+    def __call__(self, batch_norm_func, inputs):
+        if self.communication_backend_name == 'nccl':
+            return _NcclImpl(self.comm)
         else:
-            self._backend = _MpiBackend(comm)
-
-    def check_type_forward(self, in_types):
-        n_in = type_check.eval(in_types.size())
-        if n_in != 3 and n_in != 5:
-            raise type_check.InvalidType(
-                '%s or %s' % (in_types.size() == 3, in_types.size() == 5),
-                '%s == %s' % (in_types.size(), n_in))
-        x_type, gamma_type, beta_type = in_types[:3]
-        M = type_check.eval(gamma_type.ndim)
-        type_check.expect(
-            x_type.dtype.kind == 'f',
-            x_type.ndim >= gamma_type.ndim + 1,
-            x_type.shape[1:1 + M] == gamma_type.shape,
-            gamma_type.dtype.kind == 'f',
-            gamma_type.dtype == beta_type.dtype,
-            gamma_type.shape == beta_type.shape,
-        )
-        if len(in_types) == 5:
-            mean_type, var_type = in_types[3:]
-            type_check.expect(
-                mean_type.dtype == x_type.dtype,
-                mean_type.shape == gamma_type.shape,
-                var_type.dtype == x_type.dtype,
-                var_type.shape == gamma_type.shape,
-            )
-
-    def forward(self, inputs):
-        xp = backend.get_array_module(*inputs)
-        x, gamma, beta = inputs[:3]
-        if chainer.configuration.config.train:
-            if self.running_mean is None:
-                self.running_mean = xp.zeros_like(gamma, dtype=x.dtype)
-                self.running_var = xp.zeros_like(gamma, dtype=x.dtype)
-            else:
-                self.running_mean = xp.array(self.running_mean)
-                self.running_var = xp.array(self.running_var)
-        elif len(inputs) == 5:
-            self.fixed_mean = inputs[3]
-            self.fixed_var = inputs[4]
-
-        head_ndim = gamma.ndim + 1
-        expander = (None, Ellipsis) + (None,) * (x.ndim - head_ndim)
-        gamma = gamma[expander]
-        beta = beta[expander]
-
-        # cuDNN only supports these tensor dimensions because they are
-        # the most commonly used. If there is a need to support other
-        # dimensions with cuDNN, we could consider reshaping the input
-        # into a 2-dim array with channels as second dim and m=<product
-        # of all dimensions except the 2nd dimension> as the first
-        # dimension.
-        cudnn_dim_ok = x.ndim == 2 or (x.ndim == 4 and head_ndim == 2)
-        # TODO(bkvogel): Check for float16 support again in next cuDNN version.
-        # cuDNN v5 batch normalization does not seem to support float16.
-        self._can_use_cudnn = cudnn_dim_ok and x[0].dtype != numpy.float16
-
-        cudnn_updated_running_stats = False
-
-        if chainer.configuration.config.train:
-            axis = (0,) + tuple(range(head_ndim, x.ndim))
-            mean, var = self._backend.forward(axis, gamma, x, xp)
-            var += self.eps
-        else:
-            mean = self.fixed_mean
-            var = self.fixed_var + self.eps
-        self.std = xp.sqrt(var, dtype=var.dtype)
-        if xp is numpy:
-            self.x_hat = _xhat(x, mean, self.std, expander)
-            y = gamma * self.x_hat
-            y += beta
-            y = y.astype(x.dtype)
-        else:
-            self.x_hat, y = cuda.elementwise(
-                'T x, U mean, U std, U gamma, U beta', 'T x_hat, T y',
-                '''
-                x_hat = (x - mean) / std;
-                y = gamma * x_hat + beta;
-                ''',
-                'bn_fwd')(x, mean[expander], self.std[expander], gamma,
-                          beta)
-
-        if chainer.configuration.config.train and \
-                (not cudnn_updated_running_stats):
-            # Note: If in training mode, the cuDNN forward training function
-            # will do this for us, so
-            # only run following code if cuDNN was not used.
-            # Update running statistics:
-            m = x.size // gamma.size
-            adjust = m / max(m - 1., 1.)  # unbiased estimation
-            if xp is numpy:
-                self.running_mean *= self.decay
-                temp_ar = xp.array(mean)
-                temp_ar *= (1 - self.decay)
-                self.running_mean += temp_ar
-                del temp_ar
-                self.running_var *= self.decay
-                temp_ar = xp.array(var)
-                temp_ar *= (1 - self.decay) * adjust
-                self.running_var += temp_ar
-                del temp_ar
-            else:
-                cuda.elementwise(
-                    'U mean, U var, U decay, U adjust',
-                    'U r_mean, U r_var',
-                    '''
-                    r_mean = r_mean * decay + mean * (1 - decay);
-                    r_var = r_var * decay + var * (1 - decay) * adjust;
-                    ''',
-                    'update_mean_var')(mean, var, self.decay, adjust,
-                                       self.running_mean, self.running_var)
-
-        return y,
-
-    def backward(self, inputs, grad_outputs):
-        x, gamma = inputs[:2]
-        gy = grad_outputs[0]
-        head_ndim = gamma.ndim + 1
-        expander = (None, Ellipsis) + (None,) * (x.ndim - head_ndim)
-        m = gamma.dtype.type(x.size // gamma.size)
-        axis = (0,) + tuple(range(head_ndim, x.ndim))
-        xp = backend.get_array_module(x)
-        if len(inputs) == 5:
-            # This case is unlikely to be used in practice and so does not
-            # need to be optimized for performance.
-            mean = inputs[3]
-            var = inputs[4]
-            std = xp.sqrt(var, dtype=var.dtype)
-            gs = gamma / std
-            gbeta = gy.sum(axis=axis)
-            x_hat = _xhat(x, mean, std, expander)
-            ggamma = (gy * x_hat).sum(axis=axis)
-            gmean = -gs * gbeta
-            gvar = -0.5 * gamma / var * ggamma
-            gx = gs[expander] * gy
-            gx = gx.astype(x.dtype, copy=False)
-            return gx, ggamma, gbeta, gmean, gvar
-
-        # Note: If length of inputs is not 5, we must be in train mode.
-        assert chainer.configuration.config.train
-        gbeta, ggamma = self._backend.backward(axis, gamma, gy,
-                                               self.x_hat, x, xp)
-
-        if xp is numpy:
-            gx = (gamma / self.std)[expander] * (
-                gy - (self.x_hat * ggamma[expander] + gbeta[expander]) / m)
-            gx = gx.astype(x.dtype, copy=False)
-        else:
-            inv_m = numpy.float32(1) / m
-            gx = cuda.elementwise(
-                'T gy, T x_hat, U gamma, U std, U ggamma, U gbeta, \
-                U inv_m',
-                'T gx',
-                'gx = (gamma / std) * (gy - (x_hat * ggamma + gbeta) * \
-                inv_m)',
-                'bn_bwd')(gy, self.x_hat, gamma[expander],
-                          self.std[expander], ggamma[expander],
-                          gbeta[expander], inv_m)
-
-        return gx, ggamma, gbeta
+            return _MpiImpl(self.comm)
diff --git a/chainermn/links/batch_normalization.py b/chainermn/links/batch_normalization.py
--- a/chainermn/links/batch_normalization.py
+++ b/chainermn/links/batch_normalization.py
@@ -10,7 +10,7 @@
 from chainermn.functions.batch_normalization import \
     get_communication_backend
 from chainermn.functions.batch_normalization import \
-    MultiNodeBatchNormalizationFunction
+    MultiNodeBNImplSelector
 
 
 class MultiNodeBatchNormalization(link.Link):
@@ -103,11 +103,12 @@ def __call__(self, x, finetune=False):
             else:
                 decay = self.decay
 
-            func = MultiNodeBatchNormalizationFunction(
-                self.comm, self.eps, self.avg_mean, self.avg_var, decay,
-                communication_backend=self._communication_backend)
+            func = batch_normalization.BatchNormalization(
+                self.eps, self.avg_mean, self.avg_var, decay,
+                impl_selector=MultiNodeBNImplSelector(
+                    self.comm, self._communication_backend))
 
-            ret = func(x, gamma, beta)
+            ret = func.apply((x, gamma, beta))[0]
 
             self.avg_mean[:] = func.running_mean
             self.avg_var[:] = func.running_var
@@ -130,7 +131,7 @@ def start_finetuning(self):
         self.N = 0
 
     def copy(self, mode='share'):
-        to_be_preserved = ['_communication_backend', 'comm']
+        to_be_preserved = ['comm']
         preserved = {}
         for name in to_be_preserved:
             preserved[name] = getattr(self, name)


Test patch: