[RFC]TECompiler: Staged refactor and removal of compile engine #7518

csullivan · 2021-02-24T19:23:08Z

This is a staged removal of the use of compile engine. The motivation is to bring TIR compilation into the main flow of the compiler rather than producing and compiling it via a callback into the compile engine. By replacing Relay primitive function calls with TIR prim function calls that contain the lowered TIR we enable,

An intermediate stage in the lowering process where Relay and TIR coexist.
The ability to add passes at this intermediate stage,
- For example memory planning which can infer user provided information from TE and the resulting TIR.

We are starting with a proof of concept by refactoring the GraphRuntimeCodegen to use an introduced TIR/TE compiler instead of the compile engine directly. In the new flow,

The TE/TIR compiler lowers TE in the LowerTensorExpr pass
Replaces relay.Function(attr:primitive) with a PrimFnCall that contains the lowered TIR
Runs GraphPlanMemory
Finally runs GraphRuntimeCodegen::VisitExpr to lower to graph JSON

[RFC][Relay] TECompiler: Rewrite existing compile engine to match updated compiler flow

jroesch · 2021-03-24T08:08:01Z

Need to port fix from #7703 but otherwise ready for review.

jroesch · 2021-03-24T08:09:38Z

Modulo some left over polish work and documentation I think this is ready for review @icemelon9 @comaniac @csullivan @tkonolige @rkimball @junrushao1994 @areusch @mehrdadh

jroesch · 2021-03-24T08:19:18Z

cc @mbaret

src/relay/backend/graph_runtime_codegen.cc

src/relay/backend/te_compiler_cache.cc

src/relay/ir/function.cc

src/runtime/graph/graph_runtime.cc

tests/python/relay/test_backend_graph_runtime.py

src/driver/driver_api.cc

areusch

did an initial pass

src/driver/driver_api.cc

src/relay/backend/graph_plan_memory.cc

areusch · 2021-03-29T23:11:04Z

