Unify Python and C++ TIR lower API #8110

CircleSpin · 2021-05-21T22:52:22Z

Currently, the Lower API is bifurcated. There is an implementation in Python and an implementation in C++. Both are used in the codebase despite having different signatures and different implementations. This PR unifies the C++ and Python API into one C++ backend, and allows the C++ API to be called from Python through the FFI.

Unfortunately, the Python version of the API relied on duck typing quite a bit, which made it difficult to fully support the Python API. To allow the C++ backend to fully replace the Python version of lower, @electriclilies had to write three different versions of Lower: LowerModule, LowerPrimFunc, and LowerSchedule, as well as make other changes.

The major changes to the C++ API are:

1, Implementing versions of Lower for Modules, PrimFuncs and Schedules
2. Adding user-defined passes to the C++ version
3. Overload LowerSchedule to take in an Array for the arguments. The Python version accepts a List[Union[Buffer, Tensor, Var]] for the args. Unfortunately, the TVM object system does not allow union types, so Lily had to allow LowerSchedule to take in an Array for the args and raise an error if something other than a Buffer, Tensor or Var is passed in
4. Removing SchedulePostProcForTensorCore (this is deadcode, it was used in one tutorial that has since been deleted and causes errors if called in certain cases)

Additionally, Lily moved the python version of get_binds into C++ and exposed it through the FFI. Lily also renamed form_irmodule to schedule_to_module and moved that into C++. Finally, this refactor was a bit painful. @electriclilies will write a discussion post to go through why it was difficult and what can be changed to avoid issues like this in the future.

jroesch · 2021-05-25T02:04:45Z

src/driver/driver_api.cc

+      } else {
+        ICHECK(false) << "driver.lower expects the first argument to be a te::Schedule or "
+                      << "IRModule";
+        throw;


Don't add a throw here, instead return an empty IRModule(...)

jroesch · 2021-05-25T02:06:14Z

@CircleSpin @electriclilies ping me again when tests are green, left one comment that jumped out at me, otherwise looks good if we can get it to pass the test suites

jroesch · 2021-05-25T02:07:17Z

cc @manupa-arm @giuseros this is worth paying attention to, we have been working on cleaning up the internal APIs and bringing everything to C++.

…ts more types in for args

manupak

Thanks for this work! much appreciated. :)

Just few clarifcation questions about the overall design.

manupak · 2021-05-26T09:20:32Z

include/tvm/driver/driver_api.h

