From 1e4445f35b4a17750670a8936d11fb025b7abb69 Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Mon, 8 Aug 2022 10:19:37 -0700 Subject: [PATCH] [FIX,ROOFLINE] Only save tir functions for roofline Only collect TIR PrimFuncs in roofline's SaveLoweredTIR. SaveLoweredTIR was saving the full Relay main function leading which could be excessively large. Also improve the logic to only save functions right before MakePackedAPI. --- python/tvm/utils/roofline/__init__.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/python/tvm/utils/roofline/__init__.py b/python/tvm/utils/roofline/__init__.py index a54f5ed41d0a..0affb0704997 100644 --- a/python/tvm/utils/roofline/__init__.py +++ b/python/tvm/utils/roofline/__init__.py @@ -56,14 +56,11 @@ class SaveLoweredTIR: def __init__(self): self.functions = {} - self.done = False - - def run_after_pass(self, mod, info): - if not self.done: - if info.name == "tir.MakePackedAPI": - self.done = True - else: - for v, func in mod.functions.items(): + + def run_before_pass(self, mod, info): + if info.name == "tir.MakePackedAPI": + for v, func in mod.functions.items(): + if isinstance(func, tir.PrimFunc): self.functions[v] = func