You are viewing a plain text version of this content. The canonical link for it is here.
Posted to dev@tvm.apache.org by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai> on 2020/11/12 02:04:07 UTC

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler


The current approach used by auto_scheduler to extract tuning tasks leverages Relay op strategy. In short, auto_scheduler registers an implementation in Relay op strategy as AutoTVM, but instead of using the TOPI schedule function, auto_scheduler creates an empty schedule and extracts the lowered TE compute function as a tuning task (ref: https://github.com/apache/incubator-tvm/blob/main/python/tvm/relay/op/op.py#L147).

However, an obvious issue of this approach is that the scope of a tuning task is limited by Relay compile engine and op strategy. Specifically, each primitive Relay function can only have at most one complicated op (i.e., reduce ops like conv2d). Relay compile engine will mark that op as the anchor op (ref: https://github.com/apache/incubator-tvm/blob/main/src/relay/backend/compile_engine.cc#L231), and use the TOPI schedule of that op to schedule an entire Relay function (ref: https://github.com/apache/incubator-tvm/blob/main/src/relay/backend/compile_engine.cc#L152).

Here is a motivating example:

```
def @main(%data: Tensor[(1, 3, 224, 224), float32], %weight1: Tensor[(32, 3, 3, 3), float32], %weight2: Tensor[(32, 32, 3, 3), float32]) {
  %3 = fn (%data1: Tensor[(1, 3, 224, 224), float32], %weight11: Tensor[(32, 3, 3, 3), float32], %weight21: Tensor[(32, 32, 3, 3), float32], Primitive=1) {
    %0 = nn.conv2d(%data1, %weight11, padding=[1, 1, 1, 1], kernel_size=[3, 3]);
    %1 = nn.relu(%0);
    %2 = nn.conv2d(%1, %weight21, padding=[1, 1, 1, 1], kernel_size=[3, 3]);
    nn.relu(%2)
  };
  %3(%data, %weight1, %weight2)
}
```

As can be seen, we manually set `%3` to primitive so that it won't be partitioned to two separate functions after the `FuseOps` pass. If we simply build this function, we will get the follow error message:

```
Check failed: !anchor_op_.defined() || anchor_op_pattern_ < kCommReduce == false: Cannot apply TOPI schedule to a primitive function with two complicated ops anchor=Op(nn.conv2d) current=Op(nn.conv2d)
```

As a result, the goal of this RFC is to propos a mechanism that is able to make the above Relay function as an auto_scheduler tuning task, and we can also build it with the tuning logs.

The proposed mechanism is:
1. Add a mode, `use_topi_schedule`, to Relay compile engine. When `use_topi_schedule=true`, it performs as it is. When `use_topi_schedule=false`, we do not check if this function has more than one reduce ops but simply invokes `auto_schedule_topi` for an entire TE compute.
2. Propagate the flag `use_topi_schedule` all the way to `GraphRuntimeCodegen` and `relay.Build`.
	1. In `auto_scheduler.extract_tasks`, we set `use_topi_schedule=false` so that it can extract tasks.
	2. In `relay.build`, we use `auto_scheduler.DispatchContext.current` to judge whether we should query auto_scheduler schedule for an entire function, or query TOPI schedule of the anchor op.

The draft PR is available [here](https://github.com/apache/incubator-tvm/pull/6903). Note that since we now extract auto_scheduler tasks directly via compile engine, we completely removed auto_scheduler related logics from Relay op strategy.

I also provide a running script [here](https://gist.github.com/comaniac/cc10a341b7d1c2cd504a5cd5456f6b44) if you are willing to play with more Relay functions.

Comments and suggestions are welcome :)

cc @merrymercy @tqchen @jcf94 @zhiics @haichen





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/1) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/640548771452b0c3aad1deffe452c145260627e991f2e7ecfa16e3519a3e407b).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

This is a good question. This is possible for the current implementation, because we use Relay op strategy to define auto_scheduler tasks as well. In other words, we use Relay FuseOps to define the task scope, and should be able to choose to use TOPI (AutoTVM) or auto_scheduler schedule for each task. However, as you pointed out, this option becomes unavailable after this RFC as we separate their task extraction and compilation paths.

An alternative solution might be keeping all approaches in the same place. Specifically, we still keep the current task extraction (i.e., having `strategy.add_auto_scheduler` in Relay op strategy). When `use_topi_schedule=True`, we can still extract auto_scheduler tasks as for now. The case of `use_topi_schedule=False` would become totally optional.

What do you think about this proposal, or do you have any suggestion on this? Also cc @merrymercy.





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/12) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/a2806691c9a958baeb4fc45c0b6770e2c688d1c2c9916b21ddcbd8da83a4e7a3).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by Haichen Shen via Apache TVM Discuss <no...@discuss.tvm.ai>.

I have one question about `use_topi_schedule`. I assume that after we set it to False, it will always use the Ansor scheduler to schedule the ops. Will there be a case that we want have a mix of topi schedule and ansor schedule?





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/11) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/4097abc18668da96b5bf484a0ba0c11a360307236bb219d4f5dbec864b7ca2dd).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

