[jit] Add options to Operator to enable registration of alias analysis passes#18589
[jit] Add options to Operator to enable registration of alias analysis passes#18589bwasti wants to merge 14 commits into
Conversation
| #include <ATen/ATen.h> | ||
| #include <ATen/core/function_schema.h> | ||
|
|
||
| #include <torch/csrc/jit/passes/alias_analysis.h> |
There was a problem hiding this comment.
this might be bad, what do you think @zdevito ?
zdevito
left a comment
There was a problem hiding this comment.
Looks pretty good. Small nits.
| op_(std::make_shared<Operation>(std::move(op))), | ||
| options_(std::move(options)) {} | ||
|
|
||
| struct Options { |
There was a problem hiding this comment.
I would not make this struct part of the Operator class. I want to move away from exposing users directly to the Operator class (and instead have them directly call .op with the op implementation and options). So I think Options should be its own class. Putting it in its own header will help separate out the depenendencies (i.e. the options rely on alias_analysis, but the operator header does not rely directly on it).
|
|
||
| struct Options { | ||
| Options() {}; | ||
| using AliasAnalysisFunction = void (AliasDb::*)(Node*); |
There was a problem hiding this comment.
Making AliasAnalysisFunction a pointer rather than an enum doesn't buy us anything here, but it does strongly couple it to the alias_analysis.h header. How about just an enum in a separate header:
enum AliasAnalysisKind {
CREATOR,
EXTRACTOR,
}
struct Options {
};
We can add more wildly customizable versions by patching options when necessary.
There was a problem hiding this comment.
I agree. The analyzeFoo functions in alias analysis are not really designed for outside consumption and are probably not clear to users. I think we could have an enum that covers the most common cases:
- Purely functional (no aliases, no mutation on inputs).
- View (output aliases input)
- Free for all (all inputs may be written to, all outputs may alias any input)
- Look at the schema (reads a user-provided schema, allows people to do whatever fancy thing).
There was a problem hiding this comment.
reads a user-provided schema
@suo Not sure I follow how this would work
There was a problem hiding this comment.
To be clear, you don't have to put all of those options in this PR (just wanted to enumerate some common cases we'll want to add later).
For "Look at the schema"—people defining a custom op can pass a JIT schema if they want; so they should be able to instruct the alias analysis pass to use that schema to populate alias info (similar to how we do with builtin ops).
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
| const std::string& schemaOrName, | ||
| Implementation&& implementation) { | ||
| Implementation&& implementation, | ||
| OperatorOptions options = OperatorOptions()) { |
There was a problem hiding this comment.
These should be exposed in RegisterOperators::op() as well since that's the suggested way to register ops in the docs/tutorials
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
…ias analysis passes" [jit] Add options to Operator to enable registration of alias analysis passes gh-metadata: pytorch pytorch 18589 gh/bwasti/3/head
|
Looks like this broke the build? |
…es (pytorch#18589) Summary: Pull Request resolved: pytorch#18589 ghimport-source-id: dab203f Differential Revision: D14901379 Pulled By: bwasti fbshipit-source-id: d92a497e280f1b0a63b11a9fd8ae9b48bf52e6bf
…es (pytorch#18589) Summary: Pull Request resolved: pytorch#18589 ghimport-source-id: dab203f Differential Revision: D14901379 Pulled By: bwasti fbshipit-source-id: d92a497e280f1b0a63b11a9fd8ae9b48bf52e6bf
Stack from ghstack:
Differential Revision: D14901379