[TensorExpr] Insert allocations for temporary buffer at the innermost valid scope. #36836
Conversation
… valid scope. [ghstack-poisoned]
facebook-github-bot
added
the
oncall: jit
💊 Build failures summary and remediationsAs of commit ae67e49 (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Follow this link to opt-out of these comments for your Pull Requests.Please report bugs/suggestions on the GitHub issue tracker. This comment has been revised 3 times. |
| class ContainedStmtsFinder : public IRVisitor { | ||
| public: | ||
| // Simply list all Stores and LetStmts that are children of the given stmt | ||
| std::unordered_set<Stmt*> findContainedStmts(Stmt* s) { |
|
|
||
| private: | ||
| void visit(const Store* v) override { | ||
| if (stores_[v->buf()].insert(last_stmt_).second) { |
There was a problem hiding this comment.
Nit picking: it is a bit weird that the condition evaluation has an important side-effect. it might be better to separate ".insert" into its own statement.
There was a problem hiding this comment.
It's actually an idiom I've seen and used a lot previously.
|
|
||
| TORCH_API Stmt* FlattenIndexes(Stmt* s); | ||
|
|
||
| struct BufUse { |
There was a problem hiding this comment.
Nit: a class here somewhat implies that users might want to know this. But actually this class structure is very much internal to our development. It might be better to move to a internal header file, or implementation files.
There was a problem hiding this comment.
I would defer this until we decide on how our dep-analysis should look like. I added the corresponding TODO comment.
| return nullptr; | ||
| } | ||
|
|
||
| Block* findLowestContainingBlock(const std::vector<BufUse>& uses) { |
There was a problem hiding this comment.
Minor: add a TODO to see if we should improve the performance of this block for larger programs. Many systems I knew look at all the parents in reverse order, and find the lowest one.
…e innermost valid scope." Differential Revision: [D21099913](https://our.internmc.facebook.com/intern/diff/D21099913) [ghstack-poisoned]
|
@ZolotukhinM merged this pull request in b8e2d79. |
4 participants
Stack from ghstack:
Differential Revision: D21099913