Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Inductor] support masked vectorization for the tail_loop #126526

Open
wants to merge 39 commits into
base: gh/jiayisunx/10/base
Choose a base branch
from
Open

[Inductor] support masked vectorization for the tail_loop #126526

wants to merge 39 commits into from
+281 −71

Conversation

jiayisunx
Copy link
Collaborator

@jiayisunx jiayisunx commented May 17, 2024 •
edited by CaoE
Loading

Stack from ghstack (oldest at bottom):

Currently the tail_loop always uses the scalar kernel. This PR supports masked vectorization for the tail_loop to improve the performance.

Example:

import torch
import torch.nn as nn

class GN(nn.Module):
    def __init__(self, num_groups, num_channels):
        super(GN, self).__init__()
        self.gn = nn.GroupNorm(num_groups, num_channels)

    def forward(self, x):
        return self.gn(x)

input = torch.randn(2, 960, 96, 96).to(memory_format=torch.channels_last)
m = GN(32, 960).eval()
compiled_m = torch.compile(m)

with torch.no_grad():
    for _ in range(3):
        compiled_m(input)

Generated code:

  • Before:
cpp_fused_native_group_norm_0 = async_compile.cpp_pybinding(['const float*', 'const float*', 'const float*', 'float*', 'float*', 'float*'], '''
#include "/tmp/torchinductor_jiayisun/ky/cky2bufythacofebk7ujv36e4pxyqcqbpsy5r4vojoprjiwcwfxf.h"
extern "C"  void kernel(const float* in_ptr0,
                       const float* in_ptr1,
                       const float* in_ptr2,
                       float* out_ptr0,
                       float* out_ptr1,
                       float* out_ptr2)
{
    #pragma omp parallel num_threads(112)
    {
        int tid = omp_get_thread_num();
        {
            #pragma omp for collapse(2)
            for(long x0=static_cast<long>(0L); x0<static_cast<long>(2L); x0+=static_cast<long>(1L))
            {
                for(long x1=static_cast<long>(0L); x1<static_cast<long>(32L); x1+=static_cast<long>(1L))
                {
                    {
                        Welford<float> tmp_acc0 = Welford<float>();
                        Welford<at::vec::Vectorized<float>> tmp_acc0_vec = Welford<at::vec::Vectorized<float>>();
                        static WeightRecp<at::vec::Vectorized<float>> weight_recps(static_cast<long>(17280L));
                        for(long x2=static_cast<long>(0L); x2<static_cast<long>(9216L); x2+=static_cast<long>(1L))
                        {
                            for(long x3=static_cast<long>(0L); x3<static_cast<long>(16L); x3+=static_cast<long>(16L))
                            {
                                auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(x3 + (30L*x1) + (960L*x2) + (8847360L*x0)), 16);
                                tmp_acc0_vec = welford_combine(tmp_acc0_vec, tmp0, &weight_recps);
                            }
                            #pragma omp simd simdlen(8) 
                            for(long x3=static_cast<long>(16L); x3<static_cast<long>(30L); x3+=static_cast<long>(1L))
                            {
                                auto tmp0 = in_ptr0[static_cast<long>(x3 + (30L*x1) + (960L*x2) + (8847360L*x0))];
                                tmp_acc0 = welford_combine(tmp_acc0, tmp0);
                            }
                        }
                        tmp_acc0 = welford_combine(tmp_acc0, welford_vec_reduce_all(tmp_acc0_vec));
                        out_ptr0[static_cast<long>(x1 + (32L*x0))] = static_cast<float>(tmp_acc0.mean);
                        out_ptr1[static_cast<long>(x1 + (32L*x0))] = static_cast<float>(tmp_acc0.m2);
                    }
                }
            }
        }
        {
            #pragma omp for collapse(2)
            for(long x0=static_cast<long>(0L); x0<static_cast<long>(2L); x0+=static_cast<long>(1L))
            {
                for(long x1=static_cast<long>(0L); x1<static_cast<long>(9216L); x1+=static_cast<long>(1L))
                {
                    for(long x2=static_cast<long>(0L); x2<static_cast<long>(960L); x2+=static_cast<long>(16L))
                    {
                        auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(x2 + (960L*x1) + (8847360L*x0)), 16);
                        auto tmp1 =
                        [&]
                        {
                            __at_align__ std::array<float, 16> tmpbuf;
                            #pragma GCC unroll 16
                            for (long x2_inner = 0; x2_inner < 16; x2_inner++)
                            {
                                tmpbuf[x2_inner] = out_ptr0[static_cast<long>((32L*x0) + (c10::div_floor_integer((x2 + x2_inner), 30L)))];
                            }
                            return at::vec::Vectorized<float>::loadu(tmpbuf.data(), 16);
                        }
                        ()
                        ;
                        auto tmp3 =
                        [&]
                        {
                            __at_align__ std::array<float, 16> tmpbuf;
                            #pragma GCC unroll 16
                            for (long x2_inner = 0; x2_inner < 16; x2_inner++)
                            {
                                tmpbuf[x2_inner] = out_ptr1[static_cast<long>((32L*x0) + (c10::div_floor_integer((x2 + x2_inner), 30L)))];
                            }
                            return at::vec::Vectorized<float>::loadu(tmpbuf.data(), 16);
                        }
                        ()
                        ;
                        auto tmp12 = at::vec::Vectorized<float>::loadu(in_ptr1 + static_cast<long>(x2), 16);
                        auto tmp14 = at::vec::Vectorized<float>::loadu(in_ptr2 + static_cast<long>(x2), 16);
                        auto tmp2 = tmp0 - tmp1;
                        auto tmp4 = static_cast<float>(276480.0);
                        auto tmp5 = at::vec::Vectorized<float>(tmp4);
                        auto tmp6 = tmp3 / tmp5;
                        auto tmp7 = static_cast<float>(1e-05);
                        auto tmp8 = at::vec::Vectorized<float>(tmp7);
                        auto tmp9 = tmp6 + tmp8;
                        auto tmp10 = tmp9.rsqrt();
                        auto tmp11 = tmp2 * tmp10;
                        auto tmp13 = tmp11 * tmp12;
                        auto tmp15 = tmp13 + tmp14;
                        tmp15.store(out_ptr2 + static_cast<long>(x2 + (960L*x1) + (8847360L*x0)));
                    }
                }
            }
        }
    }
}
''')


