Replace sympy.Mod with custom Mod object#96270
Replace sympy.Mod with custom Mod object#96270Chillee wants to merge 2 commits intogh/chillee/195/basefrom
Conversation
[ghstack-poisoned]
| self.env[node] = self.run_node(node) | ||
| except Exception as e: | ||
| if self.extra_traceback: | ||
| if hasattr(self, "extra_traceback") and self.extra_traceback: |
There was a problem hiding this comment.
This was breaking for me 🤔
| isinstance(base, sympy.Integer) | ||
| and isinstance(modulus, sympy.Integer) | ||
| ): | ||
| return base % modulus |
There was a problem hiding this comment.
Upstream mod has a number of seemingly non-objectionable other simplifications, which may be worth taking. See https://github.com/sympy/sympy/blob/master/sympy/core/mod.py
At a brief glance, here are rules we might like:
- Mod by two can use even/odd assumptions
- Simplifying nested Mod
| ): | ||
| return base % modulus | ||
|
|
||
| if isinstance(base, sympy.Add): |
There was a problem hiding this comment.
sympy has a comment like
# (x + y + 2) % x -> Mod(y + 2, x)
which I thought was pretty useful for understanding what their Add rule is simplifying for
| if (isinstance(term, sympy.Integer) and term < 0) or ( | ||
| isinstance(term, sympy.Mul) | ||
| and isinstance(term.args[0], sympy.Integer) | ||
| and term.args[0] < 0 |
There was a problem hiding this comment.
This condition looks wrong. Why can't we use sympy assumptions to determine if the term is negative or not?
There was a problem hiding this comment.
This was just mimicked from the existing ModularIndexing code. cc: @ngimel
There was a problem hiding this comment.
I'm not against using sympy assymptions, I think when I just tried to put comparison in ModularIndexing it errored out, but I don't remember.
This exemption shouldn't be necessary really, in ModularIndexing it exists only to work around triton bug. (e.g. with valid simplifications (x - y) % x will turn in to -y % x and iirc triton has problems with negative FloorDivs or modulos, but we probably don't need to care about it in symbolic_shapes
ezyang
left a comment
There was a problem hiding this comment.
Not convinced by the Triton workaround code
cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 desertfire [ghstack-poisoned]
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Stack from ghstack (oldest at bottom):
cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @desertfire