Skip to content

Fix AutoTVM int8 vnni dense task extraction problem #11699

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mengceng15 wants to merge 2 commits into from
Closed

Fix AutoTVM int8 vnni dense task extraction problem #11699

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a5847eb
fix autotvm int8 dense vnni extraction
mengceng15 Jun 13, 2022
26337c5
pylint
mengceng15 Jun 14, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions python/tvm/autotvm/task/relay_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,10 @@ def extract_from_multiple_program(mods, params, target, target_host=None, ops=No
assert isinstance(
mod, tvm.IRModule
), "only support relay Module or Function to be tuned"

with target:

@elvin-n elvin-n Jun 14, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change have to break backward compatibility. The current ideology of AlterOpLayout and autotvm assumes that tuning does not get converted layout, just plain ones and it is responsibility of schedule to create transformed compute even for flattened data and it is responsibility of AlterOplLyout to understand that flattened data in log stands for blocked one and substitute workload to apply tuned statistics to blocked configuration.

For example here is how conv2d compute get 4d tensor, figure out that it is autotvm session and override them by 5d tensors

And here is a place where AlterOpLayout understnabds that it should correct workload to handle flattened initial tensors to packed one when AlterOpLayout is executed during network compilation.

Approach looks hacky but how it is implemented now and change in the tuning now will make previous tuned kernels not applicable for newer versions

BTW, I tried your change and it broke even tuning in cpu fp32 mode

@mengceng15 mengceng15 Jun 14, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@elvin-n Thanks for the detailed explanation. I just realize that this change could introduce serious compatibility problems due to the autotvm tuning design.

Initially, I followed @masahi 's method to apply alter op layout before the task extraction. Everything worked well. However, it could be a bit confusing to the user. Seems there is no document or hint mentioning that this workaround should be used. My intention to raise this PR was to clean up the confusion.

@masahi masahi Jun 14, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The "right" fix is to add autotvm.GLOBAL_SCOPE.in_tuning thing in the VNNI dense op. I didn't want to do that in my PR, but you are welcome to do so if this would solve your problem.

@mengceng15 mengceng15 Jun 14, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@masahi Thank you for the suggestion.

mod = tvm.relay.transform.AlterOpLayout()(mod)

relay.backend.te_compiler.get().clear()
# wrap build call in thread to avoid multiprocessing problems
build_thread = threading.Thread(target=_lower, args=(mod, target, param))
Expand Down
16 changes: 16 additions & 0 deletions tests/python/relay/test_autotvm_task_extraction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,22 @@ def get_net(batch, in_dim, out_dim, dtype, out_dtype):
assert len(tasks) == 1 and tasks[0].name == "dense_int8.cuda"


def test_task_extraction_for_dense_int8_vnni():
target = "llvm -mcpu=cascadelake"

def get_net(batch, in_dim, out_dim, dtype_input, dtype_weight, out_dtype):
data = tvm.relay.var("data", shape=[batch, in_dim], dtype=dtype_input)
weight = tvm.relay.var("weight", shape=[out_dim, in_dim], dtype=dtype_weight)
out = relay.nn.dense(data, weight, out_dtype=out_dtype)
mod, params = relay.testing.create_workload(out)
return mod, params

mod, params = get_net(1, 16, 32, "uint8", "int8", "int32")
tasks = autotvm.task.extract_from_program(mod, target=target, params=params)
assert len(tasks) == 1 and tasks[0].name == "dense_vnni.x86"


if __name__ == "__main__":
test_task_extraction()
test_task_extraction_for_dense_int8_cuda()
test_task_extraction_for_dense_int8_vnni()