async_compile.wait(globals())
del async_compile

def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    assert_size_stride(arg0_1, (960, ), (1, ))
    assert_size_stride(arg1_1, (960, ), (1, ))
    assert_size_stride(arg2_1, (2, 960, 96, 96), (8847360, 1, 92160, 960))
    buf0 = empty_strided_cpu((2, 32, 1, 1), (32, 1, 64, 64), torch.float32)
    buf1 = empty_strided_cpu((2, 32, 1, 1), (32, 1, 64, 64), torch.float32)
    buf3 = empty_strided_cpu((2, 960, 96, 96), (8847360, 1, 92160, 960), torch.float32)
    cpp_fused_native_group_norm_0(arg2_1, arg0_1, arg1_1, buf0, buf1, buf3)
    del arg0_1
    del arg1_1
    del arg2_1
    return (buf3, )
  • After:
cpp_fused_native_group_norm_0 = async_compile.cpp_pybinding(['const float*', 'const float*', 'const float*', 'float*', 'float*', 'float*'], '''
#include "/tmp/torchinductor_jiayisun/em/cemtujj65j5txpqlxc7w4pcunpmvz3qtiudkc5ocxxhcmdlknw2m.h"
extern "C"  void kernel(const float* in_ptr0,
                       const float* in_ptr1,
                       const float* in_ptr2,
                       float* out_ptr0,
                       float* out_ptr1,
                       float* out_ptr2)
{
    #pragma omp parallel num_threads(112)
    {
        int tid = omp_get_thread_num();
        {
            #pragma omp for collapse(2)
            for(long x0=static_cast<long>(0L); x0<static_cast<long>(2L); x0+=static_cast<long>(1L))
            {
                for(long x1=static_cast<long>(0L); x1<static_cast<long>(32L); x1+=static_cast<long>(1L))
                {
                    {
                        Welford<float> tmp_acc0 = Welford<float>();
                        Welford<at::vec::Vectorized<float>> tmp_acc0_vec = Welford<at::vec::Vectorized<float>>();
                        Welford<at::vec::Vectorized<float>> masked_tmp_acc0_vec = Welford<at::vec::Vectorized<float>>();
                        static WeightRecp<at::vec::Vectorized<float>> wrecps0(static_cast<long>(17280L));
                        for(long x2=static_cast<long>(0L); x2<static_cast<long>(9216L); x2+=static_cast<long>(1L))
                        {
                            for(long x3=static_cast<long>(0L); x3<static_cast<long>(16L); x3+=static_cast<long>(16L))
                            {
                                auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(x3 + (30L*x1) + (960L*x2) + (8847360L*x0)), 16);
                                tmp_acc0_vec = welford_combine(tmp_acc0_vec, tmp0, &wrecps0);
                            }
                            for(long x3=static_cast<long>(16L); x3<static_cast<long>(30L); x3+=static_cast<long>(14L))
                            {
                                auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(x3 + (30L*x1) + (960L*x2) + (8847360L*x0)), 14);
                                masked_tmp_acc0_vec = welford_combine(masked_tmp_acc0_vec, tmp0, 14, &wrecps0);
                            }
                        }
                        tmp_acc0 = welford_combine(tmp_acc0, welford_vec_reduce_all(masked_tmp_acc0_vec));
                        tmp_acc0 = welford_combine(tmp_acc0, welford_vec_reduce_all(tmp_acc0_vec));
                        out_ptr0[static_cast<long>(x1 + (32L*x0))] = static_cast<float>(tmp_acc0.mean);
                        out_ptr1[static_cast<long>(x1 + (32L*x0))] = static_cast<float>(tmp_acc0.m2);
                    }
                }
            }
        }
        {
            #pragma omp for collapse(2)
            for(long x0=static_cast<long>(0L); x0<static_cast<long>(2L); x0+=static_cast<long>(1L))
            {
                for(long x1=static_cast<long>(0L); x1<static_cast<long>(9216L); x1+=static_cast<long>(1L))
                {
                    for(long x2=static_cast<long>(0L); x2<static_cast<long>(960L); x2+=static_cast<long>(16L))
                    {
                        auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(x2 + (960L*x1) + (8847360L*x0)), 16);
                        auto tmp1 =
                        [&]
                        {
                            __at_align__ std::array<float, 16> tmpbuf;
                            #pragma GCC unroll 16
                            for (long x2_inner = 0; x2_inner < 16; x2_inner++)
                            {
                                tmpbuf[x2_inner] = out_ptr0[static_cast<long>((32L*x0) + (c10::div_floor_integer((x2 + x2_inner), 30L)))];
                            }
                            return at::vec::Vectorized<float>::loadu(tmpbuf.data(), 16);
                        }
                        ()
                        ;
                        auto tmp3 =
                        [&]
                        {
                            __at_align__ std::array<float, 16> tmpbuf;
                            #pragma GCC unroll 16
                            for (long x2_inner = 0; x2_inner < 16; x2_inner++)
                            {
                                tmpbuf[x2_inner] = out_ptr1[static_cast<long>((32L*x0) + (c10::div_floor_integer((x2 + x2_inner), 30L)))];
                            }
                            return at::vec::Vectorized<float>::loadu(tmpbuf.data(), 16);
                        }
                        ()
                        ;
                        auto tmp12 = at::vec::Vectorized<float>::loadu(in_ptr1 + static_cast<long>(x2), 16);
                        auto tmp14 = at::vec::Vectorized<float>::loadu(in_ptr2 + static_cast<long>(x2), 16);
                        auto tmp2 = tmp0 - tmp1;
                        auto tmp4 = static_cast<float>(276480.0);
                        auto tmp5 = at::vec::Vectorized<float>(tmp4);
                        auto tmp6 = tmp3 / tmp5;
                        auto tmp7 = static_cast<float>(1e-05);
                        auto tmp8 = at::vec::Vectorized<float>(tmp7);
                        auto tmp9 = tmp6 + tmp8;
                        auto tmp10 = tmp9.rsqrt();
                        auto tmp11 = tmp2 * tmp10;
                        auto tmp13 = tmp11 * tmp12;
                        auto tmp15 = tmp13 + tmp14;
                        tmp15.store(out_ptr2 + static_cast<long>(x2 + (960L*x1) + (8847360L*x0)));
                    }
                }
            }
        }
    }
}
''')


