Repository: pytorch
Issue: Jit fuser half precision support
Labels: oncall: jit

## 🐛 Bug

I have noticed that in order to generate fused kernels for half precision I am required to turn the profiling executor off 
```py
torch._C._jit_set_profiling_executor(False)
torch._C._jit_set_profiling_mode(False)
```
This finding is based on this unit test https://github.com/pytorch/pytorch/blob/7a728a86c42bf0195f83b34ba7b2e2dbdbe9f8ce/test/test_jit_fuser.py#L136

Are there plans to support half precision with the profiling executor as these settings appears to result in sub-optimal fusion groups, with slower performance. See the examples below:

Many thanks

## To Reproduce

```py
import torch 

@torch.jit.script
def mul_cat_mul_add(a, b, c):
    return torch.cat([a * b, b]) * c + c

a = torch.randn(1024,1024).cuda().half()
b = torch.randn(1024,1024).cuda().half()
c = torch.randn(2048,1024).cuda().half()
```

With profiling executor turned ON, no fusion groups are produced for `half` inputs

```py
torch._C._jit_set_profiling_executor(True)
torch._C._jit_set_profiling_mode(True)
torch._C._jit_set_num_profiled_runs(1)

mul_cat_mul_add.graph_for(a, b, c)
mul_cat_mul_add.graph_for(a, b, c)
```
```
graph(%a.1 : Tensor,
      %b.1 : Tensor,
      %c.1 : Tensor):
  %3 : int = prim::Constant[value=0]()
  %4 : int = prim::Constant[value=1]()
  %5 : Tensor = aten::mul(%a.1, %b.1) # /tmp/ipykernel_35787/878799905.py:11:22
  %6 : Tensor[] = prim::ListConstruct(%5, %b.1)
  %7 : Tensor = aten::cat(%6, %3) # /tmp/ipykernel_35787/878799905.py:11:11
  %8 : Tensor = aten::mul(%7, %c.1) # /tmp/ipykernel_35787/878799905.py:11:11
  %9 : Tensor = aten::add(%8, %c.1, %4) # /tmp/ipykernel_35787/878799905.py:11:11
  return (%9)
```
With profiling executor turned OFF, fusion groups are created using legacy fuser, which results in multiple fused kernels (this is sub-optimal compared to using float inputs where the kernel is fused end-to-end)

```py
torch._C._jit_set_profiling_executor(False)
torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_num_profiled_runs(1)

mul_cat_mul_add.graph_for(a, b, c)
mul_cat_mul_add.graph_for(a, b, c)
```
```
graph(%a.1 : Half(*, *, requires_grad=0, device=cuda:0),
      %b.1 : Half(*, *, requires_grad=0, device=cuda:0),
      %c.1 : Half(*, *, requires_grad=0, device=cuda:0)):
  %12 : Half(*, *, requires_grad=1, device=cuda:0) = prim::FusionGroup_0(%b.1, %a.1)
  %10 : Half(*, *, requires_grad=1, device=cuda:0) = prim::FusionGroup_1(%c.1, %12)
  return (%10)
with prim::FusionGroup_0 = graph(%1 : Half(*, *, requires_grad=0, device=cuda:0),
      %3 : Half(*, *, requires_grad=0, device=cuda:0)):
  %4 : Half(*, *, requires_grad=0, device=cuda:0) = aten::mul(%3, %1) # /tmp/ipykernel_35787/878799905.py:11:22
  %2 : Half(*, *, requires_grad=1, device=cuda:0) = prim::FusedConcat[dim=0](%4, %1)
  return (%2)
with prim::FusionGroup_1 = graph(%1 : Half(*, *, requires_grad=0, device=cuda:0),
      %4 : Half(*, *, requires_grad=1, device=cuda:0)):
  %2 : int = prim::Constant[value=1]()
  %5 : Half(*, *, requires_grad=1, device=cuda:0) = aten::mul(%4, %1) # /tmp/ipykernel_35787/878799905.py:11:11
  %3 : Half(*, *, requires_grad=1, device=cuda:0) = aten::add(%5, %1, %2) # /tmp/ipykernel_35787/878799905.py:11:11
  return (%3)
```
In addition to this, we can end up with scenarios where having **more** ops is actually significantly **faster**, where one of the ops is not able to pair with other ops to form a fusion group. Profiling results for the following kernel are slower than the above, despite it having the addition `mul` removed.

