[TIR] Support AllocateConst nodes in TensorIR scheduling flow#12489
Merged
Hzfengsy merged 10 commits intoapache:mainfrom Aug 22, 2022
Merged
[TIR] Support AllocateConst nodes in TensorIR scheduling flow#12489Hzfengsy merged 10 commits intoapache:mainfrom
Hzfengsy merged 10 commits intoapache:mainfrom
Conversation
masahi
commented
Aug 18, 2022
masahi
commented
Aug 18, 2022
masahi
commented
Aug 18, 2022
| auto block = block_realize->block; | ||
| block.CopyOnWrite()->body = | ||
| tir::AllocateConst(var, dtype, extents, constant_map[var], block->body); | ||
| n->body = BlockRealize(block_realize->iter_values, block_realize->predicate, block); |
Member
Author
There was a problem hiding this comment.
Please note this change. This places AllocateConst at the beginning of the body of BlockRealize. I found that putting BlockRealize as the body of AllocateConst leads to many issues since many places in TIR code assume that the body of a primfunc starts with BlockRealize.
Member
|
Thanks @masahi |
xinetzone
pushed a commit
to daobook/tvm
that referenced
this pull request
Nov 25, 2022
…#12489) * [TIR] Support AllocConstantNode in CreatePrimFunc * Handle AllocConstantNode in LeafBlockRemovalPlan * Properly handle AllocConstNode in BufferAllocationLocator * handle AllocateConst in EstimateFlops * remove NDArray printing * doc update * add test * cpplint * Removed dependency on link-params attribute from target * Restored NDArray printing to unbreak test
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TIR-level constants, represented by
AllocateConstnodes, were introduced in #8509. They are generated iflink-params = 1.At Relay level,
link_params = Truemakes constant tensors kept bound to the function body, rather than lifting them to function parameters duringFuseOps. So we end up with the following Relay prim func, for example:(if
link-params = False, which is the default, the weight also becomes the parameter of the function).I found that TensorIR related components below do not support such TIR-level constants currently:
CreatePrimFunc: It assumes that all tensors will be passed as parameters. This doesn't hold iflink-params = 1.PlanAndUpdateBufferAllocationLocation: Handling forAllocateConstnodes is incorrect. It creates a duplicated allocation for constant already allocated byAllocateConstnode, whichStorageRewritecomplains because there are two allocations of the same-name constant.LeafBlockRemovalPlan(used bycompute_atandcompute_inline): Assumes that the body ofBlockbegins withSeqStmt. This may not be the case since I placedAllocateConstat the beginning of a body, like belowThe most important change is the introduction of
CreatePrimFuncWithConstantsfunction. It reusestir::BindParamspass added in #8509 to create a PrimFunc withAllocateConstnodes. Since this new function takes an array ofruntime::NDArrayas an additional argument, I had to change all places whereCreatePrimFuncis used.cc @junrushao1994 @Hzfengsy @vinx13 @csullivan