async_compile.wait(globals())
del async_compile

def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    assert_size_stride(arg0_1, (960, ), (1, ))
    assert_size_stride(arg1_1, (960, ), (1, ))
    assert_size_stride(arg2_1, (2, 960, 96, 96), (8847360, 1, 92160, 960))
    buf0 = empty_strided_cpu((2, 32, 1, 1), (32, 1, 64, 64), torch.float32)
    buf1 = empty_strided_cpu((2, 32, 1, 1), (32, 1, 64, 64), torch.float32)
    buf3 = empty_strided_cpu((2, 960, 96, 96), (8847360, 1, 92160, 960), torch.float32)
    cpp_fused_native_group_norm_0(arg2_1, arg0_1, arg1_1, buf0, buf1, buf3)
    del arg0_1
    del arg1_1
    del arg2_1
    return (buf3, )

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

@jiayisunx
Update
473ce1e 
[ghstack-poisoned]
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented May 17, 2024 •
edited
Loading

🔗 Helpful Links

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

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

❌ 1 New Failure, 3 Unrelated Failures

As of commit f01286e with merge base ae708e9 (image):

NEW FAILURE - The following job has failed:

FLAKY - The following job failed but was likely due to flakiness present on trunk:

BROKEN TRUNK - The following job failed but was present on the merge base:

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

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

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

