Repository: pytorch
Issue: [JIT] batch operators in training/inference 
Labels: oncall: jit

## 🚀 Feature

Batching operators is a well known optimization that can lead to large speed ups in inference and training. https://github.com/pytorch/pytorch/issues/58628 describes one particular variant of operator batching, where the same inputs are reused across multiple operators, and it shows a nice speed on GPU. We have also have a [batch_mm](https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/batch_mm.cpp#L25) pass which looks for a particular pattern of adds and mms and batches them together. 

An additional batching we would like to add is to take a series of adjacent addmms and batch them in one addmm. 

The following code:
```
def foo(a, b, c, d, e, f):
    x = torch.addmm(a, b, c)
    y = torch.addmm(d, e, f)
    return x, y
```
Can be rewritten as 
```
def foo(a, b, c, d, e, f):
    x, y = torch.chunk(torch.baddmm(torch.cat(a, d), torch.cat(b, e), torch.cat(e, f)), 2)
    return x, y
```
Provided shapes/dtypes/devices of the tensors are compatible.  We had an internal model which sped up training by 28% with this batching on.

We probably want to do some benchmarking both for cpu and gpu to determine how generally profitable this is, or design the batching such that we do some amount of benchmarking at runtime. 

There are a couple options here in implementation:
- before autodiff or after autodiff 
- as a runtime operator which might do things like runtime benchmarking and check correctness preconditions, or implement it complete in terms of existing operators and generate our correctness conditions in the graph

as a pure graph transformation: 

for now, i'm going to describe after autodiff, just because that is what's set up today, but we could extend this to work pre autodiff, we would just have to a few small changes. 
``` 
import torch

@torch.jit.script
def foo(a, b, c, d, e, f):
    x = torch.addmm(a, b, c)
    y = torch.addmm(d, e, f)
    return x, y

inps = [torch.rand([20, 20]) for i in range(6)]
foo(*inps)
foo(*inps)
```
> PYTORCH_JIT_LOG_LEVEL='>>profiling_graph_executor_impl' python test/temp.py
Gives
```
After RemoveProfileNodesAndSpecializeTypes, before BatchMM
graph(%a.1 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu),
       %b.1 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu),
       %c.1 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu),
       %d.1 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu),
      %e.1 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu),
       %f.1 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu)):
   %6 : int = prim::Constant[value=1]()
   %9 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu) = aten::mm(%b.1, %c.1) # <string>:3:24
   %12 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu) = aten::add(%a.1, %9, %6) # <string>:3:17
   %15 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu) = aten::mm(%e.1, %f.1) # <string>:3:24
   %18 : Float(20, 20, strides=[20, 1], requires_grad=0, device=cpu) = aten::add(%d.1, %15, %6) # <string>:3:17
   %21 : (Tensor, Tensor) = prim::TupleConstruct(%12, %18)
   return (%21)
```
Annoyingly, the `addmm` are broken up, but we can turn off that decomposition, or just focus on batching `mm`s for now. (Side note: we should really only be decomposing addmm when it creates fusion opportunitiees, see: https://github.com/pytorch/pytorch/issues/56385). 


Anyway, we can add a pass [here](https://github.com/pytorch/pytorch/blob/3979cb0656fe2d8b0445768a769bd624b10778b5/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp#L349) which attempts to fuse adjacent `addmms` into one `baddmm`. We are going to be looking for addmms on the same block which can be moved adjacent to each other, and shapes/dtypes/devices which allow for the transformation. We will also have to set up a guard which ensures that the runtime execution of tensors follows the types we saw, following similar logic to [here](https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/tensorexpr_fuser.cpp#L358). We can do the same thing before autodiff, we just have to add remove profile nodes and specialize types as we are doing for inference currently. Depending on our benchmarking, if we are guarding on complete shapes we can AOT benchmark if making the transformation is profitable.

as a runtime operator:

most of the logic is the same, except now we'd be registering a custom operator similar to [here](https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/batch_mm.cpp#L110). With a custom operator, we can add the correctness tests within the operator, and if the tensors are incompatible to be batched together run each operator separately. Additionally, we can do some runtime testing if the batching is profitable, and run each operator separately in the case that it is not. 

This details the issue, but definitely there will many additional questions whoever implements has and I'm happy to answer more. I will also edit more pointers and details into the issue.


cc @gmagogsfm