[te] Introduce likely/unlikely CompareSelect hint (#51751)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51751
Similar in spirit to the `__builtin_expect` C intrinsic, it's useful
to be able to hint the expected branch direction in a tensor expression. Using
this flag has a few effects on codegen:
- The CompareSelect is generated using conditional branches, rather than selects
- The conditional branches are strongly hinted (like, 100000:1) in the indicated direction
- A vectorized hinted CompareSelect computes its condition in parallel with a
mask "reduction" (e.g. a bitcast from `<i1 x 8>` to `<i*>`). In AVX terms
this sequence might look like:
```
vpcmpgtd %ymm0, %ymm1, %ymm2
vmovmskps %ymm2, %eax
```
The motivating case for this addition is an attempt I'm making to replicate
fast transcendentals using tensor expressions. Floating-point numbers have
lots of special cases (denormals, inf, nan) that need special handling, and
it's convenient to be able to punt that handling off to a slow path while
keeping the fast path nice and tight.
ghstack-source-id: 121366315
Test Plan:
I'm not sure how to test this (except I can tell you it works for
the `log` implementation I'm working on right now). It would be nice to plumb
the LLIR/ASM output through programmatically so it can be used in FileCheck.
Maybe I'll do that in another diff?
Reviewed By: asuhan
Differential Revision: D26246401
fbshipit-source-id: 900f7fa0520010fb9931d6e3efc8680a51f8d844