Optimization to select(pred, 1, 0) -> convert(pred). Is it beneficial?

[apivovarov]

I'd like to check if adding the following optimization to HLO algebraic_simplifier.cc would be beneficial?

select(pred, 1, 0) -> convert(pred)                    (1)

select(pred, 1, 0) -> broadcast(convert(pred))         (2)

Optimization:

If on_true array isAll(1) and on_false array isAll(0) then we can:

1. replace select(pred, 1, 0) with convert(pred) (if pred.shape == on_true.shape)
2. replace select(pred, 1, 0) with broadcast(convert(pred)) otherwise

Assumption is that convert (or convert + broadcast) computation(s) is cheaper than (or equal to) select computation on most of CPU/GPU/TPU performance models.

Operation semantics

Select op semantic: select(pred, on_true, on_false)
Desc: Constructs an output array from elements of two input arrays, based on the values of a predicate array.

• on_true.shape is the same as on_false.shape
• pred.shape.rank is the same as on_true.shape.rank

@akuegel @lipracer @ddunl What you think?

[akuegel]

Sorry for forgetting to answer here. If you can provide microbenchmarks for GPU so that we can see how much this optimization is worth, that would be helpful. In case it is not a good optimization for all backends, we could add an AlgebraicSimplifier option to enable/disable it based on backend.

[apivovarov]

Thank you, Adrian, for your reply. We actually have found more complex pattern to simplify - Causal Mask HLO for jnp.tril(jnp.ones()) can be simplified #9709. What you think?