[Inductor] support masked vectorization for the tail_loop

[jiayisunx]

Stack from ghstack (oldest at bottom):

• -> [Inductor] support masked vectorization for the tail_loop #126526

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

Generated code:

• Before:

    {
        #pragma GCC ivdep
        for(long x0=static_cast<long>(0L); x0<static_cast<long>(2L); x0+=static_cast<long>(1L))
        {
            #pragma GCC ivdep
            for(long x1=static_cast<long>(0L); x1<static_cast<long>(3L); 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>(67L));
                    for(long x2=static_cast<long>(0L); x2<static_cast<long>(36L); 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<bfloat16>::loadu(in_ptr0 + static_cast<long>(x3 + (30L*x1) + (90L*x2) + (3240L*x0)), 16);
                            auto tmp1 = at::vec::convert<float>(tmp0);
                            tmp_acc0_vec = welford_combine(tmp_acc0_vec, tmp1, &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) + (90L*x2) + (3240L*x0))];
                            auto tmp1 = c10::convert<float>(tmp0);
                            tmp_acc0 = welford_combine(tmp_acc0, tmp1);
                        }
                    }
                    tmp_acc0 = welford_combine(tmp_acc0, welford_vec_reduce_all(tmp_acc0_vec));
                    out_ptr0[static_cast<long>(x1 + (3L*x0))] = static_cast<float>(tmp_acc0.mean);
                    out_ptr1[static_cast<long>(x1 + (3L*x0))] = static_cast<float>(tmp_acc0.m2);
                }
            }
        }
    }

• After:

    {
        #pragma GCC ivdep
        for(long x0=static_cast<long>(0L); x0<static_cast<long>(2L); x0+=static_cast<long>(1L))
        {
            #pragma GCC ivdep
            for(long x1=static_cast<long>(0L); x1<static_cast<long>(3L); 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>(36L), static_cast<long>(16L), static_cast<long>(14L));
                    for(long x2=static_cast<long>(0L); x2<static_cast<long>(36L); 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<bfloat16>::loadu(in_ptr0 + static_cast<long>(x3 + (30L*x1) + (90L*x2) + (3240L*x0)), 16);
                            auto tmp1 = at::vec::convert<float>(tmp0);
                            tmp_acc0_vec = welford_combine(tmp_acc0_vec, tmp1, &weight_recps);
                        }
                        for(long x3=static_cast<long>(16L); x3<static_cast<long>(30L); x3+=static_cast<long>(14L))
                        {
                            auto tmp0 = at::vec::Vectorized<bfloat16>::loadu(in_ptr0 + static_cast<long>(x3 + (30L*x1) + (90L*x2) + (3240L*x0)), 14);
                            auto tmp1 = at::vec::convert<float>(tmp0);
                            tmp_acc0_vec = welford_combine(tmp_acc0_vec, tmp1, 14, &weight_recps);
                        }
                    }
                    tmp_acc0 = welford_combine(tmp_acc0, welford_vec_reduce_all(tmp_acc0_vec));
                    out_ptr0[static_cast<long>(x1 + (3L*x0))] = static_cast<float>(tmp_acc0.mean);
                    out_ptr1[static_cast<long>(x1 + (3L*x0))] = static_cast<float>(tmp_acc0.m2);
                }
            }
        }
    }

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Upgrade MacOS runner to 14

✅ You can merge normally! (6 Unrelated Failures)

As of commit b72e62a with merge base bf2909b ():

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

• trunk / macos-py3-arm64-mps / test (mps, 1, 1, macos-m1-14) (gh) (matched macos rule in flaky-rules.json)
  File doesn't exist

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

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

• periodic / linux-focal-cuda11.8-py3.9-gcc9 / test (multigpu, 1, 1, linux.g5.12xlarge.nvidia.gpu) (gh) (trunk failure)
  distributed/_composable/fsdp/test_fully_shard_training.py::TestFullyShardNDTraining::test_2d_mlp_with_nd_mesh
• periodic / linux-focal-rocm6.1-py3.8 / test (distributed, 1, 2, linux.rocm.gpu) (gh) (trunk failure)
  distributed/_composable/fsdp/test_fully_shard_training.py::TestFullyShard2DTraining::test_tp_with_fsdp_offloading

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

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#126884)
  eca_halonext26ts
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.24xl.spr-metal, unstable) (gh) (#126993)
  Process completed with exit code 1.
• pull / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 5, 5, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) ()
  inductor/test_efficient_conv_bn_eval.py::EfficientConvBNEvalCudaTests::test_basic_cuda

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

[jgong5]

Why do we want to limit this? Can we treat it same as normal vec check? We are simplifying the vec checker and target removing it entirely some day. Not good to further complicate it.

[jiayisunx]

Currently, some operations on integer types do not support masked vectorization well, so these data types cannot yet support masked vectorization for the tail_loop. I will try to solve these issues in the near future.

[jgong5]

Please add a comment here that it will be removed in the near future after we support all data types.

[jiayisunx]

Done.

[jgong5]

BTW, do you mind stack a PR now to support all data types?

[jgong5]

This essentially moves the compile time checks to runtime. I don't think it is the right thing to do.

[jiayisunx]

I have modified this part, please review it again, thanks!

[jgong5]

Can we avoid code dedup between tail version and main version?

[jiayisunx]

Done.

[jgong5]

ditto

[jiayisunx]

Done.

[jgong5]

self.ranges[-1] holds for either of the conditions, right?

[jgong5]

Also, is the assumption that self.tiling_idx == len(sef.ranges) - 1 so that the vectorization happens on the inner-most loop?

[CaoE]

Yes, vectorization happens on the inner-most loop.

[CaoE]

According to what has been observed so far, self.ranges[-1] holds for either of the conditions.

[jgong5]

Then, why can't we just do self.ranges[-1]?

[jgong5]

Yes, vectorization happens on the inner-most loop.

To be safe, can we add an assertion here?

[jiayisunx]

added, please review it again, thanks!

[jgong5]

Suggested change

	self.reduction_outer_size = reduction_size / reduction_inner_size
	self.reduction_outer_size = reduction_size // reduction_inner_size

[jiayisunx]

Done, thanks!

[leslie-fang-intel]

• I think we need the UTs to check the vecmask used for tail loop.

• Will it be clearer to add a new subclass maybe CppVecMaskKernel(CppVecKernel) and lift up common code from CppVecKernel as a method to be overwritten or reuse?

[leslie-fang-intel]

Why we change the type from T::value_type to T?

[CaoE]

We use T , i.e., vec type instead of scalar type, as all element values of weight(vec type) may no longer be identical in masked vec welford reduce since weight may be masked.

[jiayisunx]

• I think we need the UTs to check the vecmask used for tail loop.

Added, thanks!

• Will it be clearer to add a new subclass maybe CppVecMaskKernel(CppVecKernel) and lift up common code from CppVecKernel as a method to be overwritten or reuse?

I don't have a strong opinion, but adding a new subclass might introduce some code duplication, @jgong5, do you have any opinion?

[jgong5]

I don't have a strong opinion, but adding a new subclass might introduce some code duplication, @jgong5, do you have any opinion?

I guess what @leslie-fang-intel meant was to factor out some functions from CppVecKernel to be overridden by CppVecMaskKernel, e.g., how the "load", "store" and "reduction" lines are generated. This can avoid code duplication you mentioned.