[TOPI] Custom schedule for standalone transpose in cuda by tkonolige · Pull Request #8030 · apache/tvm

tkonolige · 2021-05-12T18:56:00Z

This PR adds an optimized schedule for transpose if the transpose is not fused into anything else.

@altanh @junrushao1994

altanh

Some minor nits about code organization and the test, but otherwise looks good and is a welcome addition. For other reviewers: offline testing of this transpose schedule (when applicable) showed about a 2x improvement over the baseline injective transpose schedule on BERT Large training workloads. I do wonder how we could select whether or not to opt in to this vs trying to tune the original injective schedule, but I think in most cases this new kernel will probably be faster.

altanh · 2021-05-14T18:58:59Z

+    Dispatches to and optimized schedule if the transpose is standalone (not fused).
+    """
+    warp_size = int(Target.current(allow_none=False).thread_warp_size)
+    if (


is there a more principled way to do this? like maybe with an OpStrategy or something

As far as I can tell, there is not a better way to do this. There is a way to add implementations based on input sizes, but these are not on a per-target basis. If you know a better way, let me know.

altanh · 2021-05-14T19:00:32Z



-def schedule_cuda_transpose(s, out):
+def schedule_transpose(outs):


feels a bit weird to have this in sparse.py

moved to transform.py

altanh · 2021-05-14T19:02:16Z

+    r = np.random.rand(*shape)
+    tvm.testing.assert_allclose(ex.evaluate()(r).asnumpy(), np.transpose(r))
+
+    # make sure schedule does not fire here


is this a TODO? Also I wonder if it would be good to parametrize the test shape by warp size (rather than hard coding) for future proofing

It is more like a wish. Ideally we would be able to know which schedules were used, but there is to way to introspect on what was used. I've updated the comment to reflect this.

altanh · 2021-05-14T19:03:36Z


    # filter out non-packed conv2d task
-    tasks = list(filter(lambda t: len(t.args[0][1]) > 4, tasks))
+    tasks = list(filter(lambda t: len(t.args[0][1]) > 4 and "conv" in t.name, tasks))


what happened here, did this transpose change introduce a new task or something?

actually, no, but this check makes sure anyways.

Isn't the new added schedule not tunable? Or is there any concern of adding knobs?

We may want to tune it in the future.

altanh

LGTM, the only thing I'm wondering about is if someone (for whatever reason) really wanted to tune the default injective schedule for transpose, is there any way to allow that?

cc @comaniac for additional review (feel free to tag more relevant reviewers)

comaniac · 2021-05-18T22:06:46Z

LGTM, the only thing I'm wondering about is if someone (for whatever reason) really wanted to tune the default injective schedule for transpose, is there any way to allow that?

cc @comaniac for additional review (feel free to tag more relevant reviewers)

There's no reason to tune inject schedule and you basically cannot do it because injective schedule doesn't have AutoTVM knobs for tuning.

comaniac · 2021-05-18T22:13:43Z

+    # We want to make sure schedule does not fire here, but there is no way of
+    # inspecting which schedules were used.


Like this comment mentions, there is no way of inspecting which schedules were used, so it seems to me that the difference between this test and test_transpose is the workload in this test includes add to test the case of fusion. Accordingly, could we just extend test_transpose?

We could. I like to keep it separate so the intention is known.

Fair enough. Then it might be better to name it test_transpose_fuse or something like that (nit).

switched to test_transpose_unfused_schedule

comaniac · 2021-05-18T22:16:18Z


    # filter out non-packed conv2d task
-    tasks = list(filter(lambda t: len(t.args[0][1]) > 4, tasks))
+    tasks = list(filter(lambda t: len(t.args[0][1]) > 4 and "conv" in t.name, tasks))


Isn't the new added schedule not tunable? Or is there any concern of adding knobs?

comaniac

I don't have other comments. LGTM.

* [TOPI] Custom schedule for standalone transpose in cuda * check if input is not Any * fix vta test * check input shape * fix injective * move transpose out of sparse.py * update comments, use warp size * missspelled transform * formatting * rename test * comment * fix tests

tkonolige added 5 commits May 12, 2021 11:55

[TOPI] Custom schedule for standalone transpose in cuda

b5478be

check if input is not Any

b57aade

fix vta test

09a65f2

check input shape

95bf430

fix injective

c0d051f

altanh reviewed May 14, 2021

View reviewed changes

tkonolige added 3 commits May 17, 2021 11:23

move transpose out of sparse.py

cdd9f5e

update comments, use warp size

5804cc7

missspelled transform

f4982a2

altanh approved these changes May 18, 2021

View reviewed changes

formatting

410f312

comaniac reviewed May 18, 2021

View reviewed changes

comaniac approved these changes May 18, 2021

View reviewed changes

tkonolige added 3 commits May 18, 2021 15:31

rename test

804a8cd

comment

30fcc8e

fix tests

676a342

areusch merged commit 28ea03c into apache:main May 20, 2021

masahi mentioned this pull request May 21, 2021

[Target] Add a default warp size 1 for vulkan and opencl #8109

Merged

junrushao mentioned this pull request Nov 1, 2021

Apache TVM v0.8 Release Note Candidate #9416

Closed



		def schedule_cuda_transpose(s, out):
		def schedule_transpose(outs):

		# We want to make sure schedule does not fire here, but there is no way of
		# inspecting which schedules were used.

Conversation

tkonolige commented May 12, 2021

Uh oh!

altanh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

altanh left a comment

Choose a reason for hiding this comment

Uh oh!

comaniac commented May 18, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

comaniac left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants