triton/dynamo/inductor cache location and timer stats

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 - **Compute**
   - torch: update to `2.13.0` for CUDA, ROCm, IPEX
+  - torch: explicitly set inductor and triton cache locations
+  - torch: log triton/dynamo/inductor timer stats
   - cuda: update to `13.2`
   - sdnq quantization optimizations
   - sdnq attention optimizations

modules/devices.py

@@ -432,7 +432,6 @@ def test_triton(early: bool = False):
         if triton_version is None:
             try:
                 import torch._inductor.triton as torch_triton
-
                 triton_version = torch_triton.__version__
             except Exception:
                 pass
@@ -626,7 +625,23 @@ def set_cuda_params():
     except Exception:
         tunable = [False, False]
     log.info(f'Torch parameters: backend={backend} device={device_name} config={opts.cuda_dtype} dtype={dtype} fp16={"pass" if fp16_ok else "fail"} bf16={"pass" if bf16_ok else "fail"} triton={"pass" if triton_ok else "fail"} optimization="{opts.cross_attention_optimization}"')
-    log.info(f'Torch compute: context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable}')
+    try:
+        num_threads = torch._inductor.config.compile_threads # pylint: disable=protected-access
+    except Exception:
+        num_threads = None
+    log.info(f'Torch compute: context={inference_context.__name__} nohalf={opts.no_half} nohalfvae={opts.no_half_vae} upcast={opts.upcast_sampling} deterministic={opts.cudnn_deterministic} tunable={tunable} threads={num_threads}')
+
+    try:
+        from torch._inductor.runtime.runtime_utils import cache_dir
+        inductor_cache = cache_dir()
+    except Exception:
+        inductor_cache = os.getenv("TORCHINDUCTOR_CACHE_DIR", None)
+    try:
+        from triton import knobs
+        triton_cache = knobs.cache.dir
+    except Exception:
+        triton_cache = os.getenv("TRITON_CACHE_DIR", None)
+    log.info(f'Torch cache: inductor="{inductor_cache}" triton="{triton_cache}"')
 
 
 def randn(seed, shape=None):

modules/loader.py

@@ -56,6 +56,9 @@ except Exception as e:
     report(f'scipy=={scipy.__version__ if scipy is not None else None}', e)
 timer.startup.record("scipy")
 
+inductor_cache = os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "inductor"))
+triton_cache = os.environ.setdefault("TRITON_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "triton"))
+
 try:
     import atexit
     import torch._inductor.async_compile as ac

modules/processing.py

@@ -416,6 +416,27 @@ def process_samples(p: StableDiffusionProcessing, samples):
     return out_images, out_infotexts
 
 
+def print_stats():
+    log.debug(f'Processed: timers={timer.process.dct()}')
+    log.debug(f'Processed: memory={memstats.memory_stats()}')
+
+    if timer.compiler.get_total() > 0.0001:
+        log.debug(f'Processed: compile={timer.compiler.dct(min_time=0)}')
+        timer.compiler.reset()
+
+    if shared.opts.sdnq_dequantize_compile:
+        from modules.sdnq.timers import update_sdnq_attention_timers
+        update_sdnq_attention_timers()
+        if timer.autotune.get_total() > 0.0001:
+            log.debug(f'Processed: autotune={timer.autotune.dct(min_time=0)}')
+
+    if devices.triton_ok:
+        from modules.sd_models_compile import update_compile_times
+        update_compile_times()
+        if timer.dynamo.get_total() > 0.0001:
+            log.debug(f'Processed: dynamo={timer.dynamo.dct(min_time=1.0)}')
+
+
 def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     if type(p.prompt) == list:
         assert len(p.prompt) > 0
@@ -570,10 +591,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         p.scripts.postprocess(p, results)
     timer.process.record('post')
     p.ops = list(set(p.ops))
+
     if not p.disable_extra_networks:
         log.info(f'Processed: images={len(output_images)} its={(p.steps * len(output_images)) / (t1 - t0):.2f} ops={p.ops}')
-        log.debug(f'Processed: timers={timer.process.dct()}')
-        log.debug(f'Processed: memory={memstats.memory_stats()}')
+        print_stats()
 
     if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
         devices.torch_gc(force=True, reason='final')

modules/sd_hijack_te.py

@@ -80,6 +80,8 @@ def hijack_encode_prompt(*args, **kwargs):
         errors.display(e, 'Encode prompt')
     t1 = time.time()
     timer.process.add('te', t1-t0)
+    if t1 - t0 > 10:
+        log.warning(f'Encode: time={t1-t0:.3f} long encode prompt')
     shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     shared.state.end(jobid)
     # from modules import memstats

modules/sd_models_compile.py

@@ -370,3 +370,38 @@ def openvino_post_compile(op="base"): # delete unet after OpenVINO compile
             if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_refiner, "unet"):
                 shared.sd_refiner.unet.apply(sd_models_utils.convert_to_faketensors)
                 devices.torch_gc(force=True)
