[GitHub] [tvm] xqdan commented on a change in pull request #8044: [TensorIR][M1c] Lower and build TensorIR

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

xqdan commented on a change in pull request #8044:
URL: https://github.com/apache/tvm/pull/8044#discussion_r632895555



##########
File path: python/tvm/driver/build_module.py
##########
@@ -119,32 +125,39 @@ def form_irmodule(sch, args, name, binds):
     return tvm.IRModule({name: func})
 
 
-def lower(sch, args, name="main", binds=None, simple_mode=False):
+def lower(
+    inputs: Union[schedule.Schedule, PrimFunc, IRModule],
+    args: Optional[List[Union[Buffer, tensor.Tensor, Var]]] = None,
+    name: str = "main",
+    binds: Optional[Mapping[tensor.Tensor, Buffer]] = None,
+    simple_mode: bool = False,
+) -> IRModule:
     """Lowering step before build into target.
 
     Parameters
     ----------
-    sch : tvm.te.schedule.Schedule
-        The schedule to be built
+    input : Union[schedule.Schedule, PrimFunc, IRModule]
+        The TE schedule or TensorIR PrimFunc/IRModule to be built
 
-    args : list of Buffer or Tensor or Var
-        The argument lists to the function.
+    args : Optional[List[Union[Buffer, tensor.Tensor, Var]]]
+        The argument lists to the function for TE schedule.
+        It should be None if we want to lower TensorIR.
 
-    name : str, optional
+    name : str
         The name of result function.
 
-    binds : dict of :any:`Tensor` to :any:`Buffer`, optional
+    binds : Optional[Mapping[tensor.Tensor, Buffer]]
         Dictionary that maps the Tensor to Buffer which specified the data layout
         requirement of the function. By default, a new compact buffer is created
         for each tensor in the argument.
 
-    simple_mode : bool, optional
+    simple_mode : bool
         Whether only output simple and compact statement, this will skip
         LoopPartition, api wrapper generation and Unrolling.
 
     Returns
     -------
-    m : IRModule or Stmt
+    m : IRModule
        The result IRModule, if simple_mode=False
        Then the Stmt before make api is returned.

Review comment:
       We may need update Stmt related comment here




-- 
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