Clean up IRModule attrs and LowerTEPass #8914

electriclilies · 2021-09-02T23:01:23Z

In this PR, I add InferType to the end of LowerTEPass so we don't have to call it after we call LowerTEPass. I also remove per_target_modules_ from the interpreter.

To do this, I did have to fix a few bugs, namely:

Make sure we copy the attributes of the IRModule when we soft-copy the IRModule
In the type inferencer, we lookup GlobalVars by name, not by pointer address
In the type inferencer, if we encounter a tir::PrimFunc, we skip the function instead of failing (this allows us to call InferType on the unified module which contains both relay::Functions and tir::PrimFuncs)

This is follow up work from #8886

(Also note that although the branch is called remove-lowered-output, I don't actually remove LoweredOutput in this PR. I started to but there was some other cleanup I needed to do first)

electriclilies · 2021-09-03T19:21:55Z

@Mousius @mbs-octoml @mikepapadim @leandron @manupa-arm Could you take a look? Thanks!

Mousius · 2021-09-04T20:51:42Z

src/ir/module.cc

@@ -43,7 +43,8 @@ namespace tvm {

 IRModule::IRModule(tvm::Map<GlobalVar, BaseFunc> functions,
                   tvm::Map<GlobalTypeVar, TypeData> type_definitions,
-                   std::unordered_set<String> import_set, parser::SourceMap source_map) {
+                   std::unordered_set<String> import_set, parser::SourceMap source_map,
+                   DictAttrs attrs) {


Should we add attrs to SEqualReduce (return equal(attrs, other->attrs);) and SHashReduce (hash_reduce(attrs)) ?

Yes, I'll do this. FunctionNode has attrs in the hash and equal functions, so we should be consistent with that.

Mousius · 2021-09-04T21:16:28Z

src/relay/ir/transform.cc

  IRModule updated_mod =
-      IRModule(mod->functions, mod->type_definitions, mod->Imports(), mod->source_map);
+      IRModule(mod->functions, mod->type_definitions, mod->Imports(), mod->source_map, mod->attrs);


Is this getting to the point where a copy constructor would be more appropriate so it'd just be IRModule(mod) rather than having to replicate the arguments?

I was thinking about something like that, but unfortunately the copy constructor will only increment refs on the shared pointers of the IRModule fields, so it won't work.
I can make a helper function that returns IRModule(mod->functions, ..., mod->attrs), though.

Mousius · 2021-09-04T21:38:05Z

src/relay/backend/te_compiler.cc

-  // be called afterwards
-  return tvm::transform::CreateModulePass(pass_func, 1, "LowerTE", {});
+  return tvm::transform::Sequential(
+      {tvm::transform::CreateModulePass(pass_func, 1, "LowerTE", {}), InferType()});


opt_level should be 0 for LowerTE so it's always ran - it wasn't checked before this change but Sequential checks it when it decides whether to run the pass or not (that's why you're seeing weird graph executor codegen failures as the test runs with opt_level=0)

As a side note, we should keep the test to ensure things compile successfully at opt_level=0 even after the compile engine stuff is removed 😸

Cool, thanks for catching this! I'll add a note to that test about not removing it.

manupak

Thanks Lily for the change.

The changes in the scope of InfertType and removal of per_target_module seems good modulo the existing comments.

However, the addition of attrs to IRModule is an important change and a good one.
I feel we should make sure to add tests that it is passed on to lowered_funcs eventually here or in a separate PR.

(If we are going with a seperate PR, we should not introduce attrs here, just yet).

manupak · 2021-09-05T04:40:54Z

include/tvm/ir/module.h

@@ -122,6 +122,7 @@ class IRModuleNode : public Object {
    v->Visit("global_var_map_", &global_var_map_);
    v->Visit("global_type_var_map_", &global_type_var_map_);
    v->Visit("source_map", &source_map);
+    v->Visit("attrs", &attrs);


I think this is a very important change and should be called out in the PR title or the main description.

What happens to these when they are split to Map<Target, IRModule> in lowered_funcs or per_target_module in the intepretter ? Are they copied in ?

Will it be possible to add a test to make sure attrs are passed onto TIR lowering ?

After having checked :

tvm/src/relay/backend/te_compiler.cc

Lines 890 to 906 in 9f52e7e

Map<Target, IRModule> GetPerTargetModules(IRModule mod) {

std::unordered_map<Target, IRModule, backend::TargetStrHash, backend::TargetStrEqual>

per_target_modules;

for (const auto& kv : mod->functions) {

const GlobalVar& var = kv.first;

const BaseFunc& func = kv.second;

if (func->IsInstance<tir::PrimFuncNode>()) {

// Extract target

Optional<Target> target = func->GetAttr<Target>(tvm::attr::kTarget);

ICHECK(target) << "Target should be set at this point";

// Put the function in per_target_modules

if (!per_target_modules.count(target.value())) {

// Initialize the IRModule for this target and add the function

IRModule target_module;

target_module->Add(var, func);

per_target_modules[target.value()] = target_module;

I dont think it is transferred.

Hi @manupa-arm, attrs were actually added to IRModule in #8750 so anything we do now is an iterative improvement on that. Given this series of PRs aims to remove the Map<Target, IRModule> entirely (this one removes the per_target_module from the interpreter) I don't think we need to ensure the copy happens here but I agree we should have some test coverage when the unified IRModule is lowered to ensure it contains all the attributes we've accrued - this should be a follow up when we change the interface from Map<Target, IRModule> to IRModule - does that sound reasonable to you?

Would there be a PR to remove lowered_funcs as well ? @electriclilies

Its not very clear to me how we can avoid per_target_module throughout the lowering process -- we could push it way down.

Unified IRModule --> per_target_modules (lowered_funcs) --> runtime.Modules

Unless Im missing something here, there will always be a stage that IRModule contains things that gets lowered to a single runtime.Module.

Most likely, although I have not dug into that part of the codebase in depth yet so I can't say for sure.
The two options that I think are most likely are

build consuming the IRModule directly, traversing the functions in the IRModule and dealing with each directly (which is what you just mentioned)

Right before build is invoked, separating the functions in the module by Target (although we wouldn't store them in a Map<Target, IRModule>)

So to summarize, we'll either completely remove the data structure that stores functions separated by target, or just push it all the way down to right before build is called.

So the attrs of IRModule are globally valid to all targets. There I feel conveying them to the codegen might still be beneficial.

For 1, we would need to expand the build API to pass the attrs

For 2, we can embed them to each IRModule.

The attrs allows a channel that pass through all of the unified lowering process up until the concept of IRModule cease to exist.

If you agree, I feel we should pass them in to the per target IRModule, then later changed in either way we decide to proceed.

OK, I can add them to the per target IRModules and the we can figure out what to do with it later.

+1 to:

preserving module attrs through all passes on the assumption they apply to all targets

xfering them when we project a cross-target IRModule to a specific-target IRModule just before the transition into lowering phases

We should probably have a convention that IRModule attrs should be for describing properties of the code and not of the compilation? Since they will appear everywhere it would be tempting to start adding Target-like things there.

Does anyone have thoughts on whether we should try to line up the IRModule and runtime::Module worlds?

cross-target single-target IR IRModule ? Executable ? runtime::Module

That convention sounds good to me, but one thing we should be cautious of is that the attributes of functions do contain properties of the compilation flow (like targets), and this is inconsistent.

@mbs-octoml By lining up the IRModule and runtime::Module worlds do you mean also moving to a world where runtime::Modules are cross-target? I'm honestly not sure if that's something we want to do. @areusch do you have any thoughts about this?

electriclilies · 2021-09-08T17:21:39Z

@manupa-arm @Mousius This is green, can you take another look and approve? Thanks!

manupak

LGTM. Thanks for the changes Lily!

manupak · 2021-09-09T08:10:15Z

@Mousius PTAL when you have some time? and explicitly approve, possibly

electriclilies · 2021-09-10T15:16:53Z

@manupa-arm can you merge? #8974 is waiting on this PR! Thanks
(Actually, nevermind, I missed a few new comments)

manupak · 2021-09-10T15:33:50Z

Hi Lily,

After addressing those comments, could you update the description on the PR as to what it should say in the commit message ?
I ll copy that and when I do the merge.

electriclilies · 2021-09-10T20:35:21Z

@manupa-arm I changed the title and responded to the comment! Ready to merge

Mousius

This looks good to me @manupa-arm, sorry for the delay @electriclilies!

masahi · 2021-09-13T10:03:42Z

Thanks @electriclilies @manupa-arm @Mousius @mbs-octoml

)

electriclilies requested review from anijain2305, areusch, comaniac, jroesch, junrushao, jwfromm, kazum, kparzysz-quic, liangfu, MarisaKirisame, masahi, mbrookhart, slyubomirsky, tmoreau89, tqchen, vinx13, wweic, yzhliu, zhiics and ZihengJiang as code owners September 2, 2021 23:01

electriclilies changed the title ~~[DRAFT] Cleanup LowerTE + surrounding machinery more~~ [DRAFT] Run InferType in LowerTEPass and remove per_target_modules_ from the Interpreter Sep 2, 2021

electriclilies force-pushed the remove-lowered-output branch from e5e2516 to 7ee49c9 Compare September 3, 2021 18:36

electriclilies changed the title ~~[DRAFT] Run InferType in LowerTEPass and remove per_target_modules_ from the Interpreter~~ Run InferType in LowerTEPass and remove per_target_modules_ from the Interpreter Sep 3, 2021

Mousius requested changes Sep 4, 2021

View reviewed changes

manupak requested changes Sep 5, 2021

View reviewed changes

electriclilies requested a review from merrymercy as a code owner September 7, 2021 22:41

electriclilies force-pushed the remove-lowered-output branch from 2c8d8c2 to 6d00a9c Compare September 7, 2021 22:54

manupak approved these changes Sep 8, 2021

View reviewed changes

mikepapadim mentioned this pull request Sep 9, 2021

[Relay] Remove memory planing from LowerTEPass #8974

Merged

electriclilies changed the title ~~Run InferType in LowerTEPass and remove per_target_modules_ from the Interpreter~~ Clean up IRModule attrs and LowerTEPass Sep 10, 2021

Clean up LowerTEPass and pass IRModule Attrs through passes

7c77a7a

electriclilies force-pushed the remove-lowered-output branch from 3d2ba1b to 7c77a7a Compare September 11, 2021 00:01

Mousius approved these changes Sep 13, 2021

View reviewed changes

masahi merged commit bedf7cf into apache:main Sep 13, 2021

ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021

Clean up LowerTEPass and pass IRModule Attrs through passes (apache#8914

f98d426

)

junrushao mentioned this pull request Nov 1, 2021

Apache TVM v0.8 Release Note Candidate #9416

Closed

ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022

Clean up LowerTEPass and pass IRModule Attrs through passes (apache#8914

2b3ec0d

)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Clean up IRModule attrs and LowerTEPass #8914

Clean up IRModule attrs and LowerTEPass #8914

electriclilies commented Sep 2, 2021 •

edited

Loading

electriclilies commented Sep 3, 2021 •

edited

Loading

Mousius Sep 4, 2021

electriclilies Sep 7, 2021

Mousius Sep 4, 2021

electriclilies Sep 7, 2021

Mousius Sep 4, 2021

electriclilies Sep 7, 2021

manupak left a comment •

edited

Loading

manupak Sep 5, 2021

manupak Sep 5, 2021

Mousius Sep 7, 2021

manupak Sep 7, 2021 •

edited

Loading

manupak Sep 7, 2021 •

edited

Loading

electriclilies Sep 7, 2021

manupak Sep 7, 2021

electriclilies Sep 7, 2021

mbs-octoml Sep 10, 2021

electriclilies Sep 10, 2021

electriclilies commented Sep 8, 2021

manupak left a comment

manupak commented Sep 9, 2021 •

edited

Loading

electriclilies commented Sep 10, 2021 •

edited

Loading

manupak commented Sep 10, 2021

electriclilies commented Sep 10, 2021

Mousius left a comment

masahi commented Sep 13, 2021

	Map<Target, IRModule> GetPerTargetModules(IRModule mod) {
	std::unordered_map<Target, IRModule, backend::TargetStrHash, backend::TargetStrEqual>
	per_target_modules;
	for (const auto& kv : mod->functions) {
	const GlobalVar& var = kv.first;
	const BaseFunc& func = kv.second;
	if (func->IsInstance<tir::PrimFuncNode>()) {
	// Extract target
	Optional<Target> target = func->GetAttr<Target>(tvm::attr::kTarget);
	ICHECK(target) << "Target should be set at this point";

	// Put the function in per_target_modules
	if (!per_target_modules.count(target.value())) {
	// Initialize the IRModule for this target and add the function
	IRModule target_module;
	target_module->Add(var, func);
	per_target_modules[target.value()] = target_module;

Clean up IRModule attrs and LowerTEPass #8914

Clean up IRModule attrs and LowerTEPass #8914

Conversation

electriclilies commented Sep 2, 2021 • edited Loading

electriclilies commented Sep 3, 2021 • 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

manupak left a comment • 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 Sep 7, 2021 • edited Loading

Choose a reason for hiding this comment

manupak Sep 7, 2021 • 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

electriclilies commented Sep 8, 2021

manupak left a comment

Choose a reason for hiding this comment

manupak commented Sep 9, 2021 • edited Loading

electriclilies commented Sep 10, 2021 • edited Loading

manupak commented Sep 10, 2021

electriclilies commented Sep 10, 2021

Mousius left a comment

Choose a reason for hiding this comment

masahi commented Sep 13, 2021

electriclilies commented Sep 2, 2021 •

edited

Loading

electriclilies commented Sep 3, 2021 •

edited

Loading

manupak left a comment •

edited

Loading

manupak Sep 7, 2021 •

edited

Loading

manupak Sep 7, 2021 •

edited

Loading

manupak commented Sep 9, 2021 •

edited

Loading

electriclilies commented Sep 10, 2021 •

edited

Loading