[Runtime][WIP] Add prototype Relay AoT compiler directly into TVM #6219

slyubomirsky · 2020-08-06T02:20:17Z

In reference to this RFC, this PR is intended to incorporate the existing external Relay ahead-of-time (AoT) compiler, which was primarily written by @MarisaKirisame, into TVM. To start, I am simply including most of the files from the AoT compiler repo nearly verbatim, though the interfaces should be changed to better adhere to the high-level vision for TVM (especially since the initial code comes from a research prototype).

The prototype AoT compiler operates by translating Relay ASTs directly into C++ code and using TVM's JIT compiler to register all primitive functions (i.e., the C++ code calls into TVM's operator cache to handle operators). This results in producing a C++ file and requires calling the system's C++ compiler (in the prototype, assuming it to be clang).

I would be curious to hear others' thoughts (e.g., @jroesch @weberlo @tqchen) about how this compiler can be better integrated into TVM's systems. Based on the discussion in the RFC, it sounds that the interface should be made to take an IRModule and produce a runtime module that can call the compiled functions. Ideally the system could be made modular to allow for target languages other than C++

tqchen · 2020-08-06T15:27:26Z

Thanks @slyubomirsky . Some highlevel comments.

First of all, AOT should not be part of the runtime. Runtime contains the minimum set of things we need to execute the program. Most of the AOT logic are actually part of target translation phase. Taking the current organization of relay into account, they should be moved to relay/backend(eventually a better place might be target).

Of course there are additional runtime features and wrappers needed to run the program compiled from AOT. From the interface point of view, we should remove these wrapper code and completely rely on the PackedFunc and runtime.Module interface.

So the AOT compilation should take in an IRModule and output a runtime.Module, which contains the functions necessary to run the generated program. Ideally the runtime.Module should contain similar interface with other compiled programs, such as the vm and the graph runtime

slyubomirsky · 2020-08-06T22:52:58Z

Thank you for the suggestions! I will move the files to backend. I will see about getting PackedFuncs for operators directly instead of using the JIT to register them (this might fix some of the other weird bugs we've seen in the research prototype).

slyubomirsky · 2020-08-07T03:40:33Z

Hm, if the CI fails because TVM_HOME can't be assumed to be defined, is there any way for the produced C++ files to refer to TVM data structures definitions like NDArrays? I will see if there is any way to reference the compiled TVM so file from aot.py

edit: I guess I can use find_include_path() and find_lib_path() from libinfo

weberlo · 2020-08-07T21:02:43Z

python/tvm/relay/backend/aot/aot.py

+    if lib_path is None:
+        lib_path = os.curdir
+
+    debug_source_path = os.path.join(lib_path, 'source.cc')


can you put this functionality behind a debug flag?

weberlo · 2020-08-07T21:02:54Z

python/tvm/relay/backend/aot/aot.py

+    with open(debug_source_path, 'w') as source_file:
+        source_file.write(source)
+
+    # with tempfile.TmporaryDirectory() as tmpdir:


manupak · 2020-08-11T12:01:48Z

@slyubomirsky Thanks for the work! -- Some very high-level comments.

IIUC, this is bypassing the TIR lowering as it stands today. Thus, possibly losing the benefits TIR scheduling based optimizations in AOT compilation path. I just wanted to know what the roadmap looks like if we are to re-target the AOT codegen at the TIR level. Would it be an incremental work on top of this ? or would it require a complete re-write ?

MarisaKirisame · 2020-08-11T17:22:11Z

@manupa-arm the primitive function is still lowered to TIR - we only compile Relay fragment to C++. This is in accordance to how Relay had work for Interpreter/VM - Relay fragment get handle separately where primtive function get lowered to TIR and handled by TVM.

If you are talking about lowering everything to TIR, the biggest problem is the design implication. It will make TIR less like fortran and way more like SML.

manupak · 2020-08-11T20:20:40Z

@MarisaKirisame , thanks for the clarification!

manupak

Some minor comments, see if you agree.

manupak · 2020-08-12T09:05:49Z

python/tvm/relay/backend/aot/aot.py

+        return res if not record_time else (res, end - begin)
+    return _wrapper
+
+def compile_prog(func, mod, ctx, tgt, name='default', record_time=False):


Would it be a good idea to make the func the "main" inside mod ? So that we only needs to pass the mod with a "main".

Yes, this was definitely a flaw in the research prototype that we never took the time to correct. I agree that using the main function would be a much more sensible convention

manupak · 2020-08-12T09:34:21Z

python/tvm/relay/backend/aot/aot.py

+        f"-L{TVM_PATH}/build"
+    ]
+
+    if system == 'Darwin':


The compiler config could be a class variable -- a dictionary to be precise.
Maybe look that up for the flags and refactor the rest as they are mostly same.
Maybe a comment explaining the special cased flags for "Darwin" could be useful.

Yes, I think this should be written in a more maintainable manner. I will see if it can be made to programmatically match up with TVM's own C++ build configuration, for example

manupak · 2020-08-12T10:25:28Z

python/tvm/relay/backend/aot/aot.py

