Extend quantiser support so as to accelerate more binary models.

[AdamHillier]

What do these changes do?

This PR extends the LCE converter with patterns that support converting tf.where-style binary quantisers. It also adds support for binary input/output of LceQuantize/LceDequantize.

Note that in larq/larq#677 we are moving towards having the Larq quantisers implemented with tf.where instead of the current tf.sign implementation. The main change is that this PR adds support for converting tf.where-style quantisers.

The second change to add support for boolean input, meaning that a wider variety of binary quantisers will be accelerated by LCE. For example, the following wacky quantiser will now convert successfully:

lambda x: tf.where(tf.logical_or(tf.abs(x) < 0.5, tf.abs(x) > 1.0), 1.0, -1.0)

With these changes any binary quantiser that can be implemented with tf.where(boolean_condition, 1, -1) can be converted into an LceQuantize op (and consequently, a subsequent convolution can be converted to a BConv2D and thus accelerated). The boolean_condition will be implemented with TFL ops, as in the example above, but since the quantisation in general is so quick compared to the binary convolution this doesn't present much of a performance issue.

How Has This Been Tested?

MLIR test cases have been added; the end2end tests have been extended.

Benchmark Results

N/A.

Related issue number

larq/larq#677

[Tombana]

Looks great!

[Tombana]

LGTM.
Just a note to ourselves: when merging this into our private fork it will probably not cause any automatic merge conflicts, but we'll need to update the micro quantization kernels (should be fairly trivial, just by copying the changes from the public quantization kernel).

[lgeiger]

Nicely done!

It took me a while to wrap my head around boolean input for lq.quantize, but I think it makes a lot of sense and is much easier to maintain than passing an additional threshold attribute to the op or doing something similar.

I just have some additional comments regarding broadcasting and the matching of tf.where.

[lgeiger]

This is great, thank you so much for figuring this out!

I just have a few questions to understand the code, other than that this looks great to me 🚀