src/relay/backend/graph_runtime_codegen.cc

    targets_ = targets;
  }

  LoweredOutput Codegen(relay::Function func) {
-    auto pf = GetPackedFunc("relay.backend.GraphPlanMemory");
-    storage_device_map_ = (*pf)(func);
+    // Jared: why do we do this? just call C++ API.


at some point we should support overriding this. it seems like you'd need to build a C++ impl right now to do that, so I don't know that we need this hook now. but, we should plan to put some customization back in the future?

I think we should ask what do we allow customization of what and where ? I think allowing people to wholesale replace the lowering/memory/planning/etc in an ad-hoc way "seems bad", which effectively is all of these hooks (that I'm removing) do right now. They allow you to just arbitrarily overload parts of the compiler with custom code which is pretty non-compositional and hard to reason about since in certain cases, i.e Python loaded or not the behavior of the compiler is completely different.

I don't think it makes sense to allow people to overload GraphPlanMemory with its current implementation. If GraphPlanMemory were to support memory pools and tensor pinning, then I think there will be cases where people want to implement custom algorithms to place tensors in different memory pools.

I also think the reason this "seems bad" is that there is no good interface/data structure defined here. There's no reason we can't define one--in fact, I think we are hurting compiler readability by not doing so. Here we are keeping the interface to a sort of Pythonic "informal set of lists" in place of a proper data structure. Here and downstream of here, there are many instances where both TVM code and people trying to integrate with TVM code would benefit from effectively exporting StorageToken as a user-facing data structure and clearly defining its meaning.

Yeah I think we need to move up the stack to address this, because everyone is only thinking about their own memory planning needs and not really looking at the large picture.

src/relay/backend/te_compiler_cache.cc

tests/python/relay/test_backend_graph_runtime.py

src/relay/backend/graph_executor_codegen.cc

src/relay/backend/te_compiler.cc

src/runtime/graph_executor/graph_executor.cc

src/target/llvm/llvm_module.cc

tests/python/relay/test_backend_graph_executor.py

src/relay/backend/te_compiler.h

src/relay/backend/te_compiler_cache.h

src/relay/backend/te_compiler.h

src/relay/backend/te_compiler_cache.cc

src/relay/backend/graph_executor_codegen.cc

include/tvm/relay/attrs/annotation.h

python/tvm/auto_scheduler/relay_integration.py

python/tvm/micro/model_library_format.py

src/relay/backend/graph_executor_codegen.cc

src/relay/backend/te_compiler.cc

src/runtime/graph_executor/graph_executor.cc

src/target/llvm/llvm_module.cc

tests/python/relay/test_op_fast_math.py

tests/python/unittest/test_micro_model_library_format.py

Duplicate the CompileEngine interface. Refactor the graph_runtime_codegen to invoke the new LowerTE pass More changes Things appear to be working Some tracing to get Relay code to flow through too. Disable some assertions as exp. Tweak printing for now Fix a few bugs: (#13) 1. Don't add relay main function to list of lowered TIR functions 2. Don't skip visiting call to relay function in graph runtime codegen Remove debug prints. Start refactoring Split out shared data structures Fix implicit duplicate decl of IsDynamic Clean up handling of name + global prim fn Clean up the code and debug issue introduced by previous hack Clean up the debugging Do C++ lint clean up Update src/relay/backend/graph_executor_codegen.cc Co-authored-by: Chris Sullivan <csullivan@octoml.ai> Clean up handling of external functions Add more error messages More clean up Update src/runtime/graph_executor/graph_executor.cc Co-authored-by: Chris Sullivan <csullivan@octoml.ai> Update src/runtime/graph_executor/graph_executor.cc Co-authored-by: Chris Sullivan <csullivan@octoml.ai> Update src/relay/backend/te_compiler.h Co-authored-by: Haichen Shen <shenhaichen@gmail.com> Update src/relay/backend/te_compiler.h Co-authored-by: Haichen Shen <shenhaichen@gmail.com> Fix CR More CR Format Fix lowering path for C++ Fix tests Remove uncessary change Clean up a few more things CI fix Fix the default context Fix Fix broken test cases Update Fix WIP Clean up storage data structures WIP WIP Fix build errors Remove TVMLower Fix lint Lint again fix black Move UpdateMainWorkspaceSize into te_compiler.cc Fix link errors Formatting Change UpdateMainWorkspaceSize to return Map<String, FunctionInfo> Workaround for GCC 5 error caused by enums in maps (GCC 5 is on i386 CI) Testing how functions should be named Lint Change how function metadata is updated Attempt to update aot_executor_codegen to use new StaticMemoryPlan instead of storage_device_map Pass memory plan through LowerTE into UpdateMainWorkspaceSize so that we don't need to run GraphPlanMemory an extra time Fix return in UpdateMainWorkspaceSize Lint Try to fix UpdateMainWorkspaceSize Fix construction of static memory plan Clean up code while debugging Adding UpdateWorkspaceSize back Add closure + call to UpdateFunctionMetadata (WIP) UpdateFunctionMetadata builds; weird error with device ctx map though. Not sure if it came from this change or something else Add some debugging of UpdateMainWorkspaceSize Starting to move UpdateFunctionMetadata call to use process_fn infra UWhat target should be passed to UpdateFunctionMetadata? UpdateFunctionMetadata is not workinggg Added some comments about UpdateFunctionMetadata for Jared Fix the creation of function metadata Try another stab at cleaning up the information Fix Port StorageInfo and StaticMemoryPlan data structure (apache#8297) Restoring reshape opt Fix tests Caught a nasty typo from Lily, Map::Set does not mutate Format Disable stupid Google style warning Rebase cleanup Formatting Add docstring for storage info Black Post rebase fix Remove prints Disable assert that doesn't make sense for now Fix lint Add copying attrs from relay node to graph node; still need to figure out how to do this in the case of global vars Work with Lily to fix graph attrs Try to figure out where extra arguments are coming from; fix merge passes the profiling test Clean up Fix profile test Remove debugging Add attributes for BYOC uTVM case Format Dumb typo Another fix for byoc Format Fix last 3 failing tests Format Fix final two test cases Format Fix lint Fix again Fix Fix auto scheduler code Fix issue Address CR comment Format

We talked

Duplicate the CompileEngine interface. Refactor the graph_runtime_codegen to invoke the new LowerTE pass More changes Things appear to be working Some tracing to get Relay code to flow through too. Disable some assertions as exp. Tweak printing for now Fix a few bugs: (apache#13) 1. Don't add relay main function to list of lowered TIR functions 2. Don't skip visiting call to relay function in graph runtime codegen Remove debug prints. Start refactoring Split out shared data structures Fix implicit duplicate decl of IsDynamic Clean up handling of name + global prim fn Clean up the code and debug issue introduced by previous hack Clean up the debugging Do C++ lint clean up Update src/relay/backend/graph_executor_codegen.cc Co-authored-by: Chris Sullivan <csullivan@octoml.ai> Clean up handling of external functions Add more error messages More clean up Update src/runtime/graph_executor/graph_executor.cc Co-authored-by: Chris Sullivan <csullivan@octoml.ai> Update src/runtime/graph_executor/graph_executor.cc Co-authored-by: Chris Sullivan <csullivan@octoml.ai> Update src/relay/backend/te_compiler.h Co-authored-by: Haichen Shen <shenhaichen@gmail.com> Update src/relay/backend/te_compiler.h Co-authored-by: Haichen Shen <shenhaichen@gmail.com> Fix CR More CR Format Fix lowering path for C++ Fix tests Remove uncessary change Clean up a few more things CI fix Fix the default context Fix Fix broken test cases Update Fix WIP Clean up storage data structures WIP WIP Fix build errors Remove TVMLower Fix lint Lint again fix black Move UpdateMainWorkspaceSize into te_compiler.cc Fix link errors Formatting Change UpdateMainWorkspaceSize to return Map<String, FunctionInfo> Workaround for GCC 5 error caused by enums in maps (GCC 5 is on i386 CI) Testing how functions should be named Lint Change how function metadata is updated Attempt to update aot_executor_codegen to use new StaticMemoryPlan instead of storage_device_map Pass memory plan through LowerTE into UpdateMainWorkspaceSize so that we don't need to run GraphPlanMemory an extra time Fix return in UpdateMainWorkspaceSize Lint Try to fix UpdateMainWorkspaceSize Fix construction of static memory plan Clean up code while debugging Adding UpdateWorkspaceSize back Add closure + call to UpdateFunctionMetadata (WIP) UpdateFunctionMetadata builds; weird error with device ctx map though. Not sure if it came from this change or something else Add some debugging of UpdateMainWorkspaceSize Starting to move UpdateFunctionMetadata call to use process_fn infra UWhat target should be passed to UpdateFunctionMetadata? UpdateFunctionMetadata is not workinggg Added some comments about UpdateFunctionMetadata for Jared Fix the creation of function metadata Try another stab at cleaning up the information Fix Port StorageInfo and StaticMemoryPlan data structure (apache#8297) Restoring reshape opt Fix tests Caught a nasty typo from Lily, Map::Set does not mutate Format Disable stupid Google style warning Rebase cleanup Formatting Add docstring for storage info Black Post rebase fix Remove prints Disable assert that doesn't make sense for now Fix lint Add copying attrs from relay node to graph node; still need to figure out how to do this in the case of global vars Work with Lily to fix graph attrs Try to figure out where extra arguments are coming from; fix merge passes the profiling test Clean up Fix profile test Remove debugging Add attributes for BYOC uTVM case Format Dumb typo Another fix for byoc Format Fix last 3 failing tests Format Fix final two test cases Format Fix lint Fix again Fix Fix auto scheduler code Fix issue Address CR comment Format Co-authored-by: Jared Roesch <roeschinc@gmail.com>

jroesch changed the title ~~[WIP] Staged refactor and removal of compile engine~~ [RFC][WIP] TECompiler: Staged refactor and removal of compile engine Feb 24, 2021

jroesch mentioned this pull request Feb 24, 2021

[RFC] Tracking Issue for TECompiler #7526

Closed

7 tasks

jroesch force-pushed the remove-compile-engine branch from 077cc4e to 00178dc Compare March 19, 2021 00:39

jroesch changed the title ~~[RFC][WIP] TECompiler: Staged refactor and removal of compile engine~~ [RFC]TECompiler: Staged refactor and removal of compile engine Mar 24, 2021

csullivan marked this pull request as ready for review March 24, 2021 16:03

csullivan mentioned this pull request Mar 24, 2021

[Bug] Missing broadcast_to before batch_matmul for CuBLAS #7730

Closed

tkonolige reviewed Mar 26, 2021

View reviewed changes

areusch previously requested changes Mar 29, 2021

View reviewed changes

jroesch force-pushed the remove-compile-engine branch from 41a090e to 8141aa9 Compare April 7, 2021 17:12

csullivan commented Apr 7, 2021

View reviewed changes

icemelon reviewed Apr 7, 2021

View reviewed changes

jroesch force-pushed the remove-compile-engine branch from 6dab6be to 6fb66c1 Compare April 14, 2021 17:00

areusch added the status: need update need update based on feedbacks label Apr 30, 2021

jroesch force-pushed the remove-compile-engine branch from 05409d0 to a561718 Compare May 3, 2021 20:36

jroesch force-pushed the remove-compile-engine branch from a561718 to 9192543 Compare June 15, 2021 18:10

csullivan force-pushed the remove-compile-engine branch from 9192543 to cdff959 Compare June 15, 2021 20:51

areusch mentioned this pull request Jun 22, 2021

[AOT] Name mangling in AOT #8014

Merged

jroesch force-pushed the remove-compile-engine branch 2 times, most recently from 6452ca4 to f82e60d Compare June 29, 2021 21:07

areusch mentioned this pull request Jun 29, 2021

Decoupling AOT from graph memory planner #8096

Merged

jroesch force-pushed the remove-compile-engine branch from 09a0b14 to a86752c Compare June 29, 2021 23:37

tkonolige reviewed Jun 30, 2021

View reviewed changes

src/relay/backend/graph_executor_codegen.cc Outdated Show resolved Hide resolved

jroesch force-pushed the remove-compile-engine branch from e31ca8a to fe91baf Compare July 6, 2021 17:19