Repository: pytorch
Issue: NVFuser Examples to Help with Integration with Lazy Tensor Core
Labels: triaged,LazyTensor_nvfuser_integration,module: lazy

## 🚀 Feature
These are 3 Fusion examples (forward and backward) from Bert to illustrate Fusions that should be supported with Lazy Tensor Cores.  I am **only going through the forward examples**, for now.

You can find them here: https://github.com/kevinstephano/FusionMicroBenchmarks/tree/main/LazyTensorCore

## Example 1 - Bias+Gelu

https://github.com/kevinstephano/FusionMicroBenchmarks/blob/main/LazyTensorCore/bias_gelu.py

```
PYTORCH_JIT_LOG_LEVEL=">>>graph_fuser" LTC_TS_CUDA=1 python bias_gelu.py
```

Currently, LTC is decomposing Gelu into elemental components that are more optimized than the equivalent implementation in Aten.  Is there a way to turn off this elemental breakdown in LTC?

Graph from LTC:
```
[DUMP graph_fuser.cpp:1924] with prim::CudaFusionGroup_0 = graph(%14 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %18 : Float(1, 1, 4096, strides=[4096, 4096, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %19 : Float(1, 1, 4096, strides=[4096, 4096, 1], requires_grad=0, device=cuda:0)):
[DUMP graph_fuser.cpp:1924]   %4 : float = prim::Constant[value=0.5]()
[DUMP graph_fuser.cpp:1924]   %7 : float = prim::Constant[value=1.]()
[DUMP graph_fuser.cpp:1924]   %12 : float = prim::Constant[value=0.70710676908493042]()
[DUMP graph_fuser.cpp:1924]   %16 : int = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %20 : Float(1, 1, 4096, strides=[4096, 4096, 1], requires_grad=0, device=cuda:0) = aten::mul(%18, %19)
[DUMP graph_fuser.cpp:1924]   %17 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::add(%14, %20, %16)
[DUMP graph_fuser.cpp:1924]   %13 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::mul(%17, %12)
[DUMP graph_fuser.cpp:1924]   %11 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::erf(%13)
[DUMP graph_fuser.cpp:1924]   %9 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::add(%11, %7, %16)
[DUMP graph_fuser.cpp:1924]   %5 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::mul(%17, %4)
[DUMP graph_fuser.cpp:1924]   %2 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::mul(%5, %9)
[DUMP graph_fuser.cpp:1924]   return (%2, %17)
```

Graph from TorchScript:
```
[DUMP graph_fuser.cpp:1924] with prim::CudaFusionGroup_0 = graph(%3 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %4 : Float(1, 1, 4096, strides=[4096, 4096, 1], requires_grad=0, device=cuda:0)):
[DUMP graph_fuser.cpp:1924]   %1 : bool = prim::Constant[value=0]()
[DUMP graph_fuser.cpp:1924]   %5 : int = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %out1.2 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::add(%3, %4, %5) # bias_gelu.py:13:15
[DUMP graph_fuser.cpp:1924]   %out2.2 : Float(256, 128, 4096, strides=[524288, 4096, 1], requires_grad=0, device=cuda:0) = aten::gelu(%out1.2, %1) # /opt/pytorch/pytorch/torch/nn/functional.py:1559:11
[DUMP graph_fuser.cpp:1924]   return (%out2.2, %out1.2)
```

## Example 2 - Div+Mask+Softmax+Dropout

https://github.com/kevinstephano/FusionMicroBenchmarks/blob/main/LazyTensorCore/div_mask_softmax_dropout.py

```
LTC_SAVE_TENSORS_FILE=mytest PYTORCH_JIT_LOG_LEVEL=">>>graph_fuser" LTC_TS_CUDA=1 python div_mask_softmax_dropout.py
```

Only the Div+Mask are showing up in a the fusion.  This could be an issue in NVFuser as well in seeing a different form of the graph.  One issue in the output from LTC is that dropout is missing from its trace.

