Repository: pytorch
Issue: [JIT] Provide a way to access a flattened optimized graph from Profiling Executor
Labels: oncall: jit,module: bootcamp,triaged

## 🚀 Feature
Currently when a profiling executor inserts a type specialization, it generates a fallback function for the case when a type doesn't match specialized. In that case the execution goes into the fallback which has its own instance of executor, which then can lead to another respecialization and creating yet another fallback function. This cascading mechanism allows us to have several specializations, but the problem is that it is hard to get access to graphs from deeper specializations. The requested feature is to provide a way to access a graph where all specializations are present.

## To reproduce
Consider the following test case:
```python
import torch

def foo(a, b):
    c = torch.matmul(a, b)
    x = c * a * b
    y = torch.matmul(c, x)
    return y

torch._C._jit_override_can_fuse_on_cpu(True)

f = torch.jit.script(foo)
x = torch.rand(4, 4)
for _ in range(10):
    f(x, x)

print("After specialization for (4, 4) size")
print(torch.jit.last_executed_optimized_graph())


x = torch.rand(10, 10)
for _ in range(10):
    f(x, x)

print("After another specialization for (10, 10) size")
print(torch.jit.last_executed_optimized_graph())
```

In this test we first run a function with inputs of (4,4) size and create a fusion group for that size and then run the same function with inputs of (10,10) size and create a second fusion group for that new size.

We could see both fusion groups if we run the test with debugging logs enabled:
```
PYTORCH_JIT_LOG_LEVEL=">>kernel" python t.py
```

However, if we look at graphs printed by `torch.jit.last_executed_optimized_graph()`, we only would see (4,4) specialization in them, because the second specialization is hidden inside the fallback function:
```
After another specialization for (10, 10) size
graph(%a.1 : Tensor,
      %b.1 : Tensor):
  %c.1 : Tensor = aten::matmul(%a.1, %b.1) # t.py:4:8
  %24 : Float(4, 4), %25 : Float(4, 4), %26 : Float(4, 4), %27 : bool = prim::TypeCheck[types=[Float(4, 4), Float(4, 4), Float(4, 4)]](%b.1, %c.1, %a.1)
  %28 : Tensor = prim::If(%27)
    block0():
      %x.6 : Float(4, 4) = prim::TensorExprGroup_0(%24, %25, %26)
      -> (%x.6)
    block1():
      %32 : Function = prim::Constant[name="fallback_function", fallback=1]()
      %33 : (Tensor) = prim::CallFunction(%32, %b.1, %c.1, %a.1)
      %34 : Tensor = prim::TupleUnpack(%33)
      -> (%34)
  %y.1 : Tensor = aten::matmul(%c.1, %28) # t.py:6:8
  return (%y.1)
with prim::TensorExprGroup_0 = graph(%1 : Float(4, 4),
      %3 : Float(4, 4),
      %4 : Float(4, 4)):
  %5 : Float(4, 4) = aten::mul(%3, %4) # t.py:5:8
  %x.2 : Float(4, 4) = aten::mul(%5, %1) # t.py:5:8
  return (%x.2)
```

## Desired behavior
We could provide a feature to inline the specialized graph from "fallback_function" and thus get a flattened optimized graph with all optimizations. It would look something like this:
```
graph(%a.1 : Tensor,
      %b.1 : Tensor):
  %c.1 : Tensor = aten::matmul(%a.1, %b.1)
  %24 : Float(4, 4), %25 : Float(4, 4), %26 : Float(4, 4), %27 : bool = prim::TypeCheck[types=[Float(4, 4), Float(4, 4), Float(4, 4)]](%b.1, %c.1, %a.1)
  %28 : Tensor = prim::If(%27)
    block0():
      # Specialization 1: size = (4, 4)
      %x.6 : Float(4, 4) = prim::TensorExprGroup_0(%24, %25, %26)
      -> (%x.6)
    block1():
      %t0 : Float(10, 10), %t1 : Float(10, 10), %t2 : Float(10, 10), %t3 : bool = prim::TypeCheck[types=[Float(10, 10), Float(10, 10), Float(10, 10)]](%b.1, %c.1, %a.1)
      %t4 : Tensor = prim::If(%t3)
        block0():
          # Specialization 2: size = (10, 10)
          %t5 : Float(10, 10) = prim::TensorExprGroup_1(%t0, %t1, %t2)
          -> (%t5)
        block1():
          # None of the specializations worked, use the default not-optimized version
          %t6 : Tensor = aten::mul(%c.1, %a.1)
          %t7 : Tensor = aten::mul(%t6, %b.1)
          -> (%t7)
      -> (%t4)
  %y.1 : Tensor = aten::matmul(%c.1, %28)
  return (%y.1)
with prim::TensorExprGroup_0 = graph(%1 : Float(4, 4),
      %3 : Float(4, 4),
      %4 : Float(4, 4)):
  %5 : Float(4, 4) = aten::mul(%3, %4)
  %x.2 : Float(4, 4) = aten::mul(%5, %1)
  return (%x.2)
with prim::TensorExprGroup_1 = graph(%1 : Float(10, 10),
      %3 : Float(10, 10),
      %4 : Float(10, 10)):
  %5 : Float(10, 10) = aten::mul(%3, %4)
  %x.2 : Float(10, 10) = aten::mul(%5, %1)
  return (%x.2)
```

cc @gmagogsfm