Add optimized op_where

[swolchok]

It materializes separate kernels for the cases where the two input data tensors have the same dtype and the third one has dtype bool.

[swolchok]

Stack from ghstack (oldest at bottom):

• -> Add optimized op_where #8866
• Deploy BroadcastIndexesRange #8865
• add BroadcastIndexesRange #8864
• portable arg{max,min}: optimize update check #8863

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/8866

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit e52508d with merge base cca6917 ():

NEW FAILURE - The following job has failed:

• pull / unittest-editable / linux / linux-job (gh)
  backends/xnnpack/test/ops/test_conv1d.py::TestConv1d::test_qs8_conv1d_batchnorm_seq

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[digantdesai]

Love how clean it reads :)

Couple of high level comments,

• Not too familiar with how people use this op i.e. any common case we see often in the wild, but are there any short-circuits we can do to avoid coming here?
• Similarly, depending on the condition, if we know we are biased towards A vs. B, we can first copy A to the result and then go through data_cond and pick B.
• Lastly, I am assuming we want to do SIMD later, if we do then we can use predicates, that should make it less "branchy". It may not help much with load/stores though.

[swolchok]

if we know we are biased towards A vs. B

we can't possibly know this in general.

SIMD later

not currently on my agenda, but I may have to come back.

[digantdesai]

Looks good to me.

[swolchok]

unittest-editable failure is a known flake, so noting that we have green CI. I think this will cause -Wunused -Werror builds to fail though, so I need to fix that before merging.

[swolchok]

I think this will cause -Wunused -Werror builds to fail though

Checked, works fine, I misread my own code.