electriclilies commented on a change in pull request #8110:
URL: https://github.com/apache/tvm/pull/8110#discussion_r646927080
##########
File path: src/driver/driver_api.cc
##########
@@ -128,63 +128,192 @@ transform::Pass Filter(FCond fcond) {
return tir::transform::CreatePrimFuncPass(fpass, 0, "Filter", {});
}
-IRModule lower(te::Schedule sch, const Array<te::Tensor>& args, const
std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds) {
- Array<ObjectRef> out_arg_list;
+Array<tvm::transform::Pass> CreatePassList(bool simple_mode, bool
legacy_te_pass) {
auto pass_ctx = transform::PassContext::Current();
- sch = sch.normalize();
-
- // Before TIR transformation.
- auto bounds = te::InferBound(sch);
- auto stmt = te::ScheduleOps(sch, bounds, false);
- bool compact = te::VerifyCompactBuffer(stmt);
-
- Map<te::Tensor, tir::Buffer> out_binds;
- GetBinds(args, compact, binds, &out_binds, &out_arg_list);
-
- // build the function
- tir::PrimFunc f = te::SchedulePostProcToPrimFunc(out_arg_list,
std::move(stmt), out_binds);
- f = WithAttr(std::move(f), "global_symbol", runtime::String(name));
-
- bool noalias = pass_ctx->GetConfig<Bool>("tir.noalias", Bool(true)).value();
bool disable_vectorize = pass_ctx->GetConfig<Bool>("tir.disable_vectorize",
Bool(false)).value();
bool instrument_bound_checkers =
pass_ctx->GetConfig<Bool>("tir.instrument_bound_checkers",
Bool(false)).value();
- if (noalias) {
- f = WithAttr(std::move(f), "tir.noalias", Bool(true));
+ // Get any user-added passes
+ auto add_lower_pass =
+ pass_ctx->GetConfig<Array<Array<ObjectRef>>>("tir.add_lower_pass",
Array<Array<ObjectRef>>())
+ .value();
+
+ auto user_lower_phase0 = Array<tvm::transform::Pass>();
Review comment:
I'm not sure, it would be good to get some clarification from someone
else on this, maybe @junrushao1994 knows?
##########
File path: include/tvm/driver/driver_api.h
##########
@@ -42,17 +43,64 @@
#include <vector>
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:
I'm not completely sure. This logic was already in the code base, so I
just duplicated it and made it more explicit. (For context, when I started this
refactor, Jared said he wanted to get it in quickly so I should just try to
naively duplicate the existing logic, but he hasn't reviewed it yet so I'm not
sure what the timeline is now).
I can try to remove it, but I'm not sure what the best way to go about this
is since there are a few tests that call lower directly..
##########
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:
The word input is a python built in function so you can't name variables
`input.` Initially the name of the parameter to the function was `inputs` and
the name of the parameter in the documentation was `input`. I just changed the
documentation to match the function signature. I don't like that they are
different. I agree that `inputs` is not an ideal name though and I'm open to
suggestions for other names
##########
File path: include/tvm/driver/driver_api.h
##########
@@ -42,17 +43,64 @@
#include <vector>
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<te::Tensor>& args, const
std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>&
binds);
+TVM_DLL IRModule LowerSchedule(te::Schedule sch, const Array<te::Tensor>& args,
+ const std::string& name,
+ const std::unordered_map<te::Tensor,
tir::Buffer>& 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<ObjectRef>& args,
+ const std::string& name,
+ const std::unordered_map<te::Tensor,
tir::Buffer>& 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<ObjectRef>& args,
const std::string& name,
Review comment:
Initially this was a function called `form_irmodule` in python. I
translated it to C++ and renamed it as `ScheduleToModule`. Unfortunately,
form_irmodule is called by some tests. To preserve that behavior I had to split
the functions apart again and register them separately in the FFI.
The difference between these is that `ScheduleToModule` just converts the
schedule to a module that hasn't yet been lowered, whereas `LowerSchedule`
converts the schedule into a module and then applies the passes.
--
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]