+                       bool simple_mode = false);
+
+/*!
+ * \brief Build an IRModule given a module, args and binds


typo : should be a PrimFunc?

manupak · 2021-05-26T09:23:37Z

src/driver/driver_api.cc

+      pass_ctx->GetConfig<Array<Array<ObjectRef>>>("tir.add_lower_pass", Array<Array<ObjectRef>>())
+          .value();
+
+  auto user_lower_phase0 = Array<tvm::transform::Pass>();


Clarification : What is the importance of having phases and what do they represent ?

I'm not sure, it would be good to get some clarification from someone else on this, maybe @junrushao1994 knows?

It would be nice to take the refactoring opportunity to add some documentation on this. Though completely understand you have mirrored the previous lower impl.

I'm happy to add documentation, I just don't actually know what the answer is to this question

manupak · 2021-05-26T09:24:44Z

src/driver/driver_api.cc

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


Clarification : what is legacy_te_pass mean here?

Legacy te pass means that the object we are lowering is a te::Schedule. I've copied the logic and phrasing directly from the python version -- I don't know the details of why this is still here. But without it, some tests fail.

Maybe we could be more specific in the name ? e.g. using such as "bool from_te_schedule" ?

Also I see simple_mode is just enabling LoopPartition, we could use "bool enable_loop_partition" ?

manupak · 2021-05-26T09:32:33Z

src/driver/driver_api.cc

-  pass_list.push_back(tir::transform::StorageFlatten(64, instrument_bound_checkers));
-  // Phase 1
+  // PHASE 1
+  if (legacy_te_pass) {


@jroesch @CircleSpin
A general design question : What do you think of relocating the pass pipeline closer to the target registry ? So that way we could have a function (that can also be a bit parameterized based on target args) that describes pass pipeline rather than mandating the passes here.

Here, we could just query the pass pipeline based on the target and run here. I think this can remove the need to support user defined custom passes at this level.

Thoughts ?

I think doing something like this is a good idea. I also want to simplify the signature for Lower..

Right now we have a lot of function signatures for Lower because the python version allows multiple input types for each argument.

To clarify, are you saying that we would construct the pass list in the registration of driver.lower, then pass it into a function like LowerWithPassList?

Yes, OK having looked at whats being removed -- maybe this is out of scope for the PR.

My original question was about the motivation of custom passes -- specifically at the four possible locations they get inserted (known as phases). If it has some meaning, might worth putting a comment.

Then, the next one was given the feature of adding custom passes -- I was thinking we could rather have custom pass pipeline registered/provided with the proximity for the target. Thus, some could create a new target and may be re-use , re-organize the passes that needs running.

Again now I feel thats out of scope for this PR -- as this essentially mimics what the python lower is doing. :)

manupak

Thanks! some further thoughts

manupak · 2021-05-27T17:37:06Z

src/driver/driver_api.cc

-  pass_list.push_back(tir::transform::StorageFlatten(64, instrument_bound_checkers));
-  // Phase 1
+  // PHASE 1
+  if (legacy_te_pass) {


Yes, OK having looked at whats being removed -- maybe this is out of scope for the PR.

My original question was about the motivation of custom passes -- specifically at the four possible locations they get inserted (known as phases). If it has some meaning, might worth putting a comment.

Then, the next one was given the feature of adding custom passes -- I was thinking we could rather have custom pass pipeline registered/provided with the proximity for the target. Thus, some could create a new target and may be re-use , re-organize the passes that needs running.

Again now I feel thats out of scope for this PR -- as this essentially mimics what the python lower is doing. :)

manupak · 2021-05-27T17:40:44Z

src/driver/driver_api.cc

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


Maybe we could be more specific in the name ? e.g. using such as "bool from_te_schedule" ?

Also I see simple_mode is just enabling LoopPartition, we could use "bool enable_loop_partition" ?

electriclilies · 2021-05-28T01:59:17Z

@manupa-arm Yeah I'm not exactly sure what the meaning of each "phase" is-- again, I'm just trying to duplicate the python code in C++. @junrushao1994 Could you weigh in about the significance of each phase?

Also, thanks for the naming suggestions, I'm going to go through and clean stuff up later and I will take them into account then! Ideally I'd like to remove flags if possible

electriclilies

Thanks Tristan!

electriclilies · 2021-06-07T20:29:38Z

include/tvm/driver/driver_api.h

 /*!
 * \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.


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

electriclilies · 2021-06-07T20:39:04Z

include/tvm/driver/driver_api.h

+ * \param binds Buffer assignments.
+ * \return The result module.
+ */
+IRModule ScheduleToModule(te::Schedule sch, const Array<ObjectRef>& args, const std::string& name,


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.

electriclilies · 2021-06-07T20:42:08Z

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]


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

electriclilies · 2021-06-07T20:43:58Z

src/driver/driver_api.cc

+      pass_ctx->GetConfig<Array<Array<ObjectRef>>>("tir.add_lower_pass", Array<Array<ObjectRef>>())
+          .value();
+
+  auto user_lower_phase0 = Array<tvm::transform::Pass>();


I'm not sure, it would be good to get some clarification from someone else on this, maybe @junrushao1994 knows?

csullivan

Absolutely phenomenal improvement, thank you for the very well thought out refactor!

csullivan · 2021-06-07T20:10:58Z

include/tvm/driver/driver_api.h

@@ -42,17 +43,64 @@
 #include <vector>

 namespace tvm {
+
+/*!
+ * \brief Build an IRModule given a module, args and binds


Suggested change

* \brief Build an IRModule given a module, args and binds

* \brief Build an IRModule given a module

Updated the comment to reflect the API. Please check the docs for the other APIs that don't take args/binds as well.

include/tvm/driver/driver_api.h

python/tvm/autotvm/feature.py

csullivan · 2021-06-07T20:24:54Z

python/tvm/driver/build_module.py

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


Probably we should update the doc string to match the return variable names. Actually the reverse would be preferable, binds, arg_list = ffi.get_binds(...) but I understand you are limited here by the object system.

Yup can do this

src/driver/driver_api.cc

csullivan · 2021-06-07T21:12:16Z

src/relay/backend/compile_engine.cc

+      cache_node->funcs =
+          tvm::LowerSchedule(cfunc->schedule, all_args, cache_node->func_name, binds);


We shouldn't need to route through the relay.backend.lower packed function any longer right? I'd suggest either removing it and the above conditional (line 766) or also reproducing its functionality in c++ if we do need it as part of this PR. This way we avoid an unnecessary round trip through python.

Yeah I don't think we do. I'll remove it (I didn't look too closely at what this call was doing, oops!)

Co-authored-by: Chris Sullivan <csullivan@octoml.ai>

src/driver/driver_api.cc

tkonolige · 2021-06-07T23:51:00Z

include/tvm/driver/driver_api.h

 /*!
 * \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.


It is probably fine for now.

tkonolige · 2021-06-07T23:51:58Z

include/tvm/driver/driver_api.h

+ * \param binds Buffer assignments.
+ * \return The result module.
+ */
+IRModule ScheduleToModule(te::Schedule sch, const Array<ObjectRef>& args, const std::string& name,


Can you be a little more explicit about this in the documentation?

tkonolige · 2021-06-07T23:52:57Z

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]


Maybe use inp instead? I can't really think of a good name either.

Hzfengsy

Thanks for this great work! I think it should work for TensorIR as long as pass all existing unit tests :) Generally, it looks good to me except for some nits in code style

