[GitHub] [tvm] tkonolige commented on a change in pull request #8110: Unify Python and C++ TIR lower API
tkonolige commented on a change in pull request #8110:
URL: https://github.com/apache/tvm/pull/8110#discussion_r647820965
##
File path: include/tvm/driver/driver_api.h
##
@@ -42,17 +43,67 @@
#include
namespace tvm {
+
+/*!
+ * \brief Build an IRModule given an input IRModule
Review comment:
```suggestion
* \brief Lower an IRModule
```
Could you maybe also add a little more detail on what it means to lower an
irmodule?
--
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:
us...@infra.apache.org
[GitHub] [tvm] tkonolige commented on a change in pull request #8110: Unify Python and C++ TIR lower API
tkonolige commented on a change in pull request #8110:
URL: https://github.com/apache/tvm/pull/8110#discussion_r647013617
##
File path: src/driver/driver_api.cc
##
@@ -109,6 +109,52 @@ void GetBinds(const Array& args, bool compact,
}
}
+void GetBinds(const Array& args, bool compact,
+ const std::unordered_map& binds,
+ Map* out_binds, Array*
out_arg_list) {
+ *out_binds = binds;
+
+ for (const ObjectRef& x : args) {
+if (const auto* tensor_node = x.as()) {
+ auto x_ref = GetRef(tensor_node);
+ if (out_binds->find(x_ref) == out_binds->end()) {
+auto buf =
+BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
x_ref->op->name, -1, 0, compact);
+out_binds->Set(x_ref, buf);
+out_arg_list->push_back(buf);
+ } else {
+out_arg_list->push_back((*out_binds)[x_ref]);
+ }
+} else if (x.as() || x.as()) {
+ out_arg_list->push_back(x);
+} else {
+ LOG(FATAL)
+ << "Expected type of the elements of args to be te::Tensor,
te::Buffer or tir::Var, "
+ << "but got a " << typeid(x).name();
Review comment:
```suggestion
<< "but got a " << x->GetTypeKey();
```
`GetTypeKey` has a confusing name but it is what you want.
##
File path: python/tvm/driver/build_module.py
##
@@ -136,7 +98,7 @@ def lower(
Parameters
--
-input : Union[schedule.Schedule, PrimFunc, IRModule]
+inputs : Union[schedule.Schedule, PrimFunc, IRModule]
Review comment:
Maybe use `inp` instead? I can't really think of a good name either.
##
File path: include/tvm/driver/driver_api.h
##
@@ -42,17 +43,64 @@
#include
namespace tvm {
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param mod The IRmodule to lower
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerModule(IRModule mod, bool simple_mode = false);
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param func The PrimFunc to lower
+ * \param name The name of the lowered function.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerPrimFunc(tvm::tir::PrimFunc func, const std::string&
name,
+ bool simple_mode = false);
+
/*!
* \brief Build an IRModule given a schedule, args and binds
* \param sch The schedule to lower.
* \param args The arguments to the function.
* \param name The name of the lowered function.
* \param binds Buffer assignments.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
* \return The result module.
*/
-TVM_DLL IRModule lower(te::Schedule sch, const Array& args, const
std::string& name,
- const std::unordered_map&
binds);
+TVM_DLL IRModule LowerSchedule(te::Schedule sch, const Array& args,
+ const std::string& name,
+ const std::unordered_map& binds,
+ bool simple_mode = false);
+
+/*!
+ * \brief Build an IRModule given a schedule, args and binds
+ * \param sch The schedule to lower.
+ * \param args The arguments to the function (Array of Union of Tensor, Buffer
and Vars)
+ * \param name The name of the lowered function.
+ * \param binds Buffer assignments.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerSchedule(te::Schedule sch, const Array& args,
+ const std::string& name,
+ const std::unordered_map& binds,
+ bool simple_mode = false);
+
+/*!
+ * \brief Create an IRModule out of a Schedule
+ * \param sch The schedule
+ * \param args The arguments to the function.
+ * \param name The name of the lowered function.
+ * \param binds Buffer assignments.
+ * \return The result module.
+ */
+IRModule ScheduleToModule(te::Schedule sch, const Array& args,
const std::string& name,
Review comment:
Can you be a little more explicit about this in the documentation?
##
File path: include/tvm/driver/driver_api.h
##
@@ -42,17 +43,64 @@
#include
namespace tvm {
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param mod The IRmodule to lower
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerModule(IRModule mod, bool simple_mode = false);
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param func The PrimFunc to lower
+ * \param name The name of the lowered function.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerPrimFunc(tvm::tir:
[GitHub] [tvm] tkonolige commented on a change in pull request #8110: Unify Python and C++ TIR lower API
tkonolige commented on a change in pull request #8110:
URL: https://github.com/apache/tvm/pull/8110#discussion_r646870692
##
File path: include/tvm/driver/driver_api.h
##
@@ -42,17 +43,64 @@
#include
namespace tvm {
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param mod The IRmodule to lower
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerModule(IRModule mod, bool simple_mode = false);
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param func The PrimFunc to lower
+ * \param name The name of the lowered function.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerPrimFunc(tvm::tir::PrimFunc func, const std::string&
name,
+ bool simple_mode = false);
+
/*!
* \brief Build an IRModule given a schedule, args and binds
* \param sch The schedule to lower.
* \param args The arguments to the function.
* \param name The name of the lowered function.
* \param binds Buffer assignments.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
Review comment:
Why is there a special flag to disable the loop partition pass.
Shouldn't the existing `PassContext` infrastructure handle this?
##
File path: python/tvm/driver/build_module.py
##
@@ -37,92 +37,54 @@
from tvm.tir.buffer import Buffer
from tvm.tir.expr import Var
+from . import _ffi_api as ffi
+
def get_binds(args, compact=False, binds=None):
"""Internal function to get binds and arg_list given arguments.
-
Parameters
--
args : list of Buffer or Tensor or Var
The argument lists to the function.
-
compact : bool
If the statement has already bound to a compact buffer.
-
binds : dict of :any:`Tensor` to :any:`Buffer`, optional
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.
-
Returns
---
binds: dict
The bind specification
-
arg_list: list
The list of symbolic buffers of arguments.
"""
-binds = {} if binds is None else binds.copy()
-arg_list = []
-for x in args:
-if isinstance(x, tensor.Tensor):
-any_dim = any(isinstance(i, tvm.tir.Var) for i in x.shape)
-buffer_type = "auto_broadcast" if any_dim and not compact else ""
-if x not in binds:
-buf = tvm.tir.decl_buffer(
-x.shape, dtype=x.dtype, name=x.name,
buffer_type=buffer_type
-)
-binds[x] = buf
-arg_list.append(buf)
-else:
-arg_list.append(binds[x])
-elif isinstance(x, schedule.Buffer):
-arg_list.append(x)
-elif isinstance(x, tvm.tir.Var):
-arg_list.append(x)
-else:
-raise ValueError("args must be Tensor, Buffer or Var")
-return binds, arg_list
-
-
-def form_irmodule(sch, args, name, binds):
-"""According to the given schedule, form a function.
+out_arr = ffi.get_binds(args, compact, binds)
+return out_arr[0], out_arr[1]
+
+def schedule_to_module(
+sch: schedule.Schedule,
+args: Optional[List[Union[Buffer, tensor.Tensor, Var]]] = None,
+name: str = "main",
+binds: Optional[Mapping[tensor.Tensor, Buffer]] = None,
+) -> IRModule:
+"""According to the given schedule, form a function.
Parameters
--
sch : tvm.te.schedule.Schedule
The given scheduler to form the raw body
-
args : list of Buffer or Tensor or Var
The argument lists to the function.
-
name : str
The name of result function.
Review comment:
Can you add the default name to this documentation?
##
File path: include/tvm/driver/driver_api.h
##
@@ -42,17 +43,64 @@
#include
namespace tvm {
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param mod The IRmodule to lower
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerModule(IRModule mod, bool simple_mode = false);
+
+/*!
+ * \brief Build an IRModule given a module, args and binds
+ * \param func The PrimFunc to lower
+ * \param name The name of the lowered function.
+ * \param simple_mode Disables the loop partition pass. Defaults to false.
+ * \return The result module.
+ */
+TVM_DLL IRModule LowerPrimFunc(tvm::tir::PrimFunc func, const std::string&
name,
+ bool simple_mode = false);
+
/*!
* \brief Build an IRModule given a schedule, args and binds
* \param sch The schedule to lower.
* \param args The arguments to the functio
