Repository: chainer/chainer

Issue:
F.bilinear requires huge GPU memory
The following code results `cupy.cuda.memory.OutOfMemoryError: out of memory to allocate 18014398509481984 bytes (total 18014399785863168 bytes)`
```
import chainer
import cupy
b = chainer.links.Bilinear(256, 256, 256).to_gpu()
e1 = cupy.random.randn(64, 256).astype('f')
e2 = cupy.random.randn(64, 256).astype('f')
y = b(e1, e2)
print(y)
```

How to fix: merge cupy/cupy#1218 (or do not use `cupy.einsum`).
I confirmed the code run in ~5sec with
- chainer: master(6bab773dec70f291108ab2575622805252f9a208)
- cupy: (Merge: cupy/cupy@6162f9a cupy/cupy@7f89bd0)


Discussion:
The performance regression is due to #3917.

Resolving patch:
diff --git a/chainer/functions/connection/bilinear.py b/chainer/functions/connection/bilinear.py
--- a/chainer/functions/connection/bilinear.py
+++ b/chainer/functions/connection/bilinear.py
@@ -80,7 +80,19 @@ def forward(self, inputs):
         W = inputs[2]
 
         xp = cuda.get_array_module(*inputs)
-        y = xp.einsum('ij,ik,jkl->il', e1, e2, W)
+        if xp is numpy:
+            y = numpy.einsum('ij,ik,jkl->il', e1, e2, W)
+        else:
+            i_len, j_len = e1.shape
+            k_len = e2.shape[1]
+            # 'ij,ik->ijk'
+            e1e2 = e1[:, :, None] * e2[:, None, :]
+            # ijk->i[jk]
+            e1e2 = e1e2.reshape(i_len, j_len * k_len)
+            # jkl->[jk]l
+            W_mat = W.reshape(-1, W.shape[2])
+            # 'i[jk],[jk]l->il'
+            y = e1e2.dot(W_mat)
 
         if len(inputs) == 6:
             V1, V2, b = inputs[3:]
@@ -110,9 +122,23 @@ def forward(self, inputs):
         W, gy = inputs[2], inputs[-1]
 
         xp = cuda.get_array_module(*inputs)
-        ge1 = xp.einsum('ik,jkl,il->ij', e2, W, gy)
-        ge2 = xp.einsum('ij,jkl,il->ik', e1, W, gy)
-        gW = xp.einsum('ij,ik,il->jkl', e1, e2, gy)
+        if xp is numpy:
+            gW = numpy.einsum('ij,ik,il->jkl', e1, e2, gy)
+            ge1 = numpy.einsum('ik,jkl,il->ij', e2, W, gy)
+            ge2 = numpy.einsum('ij,jkl,il->ik', e1, W, gy)
+        else:
+            kern = cuda.reduce('T in0, T in1, T in2', 'T out',
+                               'in0 * in1 * in2', 'a + b', 'out = a', 0,
+                               'bilinear_product')
+
+            e1_b = e1[:, :, None, None]  # ij
+            e2_b = e2[:, None, :, None]  # ik
+            gy_b = gy[:, None, None, :]  # il
+            W_b = W[None, :, :, :]  # jkl
+
+            gW = kern(e1_b, e2_b, gy_b, axis=0)  # 'ij,ik,il->jkl'
+            ge1 = kern(e2_b, W_b, gy_b, axis=(2, 3))  # 'ik,jkl,il->ij'
+            ge2 = kern(e1_b, W_b, gy_b, axis=(1, 3))  # 'ij,jkl,il->ik'
 
         ret = ge1.reshape(inputs[0].shape), ge2.reshape(inputs[1].shape), gW
 


Test patch:
diff --git a/tests/chainer_tests/functions_tests/connection_tests/test_bilinear.py b/tests/chainer_tests/functions_tests/connection_tests/test_bilinear.py
--- a/tests/chainer_tests/functions_tests/connection_tests/test_bilinear.py
+++ b/tests/chainer_tests/functions_tests/connection_tests/test_bilinear.py
@@ -155,6 +155,27 @@ def f(*inputs):
              cuda.to_gpu(self.ggb)), **self.check_double_backward_options)
 
 
+@attr.slow
+class TestBilinearFunctionLarge(unittest.TestCase):
+
+    def setUp(self):
+        self.e1 = _uniform(256, 256)
+        self.e2 = _uniform(256, 256)
+        self.w = _uniform(256, 256, 256)
+        self.v1 = _uniform(256, 256)
+        self.v2 = _uniform(256, 256)
+        self.b = _uniform(256)
+
+    def test_cpu(self):
+        chainer.functions.bilinear(
+            self.e1, self.e2, self.w, self.v1, self.v2, self.b)
+
+    @attr.gpu
+    def test_gpu(self):
+        chainer.functions.bilinear(*map(cuda.to_gpu, (
+            self.e1, self.e2, self.w, self.v1, self.v2, self.b)))
+
+
 class TestBilinearFunctionInvalidArgument(unittest.TestCase):
 
     def setUp(self):