[Inductor] optimize the heuristics of parallel reduction (pytorch#149614)

Fix pytorch#148639.

Summary:
Optimize the heuristics of parallel reduction: When the number of steps of the first inner loop beyond the maximum parallel depth is much larger than the number of steps of all outer loops within the maximum parallel depth, change the starting depth of parallelism to the first inner loop and recalculate the maximum parallel depth. I ran the Inductor benchmark with this PR on CPU. A timm model poolformer_m36 BF16 has about 25% performance improvement, and no performance regression is seen.

Pull Request resolved: pytorch#149614
Approved by: https://github.com/jgong5, https://github.com/leslie-fang-intel, https://github.com/jansel

Author: jiayisunx

torch/_inductor/codegen/cpp.py

@@ -5457,20 +5457,21 @@ def max_parallel_depth(self):
         start_depth = 0
         max_depth = 0
         is_reduction = self.loops[0].is_reduction
-        loop_sizes = sympy.Integer(1)
+        num_steps = sympy.Integer(1)
         for loop in self.loops:
             if loop.is_reduction != is_reduction:
                 break
-            loop_sizes = loop_sizes * loop.size
+            num_steps = num_steps * FloorDiv(loop.size, loop.steps)
             max_depth += 1
 
-        # When the range of the first inner loop is much larger than the range of all outer loops,
-        # change `start_depth` to the first inner loop and recalculate `max_depth`.
+        # When the number of steps of the first inner loop is much larger than the number of steps of
+        # all outer loops, change `start_depth` to the first inner loop and recalculate `max_depth`.
         if (
             max_depth < len(self.loops)
-            and isinstance(loop_sizes, sympy.Integer)
+            and isinstance(num_steps, sympy.Integer)
             and isinstance(self.loops[max_depth].size, sympy.Integer)
-            and loop_sizes * 300 < self.loops[max_depth].size
+            and num_steps * 300
+            < FloorDiv(self.loops[max_depth].size, self.loops[max_depth].steps)
         ):
             start_depth = max_depth
             max_depth = 0