[inductor] Add ABI shim function for torch.scatter #114027
Conversation
Summary: Scatter fallback calls `at::scatter` in the C++ wrapper codegen. This doesn't work in the ABI compatibility mode, as the latter requires a shim function. One is added in this PR. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_scatter_fallback ... ---------------------------------------------------------------------- Ran 4 tests in 54.839s ``` Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/114027
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit caa0c20 with merge base 5e30741 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Summary: Scatter fallback calls `at::scatter` in the C++ wrapper codegen. This doesn't work in the ABI compatibility mode, as the latter requires a shim function. One is added in this PR. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_scatter_fallback ... ---------------------------------------------------------------------- Ran 4 tests in 54.839s ``` Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 7d9000e Pull Request resolved: #114027
aakhundov
changed the title
aakhundov
requested review from
chenyang78 and
desertfire
and removed request for
desertfire
aakhundov
added
ciflow/trunk
There was a problem hiding this comment.
The PR LGTM.
I am having a question to @SherlockNoMad: We are handling scatter different from other fallback kernels. Particularly, we have the following class hierarchies. I am wondering if we could also use ProxyExecutor to cover scatter fallback? Thanks.
class FallbackKernel(ExternKernelAlloc)
class ExternKernelAlloc(ExternKernel):
class ScatterFallback(ExternKernel)
|
FWIW, I'd prefer having an ABI shim function for the cases we can cover and rely on |
There were some discussions about this last Friday. I think it would mostly impact CPU backend. Not sure about how it would affect CUDA backend performance. It might be small. I think @SherlockNoMad volunteered to do some perf experiments. |
In some models, we can have very small kernels and / or CUDA device sync, so non-trivial CPU overhead can become visible. |
Summary: Scatter fallback calls `at::scatter` in the C++ wrapper codegen. This doesn't work in the ABI compatibility mode, as the latter requires a shim function. One is added in this PR. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_scatter_fallback s... ---------------------------------------------------------------------- Ran 4 tests in 52.713s OK (skipped=1) ``` Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]
Summary: Scatter fallback calls `at::scatter` in the C++ wrapper codegen. This doesn't work in the ABI compatibility mode, as the latter requires a shim function. One is added in this PR. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_scatter_fallback ... ---------------------------------------------------------------------- Ran 4 tests in 54.839s ``` Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 988da0b Pull Request resolved: #114027
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
4 participants
Stack from ghstack (oldest at bottom):
Summary: Scatter fallback calls
at::scatterin the C++ wrapper codegen. This doesn't work in the ABI compatibility mode, as the latter requires a shim function. One is added in this PR.Test Plan:
Reviewers:
Subscribers:
Tasks:
Tags:
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler