Use custom larq MLIR dialect for our ops

[lgeiger]

TensorFlow added a tfl.custom op to the TFL dialect in tensorflow/tensorflow@fb7ea8f. This allows us to decouple our ops from the TF dialect which was previously necessary to allow for flatbuffer serialization.

Essentially this PR does two things: It moves our op definition to a custom larq dialect and introduces a legalize pass at the end that translates our ops to tfl.custom ops which can be correctly serialized to the TFLite flatbuffer.

This has a few advantages over the current approach of adding our ops to the TF dialect:

• Full compatibility with MLIR passes and patterns
  • Proper type checking in the IR and during transformation
    • should prevent bugs that could lead to invalid flatbuffers
    • will allow us to introduce strict verification of the correcness of the IR
  • Proper handling of MLIR traits which
    • MLIR can properly check if ops have side effects and remove trailing ones without the need for custom op cleanup patterns
    • will allow us to potentially use quantization traits in the future if necessary for easier customization of bias quantizations
• Full flexibility of attribute sersialization and deserialization (If really necessary we could even get rid of flexbuffers and roll with our own attribute serialization)
• Autogenerated docs for our IR

I marked this as a draft PR since it is built ontop of #373 and #382

Click to expand autogenerated docs

'lq' Dialect

Types and operations for Larq dialect

This dialect contains operations for Larq. This dialect will be used in
conjunction with the TensorFlow dialects for converting & optimizing
TF graphs to be deployed on Larq Compute Engine.

[TOC]

Operation definition

lq.BMaxPool2d (TF::BMaxPool2dOp)

Binary MaxPool2D op.

Computes a MaxPool2D operation and outputs bitpacked binary values, for consumption by a binary convolution.

Attributes:

Attribute	MLIR Type	Description
padding	StringAttr	padding enum
stride_width	IntegerAttr	32-bit signless integer attribute
stride_height	IntegerAttr	32-bit signless integer attribute
filter_width	IntegerAttr	32-bit signless integer attribute
filter_height	IntegerAttr	32-bit signless integer attribute

Operands:

Operand	Description
input	tensor of 32-bit float or 32-bit signless integer values

Results:

Result	Description
output	tensor of 32-bit signless integer values

lq.Bconv2d (TF::Bconv2dOp)

Computes a 2-D binary convolution by binarizing and bitpacking the input and filter.

TODO

Attributes:

Attribute	MLIR Type	Description
channels_in	IntegerAttr	32-bit signless integer attribute
dilation_height_factor	IntegerAttr	32-bit signless integer attribute
dilation_width_factor	IntegerAttr	32-bit signless integer attribute
fused_activation_function	StringAttr	fused activation enum
pad_values	IntegerAttr	32-bit signless integer attribute
padding	StringAttr	padding enum
stride_height	IntegerAttr	32-bit signless integer attribute
stride_width	IntegerAttr	32-bit signless integer attribute

Operands:

Operand	Description
input	tensor of 32-bit float or 32-bit signless integer or QI8 type values
filter	tensor of 32-bit float or 32-bit signless integer values
post_activation_multiplier	tensor of 32-bit float values
post_activation_bias	tensor of 32-bit float values

Results:

Result	Description
output	tensor of 32-bit float or 32-bit signless integer or QI8 type values

lq.Bsign (TF::BsignOp)

Returns an element-wise indication of the binary sign of a number.

y = sign(x) = -1 if x < 0; 1 if x >= 0.

Operands:

Operand	Description
x	tensor of bfloat16 type or 16-bit float or 32-bit float or 64-bit float or 32-bit signless integer or 64-bit signless integer or complex type with 64-bit float elements or complex type with 32-bit float elements values

Results:

Result	Description
y	tensor of bfloat16 type or 16-bit float or 32-bit float or 64-bit float or 32-bit signless integer or 64-bit signless integer or complex type with 64-bit float elements or complex type with 32-bit float elements values

[lgeiger]

The correctness of this definition was actually never checked before.

[lgeiger]

These passes where only needed due to the need to remove dead ops.

[lgeiger]

Not sure if there is a nice way to match all ops in a dialect

[AdamHillier]

Looks good to me 👍