pytorch/features: Make a feature logger and record triton bundling (#…

…141056)

Summary:
This modifies metrics_context to allow us to store whether a feature was
used or not.

This also starts recording this for triton bundling.

X-link: pytorch/pytorch#141056
Approved by: https://github.com/masnesral

Reviewed By: izaitsevfb

Differential Revision: D66350649

Pulled By: c00w

fbshipit-source-id: 63f5014df4f8adf7bb2bf82e645b59076b3aa697

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

@@ -870,6 +870,7 @@ class CompilationMetrics:
     inductor_fx_remote_cache_backend_type: Optional[str] = None
     inductor_fx_remote_cache_hit_keys: Optional[str] = None
     inductor_fx_remote_cache_miss_keys: Optional[str] = None
+    feature_usage: Optional[dict[str, bool]] = None
 
 
 DEFAULT_COMPILATION_METRICS_LIMIT = 64
@@ -3587,3 +3588,11 @@ def record(cls):
         finally:
             if config.record_compile_time_instruction_count:
                 cls.end()
+
+
+def set_feature_use(feature: str, usage: bool):
+    """
+    Records whether we are using a feature
+    Generally a feature is a JK.
+    """
+    get_metrics_context().set_key_value("feature_usage", feature, usage)