pytorch/feature: Record if parallel compile is enabled (#141074)

c00w · pytorchmergebot · commit 0989871ac962 · 2024-12-02T19:09:11.000Z
This gets a bit messy, but this appears to be the best spot to make a true / false decision. Note that since we're looking at whether or not it's used, if the pool doesn't warm up within the time it takes for a compile, we will mark the feature use as false. Pull Request resolved: #141074 Approved by: https://github.com/masnesral ghstack dependencies: #141059
diff --git a/torch/_inductor/async_compile.py b/torch/_inductor/async_compile.py
@@ -14,7 +14,7 @@
 
 import torch
 from torch._dynamo.device_interface import get_registered_device_interfaces
-from torch._dynamo.utils import dynamo_timed
+from torch._dynamo.utils import dynamo_timed, set_feature_use
 from torch._inductor import config
 from torch._inductor.codecache import (
     CodeCacheFuture,
@@ -219,6 +219,9 @@ def triton(self, kernel_name: str, source_code: str, device_str: str = "cuda"):
 
         kernel = TritonCodeCache.load(kernel_name, source_code)
         if self._use_process_pool():
+            set_feature_use(
+                "pytorch/inductor:enable_parallel_compile_version (post_warmup)", True
+            )
             # We want to support changing these env vars after (and while) the
             # process pool is running, so pass them to the subprocess to reset.
             env_vars = ["TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR"]
@@ -232,6 +235,9 @@ def triton(self, kernel_name: str, source_code: str, device_str: str = "cuda"):
                 ),
             )
         else:
+            set_feature_use(
+                "pytorch/inductor:enable_parallel_compile_version (post_warmup)", False
+            )
             kernel.precompile()
             return kernel
 