Thanks all for the valuable feedback. Here is the summary of the finalized RFC:

### Interface
- When extracting auto_scheduler tasks, users simply call `extract_tasks`.
- `extract_tasks` now has an optional flag `include_simple_task` (default `False`). When set, each Relay function, including the one with only injective ops, becomes an auto_scheduler task.
* When building a model, users need to use `with tvm.transform.PassContext(config={"relay.backend.use_auto_scheduler": True})` to apply auto_scheduler tuning logs.

### Changes in Compile Engine
- Compile engine checks `relay.backend.use_auto_scheduler` to determine whether to use auto_scheduler schedules. If true, then compile engine calls `auto_schedule_topi`.
- In the task extraction mode,  `auto_schedule_topi` extracts a task and returns an initial schedule. Since we are not going to actually compile the model in this mode, whether the initial schedule is valid doesn't matter.
- In the build mode, `auto_schedule_topi` queries the auto_scheduler tuning log for the given workload. If success, then it returns the tuned schedule; otherwise it returns `None`. In this case, compile engine falls back to use TOPI schedule (with AutoTVM tuning logs if provided) for this workload to make sure it can be compiled and executed. As a result, mixing the use of auto_scheduler and AutoTVM schedules is also supported.

The implementation is ready for review in the PR:

https://github.com/apache/incubator-tvm/pull/6903





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/16) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/e5b55f23fe8a5cf69415e0839345bb7b8b5123b929eccc019b028836681fb0fe).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

1. We haven't figured out the plan yet, but mixing them up is definitely a trend.

2. In order to make task extraction and schedule application align, we follow the same flow as building a model to extract tasks. Both AutoTVM and auto_scheduler leverage this approach.





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/10) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/63f7438720835b184e8de2d96a2e2cbe3ef15dc29f702e95059ce1f4344e661a).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by Tristan Konolige via Apache TVM Discuss <no...@discuss.tvm.ai>.

I'm not super familiar with autotvm and auto scheduling, but I've got a couple questions:
1. What is the interaction between autoscheduler and autotvm in the future. Will we be unifying the user api for autotvm and auto scheduling? Can you mix auto scheduling and autotvm?
2. Why is the `GraphRuntimeCodegen` responsible for extracting tasks? Couldn't we just do it in a pass?





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/9) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/577abac8ffe47d8484c7f19c497e14a79a0e17eb70a1a95e5c0bfb5e63550dcf).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by tqchen via Apache TVM Discuss <no...@discuss.tvm.ai>.

```
with auto_scheduler.ApplyHistoryBest(log_file):
    with PassContext(opt_level=opt_level, config={
           "relay.CompileEngine": { use_topi_schedule: False }
    }):
        lib = relay.build(mod, target=target, params=params)
```





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/15) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/ac8450c1f7c719d1f68bd44a36916cddf923feaf1f639d61f4b760966535d1e7).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

So you meant the use case would be like the following?

```python
with auto_scheduler.ApplyHistoryBest(log_file):
    with PassContext(opt_level=opt_level, config={use_topi_schedule: False}):
        lib = relay.build(mod, target=target, params=params)
```





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/14) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/328e929aef7d8157c1d12d069a5555fe6f346a4490c205f3e780e1baa638abab).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

Thanks for the comments.

> May I ask how the graph ends up with a  `nn.conv2d + nn.relu + nn.conv2d + nn.relu`  ? Is the graph going through a BYOC kind of partitioning (sorry if the question is naive)?

There is nothing to do with BYOC. My point is that Ansor opens the door to subgraph-level scheduling, but the current Relay fusion mechanism limits a task to a subgraph with a single reduce op. With this RFC, people can investigate fusion strategies without any limitations. For example, one of our talk in the upcoming TVM conference (title: Graph-Level Scheduling Optimization with Polyhedral Analysis for Tensor Programs) will introduce our initial efforts on investigating the opportunities of fusing reduce ops. This kind of researches will be hard to proceed without this RFC.

> As for S1 vs S2, could we do both? Use an heuristic like “ignore the task without any call node” and then let the task scheduler judge if it needs to spend time on the task?

Yes we could. This issue is actually about the flexibility vs. user experience. If we filter tasks during extraction process, users will see fewer tasks (and may feel more comfortable), but advance users may want to see all tasks to perform other things such as improving the task scheduler, etc.





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/3) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/dfb458c2c796a3e0936fdbb165c337772de6f2860f8a0402f23c113dce7ab993).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by Giuseppe Rossini via Apache TVM Discuss <no...@discuss.tvm.ai>.

Hi @comaniac, 

May I ask how the graph ends up with a `nn.conv2d + nn.relu + nn.conv2d + nn.relu` ? Is the graph going through a BYOC kind of partitioning (sorry if the question is naive)?

As for S1 vs S2, could we do both? Use an heuristic like "ignore the task without any call node" and then let the task scheduler judge if it needs to spend time on the task?