LTC IR:
```
## BEGIN_GRAPH
IR {
  %0 = f32[] prim::Constant(), value=1
  %1 = f32[256,16,128,128] aten::expand(%0), size=(256, 16, 128, 128), is_scalar_expand=1
  %2 = f32[256,16,128,128] lazy_tensors::device_data(), device=GPU:0
  %3 = f32[256,16,128,128] aten::mul(%2, %1)
  %4 = f64[] prim::Constant(), value=0
  %5 = f32[256,16,128,128] lazy_tensors::device_data(), device=GPU:0
  %6 = f64[256,16,128,128] aten::div(%5, %4)
  %7 = f32[256,16,128,128] lazy_tensors::cast(%6), type=f32
  %8 = f32[256,16,128,128] aten::add(%7, %3)
  %9 = f32[256,16,128,128] aten::softmax(%8), dim=3, dtype=6, ROOT=0
}

## END_GRAPH
```

Graph from LTC:
```
[DUMP graph_fuser.cpp:1924] with prim::CudaFusionGroup_0 = graph(%0 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %4 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %5 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0)):
[DUMP graph_fuser.cpp:1924]   %2 : int = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %6 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0) = aten::mul(%4, %5)
[DUMP graph_fuser.cpp:1924]   %3 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0) = aten::add(%0, %6, %2)
[DUMP graph_fuser.cpp:1924]   return (%3)
```

Graph from TorchScript:
```
[DUMP graph_fuser.cpp:1924] with prim::CudaFusionGroup_0 = graph(%10 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %13 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %14 : float):
[DUMP graph_fuser.cpp:1924]   %2 : bool = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %1 : float = prim::Constant[value=0.10000000000000001]() # /opt/pytorch/pytorch/torch/nn/modules/dropout.py:58:32
[DUMP graph_fuser.cpp:1924]   %7 : NoneType = prim::Constant() # /opt/pytorch/pytorch/torch/nn/functional.py:1680:14
[DUMP graph_fuser.cpp:1924]   %6 : int = prim::Constant[value=-1]() # div_mask_softmax_dropout.py:19:35
[DUMP graph_fuser.cpp:1924]   %11 : int = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %out1.1 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0) = aten::div(%13, %14) # div_mask_softmax_dropout.py:17:15
[DUMP graph_fuser.cpp:1924]   %out2.1 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0) = aten::add(%out1.1, %10, %11) # div_mask_softmax_dropout.py:18:15
[DUMP graph_fuser.cpp:1924]   %ret.1 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0) = aten::softmax(%out2.1, %6, %7) # /opt/pytorch/pytorch/torch/nn/functional.py:1683:14
[DUMP graph_fuser.cpp:1924]   %out4.1 : Float(256, 16, 128, 128, strides=[262144, 16384, 128, 1], requires_grad=0, device=cuda:0) = aten::dropout(%ret.1, %1, %2) # /opt/pytorch/pytorch/torch/nn/functional.py:1169:60
[DUMP graph_fuser.cpp:1924]   return (%out4.1)
```

## Example 3 - Bias+Dropout+Add+LayerNorm

https://github.com/kevinstephano/FusionMicroBenchmarks/blob/main/LazyTensorCore/bias_dropout_add_layernorm.py

```
LTC_SAVE_TENSORS_FILE=mytest PYTORCH_JIT_LOG_LEVEL=">>>graph_fuser" LTC_TS_CUDA=1 python bias_dropout_add_layernorm.py
```

I think NVFuser is only picking up a broken up `mul` and `add` related to the 3 input `aten::add` being broken into `scalar mul` + `add` for the bias add.  The graph in LTC is actually explicitly calling `aten::native_batch_norm` for Layer Norm which we don't want.  I am guessing this is a symptom of breaking ops into elemental ops.

IR 1 from LTC:
```
## BEGIN_GRAPH
IR {
  %0 = f32[] prim::Constant(), value=1
  %1 = f32[1,1,1024] aten::expand(%0), size=(1, 1, 1024), is_scalar_expand=1
  %2 = f32[1,1,1024] lazy_tensors::device_data(), device=GPU:0
  %3 = f32[1,1,1024] aten::mul(%2, %1)
  %4 = f32[256,128,1024] lazy_tensors::device_data(), device=GPU:0
  %5 = f32[256,128,1024] aten::add(%4, %3), ROOT=0
}

## END_GRAPH
```