jiayisunx added a commit that referenced this pull request May 17, 2024
@jiayisunx
[Inductor] support masked vectorization for the tail_loop
921a8ff 
ghstack-source-id: 4e2aa6dfafd14ce90a1dc5b91cb4cbd59ad35628
Pull Request resolved: #126526
@jiayisunx jiayisunx marked this pull request as draft May 17, 2024 09:17
@jiayisunx
Update
278a763 
[ghstack-poisoned]
@jiayisunx
Update
6f1f23c 
[ghstack-poisoned]
@jiayisunx
Update
c8db188 
[ghstack-poisoned]
@jiayisunx
Update
b21a003 
[ghstack-poisoned]
jiayisunx added a commit that referenced this pull request May 20, 2024
@jiayisunx
[Inductor] support masked vectorization for the tail_loop
0627476 
ghstack-source-id: 6d8d821a710fe38916d82149fd6a4947d33c5447
Pull Request resolved: #126526
@jiayisunx
Update
e8a28af 
[ghstack-poisoned]
@jiayisunx
Update
a0061b2 
[ghstack-poisoned]
@jiayisunx
Update
b15206c 
[ghstack-poisoned]
jiayisunx added a commit that referenced this pull request May 21, 2024
@jiayisunx
[Inductor] support masked vectorization for the tail_loop
0949028 
ghstack-source-id: d468379dceab2966884ac812b9a64817151e9002
Pull Request resolved: #126526
@jiayisunx
Update
e51acda 
[ghstack-poisoned]
@jiayisunx
Update
e0184eb 
[ghstack-poisoned]
@jiayisunx
Update
cff55d9 
[ghstack-poisoned]
@jiayisunx
Update
1354b35 
[ghstack-poisoned]
@jiayisunx
Update
8aff7cf 
[ghstack-poisoned]
@jiayisunx
Update
716e01a 
[ghstack-poisoned]
@jiayisunx
Copy link
Collaborator Author

Hi @jgong5 , E made some changes to the welford reduction part to fix performance regression, could you please help review it again? thanks!

@jiayisunx jiayisunx requested a review from jgong5 July 11, 2024 09:59
@jgong5
Copy link
Collaborator

jgong5 commented Jul 12, 2024

Hi @jgong5 , E made some changes to the welford reduction part to fix performance regression, could you please help review it again? thanks!

@jiayisunx @CaoE Thanks. Do you mind share the generated code for this?

@jiayisunx
Copy link
Collaborator Author

