[inductor] Introduce CSEVariable type and use it to track if Triton variables are scalar #88347
Conversation
This fixes pytorch/torchdynamo#1515 [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/88347
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 37 PendingAs of commit dab7127: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@ngimel sorry, not sure what happened, had to close previous PR and do a new one. I copied your comment below:
I will look to address the above. |
torch/_inductor/codegen/cpp.py
Outdated
| def masked(mask, body, other): | ||
| code = BracesBuffer() | ||
| var = V.kernel.cse.newvar() | ||
| var = V.kernel.cse.newvar(is_scalar=False) |
There was a problem hiding this comment.
The arg name is_scalar is kind of confusing for cpp codegen. Variables are by default scalars in cpp and we are adding vectorization support but we manage the vector variable and scalar locally. Is it better named as is_triton_scalar and make it default to False to avoid changing cpp codegen? cc @EikanWang who is working on cpp vectorization support.
torch/_inductor/codegen/cpp.py
Outdated
| line = f"static_cast<float>({line})" | ||
| return self.cse.generate(self.loads, line) | ||
| # Not tracking whether results of loads are scalar | ||
| return self.cse.generate(self.loads, line, is_scalar_expr=False) |
torch/_inductor/codegen/cpp.py
Outdated
| self.loads, f"reduction {name} {cexpr(index)}", write=False | ||
| self.loads, | ||
| f"reduction {name} {cexpr(index)}", | ||
| is_scalar_expr=False, |
No, please hold off. |
This fixes pytorch/torchdynamo#1515 It keeps track of which variables in generated code are scalars, and does not include a mask when doing loads generated by these. It uses variable names (e.g., tmp_scalar4) to keep track of which ones are scalars, which makes generated code a little more readable, but is a little hacky. Can try to make it better, but wanted to get some feedback on overall approach first. This only really matters for triton backend, it would be nice to contain it to triton.py by didn't find a nice way of doing that. cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx lezcano [ghstack-poisoned]
|
I have rewritten this to better separate the Triton specific aspects and make it a little more generic. This required introducing a new type @ngimel after thinking about it for a bit, it seems to me pytorch/torchdynamo#1654 is related but different enough that it is better to handle it separately. I will try to put up a PR for it as well. |
fdrocha
changed the title
torch/_inductor/codegen/common.py
Outdated
| def generate( | ||
| self, buffer: IndentedBuffer, expr: typing.Union[str, CSEVariable], write=True | ||
| ) -> CSEVariable: | ||
| assert isinstance(expr, str) or isinstance(expr, CSEVariable), type(expr) |
There was a problem hiding this comment.
nit:
assert isinstance(expr, (str, CSEVariable))
…if Triton variables are scalar" This fixes pytorch/torchdynamo#1515 To fix it, we need to keep track of whether a Triton variable is a scalar (so we can not use a mask when doing indirect loads through them). This requires a way of annotating variable names generated by CSE with properties. So now CSE will use CSEVariable class to keep track of variables and let backends subclass it so they can annotate them with whatever information they want. TritonCSEVariable is such a subclass that track the `is_scalar` property. cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx lezcano [ghstack-poisoned]
The next PR in this stack is a more comprehensive test for this. I added a simple more explicit test to this one just now. |
|
Simpler test is always beneficial |
|
@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 |
…ariables are scalar (pytorch#88347) This fixes pytorch/torchdynamo#1515 To fix it, we need to keep track of whether a Triton variable is a scalar (so we can not use a mask when doing indirect loads through them). This requires a way of annotating variable names generated by CSE with properties. So now CSE will use CSEVariable class to keep track of variables and let backends subclass it so they can annotate them with whatever information they want. TritonCSEVariable is such a subclass that track the `is_scalar` property. Pull Request resolved: pytorch#88347 Approved by: https://github.com/jgong5, https://github.com/ngimel
5 participants
Stack from ghstack (oldest at bottom):
This fixes pytorch/torchdynamo#1515
To fix it, we need to keep track of whether a Triton variable is a scalar (so we can not use a mask when doing indirect loads through them). This requires a way of annotating variable names generated by CSE with properties.
So now CSE will use CSEVariable class to keep track of variables and let backends subclass it so they can annotate them with whatever information they want. TritonCSEVariable is such a subclass that track the
is_scalarproperty.cc @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @desertfire @lezcano