Hzfengsy · 2021-06-08T08:46:20Z

src/driver/driver_api.cc

+          c_binds.insert({kv.first, kv.second});
+        }
+      }
+      IRModule mod = ScheduleToModule(sch, args, name, c_binds);


Suggested change

IRModule mod = ScheduleToModule(sch, args, name, c_binds);

IRModule mod = ScheduleToModule(std::move(sch), args, name, c_binds);

Hzfengsy · 2021-06-08T08:46:56Z

src/driver/driver_api.cc

+
+IRModule LowerModule(IRModule mod, bool simple_mode) {
+  auto pass_list = CreatePassList(simple_mode, false);
+  return LowerWithPassList(mod, pass_list);


Suggested change

return LowerWithPassList(mod, pass_list);

return LowerWithPassList(std::(mod), pass_list);

Please apply std::move to all following FFI functions

Hzfengsy · 2021-06-08T08:48:54Z

src/driver/driver_api.cc

+  sch = sch.normalize();
+
+  // Before TIR transformation.
+  auto bounds = te::InferBound(sch);


Please explicitly mark the type as possible as we can. auto is not so user-friendly that others can understand the code.

Hzfengsy · 2021-06-08T08:50:10Z

src/driver/driver_api.cc

+
+  // Before TIR transformation.
+  auto bounds = te::InferBound(sch);
+  auto stmt = te::ScheduleOps(sch, bounds, false);


Please use const ref or std::move to reduce unnecessary memory copy

Suggested change

auto stmt = te::ScheduleOps(sch, bounds, false);

auto stmt = te::ScheduleOps(sch, std::move(bounds), false);

Co-authored-by: Tristan Konolige <tristan.konolige@gmail.com>

tkonolige

One small nit, but otherwise things look really good. That net -1K lines is so satisfying to see. Big thanks!

tkonolige · 2021-06-08T21:56:36Z

include/tvm/driver/driver_api.h

@@ -42,17 +43,67 @@
 #include <vector>

 namespace tvm {
+
+/*!
+ * \brief Build an IRModule given an input IRModule


Suggested change

* \brief Build an IRModule given an input IRModule

* \brief Lower an IRModule

Could you maybe also add a little more detail on what it means to lower an irmodule?

electriclilies · 2021-06-09T01:04:19Z

Thanks @tkonolige! @Hzfengsy @csullivan can you take another look and approve?

csullivan

LGTM! Thanks again for the great refactor @electriclilies, removing 4x the code than what is needed by the change == 💪 🔥 .

Hzfengsy

LTGM. Thanks a lot for the refactoring

tqchen · 2021-06-12T13:58:47Z

Thanks @CircleSpin and @electriclilies ! Thanks @Hzfengsy @tkonolige @csullivan @YuchenJin @manupa-arm for reviewing

electriclilies · 2021-06-14T16:55:54Z

Thanks @tqchen!

Jocelyn added 11 commits May 7, 2021 18:48

attempt of c->python

ae28f58

name change

1258223

build update

f57ae0e

import fix

8a4f40e

build working, import still needs work

f7e8bbd

returning null binds in driver_api.cc

e4267fb

tests pass woohoo00!

fa322f6

merge to recent tvm

51d85fb

black'd _ffi_api.py

8e30116

remove simple_mode arg from lower in build_module.py for now

ca11f46

attempt add simple_mode arg in lower c++

b117478

jwfromm requested a review from jroesch May 24, 2021 20:02

Jocelyn added 4 commits May 24, 2021 18:18

lower now can take in schedule or IRModule

41640a1

add simple_mode parameter in c++

724db08

reformat for lint

47ea819

include header details for lint

c01ca1e

jroesch reviewed May 25, 2021

View reviewed changes

Jocelyn and others added 2 commits May 25, 2021 18:07

ast lhs and rhs not matching, refactoring driver_api.cc

49d6c7f

Added user-defined passes, still failing some tests because python le…

5e43b18

…ts more types in for args

manupak reviewed May 26, 2021

View reviewed changes

manupak reviewed May 27, 2021

View reviewed changes

tests are green

59534cb

electriclilies added 5 commits May 27, 2021 19:09

Split lower api into 3

5cf0f78

got rid of legacy lower

b10d6ad

renamed lower funcs

5d5c068

remove python get_binds and rename flags

f2596c5

fix typo

d6be36e

electriclilies reviewed Jun 7, 2021

View reviewed changes

csullivan requested changes Jun 7, 2021

View reviewed changes

electriclilies and others added 5 commits June 7, 2021 14:33

respond to tristan's comments

ca4683d

update driver_api.h docs

00fcd7a

Apply suggestions from code review

7a2b404

Co-authored-by: Chris Sullivan <csullivan@octoml.ai>

remove relay.backend.lower

8988368

Merge branch 'lower_build' of github.com:CircleSpin/tvm into lower_build

a460f7e

tkonolige reviewed Jun 7, 2021

View reviewed changes

Hzfengsy requested changes Jun 8, 2021

View reviewed changes

electriclilies and others added 4 commits June 8, 2021 11:32

Respond to feedback

38a121b

remove 2nd get binds impl

55e90cf

fix lint

d3e18c2

Update src/driver/driver_api.cc

6e5a2ff

Co-authored-by: Tristan Konolige <tristan.konolige@gmail.com>

tkonolige approved these changes Jun 8, 2021

View reviewed changes

electriclilies added 3 commits June 8, 2021 17:16

clang format

4933a5e

update doc

b42fb19

Merge branch 'lower_build' of github.com:CircleSpin/tvm into lower_build

4c7049d

electriclilies added 3 commits June 8, 2021 18:56

fix typo

e462dab

black

d067a67

fix pass ctx

794c606

csullivan approved these changes Jun 11, 2021

View reviewed changes

Hzfengsy approved these changes Jun 12, 2021

View reviewed changes

tqchen approved these changes Jun 12, 2021

View reviewed changes

tqchen merged commit 9dd1286 into apache:main Jun 12, 2021

trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Jun 17, 2021

Unify Python and C++ TIR lower API (apache#8110)

4e459cf

trevor-m pushed a commit to neo-ai/tvm that referenced this pull request Jun 17, 2021

Unify Python and C++ TIR lower API (apache#8110)

042c9fe

manupak mentioned this pull request Aug 4, 2021

Additional Target Hooks RFC apache/tvm-rfcs#10

Merged

	* \brief Build an IRModule given a module, args and binds
	* \brief Build an IRModule given a module

		cache_node->funcs =
		tvm::LowerSchedule(cfunc->schedule, all_args, cache_node->func_name, binds);

	IRModule mod = ScheduleToModule(sch, args, name, c_binds);
	IRModule mod = ScheduleToModule(std::move(sch), args, name, c_binds);

	return LowerWithPassList(mod, pass_list);
	return LowerWithPassList(std::(mod), pass_list);

	auto stmt = te::ScheduleOps(sch, bounds, false);
	auto stmt = te::ScheduleOps(sch, std::move(bounds), false);

	* \brief Build an IRModule given an input IRModule
	* \brief Lower an IRModule

Unify Python and C++ TIR lower API #8110

Unify Python and C++ TIR lower API #8110

Conversation

CircleSpin commented May 21, 2021 • edited Loading

Choose a reason for hiding this comment

jroesch commented May 25, 2021

jroesch commented May 25, 2021

manupak left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

electriclilies commented May 28, 2021

electriclilies left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

csullivan left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

csullivan Jun 7, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Hzfengsy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tkonolige left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

electriclilies commented Jun 9, 2021

csullivan left a comment

Choose a reason for hiding this comment

Hzfengsy left a comment

Choose a reason for hiding this comment

tqchen commented Jun 12, 2021

electriclilies commented Jun 14, 2021

CircleSpin commented May 21, 2021 •

edited

Loading

csullivan Jun 7, 2021 •

edited

Loading