[TensorExpr] Add python bindings. #49698
Conversation
Reincarnation of #47620 by @jamesr66a. [ghstack-poisoned]
facebook-github-bot
added
cla signed
oncall: jit
💊 CI failures summary and remediationsAs of commit 2458a8d (more details on the Dr. CI page):
Extra GitHub checks: 1 failed
ci.pytorch.org: 1 failedThis 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 to the (internal) Dr. CI Users group. This comment has been revised 88 times. |
Reincarnation of #47620 by @jamesr66a. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. ghstack-source-id: 650538ae8967f4d1c54b3680aa272f2af38bda63 Pull Request resolved: #49698
Reincarnation of #47620 by @jamesr66a. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. ghstack-source-id: 952fab9e286665bcf58e2937b299c938bc9a07d2 Pull Request resolved: #49698
Reincarnation of #47620 by @jamesr66a. It's just an initial bucnh of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
ZolotukhinM
changed the title
Reincarnation of #47620 by @jamesr66a. It's just an initial bunch of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. It's just an initial bunch of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. It's just an initial bunch of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. ghstack-source-id: 42e74f0b5c05bb9358b2288f1a88a4ca4569e9b7 Pull Request resolved: #49698
|
|
||
| class TestTensorExprPyBind(JitTestCase): | ||
| def test_simple_sum(self): | ||
| with kernel_arena_scope(): |
There was a problem hiding this comment.
Not a big deal for this PR but omg I want to get rid of KernelArena before we start advertising this API. Weird low-level memory management hacks shouldn't be leaking up into Python :-p
There was a problem hiding this comment.
100% agree 😃
I think even if we keep the arena for now, there is probably a more elegant way of using it for python bindings than I did here.
There was a problem hiding this comment.
Looks like a pretty reasonable start to me. I think we all agree these APIs aren't a desirable end state (I hope) but I'm assuming no one is going to use them externally until we get them into a clean state?
(Particular things that seem clumsy are the limitations on # of args for Compute, the fact that we have to use load and/or call instead of simply indexing, KernelScope, BufferArg, and the need to create a loopnest, prepare_for_codegen, and simplify manually)
Reincarnation of #47620 by @jamesr66a. It's just an initial bunch of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
There was a problem hiding this comment.
Is it possible to put the TE-related pybinding code somewhere in jit/tensorexpr instead of in here? Maybe it's not a big deal it just feels pretty separate from the rest of torch.jit.
Good idea, I moved the TE bindings to a separate file.
I think we all agree these APIs aren't a desirable end state (I hope) but I'm assuming no one is going to use them externally until we get them into a clean state?
Yes, absolutely. This PR is just to demonstrate the API working end-to-end, making it actually usable was out of the scope :)
|
|
||
| class TestTensorExprPyBind(JitTestCase): | ||
| def test_simple_sum(self): | ||
| with kernel_arena_scope(): |
There was a problem hiding this comment.
100% agree 😃
I think even if we keep the arena for now, there is probably a more elegant way of using it for python bindings than I did here.
Reincarnation of #47620 by @jamesr66a. It's just an initial bunch of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. Differential Revision: [D25668694](https://our.internmc.facebook.com/intern/diff/D25668694) [ghstack-poisoned]
Reincarnation of #47620 by @jamesr66a. It's just an initial bunch of things that we're exposing to python, more is expected to come in future. Some things can probably be done better, but I'm putting this out anyway, since some other people were interested in using and/or developing this. ghstack-source-id: 034879a0db5ab90ad647459c25b50e7677b35218 Pull Request resolved: #49698
|
@ZolotukhinM merged this pull request in e9dc8fc. |
Stack from ghstack:
Reincarnation of #47620 by @jamesr66a.
It's just an initial bunch of things that we're exposing to python, more
is expected to come in future. Some things can probably be done better,
but I'm putting this out anyway, since some other people were interested
in using and/or developing this.
Differential Revision: D25668694