IR 2 from LTC:
```
## BEGIN_GRAPH
IR {
  %0 = pred[256,128,1024] lazy_tensors::device_data(), device=GPU:0, ROOT=0
  %1 = f32[] prim::Constant(), value=1
  %2 = f32[256,128,1024] aten::expand(%1), size=(256, 128, 1024), is_scalar_expand=1
  %3 = f32[256,128,1024] lazy_tensors::device_data(), device=GPU:0
  %4 = f32[256,128,1024] aten::mul(%3, %2)
  %5 = f32[256,128,1024] lazy_tensors::device_data(), device=GPU:0
  %6 = f32[256,128,1024] aten::add(%5, %4), ROOT=1
  %7 = f32[1,32768,1024] aten::view(%6), output_size=(1, 32768, 1024), ROOT=2
  %8 = f32[] prim::Constant(), value=0
  %9 = f32[32768] aten::expand(%8), size=(32768), is_scalar_expand=1
  %10 = f32[] prim::Constant(), value=0
  %11 = f32[32768] aten::expand(%10), size=(32768), is_scalar_expand=1
  %12 = f32[] prim::Constant(), value=0
  %13 = f32[32768] aten::expand(%12), size=(32768), is_scalar_expand=1
  %14 = f32[] prim::Constant(), value=1
  %15 = f32[32768] aten::expand(%14), size=(32768), is_scalar_expand=1
  %16 = tuple[] aten::native_batch_norm(%7, %15, %13, %11, %9), num_outputs=3, training=1, momentum=0, eps=1e-12, ROOT=5
  %17 = f32[256,128,1024] aten::view(%16.0), output_size=(256, 128, 1024), ROOT=6
  %18 = f32[] prim::Constant(), value=1
  %19 = f32[1024] lazy_tensors::device_data(), device=GPU:0
  %20 = f32[256,128,1024] aten::mul(%17, %19)
  %21 = f32[256,128,1024] aten::mul(%20, %18)
  %22 = f32[1024] lazy_tensors::device_data(), device=GPU:0
  %23 = f32[256,128,1024] aten::add(%22, %21), ROOT=7
}

## END_GRAPH
```

TS Graph from LTC:
```
[DUMP graph_fuser.cpp:1924] with prim::CudaFusionGroup_0 = graph(%0 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %4 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %5 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0)):
[DUMP graph_fuser.cpp:1924]   %2 : int = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %6 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0) = aten::mul(%4, %5)
[DUMP graph_fuser.cpp:1924]   %3 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0) = aten::add(%0, %6, %2)
[DUMP graph_fuser.cpp:1924]   return (%3)
```

TS Graph from TorchScript:
```
[DUMP graph_fuser.cpp:1924] with prim::CudaFusionGroup_0 = graph(%2 : Float(1024, strides=[1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %3 : Float(1024, strides=[1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %9 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %19 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0),
[DUMP graph_fuser.cpp:1924]       %20 : Float(1, 1, 1024, strides=[1024, 1024, 1], requires_grad=0, device=cuda:0)):
[DUMP graph_fuser.cpp:1924]   %4 : float = prim::Constant[value=9.9999999999999998e-13]()
[DUMP graph_fuser.cpp:1924]   %1 : int[] = prim::Constant[value=[1024]]()
[DUMP graph_fuser.cpp:1924]   %15 : bool = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %14 : float = prim::Constant[value=1.1111111111111112]()
[DUMP graph_fuser.cpp:1924]   %13 : float = prim::Constant[value=0.90000000000000002]()
[DUMP graph_fuser.cpp:1924]   %21 : int = prim::Constant[value=1]()
[DUMP graph_fuser.cpp:1924]   %out1.2 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0) = aten::add(%19, %20, %21) # bias_dropout_add_layernorm.py:14:15
[DUMP graph_fuser.cpp:1924]   %res.1 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0), %mask.4 : Bool(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0) = aten::native_dropout(%out1.2, %13, %14, %15) # <string>:46:23
[DUMP graph_fuser.cpp:1924]   %out3.2 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0) = aten::add(%res.1, %9, %21) # bias_dropout_add_layernorm.py:16:15
[DUMP graph_fuser.cpp:1924]   %output.1 : Float(256, 128, 1024, strides=[131072, 1024, 1], requires_grad=0, device=cuda:0), %mean.4 : Float(256, 128, 1, strides=[128, 1, 1], requires_grad=0, device=cuda:0), %rstd.4 : Float(256, 128, 1, strides=[128, 1, 1], requires_grad=0, device=cuda:0) = aten::native_layer_norm(%out3.2, %1, %2, %3, %4) # <string>:33:33
[DUMP graph_fuser.cpp:1924]   return (%output.1, %mean.4, %rstd.4, %out3.2, %mask.4)
```