[inductor] Fix block ptr store if input is constant #148679
Conversation
… block ptrs Remove formatting changes Lint Rename vars Dont reshape float32 scalars Handle constants separately
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148679
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 10ced66 with merge base c65ee72 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| assert block_ptr not in advancements, ( | ||
| "duplicate advancement for pointer '{block_ptr}' at type '{symt}'" | ||
| ) | ||
| assert ( |
There was a problem hiding this comment.
lintrunner changes
| @@ -3039,9 +3059,9 @@ def sort( | |||
| self.filter_masks(masks) | |||
| masks = sorted(masks) | |||
| assert not self._load_mask, "ops.sort not supported inside ops.masked" | |||
| assert self.persistent_reduction, ( | |||
There was a problem hiding this comment.
lintrunner changes
|
@pytorchbot label "topic: not user facing" |
| assert isinstance(value, CSEVariable) | ||
| # See `_shaped_constant`. Tensor constants are only created | ||
| # if the dtype is not float32. | ||
| if value.dtype is not torch.float32: |
There was a problem hiding this comment.
This is a lot of action at a distance. could we refactor this ?
And yes, i agree, we should add shape.
| value = triton_reshape( | ||
| str(value), [sympy.S.One], [sympy.S.One] * len(indexing.block_shape) | ||
| ) | ||
| value = f"tl.broadcast_to({value}, {V.kernel.index_to_str(indexing.block_shape)})" |
There was a problem hiding this comment.
Nice catch!
https://github.com/pytorch/pytorch/pull/151399/files recently fixed a related bug by always broadcasting block ptr stores to indexing.block_shape. In light of that change, would it make sense to reuse indexing.codegen_broadcast_and_reshape for scalars instead of having separate code paths here?
Since block ptr stores require explicit broadcasts, the input to
tl.storeneeds to be reshaped and broadcasted. Currently, it is assumed that the input to be stored is in block form (e.g.XBLOCK), however it is possible for the input to be a scalar, and so special handling is required to reshape + broadcast the scalar to the output block shape.Ideally the shape of the input would be an attribute of a
TritonCSEVariablevia shape propagation but that is not the case today. The patch in this PR determines if the input is a constant by checking the arguments to an FX store node which is not ideal. Maybe there is an alternative and simpler methodFixes #ISSUE_NUMBER
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov