Transform Larq dialect ops to Xcore

[panickal-xmos]

Hi,

I'm working on a project to transform Larq dialect ops to ones in Xcore dialect to be deployed on XMOS xcore architecture.

There are two concerns I wanted to discuss -

• There is a small patch we need to apply to be able to include and build with the compute-engine repo. This is also so that the lce_ops.td file can be included from the XMOS project. Here's the patch, https://github.com/xmos/ai_tools/blob/5f7412fe0e71e0f904dd539773f15fd51b54b899/experimental/xformer/patches/build_fix.patch
  It would be good if this patch can be included in the compute-engine repo. I'm happy to put up a PR for this. Similar changes might be needed when the compute-engine repo is moved up to Tensorflow 2.7.

• The second concern is a more general question on transforming the Larq dialect ops. Currently, we are using the lce python interface to convert a Keras model to a tflite model containing Larq ops such as LceQuantize and LceBconv2d. We would then load this model and do the necessary graph transformations to the Xcore dialect in our MLIR transformer.
  The issue is that by the time we load the model, the Larq ops have already been converted to the TFL custom ops.
  We would have to convert the custom ops back to ones in the Larq dialect by parsing the flexbuffer data and replacing them with new Larq ops. This seems quite error-prone since there is no way to ensure type-safety using flexbuffers as the medium if there are changes to the ops.
  I'm interested in knowing if there is a better way of doing this. Is there a way to avoid this conversion so that the ops remain in the Larq dialect? Would we have to add the Larq ops to the TFL flatbuffer schema?

Thanks for your time!

[lgeiger]

It would be good if this patch can be included in the compute-engine repo. I'm happy to put up a PR for this.

This sounds good. We'd be very happy to accept a PR adding this build target.

Is there a way to avoid this conversion so that the ops remain in the Larq dialect?

Unfortunately I think in the current setup there is no easy way of avoiding this conversion.
I think in order to avoid the conversion to custom TFLite ops the only way would be to add the larq ops to the TFlite flatbuffer schema. But this comes with a few technical problems: Not only would the converted flatbuffer models not be compatible with the default TFLite runtime anymore which means we would need to modify that TFLite runtime itself, but we would also need to write custom serialisers and deserialisers for the larq ops.
Another alternative would be to dump the MLIR IR directly to file in case one wants to consume it in future conversion passes, but I am not sure if this would work in general.
I guess as a compromise we could add a pass which converts the custom ops back to ops in the Larq dialect. That way we could still have a stable serialisation format that doesn't require modifications to the TFLite runtime or schema, but further downstream conversion infrastructure could rely on this pass to de-serialise the larq ops.
@panickal-xmos I haven't put much thought into ways to properly serialise larq ops yet, but I'm curious what you think.

[panickal-xmos]

Hi @lgeiger, thank you for your reply. I think your idea of adding a pass to convert the custom ops back to Larq dialect makes most sense. It's cleaner than dumping to MLIR or changing the schema. Shall I put up PRs for both the build target and custom op conversion?

[lgeiger]

Shall I put up PRs for both the build target and custom op conversion?

That would be great!

[panickal-xmos]

@lgeiger , I have put up the first PR here, #694 . Please let me know if there are changes to be made. Thanks!

[panickal-xmos]

@lgeiger I have put up an initial PR for the translation pass, #698 . Did you have something like that in mind? Wanted to get your thoughts on it and then can adapt/clean up.