Move the allocates of AoT codegen to be TVMBAWs by manupak · Pull Request #9065 · apache/tvm

manupak · 2021-09-21T16:30:08Z

This commit introduces changes to aot_executor_codegen.cc
to place tir.allocate to use storage_scope = 'global.workspace'.
The lower_tvm_builtin pass is modified slightly to generate
TVMBAW calls.

manupak · 2021-09-21T16:30:37Z

cc : @areusch @mbs-octoml @tqchen

areusch · 2021-09-21T18:15:04Z



+def test_aot_codegen_backend_alloc_workspace_calls():
+ dtype = "float32"


could you add a brief docstring explaining what this test asserts?

areusch · 2021-09-21T18:15:45Z

+ inputs = OrderedDict([("data", i_data), ("weight", w1_data)])
+ output_list = generate_ref_data(mod, inputs)
+
+ compiled_runtime_modules = compile_models(


want to assert that all the tir.allocate nodes are correctly tagged somewhere?

Ah it is a bit cumbersome to do that :), Instead I used relay in primitive form so its clear that main function should only have three allocates.

@mbs-octoml @denise-k can we add a tracking/cleanup task to make this kind of assert easier to write? And flag to cleanup this test?

@areusch roadmap item and task tracking have been created.

mbs-octoml

LGTM (FWIW). Thanks I learned from this one.

mbs-octoml · 2021-09-21T19:59:35Z

 for (auto kv : storage_device_map_) {
 for (auto sid : kv.second->storage_ids) {
+ // The buffer_var is created with storage_scope to be global.workspace to be serviced by
+ // TVMBAWs, explicitly. The reasoning being the executor allocates should be serviced by


For poor folks like me: use TVMBackendAllocWorkspace at least once.

jroesch · 2021-09-22T06:08:26Z

@csullivan @mbrookhart @junrushao1994 can you guys sign off on this one as well to make sure this won't break anyone?

jroesch · 2021-09-22T06:08:57Z

LGTM, will just ping other core compiler folks to make sure this works for them.

csullivan · 2021-09-22T16:39:08Z

+ True,
+ )
+
+ source = compiled_runtime_modules[0].lib.imported_modules[0].get_source()


The AoT codegen for main ends up as an imported module? Naively I would expect the TVMBackendAllocateWorkspace calls in the imported_modules list to be intra-op only, e.g. for the conv2d, not AoT main.

I can understand the reasoning but the current flow just creates per target IRModules just before runtime.Modules are created. Therefore all host_target (i.e. CPU) PrimFuncs end up in a single runtime.Module.

manupak

I've addressed the comments now. PTAL.

manupak · 2021-09-22T17:25:48Z

+ inputs = OrderedDict([("data", i_data), ("weight", w1_data)])
+ output_list = generate_ref_data(mod, inputs)
+
+ compiled_runtime_modules = compile_models(


Ah it is a bit cumbersome to do that :), Instead I used relay in primitive form so its clear that main function should only have three allocates.

manupak · 2021-09-22T17:27:32Z

+ True,
+ )
+
+ source = compiled_runtime_modules[0].lib.imported_modules[0].get_source()


I can understand the reasoning but the current flow just creates per target IRModules just before runtime.Modules are created. Therefore all host_target (i.e. CPU) PrimFuncs end up in a single runtime.Module.

manupak · 2021-09-23T08:43:11Z

This is failing CI on an unrelated error (I think due to : #9013). I think #9076 should fix it (Thanks @mbs-octoml ). Therefore it would be great if we can get an approval if this looks good as it is blocking our work -- so that once #9076 lands we could merge this.

jroesch · 2021-09-23T16:59:56Z

@manupa-arm @junrushao1994 and I landed the other patch, if you guys are still working today can you rebase before you log off and I will try and merge before end of day pacific time?

manupak · 2021-09-23T17:04:52Z

Thanks! Done.

This commit introduces changes to aot_executor_codegen.cc to place tir.allocate to use storage_scope = 'global.workspace'. The lower_tvm_builtin pass is modified slightly to generate TVMBAW calls. Change-Id: Iba4ba437c1431c5197bf11abb826e03807bbcf66

*Adding more comments and descriptions *Modified the test case to use primitive relay Change-Id: Ia18a169d94bded3f81af7b3081c7d1ac29c669bc

* Move the allocates of AoT codegen to be TVMBAWs This commit introduces changes to aot_executor_codegen.cc to place tir.allocate to use storage_scope = 'global.workspace'. The lower_tvm_builtin pass is modified slightly to generate TVMBAW calls. Change-Id: Iba4ba437c1431c5197bf11abb826e03807bbcf66 * Move the allocates of AoT codegen to be TVMBAWs *Adding more comments and descriptions *Modified the test case to use primitive relay Change-Id: Ia18a169d94bded3f81af7b3081c7d1ac29c669bc

manupak requested review from Hzfengsy, MarisaKirisame, ZihengJiang, anijain2305, areusch, comaniac, jroesch, junrushao, jwfromm, kparzysz-quic, masahi, mbrookhart, merrymercy, slyubomirsky, tqchen, vinx13, wweic, yzhliu and zhiics as code owners September 21, 2021 16:30

manupak mentioned this pull request Sep 21, 2021
[Bug] BuiltinLower has hard-coded heuristic for alloca which is not appropriate for all kDLCPU target devices #9022
Closed

areusch reviewed Sep 21, 2021
View reviewed changes

mbs-octoml approved these changes Sep 21, 2021
View reviewed changes

csullivan reviewed Sep 22, 2021
View reviewed changes

manupak force-pushed the move_aot_allocates_global_workspace branch from 1f2a365 to db37329Compare September 22, 2021 17:24

manupak commented Sep 22, 2021
View reviewed changes

manupak force-pushed the move_aot_allocates_global_workspace branch from db37329 to 321ba2cCompare September 22, 2021 20:33

manupak force-pushed the move_aot_allocates_global_workspace branch from 321ba2c to 14f3cf4Compare September 23, 2021 17:04

jroesch approved these changes Sep 23, 2021
View reviewed changes

denise-k mentioned this pull request Sep 23, 2021
[CLEANUP] Clean up existing structural tests when structural testing infrastructure is created #9096
Closed
1 task

manupak added 2 commits September 24, 2021 06:46

Move the allocates of AoT codegen to be TVMBAWs
bfa7944
*Adding more comments and descriptions *Modified the test case to use primitive relay Change-Id: Ia18a169d94bded3f81af7b3081c7d1ac29c669bc

manupak force-pushed the move_aot_allocates_global_workspace branch from 14f3cf4 to bfa7944Compare September 24, 2021 05:46

leandron approved these changes Sep 24, 2021
View reviewed changes

leandron merged commit 7c6a334 into apache:mainSep 24, 2021



		def test_aot_codegen_backend_alloc_workspace_calls():
		dtype = "float32"

Uh oh!

Conversation

manupak commented Sep 21, 2021

Uh oh!

manupak commented Sep 21, 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!

mbs-octoml left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jroesch commented Sep 22, 2021

Uh oh!

jroesch commented Sep 22, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

manupak left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

manupak commented Sep 23, 2021

Uh oh!

jroesch commented Sep 23, 2021

Uh oh!

manupak commented Sep 23, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants