[Unity] Allocate workspace for all functions

[yelite]

This PR makes AllocateWorkspace allocate workspace in all Relax functions that doesn't have Codegen or Composite attributes.

cc @masahi

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[masahi]

This assumes that all entry functions get the same workspace size right?

I also have a WIP branch https://github.com/masahi/tvm/tree/workspace-provider-entry-func that calculates the minimum-necessary workspace size for each function that calls another functions with kWorkSpaceSize attribute. It doesn't yet handle the case where the same function with kWorkSpaceSize is called by multiple entry functions, with different maximum workspace sizes. Since we modify the function signature for a function that needs workspace param, the param size needs to take into account all entry funcs that call it.

So the upshot is, even though different entry funcs can have its own minimum workspace size, to handle the case above we need to take the max over all functions with kWorkSpaceSize, even though a single entry func might not call all of them. I was planning to add such logic on top of my branch, but looking at this PR makes me realize that this simple patch is effectively sufficient.

Do you agree with my assessment? @yelite

[yelite]

@masahi I think what you described is ideal and the workspace param can be made as R.Object to workaround the multi-caller limitation. I am just not sure what we lose by stripping out shape information from the workspace param. If all the usages of workspace validates the actual size of workspace, and allocate if workspace is not large enough, like what attention does currently, I think it should be fine.

[masahi]

R.Object sounds like a good solution. For now this patch should be fine.