+
+
+def update_compile_times():
+    try:
+        from torch._dynamo.utils import compile_times, reset_frame_count
+        raw_str = compile_times()
+        reset_frame_count()
+    except Exception:
+        return {}
+    from modules.timer import dynamo
+    dynamo.reset()
+    lines = raw_str.strip().split('\n')
+    # parsed = []
+    for line in lines:
+        if not line or 'TorchDynamo compilation metrics' in line or 'Function, Runtimes' in line:
+            continue
+        parts = line.split(',')
+        fn = parts[0].strip()
+        try:
+            times = [float(t.strip()) for t in parts[1:] if t.strip()]
+            if times:
+                # parsed.append((fn, sum(times), len(times), max(times)))
+                dynamo.add(fn, round(sum(times), 2))
+        except ValueError:
+            continue
+    """
+    parsed.sort(key=lambda x: x[1], reverse=True)
+    results = {}
+    min_time = 0.1
+    for fn, total, count, max_val in parsed:
+        if total > min_time:
+            dynamo.ts(fn, total)
+            results[fn] = { "total": round(total, 2), "count": count, "avg": round(total / count, 2), "max": round(max_val, 2) }
+    return results
+    """

modules/sdnq/common.py

@@ -2,6 +2,7 @@
 
 import os
 import json
+import time
 import torch
 
 from modules import shared
@@ -338,10 +339,13 @@ weights_dtype_order = [
 
 use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply
 
+
 def check_torch_compile() -> bool: # dynamo can be disabled after startup
     return use_torch_compile and not torch._dynamo.config.disable # pylint: disable=protected-access
 
+
 if use_torch_compile:
+    from modules.timer import compiler as compile_timer
     if hasattr(torch._dynamo.config, "recompile_limit"):
         torch._dynamo.config.recompile_limit = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))
     if hasattr(torch._dynamo.config, "cache_size_limit"):
@@ -361,7 +365,10 @@ if use_torch_compile:
         if os.environ.get("SDNQ_COMPILE_KWARGS", None) is not None:
             for key, value in json.loads(os.environ.get("SDNQ_COMPILE_KWARGS")).items():
                 kwargs[key] = value
-        return torch.compile(fn, **kwargs)
+        t0 = time.time()
+        res = torch.compile(fn, **kwargs)
+        compile_timer.ts(fn.__name__, t0)
+        return res
 else:
     def compile_func(fn, **kwargs): # pylint: disable=unused-argument
         return fn

modules/sdnq/timers.py

@@ -0,0 +1,10 @@
+def update_sdnq_attention_timers():
+    from modules.timer import autotune
+    autotune.reset()
+    from modules.sdnq.kernels import triton_atten, triton_mm, triton_scaled_mm
+    if getattr(triton_atten.sdnq_attn_kernel, 'bench_time', None) is not None:
+        autotune.add('sdnq_attn_kernel', getattr(triton_atten.sdnq_attn_kernel, 'bench_time', 0))
+    if getattr(triton_mm.sdnq_triton_mm, 'bench_time', None) is not None:
+        autotune.add('sdnq_triton_mm', getattr(triton_mm.sdnq_triton_mm, 'bench_time', 0))
+    if getattr(triton_scaled_mm.sdnq_scaled_mm, 'bench_time', None) is not None:
+        autotune.add('sdnq_scaled_mm', getattr(triton_scaled_mm.sdnq_scaled_mm, 'bench_time', 0))

modules/timer.py

@@ -10,11 +10,11 @@ except Exception:
 
 
 class Timer:
-    def __init__(self):
+    def __init__(self, profile=False):
         self.start = time.time()
         self.records = {}
         self.total = 0
-        self.profile = False
+        self.profile = profile
 
     def elapsed(self, reset=True):
         end = time.time()
@@ -71,10 +71,14 @@ class Timer:
         return res
 
     def reset(self):
-        self.__init__()
+        self.records.clear()
+        self.__init__(self.profile)
 
 startup = Timer()
 process = Timer()
 launch = Timer()
 init = Timer()
 load = Timer()
+dynamo = Timer()
+compiler = Timer(profile=True)
+autotune = Timer(profile=True)

modules/video_models/video_load.py

@@ -155,9 +155,7 @@ def load_model(selected: models_def.Model):
             shared.sd_model = load_custom(selected.repo)
         else:
             log.debug(f'Load video: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}')
-            print('HERE1')
             sd_models.hf_prefetch_configs(selected.repo, {}, 'video')
-            print('HERE2')
             shared.sd_model = selected.repo_cls.from_pretrained(
                 pretrained_model_name_or_path=selected.repo,
                 revision=selected.repo_revision,

ui/gallery.ts

@@ -1013,7 +1013,6 @@ export async function gallerySort(key) {
   const folderNames = Array.from<string>(folderGroups.keys());
   const sortedFolderNames = currentSort.endsWith('A') ? folderNames.sort((a, b) => a.localeCompare(b)) : folderNames.sort((a, b) => b.localeCompare(a));
 
-  console.log('HERE', sortedFolderNames);
   for (const folderName of sortedFolderNames) {
     const files = folderGroups.get(folderName);
     files.sort(sortMode.func);