Thanks





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/2) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/08198256fdd8730ea14f00dfebb803035f3411d12abf9210f7ccfa0178cd95d3).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

We haven’t planned that far yet, as currently we lower a Relay function to a TE compute, which relies on Relay op strategy to map Relay ops to TOPI computes.

I’m not familiar with custom Relay ops, but it would be great if you have any suggestion that could make this RFC potentially work for custom Relay ops in the future (or even just follow up PRs). My current subgraph workloads include zero or more than one reduce ops, we can definitely seek for collaborations in this direction :)





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/8) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/25661e8aa7e6f81d808f8d4adbab207f55c7a5bbd313a1556954bce3dcb76702).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by moderato via Apache TVM Discuss <no...@discuss.tvm.ai>.

This is interesting work. I'm curious if the plan is that in the future auto_scheduler would not rely on any new custom Relay ops for tuning subgraphs like the example that is shown, i.e. it directly tunes a primitive function as the user designates?

I'm working on a similar thing of subgraph tuning. Now it runs with topi schedules + pattern matching + customer Relay ops, without auto_scheduler being involved yet. I'm looking to working on more complex subgraphs, and would love to know if there's a chance to collaborate!





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/7) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/7ee0d3e4888f88d03c3d4c979860b155629c0eb8419294787cc401144dd9e7fa).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by tqchen via Apache TVM Discuss <no...@discuss.tvm.ai>.

I agree it could be part of the PassContext, but perhaps not at the top level as opt_level, but more as a sub-level attribute, like the other attributes in loop unrolling





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/13) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/06ef3049a8fe4375310ffecd2623474e3c7338df6bd4910033bd94d70acd3b11).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by "Cody H. Yu via Apache TVM Discuss" <no...@discuss.tvm.ai>.

Here are some more details about the interface change in this RFC. The new added `use_topi_schedule` flag is propagated from the compile engine to `relay._build`. As a result, this actually doesn't expose to users. The use cases are the following:

1. Use TOPI schedule with fallback values (same as current).
    ```python
    with PassContext(opt_level=opt_level):
        lib = relay.build(mod, target=target, params=params)
    ```

2. Use TOPI schedule with AutoTVM tuning logs (same as current).
    ```python
    with autotvm.apply_history_best(log_file):
        with PassContext(opt_level=opt_level):
            lib = relay.build(mod, target=target, params=params)
    ```

3. Extract auto_scheduler tasks. It calls `GraphRuntimeCodegen(use_topi_schedule=False)` to launch the compile engine in order to lower the Relay functions to TE compute DAGs.
    ```python
    tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)
    ```

    In `extract_tasks`:

    ```python
    with transform.PassContext(opt_level=3):
        opt_mod, _ = relay.optimize(mod, target, params)
        grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target, use_topi_schedule=False)
        grc.codegen(opt_mod["main"])
    ```


4. Use auto_scheduler tuning logs. In `relay.build`, it invokes `relay._build(use_topi_schedule=False)` because it finds the `auto_scheduler.DispatchContext` is not `None`, meaning that users want to apply the auto_scheduler log.
    ```python
    with auto_scheduler.ApplyHistoryBest(log_file):
        with PassContext(opt_level=opt_level):
            lib = relay.build(mod, target=target, params=params)
    ```

As a result, the changes are hid from an end-user's point of view. On the other hand, putting this flag to the `PassContext` results in the change of case 3 and case 4:

3. In `extract_tasks`, we can still add the flag for users.

    ```python
    with transform.PassContext(opt_level=3, use_topi_schedule=False):
        opt_mod, _ = relay.optimize(mod, target, params)
        grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
        grc.codegen(opt_mod["main"])
    ```

4. Users have to manually add the flag anyways.
    ```python
    with auto_scheduler.ApplyHistoryBest(log_file):
        with PassContext(opt_level=opt_level, use_topi_schedule=False):
            lib = relay.build(mod, target=target, params=params)
    ```

IMHO, this changes the interface more than the current solution.





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/5) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/7fd1a3f875996b223bc241cc6d432f2789adb884f70db0b6e70b900badb73a44).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by Zhi via Apache TVM Discuss <no...@discuss.tvm.ai>.

This looks okay to me. But I have one comment because this sounds like we need to add one more argument to the build interface which users may not need to know the details. Another possible option is that we can bake it into `PassContext` as a config. However, I understand that this configure is really not a pass config. @tqchen do you have any better suggestion?





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/4) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/627ecdecfda1f8035585a6c5e472b10858d8dbbc7f63c8469e65c03fd974ada0).

[Apache TVM Discuss] [Development/RFC] [RFC] A general task extraction mechanism for auto_scheduler

Posted by Junru Shao via Apache TVM Discuss <no...@discuss.tvm.ai>.

CC: @Hzfengsy @spectrometerHBH if you guys are interested





---
[Visit Topic](https://discuss.tvm.apache.org/t/rfc-a-general-task-extraction-mechanism-for-auto-scheduler/8444/6) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/154aed9447749420f4b4b195cf172b442f9bf3920ecf988c8504843d7ed558eb).