Repository: pytorch
Issue: TrilinearBackward takes 98.4% of total computational time, is this to be expected?
Labels: module: performance,module: nn,triaged

I made a new neural network yesterday that seems to be running very slow for some reason, so I used the profiler on it today and got the following. When just running it on the forward pass of the model I get the following:

```
-----------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
-----------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
              model_inference         6.38%      74.870ms       100.00%        1.173s        1.173s             1  
               aten::bilinear         0.02%     215.916us        68.38%     802.342ms      66.862ms            12  
             aten::_trilinear         8.46%      99.252ms        68.33%     801.711ms      66.809ms            12  
                    aten::bmm        47.37%     555.859ms        48.33%     567.123ms     184.610us          3072  
             aten::unique_dim         4.30%      50.488ms        20.53%     240.858ms     240.858ms             1  
                 aten::select         7.48%      87.707ms         9.29%     108.992ms       4.468us         24394  
                   aten::add_         4.60%      54.026ms         4.60%      54.026ms       8.473us          6376  
             aten::as_strided         3.28%      38.487ms         3.28%      38.487ms       0.878us         43846  
                aten::permute         2.34%      27.502ms         2.93%      34.392ms       3.732us          9216  
                 aten::narrow         1.17%      13.781ms         2.92%      34.317ms       5.583us          6147  
-----------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 1.173s

```
In this case the bilinear operation seems to take a rather large percentage of the total time, which is a bit surprising since I wouldn't have expected it to be that big a part of the network.

However when I profile the full network when training (forward pass + backward), I get the following:
```
-----------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                               Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
-----------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                    model_inference         0.26%     114.160ms       100.00%       44.107s       44.107s             1  
                   aten::_trilinear        17.29%        7.628s        98.41%       43.404s        1.033s            42  
                  TrilinearBackward         0.00%     427.701us        95.71%       42.214s        3.518s            12  
                          aten::mul        52.79%       23.282s        52.87%       23.320s       5.735ms          4066  
                          aten::bmm        26.15%       11.534s        26.24%       11.575s       1.675ms          6912  
                     aten::bilinear         0.00%     284.004us         2.70%        1.192s      99.321ms            12  
                         aten::add_         1.00%     440.932ms         1.00%     441.583ms      42.960us         10279  
                   aten::unique_dim         0.11%      46.672ms         0.54%     240.012ms     240.012ms             1  
                       aten::narrow         0.20%      87.566ms         0.50%     221.255ms      11.521us         19204  
                      aten::permute         0.26%     113.801ms         0.32%     142.023ms       6.849us         20736  
-----------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 44.107s

```

Now in this case the bilinear operation (which has to be the one calling the trilinear operation), uses more than 98% of the total time, which to me suggests that the operation either has a bug or is highly unoptimized? The alternative is that the bilinear function is effectively un-useable with autograd from what I can see.

I don't have a nice small example of the code I'm currently running, but I can produce one if needed. However the essential part of my network is this propagation block, where it is the bilinear operations that takes all the compute both in the forward and backward passes:


```
class PropagationBlock(nn.Module):
    def __init__(self, xn_dim, xn_attr_dim, xe_dim, xe_attr_dim):
        super().__init__()

        self.fc1 = FullyConnectedNet([1, xe_dim],activation_fnc=torch.nn.functional.silu)
        self.fc2 = FullyConnectedNet([1, xn_dim],activation_fnc=torch.nn.functional.silu)

        self.xn_bilinear = nn.Bilinear(xn_dim, xn_attr_dim, xn_dim, bias=False)
        self.xe_bilinear = nn.Bilinear(xe_dim, xe_attr_dim, xe_dim, bias=False)

        self.nodes_to_edges = NodesToEdges(xn_dim,xn_dim)
        self.edges_to_nodes = EdgesToNodes(xe_dim,xe_dim)

        self.activation = torch.nn.functional.silu

        return

    def forward(self, xn, xn_attr, xe_attr, xe_src, xe_dst):
        eps = 1e-9

        xn = self.xn_bilinear(xn, xn_attr)
        xn = xn / (xn.std(dim=1)[:,None] + eps)

        weight = self.fc1(xe_attr)
        xe = self.nodes_to_edges(xn, xe_src, xe_dst, weight)

        xe = self.xe_bilinear(xe, xe_attr)
        xe = xe / (xe.std(dim=1)[:,None] + eps)

        weight = self.fc2(xe_attr)
        xn = self.edges_to_nodes(xe, xe_src, xe_dst, weight)

        xn = self.activation(xn)
        xn = xn / (xn.std(dim=1)[:,None] + eps)

        return xn

```

Edit: I should add that I'm running this with the following versions:
```
(e3nn_cpu) tue@tue-laptop:~/PycharmProjects/GraphNetworks$ pip list
Package               Version
--------------------- ---------
ase                   3.21.1
certifi               2020.12.5
chardet               4.0.0
cycler                0.10.0
decorator             4.4.2
e3nn                  0.2.7
googledrivedownloader 0.4
h5py                  3.2.1
idna                  2.10
isodate               0.6.0
Jinja2                2.11.3
joblib                1.0.1
kiwisolver            1.3.1
llvmlite              0.36.0
MarkupSafe            1.1.1
matplotlib            3.4.1
mpmath                1.2.1
networkx              2.5.1
numba                 0.53.1
numpy                 1.20.2
pandas                1.2.4
Pillow                8.2.0
pip                   21.0.1
pyparsing             2.4.7
python-dateutil       2.8.1
python-louvain        0.15
pytz                  2021.1
rdflib                5.0.0
requests              2.25.1
scikit-learn          0.24.1
scipy                 1.6.2
setuptools            56.0.0
six                   1.15.0
sympy                 1.8
threadpoolctl         2.1.0
torch                 1.8.1+cpu
torch-cluster         1.5.9
torch-geometric       1.7.0
torch-scatter         2.0.6
torch-sparse          0.6.9
torch-spline-conv     1.2.1
torchaudio            0.8.1
torchvision           0.9.1+cpu
tqdm                  4.60.0
typing-extensions     3.7.4.3
urllib3               1.26.4

(e3nn_cpu) tue@tue-laptop:~/PycharmProjects/GraphNetworks$ python --version
Python 3.8.5

```



cc @albanD @mruberry @jbschlosser @VitalyFedyunin @ngimel @heitorschueroff @jianyuh @nikitaved @pearu @walterddr @IvanYashchuk @xwang233 @Lezcano