Shape inference time doubled after #649

[zasdfgbnm]

Before #649:

-------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations
-------------------------------------------------------------------------------------------
LayerNormBackward_ShapeInference                       87.8 us         87.6 us         7758
LayerNormForward_ShapeInference                        29.9 us         29.8 us        22938
LayerNormBackward_NoShapeInferenceCachedBaseline       36.5 us         36.5 us        18906
LayerNormForward_NoShapeInferenceCachedBaseline        21.4 us         21.3 us        33063

After:

-------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations
-------------------------------------------------------------------------------------------
LayerNormBackward_ShapeInference                        162 us          162 us         4332
LayerNormForward_ShapeInference                        50.3 us         50.2 us        13811
LayerNormBackward_NoShapeInferenceCachedBaseline        111 us          111 us         6304
LayerNormForward_NoShapeInferenceCachedBaseline        42.4 us         42.3 us        16632

[naoyam]

Any plan for this issue?

[IvanYashchuk]

I wrote a simple script that measures the performance of fusion execution with an increasing number of inputs.

[---------------------  ---------------------]
   N_args  |  Before #649 PR  |  After #649 PR
1 threads: -----------------------------------
      2    |       14.8       |        27.9
      4    |       17.8       |        40.2
      8    |       24.0       |        73.3
      16   |       37.3       |       168.0
      32   |       63.6       |       516.5
      64   |      124.2       |      2108.0
      128  |      251.1       |      7313.7

Times are in microseconds (us).

Click for code

I ran it twice for different versions of nvFuser

from nvfuser import FusionDefinition, DataType
import nvfuser
from looseversion import LooseVersion
import torch
import torch.utils.benchmark as benchmark

nv_version = LooseVersion(nvfuser.__version__)
if nv_version < LooseVersion("0.0.17"):
    description = "Before #649 PR"
else:
    description = "After #649 PR"

results = []

for N in (2 ** x for x in range(1, 8)):
    args = [torch.randn(3, 3, device='cuda') for _ in range(N)]

    def define_input(fd: FusionDefinition):
        if nv_version >= LooseVersion("0.0.17"):
            return fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float)
        else:
            return fd.define_tensor(symbolic_sizes=[-1, -1], contiguity=[True, True], dtype=DataType.Float)

    def nvfuser_fusion_id0(fd: FusionDefinition) -> None :
        T0 = define_input(fd)
        for _ in range(N - 1):
            define_input(fd)
        S0 = fd.define_scalar(0.00000, dtype=DataType.Double)
        T_out = fd.ops.add(S0, T0)
        fd.add_output(T_out)

    with FusionDefinition() as fd:
        nvfuser_fusion_id0(fd)

    out = fd.execute(args)

    timer = benchmark.Timer(
        stmt='fd.execute(args)',
        globals={'fd': fd, 'args': args},
        num_threads=1,
        sub_label=f'{N}',
        description=description,
    )
    results.append(timer.blocked_autorange(min_run_time=1))

import pickle
if nv_version < LooseVersion("0.0.17"):
    with open('test_overhead_nvfuser_scaling_before649.pkl', 'wb') as f:
        pickle.dump(results, f)
else:
    with open('test_overhead_nvfuser_scaling_after649.pkl', 'wb') as f:
        pickle.dump(results, f)

# If file exists, load it
try:
    with open('test_overhead_nvfuser_scaling_before649.pkl', 'rb') as f:
        results_before649 = pickle.load(f)
except FileNotFoundError:
    results_before649 = [None]

try:
    with open('test_overhead_nvfuser_scaling_after649.pkl', 'rb') as f:
        results_after649 = pickle.load(f)
except FileNotFoundError:
    results_after649 = [None]

results = [*results_before649, *results_after649]
benchmark.Compare(results).print()

[IvanYashchuk]

Here are the results with #768

[------------  -----------]
           |  With #768 PR
1 threads: ----------------
      2    |       25.3
      4    |       34.6
      8    |       52.6
      16   |       86.6
      32   |      179.8
      64   |      393.0
      128  |      743.2

Times are in microseconds (us).

[zasdfgbnm]

@IvanYashchuk thank you for sharing. I will pull numbers from this script to see how the cpu overhead optimization goes.

[zasdfgbnm]

With #813, the cpu overhead is almost back. Closing this as resolved. Feel free to reopen or open a new issue if we need to optimize cpu overhead further.