Repository: pytorch
Issue:  [Inductor] [CPU] optimize thread parallel and loop collapse
Labels: triaged,oncall: pt2,oncall: cpu inductor

### Description
Current thread parallel and loop collapse implementation in `TorchInductor` are not optimal, e.x.: not consider the loop body workload but just loop times at each level. The config parameter `min_chunk_size` determines the minimum loop times one thread should at least take, regardless of the loop body workload which varies a lot.


### Real Case
The model `speech_transformer` suffers from the issue. We verified this by simply tuning `min_chunk_size` to see if the performance improves. 

<html xmlns:v="urn:schemas-microsoft-com:vml"
xmlns:o="urn:schemas-microsoft-com:office:office"
xmlns:x="urn:schemas-microsoft-com:office:excel"
xmlns="http://www.w3.org/TR/REC-html40">

<head>

<meta name=ProgId content=Excel.Sheet>
<meta name=Generator content="Microsoft Excel 15">
<link id=Main-File rel=Main-File
href="file:///C:/Users/xuanliao/AppData/Local/Temp/msohtmlclip1/01/clip.htm">
<link rel=File-List
href="file:///C:/Users/xuanliao/AppData/Local/Temp/msohtmlclip1/01/clip_filelist.xml">
<!--table
	{mso-displayed-decimal-separator:"\.";
	mso-displayed-thousand-separator:"\,";}
@page
	{margin:.75in .7in .75in .7in;
	mso-header-margin:.3in;
	mso-footer-margin:.3in;}
tr
	{mso-height-source:auto;}
col
	{mso-width-source:auto;}
br
	{mso-data-placement:same-cell;}
td
	{padding-top:1px;
	padding-right:1px;
	padding-left:1px;
	mso-ignore:padding;
	color:black;
	font-size:11.0pt;
	font-weight:400;
	font-style:normal;
	text-decoration:none;
	font-family:Calibri, sans-serif;
	mso-font-charset:0;
	mso-number-format:General;
	text-align:general;
	vertical-align:bottom;
	border:none;
	mso-background-source:auto;
	mso-pattern:auto;
	mso-protection:locked visible;
	white-space:nowrap;
	mso-rotate:0;}
.xl63
	{border:.5pt solid windowtext;}
.xl64
	{color:gray;
	border:.5pt solid windowtext;}
.xl65
	{mso-number-format:Percent;
	border:.5pt solid windowtext;}
.xl66
	{font-weight:700;
	mso-number-format:Percent;
	border:.5pt solid windowtext;}
-->
</head>

<body link="#0563C1" vlink="#954F72">



min_chunk_size | inductor time | RSD | #collapse | #parallel | increased performance
-- | -- | -- | -- | -- | --
4096 | 0.077303238 | 0.57% | 30 | 55 | /
0 | 0.063956184 | 0.48% | 31 | 98 | **17.27%**



</body>

</html>

The default `min_chunk_size` is 4096. When it is set to 0, we do not limit the loop times one thread takes. The above result shows that the performance could improve **17.27%** if we optimize thread parallel and loop collapse.


We also want to know what has happened in kernels when `min_chunk_size` changes. `Graph8, kernel0` is taken as an example.

**min_chunk_size=4096**
```
kernel_cpp_0 = async_compile.cpp('''
#include <ATen/record_function.h>
#include "/tmp/torchinductor_root/rp/crpdeql3xwpfmcyakwtqpzihz525if6mt25mozau77xvmnh7vqyu.h"
extern "C" void kernel(const long* __restrict__ in_ptr0,
                       const float* __restrict__ in_ptr1,
                       const float* __restrict__ in_ptr2,
                       float* __restrict__ out_ptr0,
                       float* __restrict__ out_ptr1,
                       float* __restrict__ out_ptr2)
{
    RECORD_FUNCTION("kernel_cpp_0", c10::ArrayRef<c10::IValue>({}));
    #pragma GCC ivdep
    for(long i0=0; i0<10; i0+=1)
    {
        #pragma GCC ivdep
        for(long i1=0; i1<22; i1+=1)
        {
            #pragma GCC ivdep
            for(long i2=0; i2<512; i2+=1)
            {
                {
                    {
                        auto tmp0 = in_ptr0[i1 + (22*i0)];
                        auto tmp4 = in_ptr2[i2 + (512*i1)];
                        auto tmp1 = in_ptr1[i2 + (512*tmp0)];
                        auto tmp2 = static_cast<float>(22.627416997969522);
                        auto tmp3 = tmp1 * tmp2;
                        auto tmp5 = tmp3 + tmp4;
                        out_ptr0[i2 + (512*i1) + (11264*i0)] = tmp5;
                        out_ptr1[i2 + (512*i1) + (11264*i0)] = tmp5;
                        out_ptr2[i2 + (512*i1) + (11264*i0)] = tmp5;
                    }
                }
            }
        }
    }
}
''')
```

**min_chunk_size=0**
```
kernel_cpp_0 = async_compile.cpp('''
#include <ATen/record_function.h>
#include "/tmp/torchinductor_root/rp/crpdeql3xwpfmcyakwtqpzihz525if6mt25mozau77xvmnh7vqyu.h"
extern "C" void kernel(const long* __restrict__ in_ptr0,
                       const float* __restrict__ in_ptr1,
                       const float* __restrict__ in_ptr2,
                       float* __restrict__ out_ptr0,
                       float* __restrict__ out_ptr1,
                       float* __restrict__ out_ptr2)
{
    RECORD_FUNCTION("kernel_cpp_0", c10::ArrayRef<c10::IValue>({}));
    #pragma omp parallel num_threads(28)
    {
        #pragma omp for  collapse(2)
        for(long i0=0; i0<10; i0+=1)
        {
            for(long i1=0; i1<22; i1+=1)
            {
                #pragma GCC ivdep
                for(long i2=0; i2<512; i2+=1)
                {
                    {
                        {
                            auto tmp0 = in_ptr0[i1 + (22*i0)];
                            auto tmp4 = in_ptr2[i2 + (512*i1)];
                            auto tmp1 = in_ptr1[i2 + (512*tmp0)];
                            auto tmp2 = static_cast<float>(22.627416997969522);
                            auto tmp3 = tmp1 * tmp2;
                            auto tmp5 = tmp3 + tmp4;
                            out_ptr0[i2 + (512*i1) + (11264*i0)] = tmp5;
                            out_ptr1[i2 + (512*i1) + (11264*i0)] = tmp5;
                            out_ptr2[i2 + (512*i1) + (11264*i0)] = tmp5;
                        }
                    }
                }
            }
        }
    }
}
''')
```


cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @soumith @ngimel