+    _LIB_COUNTER += 1
+    return lib_name, packed_name
+
+def _mk_wrapper(func, ctx, constants, record_time):


This wrapper seems to enable runtime perf. measurement; Maybe add a comment describing the functionality;
Also it would be better make record_time, default to False.
[Suggestion] make this a decorator and remove the wrapper from the implemetation and use the decorator where the performance measurement is needed

I may remove the time recording feature entirely for now, as it was something we included ad hoc for a single experiment

manupak · 2020-08-12T10:49:24Z

python/tvm/relay/backend/aot/aot.py

+    func = compiler.visit(func)
+    lib_name, packed_name = lib_and_func_name(name)
+    constants, source_code = to_source.to_source(func, compiler.gv_map, ctx, packed_name)
+    lib_name = f"librelay_aot_{_LIB_COUNTER}.so"


Is it possible for user to optionally give the paths for the artifacts : .so and .cc ?

That would be a good option, I will include that in my next revision.

manupak · 2020-08-12T11:31:11Z

[Clarification Question] How are the memory allocated for tensors -- in-between primitive functions -- ? Please point me to the code if its there -- It seems I have missed that. Do you do storage_id usage optimizations such as done in graph plan memory ?

manupak · 2020-08-12T12:09:11Z

python/tvm/relay/backend/aot/aot.py

+
+    must_run_process(["clang-format", "-i", debug_source_path])
+
+    system = os.uname()[0]


Could we make this an argument to the compile cpp_function?
In future, it would be easier to enable cross-compilation.

slyubomirsky · 2020-08-15T04:18:30Z

[Clarification Question] How are the memory allocated for tensors -- in-between primitive functions -- ? Please point me to the code if its there -- It seems I have missed that. Do you do storage_id usage optimizations such as done in graph plan memory ?

There are not, to my knowledge (@MarisaKirisame wrote most of the compiler), any memory planning optimizations in the AoT prototype, though it would certainly be a good addition. I never specifically looked into the memory allocation behavior (it was an area we ignored in the prototype altogether), but I believe allocations simply happen when the NDArray constructor is called in the generated code -- I will check that.

jroesch · 2021-08-27T18:20:29Z

I am doing triage on old PRs, going to close this, please feel free to follow up if you would like to still merge these changes. Thanks for your contributions!.

Add prototype Relay AoT compiler directly into TVM

a0a0e69

tqchen self-assigned this Aug 6, 2020

tqchen added the status: RFC label Aug 6, 2020

slyubomirsky added 3 commits August 6, 2020 16:33

Move AOT files to relay/backend

fd3579f

Add ASF headers to AoT files

8961bcb

Fix tons of pylint errors

44358e2

weberlo mentioned this pull request Aug 7, 2020

[RFC][μTVM] Bringing TVM to Bare-Metal Devices #2563

Closed

6 tasks

weberlo suggested changes Aug 7, 2020

View reviewed changes

manupak requested changes Aug 12, 2020

View reviewed changes

manupak reviewed Aug 12, 2020

View reviewed changes

tqchen changed the base branch from master to main October 11, 2020 18:21

jroesch closed this Aug 27, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Runtime][WIP] Add prototype Relay AoT compiler directly into TVM #6219

[Runtime][WIP] Add prototype Relay AoT compiler directly into TVM #6219

slyubomirsky commented Aug 6, 2020

tqchen commented Aug 6, 2020

slyubomirsky commented Aug 6, 2020

slyubomirsky commented Aug 7, 2020 •

edited

Loading

weberlo Aug 7, 2020

slyubomirsky Aug 10, 2020

weberlo Aug 7, 2020

manupak commented Aug 11, 2020

MarisaKirisame commented Aug 11, 2020

manupak commented Aug 11, 2020

manupak left a comment

manupak Aug 12, 2020

slyubomirsky Aug 15, 2020

manupak Aug 12, 2020

slyubomirsky Aug 15, 2020

manupak Aug 12, 2020 •

edited

Loading

slyubomirsky Aug 15, 2020

manupak Aug 12, 2020

slyubomirsky Aug 15, 2020

manupak commented Aug 12, 2020

manupak Aug 12, 2020

slyubomirsky commented Aug 15, 2020

jroesch commented Aug 27, 2021


		must_run_process(["clang-format", "-i", debug_source_path])

		system = os.uname()[0]

[Runtime][WIP] Add prototype Relay AoT compiler directly into TVM #6219

[Runtime][WIP] Add prototype Relay AoT compiler directly into TVM #6219

Conversation

slyubomirsky commented Aug 6, 2020

tqchen commented Aug 6, 2020

slyubomirsky commented Aug 6, 2020

slyubomirsky commented Aug 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak commented Aug 11, 2020

MarisaKirisame commented Aug 11, 2020

manupak commented Aug 11, 2020

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

manupak Aug 12, 2020 • 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

manupak commented Aug 12, 2020

Choose a reason for hiding this comment

slyubomirsky commented Aug 15, 2020

jroesch commented Aug 27, 2021

slyubomirsky commented Aug 7, 2020 •

edited

Loading

manupak Aug 12, 2020 •

edited

Loading