Fuse nodes with sizes (s0*s1*...,) and (s0, s1, s2, ...)

[vfdev-5]

Description:

• PR tries to fuse nodes with compatible sizes, for example node1: (s0, s1, s2) and node2: (s0 * s1 * s2). On main these two nodes can be fused due to different sizes. With this PR we can recompute node2 size, body etc using node1 indexing constraint and thus be able to fuse two nodes.
• this should influence only cpu device

Example:

from unittest.mock import patch
import torch
from torch._inductor.graph import GraphLowering
from torch._inductor import config


# Force multple scheduler nodes creation to fuse them
config.realize_opcount_threshold = 1


@torch.compile(fullgraph=True, dynamic=True)
def fn(x: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor) -> torch.Tensor:
    o1 = x * w1.view(1, 1, 1, -1)
    o2 = x * w2.view(1, 1, 1, -1)
    output = o1 + o2
    return output


in_nodes = []
outputs = []
run_node = GraphLowering.run_node

graph_lowering_obj = None

def run_node_alt(self, n):
    global graph_lowering_obj

    graph_lowering_obj = self
    in_nodes.append(n)
    output = run_node(self, n)
    outputs.append(output)

    return output


x = torch.rand(1, 3, 32, 32)
w1 = torch.randn(32)
w2 = torch.randn(32)

with patch.object(GraphLowering, "run_node", run_node_alt):
    fn(x, w1, w2)

print("graph_lowering_obj.buffers:", graph_lowering_obj.buffers)
print("graph_lowering_obj.scheduler:", graph_lowering_obj.scheduler.nodes)

Output on main:

