[inductor] Refactor generate_kernel_call #102018
Conversation
Summary: Refactor generate_kernel_call to support codegen call to Triton kernel [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/102018
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 283f4f5: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
torch/_inductor/codegen/wrapper.py
Outdated
| def generate_kernel_call(self, name, call_args, device_index=None, cpp=False): | ||
| self.writeline(self.wrap_kernel_call(name, call_args)) | ||
| def generate_kernel_call( | ||
| self, name, call_args, grid=None, device_index=None, cuda=None |
There was a problem hiding this comment.
what type is cuda? You are passing cuda=True, so default should probably be False, not None?
There was a problem hiding this comment.
I will change it to cuda=True as the triton backend is the default one.
| self, name: str, kernel: str, metadata: Optional[str] = None, cuda=True | ||
| ): | ||
| if cpp: | ||
| if not cuda: |
There was a problem hiding this comment.
what's the reason for changing arg from cpp to cuda?
There was a problem hiding this comment.
The main reason is because cpp means two things in inductor's codegen, the cpu backend and the cpp wrapper codegen. Because I added generate_kernel_call for the triton backend into the base class here, WrapperCodeGen, now we need a flag to distinguish the two backends, and cpp wouldn't be a good naming here. cuda may not be the best naming, but it's slightly better.
There was a problem hiding this comment.
@desertfire , I'm trying to enable Inductor to support other backends at runtime #100706. The design assumes the backend just needs to customize two meta classes, BaseScheduling and wrapper.WrapperCodeGen, by overriding some member functions. Regarding the define_kernel, I'm kind of confused by the cuda parameter because it is the logic of CudaWrapperCodeGen and it has implied cuda is True. My question is this parameter is must-have.
There was a problem hiding this comment.
Initially I didn't have cuda as a parameter, but the problem appears when we have a mix of devices (see test_multi_device). Because we only keep one wrapper codegen instance, i.e. CudaWrapperCodeGen when there is a mix of cuda and cpu tensors, we need this additional parameter to help generate correct kernel definition and kernel call.
Summary: Refactor generate_kernel_call to support codegen call to Triton kernel cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 [ghstack-poisoned]
|
@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 |
6 participants
Stack from ghstack (oldest at bottom):
Summary: Refactor generate_kernel_call to support codegen call to Triton
kernel
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10