[tvm] branch unity updated: [Unity] Only process relax functions in workspace annotation (#14896)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new d3ad96e43c [Unity] Only process relax functions in workspace annotation (#14896)
d3ad96e43c is described below

commit d3ad96e43cb8ad452fd5851517feb22b76850fbd
Author: Lite Ye <ye...@gmail.com>
AuthorDate: Fri May 19 22:37:21 2023 -0400

    [Unity] Only process relax functions in workspace annotation (#14896)
    
    Only process relax function in annotate workspace
---
 python/tvm/relax/backend/contrib/cutlass.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relax/backend/contrib/cutlass.py b/python/tvm/relax/backend/contrib/cutlass.py
index 60e25abffb..f27e0d9e50 100644
--- a/python/tvm/relax/backend/contrib/cutlass.py
+++ b/python/tvm/relax/backend/contrib/cutlass.py
@@ -412,8 +412,9 @@ def annotate_workspace(mod, _):
     """Pass to annotate a workspace requirement for each CUTLASS-offloaded function."""
     annotator = WorkspaceAnnotator(mod)
     for name, f in mod.functions.items():
-        new_f = annotator.visit_expr(f)
-        mod.update_func(name, new_f)
+        if isinstance(f, Function):
+            new_f = annotator.visit_expr(f)
+            mod.update_func(name, new_f)
     return mod