jiayisunx commented Jul 14, 2024 •
edited
Loading

Hi @jgong5 , E made some changes to the welford reduction part to fix performance regression, could you please help review it again? thanks!

@jiayisunx @CaoE Thanks. Do you mind share the generated code for this?

Sure, I updated the generated code in the top comment.

jgong5
jgong5 requested changes Jul 15, 2024
View reviewed changes
torch/_inductor/codegen/cpp.py Show resolved Hide resolved
torch/_inductor/codegen/cpp_prefix.h Outdated Show resolved Hide resolved
torch/_inductor/codegen/cpp_prefix.h Show resolved Hide resolved
torch/_inductor/codegen/cpp.py Show resolved Hide resolved
@jiayisunx
Update
1cfb750 
[ghstack-poisoned]
@jiayisunx jiayisunx mentioned this pull request Jul 15, 2024
Open
@jiayisunx
Update
552f102 
[ghstack-poisoned]
@jiayisunx jiayisunx requested a review from jgong5 July 16, 2024 07:28
CaoE pushed a commit to CaoE/pytorch that referenced this pull request Jul 18, 2024
@jiayisunx
[Inductor] support masked vectorization for the tail_loop
d273f3c 
ghstack-source-id: de3061b3de1ad8e7032111cff88ab43af9aba0af
Pull Request resolved: pytorch#126526
@jiayisunx
Update
0e31eb4 
[ghstack-poisoned]
@jiayisunx jiayisunx requested a review from jgong5 July 18, 2024 09:24
jgong5
jgong5 requested changes Jul 19, 2024
View reviewed changes
torch/_inductor/codegen/cpp.py Outdated Show resolved Hide resolved
jgong5
jgong5 reviewed Jul 19, 2024
View reviewed changes
torch/_inductor/codegen/cpp.py
@@ -2146,6 +2142,8 @@ def __init__(
tiling_factor = self.vec_isa.nelements(dtype=tiling_dtype)
self.tiling_factor = tiling_factor
self.tiling_idx = tiling_idx
self.tail_size = tail_size
self.num_elems = tail_size if tail_size else tiling_factor
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How do we handle the case when self.tail_size is symbolic? This seems problematic when we use self.tail_size as constant in the code like in GCC unroll?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked has_free_symbols in CppVecKernelChecker, for dynamic shapes, the tail loop still uses scalar kernel.

@jiayisunx
Update
7255a24 
[ghstack-poisoned]
@jiayisunx jiayisunx requested a review from jgong5 July 19, 2024 06:48
jgong5
jgong5 approved these changes Jul 19, 2024
View reviewed changes
torch/_inductor/codegen/cpp.py
@@ -2828,6 +2942,9 @@ def store(self, name, index, value, mode=None):
return self.simd_vec

def reduction(self, dtype, src_dtype, reduction_type, value):
if has_free_symbols(self.ranges):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK. Are you going to enable this for dynamic shapes in future PRs? I don't think it is very hard to do so?

@jiayisunx jiayisunx mentioned this pull request Jul 19, 2024
Open
@leslie-fang-intel leslie-fang-intel mentioned this pull request Jul 22, 2024
Open
@jiayisunx jiayisunx mentioned this pull request Jul 25, 2024
Draft
@jiayisunx
Update
eceeff5 
[ghstack-poisoned]
@jiayisunx jiayisunx requested a review from jansel July 31, 2024 05:00
@jiayisunx
Copy link
Collaborator Author

@jansel , could you please review this PR? Thanks!

@CaoE
Update
f01286e 
[ghstack-poisoned]
This was referenced Jul 31, 2024
Draft
Draft
Draft
Draft
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jgong5 jgong5 jgong5 approved these changes

@leslie-fang-intel leslie-fang-intel leslie-fang-intel approved these changes

@CaoE CaoE Awaiting requested review from CaoE

@jansel jansel Awaiting requested review from jansel

Assignees
No one assigned
Labels
ciflow/inductor ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR ciflow/trunk Trigger trunk jobs on your pull request module: inductor open source release notes: fx release notes category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
5 participants
@jiayisunx @jgong5 @CaoE @leslie-fang-intel @pytorchbot