tqchen commented on a change in pull request #8044:
URL: https://github.com/apache/tvm/pull/8044#discussion_r632955391
##########
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:
I agree, @Hzfengsy let us remove the comment of simple_mode since it no
longer applies.
--
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:
[email protected]
