[NO CP] [release/2.9] Add partitioned buffer approach for scatter add op

test/inductor/test_metrics.py

@@ -76,6 +76,7 @@ def test_parse_reduction_hint(self):
         )
 
     @config.patch("fx_graph_remote_cache", False)
+    @config.patch("partitioned_scatter_enabled", False)
     def test_atomic_add(self):
         @torch.compile
         def f(lhs, index, rhs):

torch/_inductor/config.py

@@ -819,6 +819,9 @@ def decide_worker_start_method() -> str:
 
 _micro_pipeline_tp: bool = False
 
+# Enable/disable partitioned scatter optimization for atomic add kernels
+# this will improve kernel performance at cost of memory usage.
+partitioned_scatter_enabled = os.environ.get("TORCHINDUCTOR_PARTITIONED_SCATTER_ENABLED", "1") == "1"
 
 class _collective:
     auto_select: bool = False

torch/_inductor/fx_passes/post_grad.py

@@ -63,6 +63,7 @@
 from .pre_grad import is_same_dict, save_inductor_dict
 from .reinplace import reinplace_inplaceable_ops
 from .split_cat import POST_GRAD_PATTERNS
+from .reduced_atomic_contention import partitioned_scatter_optimization_pass
 
 
 _T = TypeVar("_T")
@@ -140,6 +141,10 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
             GraphTransformObserver(gm, f"pass_pattern_{i}").apply_graph_pass(
                 patterns.apply
             )
+        if config.partitioned_scatter_enabled:
+            GraphTransformObserver(gm, "partitioned_scatter_optimization").apply_graph_pass(
+                partitioned_scatter_optimization_pass
+            )
         for pass_name in config.post_grad_fusion_options:
             # skip all patterns for group batch fusions or quantization patterns
             if pass_name in POST_GRAD_FUSIONS or pass_name in OPTIMUS_EXCLUDE_POST_GRAD: