[mlir] Support fast-math friendly constants for identity value
Add an option to the family of `getIdentity` helper functions so that it is possible to produce fast-math friendly constants. For instance, for maxf the identity value is `-inf`, however, if the related operations are lowered with fast-math (`noinf` in particular), then the value becomes `poison` and chances are the whole codegen is not going to do what we want. To avoid this problem, we add an option to `getIdentity` and friends that specifies whether a finite value needs to be produced or not. The patch is NFC for all the code but the lowering of `linalg::softmax` because we know we lower that with fast-math down the line. I didn't audit the rest of the code to check if it would make sense to set this boolean in more places. Note: It feels kind of wrong to have to know what the lowering may do, but I don't know what the right (at least short-term) solution is. Long term, we may want a special "neutral element" attribute for the respective ops. I didn't think too much about the implications for that. Differential Revision: https://reviews.llvm.org/D156471
parent
c3cfbbc4
Please register or sign in to comment