[GitHub] [tvm] trevor-m commented on a change in pull request #7412: TRT Dynamic Reshape Fix

[GitBox <gi...@apache.org>]

trevor-m commented on a change in pull request #7412:
URL: https://github.com/apache/tvm/pull/7412#discussion_r572445116



##########
File path: tests/python/contrib/test_tensorrt.py
##########
@@ -631,6 +632,106 @@ def get_graph(x_shape, new_shape):
     run_and_verify_func(get_graph((1, 1, 2, 3), (1, 6)))
 
 
+class AreOpsOnGraph(ExprVisitor):
+    """
+    Visits the Graph recursively and checks if it contains ops in the op_list
+    """
+
+    def __init__(self, op_list):
+        ExprVisitor.__init__(self)
+        self.op_list = op_list
+        self.on_graph = False
+
+    def visit_call(self, call):
+        if isinstance(call.op, tvm.tir.op.Op):
+            if str(call.op) in self.op_list:
+                self.on_graph = True
+
+        return super().visit_call(call)
+
+    def are_ops_on_graph(self, subgraph) -> bool:
+        """
+        This function recursively visits the graph and checks if op_list ops are ongraph"
+        """
+        self.visit(subgraph)
+        return self.on_graph
+
+
+def are_ops_on_trt(mod, op_list):
+    for subgraph in mod.get_global_vars():
+        name = subgraph.name_hint
+        op_on_trt = False
+        op_on_tvm = True
+        if name == "main":
+            op_on_tvm = AreOpsOnGraph(op_list).are_ops_on_graph(mod[name].body)
+        elif mod[name].attrs and mod[name].attrs["Compiler"] == "tensorrt":
+            op_on_trt = AreOpsOnGraph(op_list).are_ops_on_graph(mod[name].body)
+        elif mod[name].attrs and mod[name].attrs["Compiler"] != "tensorrt":

Review comment:
       Do we need this third condition? I think it could just be `else:` if we do. 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org