```py
@torch.jit.script
def mul_cat_mul(a, b, c):
    return torch.cat([a * b, b]) * c

mul_cat_mul.graph_for(a, b, c)
mul_cat_mul.graph_for(a, b, c)
```
```
graph(%a.1 : Half(*, *, requires_grad=0, device=cuda:0),
      %b.1 : Half(*, *, requires_grad=0, device=cuda:0),
      %c.1 : Half(*, *, requires_grad=0, device=cuda:0)):
  %9 : Half(*, *, requires_grad=1, device=cuda:0) = prim::FusionGroup_0(%b.1, %a.1)
  %7 : Half(*, *, requires_grad=1, device=cuda:0) = aten::mul(%9, %c.1) # /tmp/ipykernel_35787/878799905.py:7:11
  return (%7)
with prim::FusionGroup_0 = graph(%1 : Half(*, *, requires_grad=0, device=cuda:0),
      %3 : Half(*, *, requires_grad=0, device=cuda:0)):
  %4 : Half(*, *, requires_grad=0, device=cuda:0) = aten::mul(%3, %1) # /tmp/ipykernel_35787/878799905.py:7:22
  %2 : Half(*, *, requires_grad=1, device=cuda:0) = prim::FusedConcat[dim=0](%4, %1)
  return (%2)
```

## Expected behavior

Ideally we would have equivalent support that we see with `float` inputs where the entire kernel is fused end-to-end in a single fusion group, with slower performance.

```py
torch._C._jit_set_profiling_executor(True)
torch._C._jit_set_profiling_mode(True)
torch._C._jit_set_num_profiled_runs(1)

a = torch.randn(1024,1024).cuda().float()
b = torch.randn(1024,1024).cuda().float()
c = torch.randn(2048,1024).cuda().float()

mul_cat_mul_add.graph_for(a, b, c)
mul_cat_mul_add.graph_for(a, b, c)
```
```
graph(%a.1 : Tensor,
      %b.1 : Tensor,
      %c.1 : Tensor):
  %34 : Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0), %35 : Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0), %36 : Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0), %37 : bool = prim::TypeCheck[types=[Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0), Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0), Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0)]](%c.1, %b.1, %a.1)
  %38 : Tensor = prim::If(%37)
    block0():
      %20 : Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0) = prim::TensorExprGroup_0(%34, %35, %36)
      -> (%20)
    block1():
      %47 : Function = prim::Constant[name="fallback_function", fallback=1]()
      %48 : (Tensor) = prim::CallFunction(%47, %c.1, %b.1, %a.1)
      %49 : Tensor = prim::TupleUnpack(%48)
      -> (%49)
  return (%38)
with prim::TensorExprGroup_0 = graph(%c.1 : Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0),
      %b.1 : Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0),
      %a.1 : Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0)):
  %2 : int = prim::Constant[value=1]()
  %7 : int = prim::Constant[value=0]()
  %13 : Float(1024, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0) = aten::mul(%a.1, %b.1) # /tmp/ipykernel_35787/878799905.py:11:22
  %11 : Tensor[] = prim::ListConstruct(%13, %b.1)
  %8 : Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0) = aten::cat(%11, %7) # /tmp/ipykernel_35787/878799905.py:11:11
  %5 : Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0) = aten::mul(%8, %c.1) # /tmp/ipykernel_35787/878799905.py:11:11
  %3 : Float(2048, 1024, strides=[1024, 1], requires_grad=0, device=cuda:0) = aten::add(%5, %c.1, %2) # /tmp/ipykernel_35787/878799905.py:11:11
  return (%3)
```

## Environment

```
PyTorch version: 1.10.0a0+gitc545b09    
Is debug build: False                                                                                                                                                     
CUDA used to build PyTorch: 11.2                                                                                                                                          
ROCM used to build PyTorch: N/A                                                                                                                                                                                                                                                                                                                     OS: Ubuntu 18.04.6 LTS (x86_64)                                                                                                                                           
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0                                                                                                                          Clang version: Could not collect                                                                                                                                          
CMake version: version 3.21.2                                                                                                                                             
Libc version: glibc-2.27                                                                                                                                                                                                                                                                                                                            Python version: 3.8.12 (default, Sep 10 2021, 00:16:05)  [GCC 7.5.0] (64-bit runtime)                                                                                     Python platform: Linux-5.4.0-80-generic-x86_64-with-glibc2.27                                                                                                             Is CUDA available: True                                                                                                                                                   
CUDA runtime version: 11.2.152                                                                                                                                            
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti
GPU 7: GeForce RTX 2080 Ti
GPU 8: GeForce RTX 2080 Ti
GPU 9: GeForce RTX 2080 Ti

Nvidia driver version: 460.32.03
cuDNN version: Probably one of the following:
```