graph_lowering_obj.buffers: [ComputedBuffer(name='buf0', layout=FixedLayout('cpu', torch.float32, size=[1, s1, s0, s0], stride=[s0**2*s1, s0**2, s0, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  def inner_fn(index):
      _, i1, i2, i3 = index
      tmp0 = ops.load(arg3_1, i3 + i1 * s0**2 + i2 * s0)
      tmp1 = ops.load(arg1_1, i3)
      tmp2 = tmp0 * tmp1
      return tmp2
  ,
  ranges=[1, s1, s0, s0],
  origin_node=mul,
  origins={mul}
)), ComputedBuffer(name='buf1', layout=FixedLayout('cpu', torch.float32, size=[1, s1, s0, s0], stride=[s0**2*s1, s0**2, s0, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  def inner_fn(index):
      _, i1, i2, i3 = index
      tmp0 = ops.load(arg3_1, i3 + i1 * s0**2 + i2 * s0)
      tmp1 = ops.load(arg4_1, i3)
      tmp2 = tmp0 * tmp1
      return tmp2
  ,
  ranges=[1, s1, s0, s0],
  origin_node=mul_1,
  origins={mul_1}
)), ComputedBuffer(name='buf2', layout=FixedLayout('cpu', torch.float32, size=[1, s1, s0, s0], stride=[s0**2*s1, s0**2, s0, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  def inner_fn(index):
      _, i1, i2, i3 = index
      tmp0 = ops.load(buf0, i3 + i1 * s0**2 + i2 * s0)
      tmp1 = ops.load(buf1, i3 + i1 * s0**2 + i2 * s0)
      tmp2 = tmp0 + tmp1
      return tmp2
  ,
  ranges=[1, s1, s0, s0],
  origin_node=add,
  origins={add}
))]
graph_lowering_obj.scheduler: [FusedSchedulerNode(nodes=buf0_buf1), SchedulerNode(name='buf2')]

Output on this PR:

graph_lowering_obj.buffers: [ComputedBuffer(name='buf0', layout=FixedLayout('cpu', torch.float32, size=[1, s1, s0, s0], stride=[s0**2*s1, s0**2, s0, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  def inner_fn(index):
      _, i1, i2, i3 = index
      tmp0 = ops.load(arg3_1, i3 + i1 * s0**2 + i2 * s0)
      tmp1 = ops.load(arg1_1, i3)
      tmp2 = tmp0 * tmp1
      return tmp2
  ,
  ranges=[1, s1, s0, s0],
  origin_node=mul,
  origins={mul}
)), ComputedBuffer(name='buf1', layout=FixedLayout('cpu', torch.float32, size=[1, s1, s0, s0], stride=[s0**2*s1, s0**2, s0, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  def inner_fn(index):
      _, i1, i2, i3 = index
      tmp0 = ops.load(arg3_1, i3 + i1 * s0**2 + i2 * s0)
      tmp1 = ops.load(arg4_1, i3)
      tmp2 = tmp0 * tmp1
      return tmp2
  ,
  ranges=[1, s1, s0, s0],
  origin_node=mul_1,
  origins={mul_1}
)), ComputedBuffer(name='buf2', layout=FixedLayout('cpu', torch.float32, size=[1, s1, s0, s0], stride=[s0**2*s1, s0**2, s0, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  def inner_fn(index):
      _, i1, i2, i3 = index
      tmp0 = ops.load(buf0, i3 + i1 * s0**2 + i2 * s0)
      tmp1 = ops.load(buf1, i3 + i1 * s0**2 + i2 * s0)
      tmp2 = tmp0 + tmp1
      return tmp2
  ,
  ranges=[1, s1, s0, s0],
  origin_node=add,
  origins={add}
))]
graph_lowering_obj.scheduler: [FusedSchedulerNode(nodes=buf0_buf1_buf2)]

Context:
While working on #120411, upsampling bicubic decomposition, I saw an extra for-loop in C++ generated code summing up two buffers. Exploring the cause, it happend due to buffer number of ops goes beyond config.realize_opcount_threshold.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/120077

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (4 Unrelated Failures)

As of commit 6a8d136 with merge base bd19d6d ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_huggingface, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
  DebertaForQuestionAnswering
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_huggingface, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
  DebertaForQuestionAnswering
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (cpu_inductor_torchbench, 2, 2, linux.12xlarge) (gh)
  llava
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (dynamic_cpu_inductor_torchbench, 2, 2, linux.12xlarge) (gh)
  llava

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[peterbell10]

One concern is that retracing the body is fairly expensive, so we'll need to benchmark this to see how it impacts compile time overhead. I don't expect it to be too bad though as scheduler time doesn't usually dominate compilation.

[lezcano]

Type the function to make it more readable. Same below.

[vfdev-5]

well, if you check can_fuse_* methods here, they are not typed neither, I follow the trend here :)

[jgong5]

@vfdev-5 Do you have performance comparison with and without this PR on the three inductor benchmark suites? I guess my major concern is that whether the fusion would always bring performance benefit. It might change the parallelization scheduling and data access patterns of the node with sizes s0*s1*... cc @chuanqi129 for awareness. We may evaluate its performance impact from our side too.

[vfdev-5]

three inductor benchmark suites

@jgong5 can you please point out which benchmark suites you think of? Thanks!

I'll provide a benchmark of the case I worked on (upsampling bicubic) where this fusion leads to some speed-up.

[jgong5]

@jgong5 can you please point out which benchmark suites you think of? Thanks!

@vfdev-5 we are tracking the performance with the three benchmark suites: torchbench, huggingface and timm. @chuanqi129 can share the steps how to run them to check the performance.

[lezcano]

@vfdev-5 run the torchbench suite on the hud, please.

[chuanqi129]

@jgong5 can you please point out which benchmark suites you think of? Thanks!

@vfdev-5 we are tracking the performance with the three benchmark suites: torchbench, huggingface and timm. @chuanqi129 can share the steps how to run them to check the performance.

Hi @vfdev-5 @jgong5 , for Inductor CPU performance dashboard benchmark side, we used this Dockerfileto prepare the test environment with cmd docker build --build-arg http_proxy=${http_proxy} --build-arg https_proxy=${https_proxy} --build-arg PT_REPO=https://github.com/pytorch/pytorch --build-arg PT_COMMIT=xxxxxxx -t pt:test -f Dockerfile --target image ..
After that, we can sue the test script inductor_test.sh to launch 3 suites (torchbench, huggingface and timm) with multi-thread mode and single thread mode on cpu together by using cmd bash inductor_test.sh under pytorch source code root folder. By default, it will test float32 static shape default wrapper inference accuracy and performance, we can pass different parameter to this script to test other combinations.
If we want to test single model, we can directly use inductor_single_run.sh with proper parameters. Please feel free let me know if there any issues.

[lezcano]

If CPU benchmarks are not run on the hud, then it may be better for them to be run by someone from intel so that they are run on the relevant hardware, perhaps with core pinning, disabled throttling, w/without hyperthreading, etc... Also, we don't have access to a box with AVX512, for example. @jgong5

[vfdev-5]

[------------------------------------------------------------------------------------------------ Interpolate, cpu ------------------------------------------------------------------------------------------------]
                                                                                                                                                   |  Eager (2.3.0a0)   |  Compiled (2.3.0a0)
1 threads: ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

Downsampling to small output image 256x256
- No fusion (8ef4a43a31a)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: True, antialias: False, osize: (256, 256)     |        1.745 (+-0.073)        |         3.598 (+-0.111)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: False, antialias: False, osize: (256, 256)    |        1.749 (+-0.065)        |         3.852 (+-0.140)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: True, antialias: False, osize: (256, 256)         |        2.804 (+-0.131)        |         3.784 (+-0.141)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: False, antialias: False, osize: (256, 256)        |        2.817 (+-0.138)        |         4.134 (+-0.132)

- With fusion (PR 120077, merged as ffd0b4de1d3) -> fusion benefits are invisible
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: True, antialias: False, osize: (256, 256)     |        1.756 (+-0.060)        |         3.574 (+-0.108)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: False, antialias: False, osize: (256, 256)    |        1.753 (+-0.056)        |         3.686 (+-0.123)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: True, antialias: False, osize: (256, 256)         |        2.834 (+-0.135)        |         3.782 (+-0.088)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: False, antialias: False, osize: (256, 256)        |        2.862 (+-0.126)        |         4.130 (+-0.105)

- With extended fusion -> fusion benefits are invisible
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: True, antialias: False, osize: (256, 256)     |        1.768 (+-0.061)        |         3.601 (+-0.108)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: False, antialias: False, osize: (256, 256)    |        1.761 (+-0.075)        |         3.690 (+-0.112)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: True, antialias: False, osize: (256, 256)         |        2.876 (+-0.126)        |         3.719 (+-0.123)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: False, antialias: False, osize: (256, 256)        |        2.859 (+-0.123)        |         4.061 (+-0.131)

Upsampling to large output image 1024x1024
- No fusion (8ef4a43a31a)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: True, antialias: False, osize: (1024, 1024)   |        27.097 (+-1.031)       |         56.546 (+-1.339)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: False, antialias: False, osize: (1024, 1024)  |        26.972 (+-1.128)       |         66.710 (+-3.071)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: True, antialias: False, osize: (1024, 1024)       |        43.651 (+-1.563)       |         61.883 (+-2.002)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: False, antialias: False, osize: (1024, 1024)      |        43.777 (+-1.466)       |         80.301 (+-3.199)

- With fusion (PR 120077, merged as ffd0b4de1d3) -> visible only for ac=false and CF
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: True, antialias: False, osize: (1024, 1024)   |        26.906 (+-1.212)       |         56.570 (+-1.846)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: False, antialias: False, osize: (1024, 1024)  |        26.940 (+-1.122)       |         58.319 (+-1.924)    <---
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: True, antialias: False, osize: (1024, 1024)       |        44.541 (+-1.838)       |         61.911 (+-2.280)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: False, antialias: False, osize: (1024, 1024)      |        44.513 (+-1.792)       |         80.829 (+-2.849)

- With extended fusion -> visible only for ac=false and CF / CL
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: True, antialias: False, osize: (1024, 1024)   |        27.171 (+-1.108)       |         56.638 (+-1.361)
      Input (1, 3, 500, 400), torch.float32, torch.contiguous_format | mode: bicubic, align_corners: False, antialias: False, osize: (1024, 1024)  |        27.242 (+-1.067)       |         58.309 (+-1.500)    <---
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: True, antialias: False, osize: (1024, 1024)       |        44.792 (+-1.700)       |         58.871 (+-1.848)
      Input (1, 3, 500, 400), torch.float32, torch.channels_last | mode: bicubic, align_corners: False, antialias: False, osize: (1024, 1024)      |        44.419 (+-1.432)       |         64.399 (+-1.467)    <---


Times are in milliseconds (ms).

Observations:

• For CL, ac=false and ac=true there is no fusion of the last summation buffer into the main loop for "with fusion" option
• For CF, ac=true, there is no last summation buffer for "no fusion" and "with fusion" options -> there is no speed-up. This is due to ac=true has less ops vs ac=false and thus config.realize_opcount_threshold is not reached.

"extended fusion" refers to the following update to this PR to handle channels-last case:

-        c3 = len(vars1) == 1 or len(vars2) == 1
+        c3 = len(vars1) in (1, 2) or len(vars2) in (1, 2)

[jgong5]

Observations:

• For CL, ac=false and ac=true there is no fusion of the last summation buffer into the main loop for "with fusion" option
• For CF, ac=true, there is no last summation buffer for "no fusion" and "with fusion" options -> there is no speed-up. This is due to ac=true has less ops vs ac=false and thus config.realize_opcount_threshold is not reached.

Thanks for sharing the numbers. @chuanqi129 will check the impact on the inference performance on the three benchmark suites.

[chuanqi129]

Observations:

• For CL, ac=false and ac=true there is no fusion of the last summation buffer into the main loop for "with fusion" option
• For CF, ac=true, there is no last summation buffer for "no fusion" and "with fusion" options -> there is no speed-up. This is due to ac=true has less ops vs ac=false and thus config.realize_opcount_threshold is not reached.

Thanks for sharing the numbers. @chuanqi129 will check the impact on the inference performance on the three benchmark suites.

Have verified for FP32 static shape default wrapper inference test with 3 suites, there is no performance drop. cc: @jgong5 @vfdev-5

[jgong5]

Have verified for FP32 static shape default wrapper inference test with 3 suites, there is no performance drop. cc: @jgong5 @vfdev-

Thanks @chuanqi129 !

[vfdev-5]

@chuanqi129 thanks for the feedback, I wonder whether you have seen any perf improvements on some tests ?

[lezcano]

sgtm. @peterbell10 wanna have a last look?

Also, @vfdev-5 will you implement the more generalised version of this algorithm after this PR is merged?

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased cpp-nodes-fusion onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout cpp-nodes-fusion && git pull --rebase)

[vfdev-5]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased cpp-nodes-fusion onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout cpp-nodes-fusion && git pull --rebase)

[lezcano]

this comment is no longer relevant :)

[vfdev-5]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here